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

Merge remote-tracking branch 'origin/main' into woa-nvidia-wsl-fallback

Browse source
This commit is contained in:
Daniel Han 2026-06-09 19:11:51 -07:00
commit 53de77b007
686 changed files with 57347 additions and 17793 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -990,6 +990,7 @@ 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",

32
.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-IQ2_XXS (~570 MiB) + mmproj-F16 (~780 MiB).
# response_format JSON-schema decoding and OpenAI image_url
# (data URI) plus Anthropic source/base64 image inputs.
#
@ -791,9 +791,9 @@ 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
GGUF_VARIANT: UD-IQ2_XXS
GGUF_FILE: Qwen3-VL-2B-Instruct-UD-IQ2_XXS.gguf
MMPROJ_FILE: mmproj-F16.gguf
STUDIO_PORT: '18890'
HF_HOME: ${{ github.workspace }}/hf-cache
@ -888,13 +888,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:

10
.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
@ -843,9 +843,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
@ -1123,7 +1123,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.

2
.pre-commit-config.yaml
View file

@ -1,6 +1,6 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.15
rev: v0.15.16
hooks:
- id: ruff
args:

329
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,17 +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.
# 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-syntactic import patterns (static/dynamic/require/re-export)
# could be a real module reference. Markdown gets a separate gate (.mdx is
# 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 = [
@ -102,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"
@ -153,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])
@ -168,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 {})
@ -185,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()
@ -215,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
@ -241,30 +227,22 @@ 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")`.
@ -273,17 +251,15 @@ def classify(pkg: str, file: str, content: str) -> str | None:
# require / require.resolve
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):
return "html_script"
@ -295,8 +271,7 @@ def classify(pkg: str, file: str, content: str) -> str | None:
# 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,
@ -309,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
@ -350,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())
@ -361,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)
@ -386,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
@ -429,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] = []
@ -453,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)
@ -476,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
@ -497,42 +453,29 @@ 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
@ -540,21 +483,17 @@ def _next_real_bin(words: list[str], idx: int) -> str | None:
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.
# 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("-")
@ -572,18 +511,12 @@ def _next_real_bin(words: list[str], idx: int) -> str | 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.
"""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
@ -599,7 +532,7 @@ def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, li
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
@ -613,11 +546,8 @@ def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, li
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"):
@ -635,26 +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.
# `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 {}
@ -683,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
@ -707,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*)?|"
@ -754,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
@ -762,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 {
@ -807,11 +711,10 @@ def _read_file(path: str) -> list[str]:
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 = []
@ -822,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)
@ -841,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:]]+"
@ -894,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/") :]
@ -908,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 []
@ -997,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 {}
@ -1011,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)
@ -1074,12 +961,9 @@ def main() -> int:
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)
@ -1091,9 +975,8 @@ 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(
@ -1106,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))
@ -1115,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:

77
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
@ -68,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:
@ -91,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")
@ -118,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
@ -133,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
@ -155,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:
@ -190,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]:
@ -203,7 +161,6 @@ 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>"
scripts = _fetch_registry_scripts(name, version)
if scripts:
@ -226,9 +183,7 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
return findings
# ─────────────────────────────────────────────────────────────────────
# CLI.
# ─────────────────────────────────────────────────────────────────────
def main(argv: list[str] | None = None) -> int:

116
scripts/enforce_kwargs_spacing.py
View file

@ -20,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:
@ -142,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):
@ -164,15 +159,12 @@ def remove_redundant_passes(text: str) -> tuple[str, 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.
"""Drop blank line(s) after an import block in a small nested suite.
Inside an indented suite of <= 3 statements (function/try/if/with/etc., never
module level), when a run of consecutive ``import`` / ``from ... import``
statements is directly followed -- across one or more blank lines and nothing
else -- by another statement in the same suite, remove those blank lines so
the import sits next to the code that uses it. A comment in the gap blocks the
rule. Removing blank lines never changes the AST, so this is always
semantics-preserving.
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)
@ -226,12 +218,11 @@ _DEF_MIN_PARAMS_FOR_MULTILINE = 3 # signatures with < this many params stay one
def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]:
"""Map the line of each def keyword to (param count, has-any-default).
"""Map each def keyword line to (param count, has-any-default).
One def per line, so the line is a stable key. ``*`` / ``/`` markers are not
parameters and are not counted. A default exists if any positional default
is present or any keyword-only default is not ``None`` (a ``None`` entry in
``kw_defaults`` means a required keyword-only arg).
``*`` / ``/`` 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):
@ -250,20 +241,12 @@ def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]:
def normalize_def_trailing_comma(text: str) -> tuple[str, bool]:
"""Force a def / async-def signature one-per-line iff it has >= 3 parameters
AND at least one default value; otherwise keep it collapsible.
"""Force a def signature one-per-line iff >= 3 params AND a default; else collapsible.
Rationale: signatures with defaults read better one parameter per line, but
only once they are non-trivial (< 3 params always stay on one line). A
signature with >= 3 params and a default gets a magic trailing comma added
(ruff then wraps it one-per-line regardless of length); every other
signature has its trailing comma stripped so ruff collapses it onto one line
when it fits (and wraps a genuinely long one by length alone).
Function-definition parameter lists only, never call sites or collection
literals. Parameter counts and defaults come from the AST. Run BEFORE ruff
format. Adding or removing a def trailing comma never changes the AST, which
is re-checked before returning.
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)
@ -332,12 +315,10 @@ def normalize_def_trailing_comma(text: str) -> tuple[str, bool]:
def _split_string_token(s: str) -> tuple[str, str, str] | None:
"""Split a string literal's source into (prefix, quote, body).
"""Split a string literal source into (prefix, quote, body).
``prefix`` is the letters before the opening quote (``r``/``f``/``b``/``u``
in any case/order), ``quote`` is the opening delimiter (``'``, ``"``,
``'''`` or ``\"\"\"``) and ``body`` is everything between the delimiters.
Returns ``None`` if ``s`` is not a recognizable string literal.
``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 ("'", '"'):
@ -393,15 +374,12 @@ def _string_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 a list of ``(kind, raw_source)`` where kind is ``"str"`` or
``"f"``. Rules: bytes are left side-by-side (return ``None``); a run with no
f-string merges plain/raw/unicode pieces sharing one prefix+quote by simple
body concatenation; a run mixing an f-string with at least one plain string
(and no bytes, no raw) folds into a single f-string -- f pieces keep their
bodies verbatim and plain pieces have their braces escaped (``{`` -> ``{{``).
Runs of only f-strings are left side-by-side. The caller re-checks the file
AST and drops the change if it differs, so any subtle case (e.g. ``\\N{...}``)
that this would mis-handle is caught and skipped.
``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:
@ -455,13 +433,10 @@ def _fold_collapses(
) -> bool:
"""Whether an f-string fold at ``row[c0:c1]`` -> ``merged`` is safe to apply.
Only ``assert`` statements wrap awkwardly when a message is folded: ruff
parenthesizes the *condition* once ``assert cond, msg`` no longer fits on one
line. For every other construct (call argument, ``raise``, assignment, ...) a
folded long message wraps acceptably, so the fold is always allowed. For an
``assert`` the fold is allowed only when the statement is already one physical
line, or its estimated one-line length after folding fits the line length;
otherwise the message is left side-by-side.
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):
@ -483,15 +458,12 @@ def _fold_collapses(
def merge_adjacent_string_literals(text: str) -> tuple[str, bool]:
"""Merge a run of adjacent string literals on ONE physical line into a single
literal (the ``"a" "b"`` form ruff emits when it collapses an implicit
concatenation). Plain/raw/unicode runs merge by concatenation; a run mixing
an f-string with a plain string folds into one f-string (plain parts' braces
escaped) -- but only when the statement still fits on one line, so a long
message is left side-by-side rather than forcing the statement to re-wrap.
Runs of only f-strings, and bytes, are left side-by-side. The whole file's AST
is re-checked and the change is dropped if it would differ, so the transform
can never change meaning.
"""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))
@ -552,17 +524,11 @@ def merge_adjacent_string_literals(text: str) -> tuple[str, bool]:
def collapse_short_asserts(text: str) -> tuple[str, bool]:
"""Collapse a multi-line ``assert`` onto one line when it would fit.
An ``assert`` is often kept multi-line only by a magic trailing comma inside
a collection / call / tuple-message (``assert x == {\"a\": 1,}`` written
across lines). When the whole statement's estimated one-line length fits the
line length, strip those trailing commas (the comma before a ``)`` / ``]`` /
``}``) so ruff joins it back onto one line on the following format pass.
Run BEFORE ruff format. Skips any assert that contains a comment (a comment
forces ruff to keep it multi-line, which would oscillate). Stripping a
trailing comma is non-semantic except for a one-element tuple ``(x,)``; the
file AST is re-checked and any assert whose strip would change it is left
alone.
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)

40
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

211
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.
@ -545,12 +463,10 @@ def audit_npm_lockfile(path: Path) -> list[Finding]:
)
)
# 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(
@ -575,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*$")
@ -706,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",
@ -779,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()
@ -809,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 ())]
@ -835,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]
@ -854,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)
@ -876,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(

69
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):
@ -308,21 +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("-", "_")
# 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)
@ -349,11 +313,7 @@ Examples:
)
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.
# Default True for backwards compat; pass --no-allow-shell for untrusted notebooks.
parser.add_argument(
"--allow-shell",
dest = "allow_shell",
@ -371,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)

138
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,12 +78,10 @@ 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",
@ -145,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")
@ -198,11 +192,8 @@ def install_cells(nb: dict[str, Any]) -> list[tuple[int, str]]:
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
@ -466,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()
@ -543,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",
@ -575,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:
@ -652,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:
@ -724,11 +703,9 @@ def rule_inst_006_double_bang(install_cell: str, file: str, cell_idx: int) -> li
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
@ -736,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 (
@ -799,13 +772,9 @@ 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.
"""
"""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)
@ -868,14 +837,9 @@ def cmd_drift(args: argparse.Namespace) -> int:
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:
@ -920,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,
@ -1004,19 +967,16 @@ 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
@ -1147,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()
@ -1185,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():

6
scripts/run_ruff_format.py
View file

@ -19,10 +19,8 @@ def main(argv: list[str]) -> int:
spacing_script = HERE / "enforce_kwargs_spacing.py"
# Pre-ruff: normalize def-signature trailing commas (>=3 params with a
# default -> one-per-line; everything else collapsible) and strip the magic
# comma from a short multi-line assert, so ruff wraps signatures accordingly
# and joins the assert back onto one line.
# 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:

215
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}")
@ -873,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
@ -889,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 (
@ -903,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"
@ -919,10 +861,8 @@ 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
@ -941,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)
@ -1008,10 +947,8 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
),
)
)
# 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(
@ -1071,20 +1008,12 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
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(
@ -1127,8 +1056,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
)
)
# 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(
@ -1145,8 +1073,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
)
)
# 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(
@ -1164,11 +1092,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
)
)
# 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):
@ -1211,9 +1136,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
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",
@ -1241,12 +1165,9 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]:
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
@ -1288,12 +1209,12 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]:
def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | None]:
"""Download + extract + scan a single package. Cleans up its dir.
"""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)

261
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(
@ -312,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"
@ -326,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,}"
@ -342,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"
@ -363,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*\("
@ -374,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"
@ -388,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"
@ -400,9 +384,6 @@ RE_SHELL_DROPPER = re.compile(
)
# ---------------------------------------------------------------------------
# Finding dataclass
# ---------------------------------------------------------------------------
@dataclass
class Finding:
severity: str
@ -412,9 +393,7 @@ class Finding:
evidence: str = ""
# ---------------------------------------------------------------------------
# Checkers
# ---------------------------------------------------------------------------
def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
@ -425,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
@ -470,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(
@ -483,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:
@ -521,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))
@ -546,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:
@ -740,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:
@ -879,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:
@ -978,23 +951,18 @@ def _extract_evidence(
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))
@ -1119,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(
@ -1166,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
@ -1183,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}"
@ -1197,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)
@ -1225,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} "
@ -1262,9 +1220,8 @@ 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}",
@ -1305,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(
@ -1327,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:
@ -1342,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,
@ -1370,9 +1319,7 @@ 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_.\-+!]+)")
@ -1382,10 +1329,8 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find
"""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] = []
@ -1407,7 +1352,7 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find
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
@ -1416,14 +1361,11 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find
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"
@ -1433,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",
@ -1445,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._-]")
@ -1459,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,
@ -1495,7 +1428,7 @@ 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:
@ -1511,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)
@ -1552,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):
@ -1560,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])?)")
@ -1591,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:
@ -1624,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)]
@ -1634,9 +1562,7 @@ def get_downloaded_version(archive_path: str) -> str | None:
return None
# ---------------------------------------------------------------------------
# Display
# ---------------------------------------------------------------------------
def severity_color(sev: str) -> str:
@ -1652,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}")
@ -1682,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:
@ -1708,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
@ -1732,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)
@ -1772,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):
@ -1820,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:
@ -1850,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}",
@ -1870,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()
@ -1883,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 "."
@ -1909,7 +1822,7 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> 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(".", "_")
@ -1931,8 +1844,7 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
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"]
@ -1947,7 +1859,6 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
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:
@ -1986,7 +1897,6 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
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}")
@ -1995,9 +1905,7 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
print(f"\n Re-run without --fix to verify the scan is clean.")
# ---------------------------------------------------------------------------
# Directory scanning
# ---------------------------------------------------------------------------
def _find_requirements_files(root: str) -> list[str]:
@ -2014,7 +1922,7 @@ 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
@ -2024,18 +1932,15 @@ def _find_requirements_files(root: str) -> list[str]:
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:
@ -2134,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)
@ -2172,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"

7
scripts/stamp_studio_release.py
View file

@ -22,11 +22,8 @@ def _atomic_write_text(
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``.
"""
"""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)

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

64
scripts/verify_import_hoist.py
View file

@ -122,16 +122,13 @@ class _Builder(ast.NodeVisitor):
def __init__(self):
self.module = Scope("module", "<module>", None)
self.uses: list[tuple[Scope, str, int]] = [] # (scope, name, lineno) hard loads
self.soft_uses: list[tuple[Scope, str, int]] = [] # annotations: count as "used"
# but never as "unresolved"
# (forward refs / string annos)
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:
"""Annotation context: with `from __future__ import annotations` these are
never evaluated (strings), and even otherwise they routinely contain forward
references. Record contained names as SOFT uses so an import used only in an
annotation still counts as used, but a forward-ref name is never 'unresolved'."""
"""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):
@ -189,7 +186,7 @@ class _Builder(ast.NodeVisitor):
child = Scope("function", f"{scope.qualname}.{node.name}", scope)
self._bind_type_params(node, child)
self._bind_args(node.args, child)
# arg + return annotations: soft uses (may be strings / forward refs)
# 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)
@ -355,7 +352,7 @@ class _Builder(ast.NodeVisitor):
if isinstance(node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)):
child = Scope("comp", f"{scope.qualname}.<comp>", scope)
for i, gen in enumerate(node.generators):
# first iterable is evaluated in the enclosing scope
# 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:
@ -394,9 +391,8 @@ def _any_star(scope: Scope) -> bool:
def _resolve(scope: Scope, name: str):
"""LEGB resolution. Returns (status, bindings) where status in
"""LEGB resolution. Returns (status, bindings); status in
{'local','import','other','builtin','star','unresolved'}."""
# global / nonlocal redirection
start = scope
if name in scope.globals:
chain = [_module_of(scope)]
@ -441,7 +437,7 @@ def _legb_chain(scope: Scope) -> list[Scope]:
chain = [scope]
p = scope.parent
while p is not None:
if p.kind != "class" or p.parent is None: # module-level class never happens; keep module
if p.kind != "class" or p.parent is None: # skip class scopes, keep module
if p.kind != "class":
chain.append(p)
p = p.parent
@ -455,7 +451,7 @@ def _analyze(src: str):
tree = ast.parse(src)
b = _Builder()
b.run(tree)
# Per-scope: unresolved load names, and import targets it resolves to.
# 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]] = {}
@ -467,7 +463,7 @@ def _analyze(src: str):
tids = {bd.target for bd in binds if bd.target}
targets_by_scope.setdefault(scope.qualname, set()).update(tids)
target_by_use.setdefault((scope.qualname, name), set()).update(tids)
# soft uses (annotations): only contribute to "used", never to "unresolved"
# 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":
@ -495,7 +491,7 @@ def _analyze(src: str):
x.kind not in ("import", "importfrom") for x in bs
):
ambiguous.setdefault(scope.qualname, set()).add(n)
# scope tree isn't stored; rebuild via uses is hard. We approximate with module only.
# scope tree isn't stored; approximate with module only.
walk_scopes(module)
return {
@ -524,14 +520,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
Blocker signals (precise, no relocation false-positives):
UNRESOLVED-NEW - a load became undefined (dangling alias / removed import).
NEW-UNUSED-HOIST - a module-level import added by THIS change is resolved by
NO load. A correct hoist always wires its new import to a
reference; if the alias was left un-normalized OR renamed
to the wrong name, the hoisted import ends up unused. This
single signal catches BOTH user-described failure modes and
does NOT fire for code merely relocated to another file
(that removes the import, it doesn't add an unused one).
TARGET-CHANGED - the same (scope, name) load resolves to a different import
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)
@ -566,14 +557,13 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
)
)
# 2. HOISTED-IMPORT-UNUSED (the core botched-hoist / wrong-rename signal)
# A module-level import in AFTER that NO load resolves to, and which was
# either newly added by this change OR was actually used before. Excludes:
# - relocation (the import is REMOVED, so it's not in after at all)
# - stable pre-existing re-exports (unused before AND after, not newly added)
# 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 by something -> fine
continue # resolved -> fine
newly_added = bool(tids - before_module_targets)
was_used_before = bool(tids & before_used)
if newly_added or was_used_before:
@ -619,8 +609,7 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
findings.append(("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}"))
# 6. TARGET-MISSING (informational): a scope stopped resolving to an import
# target. Real bugs are already covered above; remaining cases are code
# relocated to another file (e.g. a moved helper). Shown for transparency.
# 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):
@ -667,13 +656,13 @@ _SELF_TESTS = {
"BLOCKER",
),
"local_var_clash": (
# _b renamed to b, but b is a LOCAL variable in f -> import silently unused
# _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 a config_copy var exists: NO false positive
# 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"
@ -728,10 +717,9 @@ def _pyflakes_undefined(path: str) -> set[str] | None:
def audit_files(paths: list[str]) -> int:
"""Single-version robustness audit. For every file: confirm the analyzer does
not crash, then cross-check its 'unresolved' names against pyflakes. Any name
the resolver flags that pyflakes does NOT call undefined is a tool FALSE
POSITIVE (a resolver gap to fix)."""
"""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] = {}

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

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,

20
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)
@ -112,8 +112,8 @@ def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[st
"""
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:
@ -128,7 +128,7 @@ def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[st
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.
"""
@ -156,15 +156,7 @@ async def get_current_subject_allow_password_change(
async def _get_current_subject(
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()

82
studio/backend/auth/storage.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
"""
SQLite storage for authentication data (user credentials + JWT secret).
"""
"""SQLite storage for auth data (user credentials + JWT secret)."""
import hashlib
import os
@ -17,42 +15,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 +84,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()
@ -176,22 +163,19 @@ def get_connection() -> sqlite3.Connection:
# ── 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:
@ -233,22 +217,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(
@ -513,7 +485,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(),),

57
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,11 +23,10 @@ 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 and validates that the result is a real HTTPS Colab
URL before returning. Falls back to http://localhost:{port} only when all
attempts fail.
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
@ -43,7 +40,7 @@ def get_colab_url(port: int = 8888) -> str:
for attempt in range(3):
try:
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10)
# A valid Colab proxy URL starts with https:// and embeds the port.
# 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:
@ -61,16 +58,13 @@ def get_colab_url(port: int = 8888) -> str:
def show_link(port: int = 8888, *, _url: "str | None" = None):
"""Display a styled clickable link to the UI.
*_url* is an optional pre-fetched Colab proxy URL. When omitted,
``get_colab_url(port)`` is called internally. Pass it from
``_show_and_embed`` to avoid a second ``eval_js`` round-trip.
*_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip.
"""
from IPython.display import display, HTML
url = _url if _url is not None else get_colab_url(port)
# Build a truncated display URL. Wrap in try/except so an unexpected URL
# shape never prevents the link from rendering.
# Truncated display URL; try/except so an odd URL shape still renders the link.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
@ -79,8 +73,7 @@ def show_link(port: int = 8888, *, _url: "str | None" = None):
except (ValueError, IndexError):
short_url = url
# Also emit a plain-text line so the URL is visible even if HTML display
# is suppressed or fails.
# Plain-text line so the URL shows even if HTML display fails.
logger.info(f"🌐 Unsloth Studio URL: {url}")
html = f"""
@ -123,12 +116,8 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
def _show_and_embed(port: int):
"""Embed the Studio inline for *port* with a branded header bar.
Fetches the Colab proxy URL once (registering the port with Colab's
reverse-proxy at the same time) then renders a header bar + full-height
iframe as a single HTML block.
Falls back to ``serve_kernel_port_as_iframe`` if ``IPython.display.HTML``
is unavailable for any reason.
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}")
@ -138,7 +127,7 @@ def _show_and_embed(port: int):
iframe_id = f"unsloth-studio-{port}"
# Truncated URL shown in the header — best-effort, falls back to full URL.
# Truncated header URL — best-effort, falls back to full URL.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
@ -167,7 +156,7 @@ def _show_and_embed(port: int):
""")
)
except Exception:
# Fallback: Colab's built-in helper (less control, but always works)
# 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%")
@ -187,9 +176,8 @@ def start(port: int = 8888):
logger.info("🦥 Starting Unsloth Studio...")
# --- Fast path: Studio is already running (cell re-run) ---
# Re-launching would either collide on the port or silently shift to a new
# port and confuse the user. Just re-show the link and iframe instead.
# 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)
@ -222,16 +210,14 @@ def start(port: int = 8888):
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
# run_server auto-increments the port when the requested one is already in
# use (e.g. Jupyter occupying 8888). Read back the actual bound port so the
# Colab proxy URL and iframe always point at the right place.
# 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
logger.info(f" Server started on port {actual_port}!")
# Poll health endpoint to confirm the server is truly reachable before
# showing the link and registering the iframe — avoids the race where
# ready_event fires but the process hasn't finished binding.
# 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
@ -252,9 +238,8 @@ def start(port: int = 8888):
_show_and_embed(actual_port)
# Keep kernel alive so the daemon server thread stays running.
# Handle KeyboardInterrupt cleanly so the user gets a readable message
# rather than a raw traceback when they interrupt the cell.
# 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)

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

45
studio/backend/core/_torchao_stub.py
View file

@ -3,18 +3,11 @@
"""Shared torchao Windows-ROCm import stub.
torchao (pulled in by transformers.quantizers) imports
torch.distributed._functional_collectives at module level, which imports
distributed_c10d.py unconditionally — that file crashes on Windows ROCm because
torch._C._distributed_c10d (the RCCL backend) is absent.
torch/distributed/__init__.py itself is guarded by `if is_available()` so
`import torch.distributed` alone is safe; the crash only comes via torchao's
import chain. Stubbing torchao short-circuits it entirely.
_StubSubpackageFinder handles any depth of torchao.xxx.yyy imports.
This logic used to be duplicated inline inside run_export_process() and
run_training_process(); it now lives here so both worker subprocesses call the
single `install_torchao_windows_rocm_stub()` entrypoint before importing
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.
"""
@ -28,12 +21,9 @@ import importlib.machinery
_STUB_SENTINEL = object()
# Metaclass for stub types so that isinstance(x, StubClass) returns False
# instead of raising TypeError ("arg 2 must be a type").
# peft/tuners/lora/torchao.py does:
# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor
# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor))
# If those names resolve to stub modules rather than types, isinstance() raises.
# 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
@ -71,8 +61,7 @@ def _make_mod_stub(mod_name):
):
if attr.startswith("__"):
raise AttributeError(attr)
# Return a stub CLASS (not a module) so that isinstance(x, attr)
# works and returns False instead of raising TypeError.
# 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
@ -114,15 +103,12 @@ class _StubSubpackageFinder(importlib.abc.MetaPathFinder):
def install_torchao_windows_rocm_stub() -> None:
"""Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash.
No-op on every other platform (Windows CUDA included — there torchao is real
and shadowing it would break torchao-based quantization paths). Must run
before any import of transformers / unsloth_zoo. Safe to call once per worker
process.
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 stay set after a user installs the HIP SDK and reverts to a
# CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip
# but still encode "rocm" in torch.__version__, so accept either.
# 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:
@ -135,8 +121,7 @@ def install_torchao_windows_rocm_stub() -> None:
except Exception:
pass
if _is_win32_rocm:
# Register the finder only on Windows ROCm -- on other platforms there
# are no stub modules seeded, so appending is a pure accumulation.
# 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 (

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

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

@ -95,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(

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

@ -132,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
@ -202,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
@ -537,15 +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

4
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"
)
@ -327,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:

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

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

@ -73,13 +73,10 @@ def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Pat
def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) -> None:
"""
Subprocess entrypoint.
Sends events to `event_queue`.
"""
"""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
@ -115,8 +112,8 @@ def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any])
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 (

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

@ -184,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:

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

@ -168,8 +168,7 @@ 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 only build them when this host allows it (desktop / explicit opt-in).
# Skip them otherwise so a recipe carried onto a hosted host cannot spawn.
# 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()
@ -258,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
@ -282,9 +281,8 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None =
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 # pyright: ignore[reportMissingImports]
model_providers = [

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

79
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
@ -46,10 +43,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
@ -110,18 +105,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")
@ -137,8 +129,7 @@ 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)
@ -159,12 +150,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():
@ -174,11 +164,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
@ -246,7 +234,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...")
@ -258,14 +246,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
@ -349,7 +335,6 @@ class ExportBackend:
else:
save_method = "merged_16bit"
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
logger.info(f"Saving merged model locally to: {save_directory}")
@ -370,7 +355,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 (
@ -451,7 +435,6 @@ class ExportBackend:
output_path: Optional[str] = None
try:
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
logger.info(f"Saving base model locally to: {save_directory}")
@ -459,7 +442,7 @@ class ExportBackend:
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,
@ -474,7 +457,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 (
@ -509,12 +491,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"
)
# Create repo
hf_api = HfApi(token = hf_token)
repo_id = PushToHubMixin._create_repo(
PushToHubMixin,
@ -524,7 +505,6 @@ class ExportBackend:
)
username = repo_id.split("/")[0]
# Create and push model card
content = MODEL_CARD.format(
username = username,
base_model = base_model,
@ -535,7 +515,6 @@ class ExportBackend:
card = ModelCard(content)
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,
@ -585,13 +564,11 @@ class ExportBackend:
output_path: 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 (
@ -610,30 +587,23 @@ 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.
# Absolute path so unsloth's relative-path internals resolve
# against the repo root cwd, not the export directory.
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")
self.current_model.save_pretrained_gguf(
model_save_path,
@ -641,18 +611,14 @@ class ExportBackend:
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".
# 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}/")
# 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 any .gguf from subdirs (e.g. model_gguf/) into abs_save_dir.
for sub in list(Path(abs_save_dir).iterdir()):
if not sub.is_dir():
continue
@ -660,13 +626,11 @@ class ExportBackend:
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"
@ -686,8 +650,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",
@ -696,7 +658,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 (
@ -750,7 +711,6 @@ class ExportBackend:
output_path: Optional[str] = None
try:
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
logger.info(f"Saving LoRA adapter locally to: {save_directory}")
@ -766,7 +726,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 (

118
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:
@ -268,9 +227,8 @@ class ExportOrchestrator:
) -> 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
@ -279,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
@ -294,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(
{
@ -315,7 +269,7 @@ 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,
@ -362,12 +316,11 @@ 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:
@ -486,14 +439,10 @@ 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.
"""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():
@ -527,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
@ -542,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
@ -553,7 +501,7 @@ class ExportOrchestrator:
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."""
"""Scan for checkpoints — runs locally, no ML imports."""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir = outputs_dir)

109
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
@ -372,23 +350,16 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
"""
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
@ -419,7 +390,7 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
)
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
@ -432,10 +403,8 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
)
# ── 1c. Stub torchao on Windows ROCm ──
# Shared with the training worker; see core/_torchao_stub.py for the full
# rationale (torchao -> torch.distributed._functional_collectives crashes on
# Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm.
# Must run before any import of transformers / unsloth_zoo.
# 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()
@ -511,7 +480,7 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
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)

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)

150
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.
@ -46,9 +74,9 @@ def anthropic_messages_to_openai(
) -> 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] = []
@ -75,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:
@ -104,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] = []
@ -137,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})
@ -181,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
@ -208,19 +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(
@ -236,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,
},
},
},
)
@ -259,22 +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 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),
},
)
)
@ -317,19 +373,19 @@ class AnthropicStreamEmitter:
return self._tool_arguments_delta(args)
events = []
# Close current text block if open.
if self._text_block_open:
events.append(self._close_block())
# Defensive: if a replacement/different tool_start arrives while a
# tool_use block is open, close the stale block before starting another.
# 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(
@ -339,7 +395,7 @@ class AnthropicStreamEmitter:
"index": self.block_index,
"content_block": {
"type": "tool_use",
"id": tool_call_id,
"id": self._open_tool_use_id,
"name": event.get("tool_name", ""),
"input": {},
},
@ -374,7 +430,11 @@ class AnthropicStreamEmitter:
# 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(
@ -382,7 +442,7 @@ class AnthropicStreamEmitter:
"tool_result",
{
"type": "tool_result",
"tool_use_id": event.get("tool_call_id", ""),
"tool_use_id": tool_use_id,
"content": event.get("result", ""),
},
)
@ -421,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:
@ -433,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",
@ -448,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,
},
},
},
)
@ -498,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"
@ -541,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
@ -561,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),
},
)
)

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

@ -90,8 +90,8 @@ class AudioCodecManager:
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)
# 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):
@ -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"),
@ -178,13 +176,10 @@ class AudioCodecManager:
# ── Decoders ─────────────────────────────────────────────────
def decode_snac(self, generated_ids: torch.Tensor, device: str) -> Tuple[bytes, int]:
"""
Decode SNAC tokens (Orpheus) into WAV bytes.
"""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)
@ -192,7 +187,7 @@ 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
# 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]
@ -229,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).
"""
@ -256,8 +248,8 @@ class AudioCodecManager:
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]
@ -279,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).
"""

4
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

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)

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

File diff suppressed because it is too large Load diff

367
studio/backend/core/inference/inference.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
"""
Core inference backend - streamlined
"""
"""Core inference backend."""
from unsloth import FastLanguageModel, FastVisionModel
from unsloth.chat_templates import get_chat_template
@ -36,29 +34,15 @@ logger = get_logger(__name__)
class HarmonyTextStreamer:
"""Streaming text decoder for gpt-oss harmony channel protocol.
"""Streaming text decoder for the gpt-oss harmony channel protocol.
gpt-oss models emit multi-channel output using special tokens like
``<|channel|>analysis<|message|>...`` and ``<|channel|>final<|message|>...``.
A plain ``TextIteratorStreamer(skip_special_tokens=True)`` strips the special
tokens but leaves the channel names concatenated with content, producing
garbled output such as ``analysisWe need to respond...assistantfinalHello!``.
This streamer decodes with ``skip_special_tokens=False`` so the full
harmony markup is visible, then uses **stateful incremental** parsing
to emit properly-formatted text:
- ``<think>`` emitted once when the ``analysis`` channel is first seen
- Analysis content streamed incrementally
- ``</think>`` emitted once when the ``final`` channel is first seen
- Final content streamed incrementally
This avoids the delta-on-transformed bug where wrapping tags shift
position as content grows.
Implements the same ``put`` / ``end`` / iterator interface as
``TextIteratorStreamer`` so ``generate_stream`` can use it as a drop-in
replacement.
gpt-oss emits multi-channel output via ``<|channel|>analysis<|message|>...``
/ ``<|channel|>final<|message|>...``. Plain skip_special_tokens streaming
glues channel names to content. This decodes with skip_special_tokens=False
and parses statefully: emit ``<think>`` on first analysis, stream analysis,
emit ``</think>`` on first final, stream final. Tracking per-channel lengths
avoids the delta-on-transformed bug where wrapping tags shift position.
Same put/end/iterator interface as TextIteratorStreamer.
"""
import re as _re
@ -87,22 +71,20 @@ class HarmonyTextStreamer:
self._is_first_put: bool = True
self._stop: bool = False
# Stateful channel tracking — avoids delta-on-transformed bugs
# Stateful channel tracking avoids delta-on-transformed bugs
self._emitted_think_open: bool = False
self._emitted_think_close: bool = False
self._analysis_emitted: int = 0 # chars of analysis content emitted
self._final_emitted: int = 0 # chars of final content emitted
# ------------------------------------------------------------------
# put / end — called from the generation thread
# ------------------------------------------------------------------
def put(self, value):
"""Receive new token IDs from model.generate()."""
import torch
if isinstance(value, torch.Tensor):
# value shape: (batch, seq) — take first batch element
# shape (batch, seq) — take first batch element
ids = value[0].tolist() if value.dim() > 1 else value.tolist()
elif isinstance(value, (list, tuple)):
ids = list(value)
@ -110,7 +92,7 @@ class HarmonyTextStreamer:
ids = [value]
if self._is_first_put and self.skip_prompt:
# First call contains the full prompt; remember its length
# First call is the full prompt; remember its length.
self._prompt_len = len(ids)
self._token_ids = list(ids)
self._is_first_put = False
@ -118,20 +100,20 @@ class HarmonyTextStreamer:
self._token_ids.extend(ids)
# Decode only the generated part (after the prompt)
# Decode only the generated part (after the prompt).
gen_ids = self._token_ids[self._prompt_len :]
raw = self.tokenizer.decode(gen_ids, skip_special_tokens = False)
self._process_incremental(raw)
def end(self):
"""Signal generation is complete."""
# Final decode to capture any remaining content
# Final decode to capture remaining content.
gen_ids = self._token_ids[self._prompt_len :]
if gen_ids:
raw = self.tokenizer.decode(gen_ids, skip_special_tokens = False)
self._process_incremental(raw)
# Close any open think tags
# Close any open think tags.
if self._emitted_think_open and not self._emitted_think_close:
self._queue.put("</think>")
self._emitted_think_close = True
@ -139,9 +121,7 @@ class HarmonyTextStreamer:
self._stop = True
self._queue.put(None) # sentinel
# ------------------------------------------------------------------
# Iterator interface — consumed by the streaming loop
# ------------------------------------------------------------------
def __iter__(self):
return self
@ -159,34 +139,19 @@ class HarmonyTextStreamer:
raise StopIteration
return val
# ------------------------------------------------------------------
# Stateful incremental harmony protocol parsing
# ------------------------------------------------------------------
def _process_incremental(self, raw: str) -> None:
"""Parse harmony channels and emit deltas per-channel.
Instead of transforming the entire raw text and computing a string
delta (which breaks when wrapping ``<think>`` tags shift position),
this tracks per-channel content lengths and emits:
- ``<think>`` once when analysis channel first appears
- analysis content deltas (computed on channel content directly)
- ``</think>`` once when final channel first appears
- final content deltas
"""
# If raw contains <|channel|> but no complete channel+message pair yet,
# buffer silently — don't emit partial channel names as text.
"""Parse harmony channels and emit per-channel deltas (tracked by length, not whole-text diff)."""
# If raw has <|channel|> but no complete channel+message pair yet, buffer.
has_channel_token = "<|channel|>" in raw
matches = list(self._HARMONY_RE.finditer(raw))
if has_channel_token and not matches:
# Partial harmony markup still building — wait for more tokens
# Partial harmony markup still building — wait for more tokens.
return
if not has_channel_token and not matches:
# No harmony protocol at all — should not happen for gpt-oss
# but handle gracefully by not emitting anything
return
for m in matches:
@ -228,11 +193,9 @@ class InferenceBackend:
self.device = get_device().value
self._audio_codec_manager = AudioCodecManager()
# Thread safety — _generation_lock serializes model.generate() calls.
# Must be a regular Lock (NOT RLock) because in async FastAPI, multiple
# requests share the same event-loop thread, so RLock reentrancy lets
# concurrent compare-mode requests race on the GPU. The lock is
# acquired by the *background generation thread*, not the event-loop.
# _generation_lock serializes model.generate(). Plain Lock (NOT RLock):
# RLock reentrancy would let concurrent compare-mode requests race on
# the GPU. Acquired by the background generation thread, not the event-loop.
import threading
self._generation_lock = threading.Lock()
@ -242,7 +205,7 @@ class InferenceBackend:
@staticmethod
def _normalize_top_k(top_k: int) -> int:
# API supports -1 as "disable top-k"; transformers expects 0 to disable.
# API uses -1 to disable top-k; transformers uses 0.
return 0 if top_k < 0 else top_k
def load_model(
@ -255,9 +218,7 @@ class InferenceBackend:
trust_remote_code: bool = False,
gpu_ids: Optional[list[int]] = None,
) -> bool:
"""
Load any model: base, LoRA adapter, text, or vision.
"""
"""Load any model: base, LoRA adapter, text, or vision."""
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
if max_seq_length <= 0:
max_seq_length = 2048
@ -265,13 +226,13 @@ class InferenceBackend:
try:
model_name = config.identifier
# Check if already loaded
# Already loaded?
if model_name in self.models and self.models[model_name].get("model"):
logger.info(f"Model {model_name} already loaded")
self.active_model_name = model_name
return True
# Check if currently loading
# Currently loading?
if model_name in self.loading_models:
logger.info(f"Model {model_name} is already being loaded")
return False
@ -322,14 +283,13 @@ class InferenceBackend:
from unsloth import FastModel
if config.is_lora and config.base_model:
# LoRA adapter: load from local adapter path.
# base_model is e.g. /home/.../Spark-TTS-0.5B/LLM
# The BiCodec weights are in the parent dir (Spark-TTS-0.5B/).
# LoRA adapter: base_model is .../Spark-TTS-0.5B/LLM;
# BiCodec weights live in the parent dir.
base_path = config.base_model
if os.path.isdir(base_path):
abs_repo_path = os.path.abspath(os.path.dirname(base_path))
else:
# base_model is an HF ID — download it
# base_model is an HF ID — download it.
from huggingface_hub import snapshot_download
local_dir = base_path.split("/")[-1]
@ -348,7 +308,7 @@ class InferenceBackend:
trust_remote_code = trust_remote_code,
)
else:
# Base model: download full HF repo, then load from /LLM subfolder
# Base model: download full HF repo, load from /LLM subfolder
from huggingface_hub import snapshot_download
hf_repo = config.path
@ -406,7 +366,7 @@ class InferenceBackend:
FastModel.for_inference(model)
model.eval()
# Create ASR pipeline (per notebook)
# ASR pipeline (per notebook)
from transformers import pipeline as hf_pipeline
whisper_pipe = hf_pipeline(
@ -435,8 +395,8 @@ class InferenceBackend:
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
# Load the external codec for TTS audio types
# (Whisper is ASR, audio_vlm is audio input — neither needs a codec)
# Load external codec for TTS audio types
# (Whisper is ASR, audio_vlm is audio input — neither needs one)
if audio_type not in ("whisper", "audio_vlm"):
model_repo_path = self.models[model_name].get("model_repo_path")
self._audio_codec_manager.load_codec(
@ -457,7 +417,7 @@ class InferenceBackend:
logger.info(f"Loading {model_type} model{adapter_info}: {model_name}")
log_gpu_memory(f"Before loading {model_name}")
# Load model - same approach for base models and LoRA adapters
# Same load path for base models and LoRA adapters
if config.is_vision:
# Vision model (or vision LoRA adapter)
model, processor = FastVisionModel.from_pretrained(
@ -470,19 +430,16 @@ class InferenceBackend:
trust_remote_code = trust_remote_code,
)
# Apply inference optimization
FastVisionModel.for_inference(model)
# FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast)
# instead of a proper Processor for some models (e.g. Gemma-3).
# In that case, load the real processor from the base model.
# FastVisionModel may return a raw tokenizer instead of a
# Processor for some models (e.g. Gemma-3); load the real one.
from transformers import ProcessorMixin
if not (
isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")
):
# For LoRA adapters, use the base model. For local merged exports,
# read export_metadata.json to find the original base model.
# LoRA adapters: use base model. Local merged exports: read base from export_metadata.json.
processor_source = config.base_model if config.is_lora else config.identifier
if not config.is_lora and config.is_local:
_meta_path = Path(config.path) / "export_metadata.json"
@ -522,7 +479,6 @@ class InferenceBackend:
trust_remote_code = trust_remote_code,
)
# Apply inference optimization
FastLanguageModel.for_inference(model)
self.models[model_name]["model"] = model
@ -530,7 +486,6 @@ class InferenceBackend:
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
# Load chat template info
self._load_chat_template_info(model_name)
self.active_model_name = model_name
@ -552,29 +507,24 @@ class InferenceBackend:
raise Exception(error_msg)
def unload_model(self, model_name: str) -> bool:
"""
Completely removes a model from the registry and clears GPU memory.
"""
"""Remove a model from the registry and clear GPU memory."""
if model_name in self.models:
try:
# If this was an audio model, clean up codecs
# Clean up codecs for audio models
if self.models[model_name].get("is_audio"):
self._audio_codec_manager.unload()
logger.info(f"Unloading model '{model_name}' from memory.")
# Delete the model entry from our registry
del self.models[model_name]
# Clear the active model if it was the one being unloaded
# Clear the active model if it was the one unloaded
if self.active_model_name == model_name:
self.active_model_name = None
# Clear GPU memory cache
clear_gpu_cache()
# Remove stale compiled cache so the next model gets a fresh one.
# On spawn-based platforms, preserve trainer files so that any
# concurrent training dataset.map() workers can still import them.
# Drop stale compiled cache for the next model. On spawn platforms,
# preserve trainer files so concurrent dataset.map() workers can import them.
import sys as _sys
from utils.cache_cleanup import clear_unsloth_compiled_cache
@ -593,27 +543,23 @@ class InferenceBackend:
return True
def revert_to_base_model(self, base_model_name: str) -> bool:
"""
Reverts the model to its pristine base state by unloading AND
deleting all adapter configurations, as instructed.
"""
"""Revert the model to its pristine base state by unloading and
deleting all adapter configurations."""
if base_model_name not in self.models:
return False
model = self.models[base_model_name].get("model")
try:
# Step 1: Unload the adapter weights if model is a PeftModel.
# Unload adapter weights if model is a PeftModel.
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(f"Unloading LoRA adapters from '{base_model_name}'...")
unwrapped_base_model = model.unload()
self.models[base_model_name]["model"] = unwrapped_base_model
model = unwrapped_base_model
# Step 2: Clear any lingering peft_config from the unwrapped model.
# After model.unload(), the base model may still carry a peft_config
# attribute. Removing it ensures PeftModel.from_pretrained() gets
# a clean base model without "multiple adapters" warnings.
# model.unload() can leave a peft_config; removing it avoids
# "multiple adapters" warnings on the next from_pretrained().
if hasattr(model, "peft_config"):
del model.peft_config
@ -636,10 +582,8 @@ class InferenceBackend:
hf_token: Optional[str] = None,
gpu_ids: Optional[list[int]] = None,
) -> Tuple[bool, Optional[str], Optional[str]]:
"""
Final Corrected Version:
Ensures the base model and the specified adapter are loaded.
This function is idempotent and handles all states correctly.
"""Ensure the base model and the given adapter are loaded.
Idempotent and handles all states correctly.
"""
try:
from utils.models import ModelConfig
@ -650,7 +594,7 @@ class InferenceBackend:
base_model_name = lora_config.base_model
# 1. Load the base model if it's not already in memory
# 1. Load the base model if not already in memory
if base_model_name not in self.models or not self.models[base_model_name].get("model"):
logger.info(f"Base model '{base_model_name}' not loaded, loading now.")
base_config = ModelConfig.from_ui_selection(base_model_name, None, is_lora = False)
@ -666,11 +610,11 @@ class InferenceBackend:
self.active_model_name = base_model_name
# 2. Determine the required adapter name from the user's selection
# 2. Derive adapter name from the user's selection
adapter_name = lora_path.split("/")[-1].replace(".", "_")
# 3. Call our robust load_adapter function to ensure this specific adapter is loaded.
# It will only load from disk if the model doesn't already have it.
# 3. Ensure this adapter is loaded (load_adapter only reads from
# disk if the model doesn't already have it).
adapter_success = self.load_adapter(
base_model_name = base_model_name,
adapter_path = lora_path,
@ -679,7 +623,7 @@ class InferenceBackend:
if not adapter_success:
return False, base_model_name, None
# 4. Return the correct, verified adapter name for the UI logic to use.
# 4. Return the verified adapter name for the UI.
return True, base_model_name, adapter_name
except Exception as e:
@ -690,12 +634,10 @@ class InferenceBackend:
return False, None, None
def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str) -> bool:
"""
Loads an adapter onto the model ONLY if it's not already attached.
"""
"""Load an adapter onto the model only if not already attached."""
model = self.models[base_model_name].get("model")
# Check if this adapter name is already part of the model's config. This is the most reliable check.
# Most reliable check: adapter name already in the model's config.
if hasattr(model, "peft_config") and adapter_name in model.peft_config:
logger.info(
f"Adapter '{adapter_name}' is already attached to the model. Skipping load."
@ -708,7 +650,7 @@ class InferenceBackend:
)
model.load_adapter(adapter_path, adapter_name = adapter_name)
# Update our internal registry ONLY after a successful load.
# Update the registry only after a successful load.
if "loaded_adapters" not in self.models[base_model_name]:
self.models[base_model_name]["loaded_adapters"] = {}
self.models[base_model_name]["loaded_adapters"][adapter_name] = adapter_path
@ -723,9 +665,7 @@ class InferenceBackend:
return False
def set_active_adapter(self, base_model_name: str, adapter_name: str) -> bool:
"""
Sets the active adapter for generation. This replaces the flawed 'enable_adapter'.
"""
"""Set the active adapter for generation."""
model = self.models[base_model_name].get("model")
try:
logger.info(f"Setting active adapter to: '{adapter_name}'")
@ -733,22 +673,16 @@ class InferenceBackend:
self.models[base_model_name]["active_adapter"] = adapter_name
return True
except Exception as e:
# This will catch the "adapter not found" error if something goes wrong.
# Catches "adapter not found" if something goes wrong.
logger.error(f"Failed to set active adapter to '{adapter_name}': {e}")
return False
def _apply_adapter_state(self, use_adapter: Optional[Union[bool, str]]) -> None:
"""
Apply adapter state before generation. Must be called under _generation_lock.
"""Apply adapter state before generation (must hold _generation_lock).
Uses PEFT's disable_adapter_layers() / enable_adapter_layers() which toggle
a boolean flag on each LoRA layer. Unsloth's fast_linear_forward checks this
flag (proj.disable_adapters) and skips LoRA computation when True.
This is non-destructive — no model unloading/reloading needed.
Args:
use_adapter: None = no change, False = disable (base model),
True = enable current adapter, str = enable specific adapter.
Toggles PEFT enable/disable_adapter_layers (non-destructive, no reload).
use_adapter: None = no change, False = base model, True = current adapter,
str = named adapter.
"""
if use_adapter is None:
return
@ -763,7 +697,7 @@ class InferenceBackend:
return
if use_adapter is False:
# Disable LoRA layers → base model output
# Disable LoRA layers -> base model output.
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(
f"Compare mode: disabling adapters on '{base}' for base model generation"
@ -773,7 +707,7 @@ class InferenceBackend:
logger.info(f"Compare mode: model '{base}' is not a PeftModel, already base")
elif use_adapter is True:
# Re-enable LoRA layers → adapter output
# Re-enable LoRA layers -> adapter output.
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(f"Compare mode: enabling adapters on '{base}' for LoRA generation")
model.base_model.enable_adapter_layers()
@ -781,7 +715,7 @@ class InferenceBackend:
logger.warning("use_adapter=true but model is not a PeftModel")
elif isinstance(use_adapter, str):
# Enable adapters and set the specific one active
# Enable adapters and set the named one active.
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(f"Compare mode: enabling adapter '{use_adapter}' on '{base}'")
model.base_model.enable_adapter_layers()
@ -795,17 +729,11 @@ class InferenceBackend:
cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
"""
Thread-safe generation with optional adapter toggling.
"""Thread-safe generation with optional adapter toggling.
The adapter toggle + model.generate() are serialized by _generation_lock
inside the background generation thread — NOT in the event-loop thread.
This prevents the RLock-reentrant race that occurs when two async SSE
handlers share the same event-loop thread.
Args:
use_adapter: Adapter control (None/False/True/str). See _apply_adapter_state.
**gen_kwargs: Forwarded to generate_chat_response.
Adapter toggle + model.generate() are serialized by _generation_lock in
the background thread, avoiding the RLock-reentrant race when two async
SSE handlers share one event-loop thread. use_adapter: see _apply_adapter_state.
"""
yield from self._generate_chat_response_inner(
cancel_event = cancel_event, _adapter_state = use_adapter, **gen_kwargs
@ -833,9 +761,8 @@ class InferenceBackend:
):
"""Run an agentic tool loop on top of ``generate_chat_response``.
Yields the same event-dict protocol used by the GGUF path so
the route layer can stream both backends through one helper.
Each event is one of:
Yields the same event-dict protocol as the GGUF path so the route
layer can stream both backends through one helper. Each event is one of:
* ``{"type": "status", "text": ...}``
* ``{"type": "content", "text": cumulative_text}``
@ -845,8 +772,11 @@ class InferenceBackend:
from core.inference.safetensors_agentic import run_safetensors_tool_loop
from core.inference.tools import execute_tool
def _single_turn(conv: list):
def _single_turn(conv: list, *, active_tools: Optional[list[dict]] = None):
# conv already has the system message -- avoid double-prepend.
# `active_tools` is supplied by run_safetensors_tool_loop so one-shot
# tools such as render_html can be removed from later same-response prompts.
turn_tools = active_tools if active_tools is not None else tools
yield from self._generate_chat_response_inner(
messages = conv,
system_prompt = "",
@ -857,7 +787,7 @@ class InferenceBackend:
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,
@ -896,15 +826,11 @@ class InferenceBackend:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""
Generate response for text or vision models.
The generation lock is acquired by the background generation thread.
"""Generate response for text or vision models (lock held by background thread).
``tools`` / ``enable_thinking`` / ``reasoning_effort`` /
``preserve_thinking`` are forwarded into
``tokenizer.apply_chat_template`` so templates that understand
these kwargs (Qwen3, Llama 3.1+, gpt-oss harmony, ...) advertise
the tool schemas and reasoning controls to the model.
``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking``
are forwarded into ``apply_chat_template`` so templates that understand them
(Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls.
"""
yield from self._generate_chat_response_inner(
messages = messages,
@ -941,9 +867,8 @@ class InferenceBackend:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""
Inner generation logic. Called by both generate_chat_response
and generate_with_adapter_control.
"""Inner generation logic, called by generate_chat_response and
generate_with_adapter_control.
_adapter_state is passed to generate_stream/vision so the background
thread can toggle adapters under the generation lock.
@ -955,15 +880,13 @@ class InferenceBackend:
model_info = self.models[self.active_model_name]
is_vision = model_info.get("is_vision", False)
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
# Unwrap processor → raw tokenizer for VLMs on the text path
# Unwrap processor -> raw tokenizer for VLMs on the text path.
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
top_k = self._normalize_top_k(top_k)
if is_vision and image:
# Vision model generation (only when an image is actually provided)
# Check that the stored processor can actually handle images.
# FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast)
# instead of a proper ProcessorMixin for some models (e.g. Gemma-3).
# Verify the stored processor can handle images; FastVisionModel may
# return a raw tokenizer instead of a ProcessorMixin (e.g. Gemma-3).
from transformers import ProcessorMixin
processor = model_info.get("processor")
@ -991,10 +914,9 @@ class InferenceBackend:
f"falling back to text-only generation (image will be ignored)."
)
# Text path: Use training pipeline approach
# Messages are already in ChatML format from eval.py
# Text path: messages are already in ChatML format from eval.py.
# Step 1: Apply get_chat_template if model is in mapper
# Step 1: apply get_chat_template if model is in mapper.
try:
from utils.datasets import (
MODEL_TO_TEMPLATE_MAPPER,
@ -1002,14 +924,12 @@ class InferenceBackend:
)
model_name_lower = self.active_model_name.lower()
# Check if model has a registered template
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
logger.info(
f"Applying chat template '{template_name}' for {self.active_model_name}"
)
# This modifies the tokenizer with the correct template
tokenizer = get_chat_template(
tokenizer,
chat_template = template_name,
@ -1021,7 +941,7 @@ class InferenceBackend:
except Exception as e:
logger.warning(f"Could not apply get_chat_template: {e}")
# Step 2: Format with tokenizer.apply_chat_template()
# Step 2: format with tokenizer.apply_chat_template().
if system_prompt:
template_messages = [{"role": "system", "content": system_prompt}] + messages
else:
@ -1046,10 +966,10 @@ class InferenceBackend:
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
except Exception as e:
logger.error(f"Error applying chat template: {e}")
# Fallback to manual formatting
# Fall back to manual formatting
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
# Step 3: Generate
# Step 3: generate
yield from self.generate_stream(
formatted_prompt,
temperature,
@ -1080,7 +1000,7 @@ class InferenceBackend:
model = model_info["model"]
processor = model_info["processor"]
# FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast)
# instead of a Processor for some models. Safe unwrap for tokenize-only ops.
# for some models. Safe unwrap for tokenize-only ops.
raw_tokenizer = getattr(processor, "tokenizer", processor)
# Extract user message
@ -1136,7 +1056,7 @@ class InferenceBackend:
return_tensors = "pt",
).to(model.device)
else:
# Text-only for vision model
# Text-only path for a vision model
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device)
@ -1233,7 +1153,7 @@ class InferenceBackend:
repetition_penalty,
cancel_event = None,
) -> Generator[str, None, None]:
"""Handle audio input (ASR) generation — accepts audio numpy array, streams text output.
"""Audio-input (ASR) generation: takes an audio numpy array, streams text.
Uses processor.apply_chat_template with audio embedded in messages (Gemma 3n pattern).
"""
@ -1245,7 +1165,7 @@ class InferenceBackend:
processor = model_info.get("processor") or model_info.get("tokenizer")
raw_tokenizer = getattr(processor, "tokenizer", processor)
# Extract last user text — default matches notebook prompt
# Last user text; default matches the notebook prompt
user_text = "Please transcribe this audio."
if messages:
for msg in reversed(messages):
@ -1253,11 +1173,11 @@ class InferenceBackend:
user_text = msg["content"]
break
# Use ASR-specific system prompt if user hasn't set a custom one
# ASR-specific default system prompt if none set
if not system_prompt:
system_prompt = "You are an assistant that transcribes speech accurately."
# Build messages in Gemma 3n format — audio goes INTO apply_chat_template
# Gemma 3n format — audio goes INTO apply_chat_template
audio_messages = [
{"role": "system", "content": [{"type": "text", "text": system_prompt}]},
{
@ -1269,7 +1189,7 @@ class InferenceBackend:
},
]
# apply_chat_template handles audio embedding + tokenization in one step
# apply_chat_template does audio embedding + tokenization in one step
inputs = processor.apply_chat_template(
audio_messages,
add_generation_prompt = True,
@ -1290,7 +1210,7 @@ class InferenceBackend:
timeout = 0.2,
)
# Notebook uses do_sample=False for ASR (greedy decoding for accuracy)
# Notebook uses do_sample=False (greedy) for ASR accuracy
generation_kwargs = dict(
**inputs,
streamer = streamer,
@ -1354,9 +1274,9 @@ class InferenceBackend:
audio_array,
cancel_event = None,
) -> Generator[str, None, None]:
"""Whisper ASR — takes audio numpy array, yields transcribed text.
"""Whisper ASR: takes an audio numpy array, yields transcribed text.
Uses the pre-built transformers pipeline (created during model loading).
Uses the pre-built transformers pipeline created at model load.
"""
model_info = self.models[self.active_model_name]
whisper_pipe = model_info.get("whisper_pipeline")
@ -1376,7 +1296,7 @@ class InferenceBackend:
yield f"Error: {str(e)}"
def _is_gpt_oss_model(self, model_name: str = None) -> bool:
"""Check if the given (or active) model uses the gpt-oss harmony protocol."""
"""Whether the given (or active) model uses the gpt-oss harmony protocol."""
from utils.datasets import is_gpt_oss_model_name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
@ -1392,10 +1312,10 @@ class InferenceBackend:
cancel_event = None,
_adapter_state = None,
) -> Generator[str, None, None]:
"""Generate streaming text response (text models only).
"""Generate a streaming text response (text models only).
_adapter_state: if not None, the background thread toggles adapters
before model.generate(), all under _generation_lock.
before model.generate(), under _generation_lock.
"""
if not self.active_model_name:
yield "Error: No active model"
@ -1403,9 +1323,9 @@ class InferenceBackend:
model_info = self.models[self.active_model_name]
model = model_info["model"]
# For VLMs the stored "tokenizer" is actually the processor.
# Unwrap to get the real tokenizer so TextIteratorStreamer's
# skip_prompt / skip_special_tokens work correctly.
# For VLMs the stored "tokenizer" is actually the processor. Unwrap to
# the real tokenizer so TextIteratorStreamer's skip_prompt /
# skip_special_tokens work correctly.
tokenizer = model_info["tokenizer"]
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
@ -1415,8 +1335,8 @@ class InferenceBackend:
from transformers import TextIteratorStreamer
import threading
# Use HarmonyTextStreamer for gpt-oss models to properly parse
# the multi-channel harmony protocol into <think> tags
# gpt-oss models: HarmonyTextStreamer parses the multi-channel
# harmony protocol into <think> tags
if self._is_gpt_oss_model():
try:
streamer = HarmonyTextStreamer(
@ -1513,11 +1433,10 @@ class InferenceBackend:
cleaned = self._clean_generated_text(output)
yield cleaned
finally:
# Only set cancel_event when we exited early (user cancel),
# NOT on normal completion. cancel_event is a shared mp.Event
# — setting it unconditionally would leave a stale cancel
# signal that could interfere with the next serialized
# generation request (e.g. in compare mode).
# Set cancel_event only on early exit (user cancel), NOT on
# normal completion. It's a shared mp.Event; setting it
# unconditionally would leave a stale cancel signal that could
# disrupt the next serialized request (e.g. compare mode).
if cancel_event is not None and not generation_complete:
cancel_event.set()
thread.join(timeout = 10)
@ -1544,10 +1463,8 @@ class InferenceBackend:
repetition_penalty: float = 1.0,
use_adapter: Optional[Union[bool, str]] = None,
) -> Tuple[bytes, int]:
"""
Generate audio from text for TTS models.
Returns (wav_bytes, sample_rate).
Blocking — generates complete audio before returning.
"""Generate audio from text for TTS models.
Returns (wav_bytes, sample_rate). Blocking — full audio before return.
"""
if not self.active_model_name:
raise RuntimeError("No active model")
@ -1661,8 +1578,8 @@ class InferenceBackend:
repetition_penalty,
):
"""Generate audio using DAC (OuteTTS). Follows Oute_TTS_(1B).ipynb exactly."""
# Monkey-patch RepetitionPenaltyLogitsProcessor with a 64-token penalty
# window (same as the OuteTTS notebook) to avoid degenerate repetition.
# Monkey-patch RepetitionPenaltyLogitsProcessor with a 64-token window
# (same as the OuteTTS notebook) to avoid degenerate repetition.
self._patch_repetition_penalty_processor()
prompt = (
@ -1689,10 +1606,9 @@ class InferenceBackend:
@classmethod
def _patch_repetition_penalty_processor(cls):
"""
Monkey-patch transformers' RepetitionPenaltyLogitsProcessor with a
64-token sliding window variant (from the OuteTTS notebook).
Only applied once per process.
"""Monkey-patch transformers' RepetitionPenaltyLogitsProcessor with a
64-token sliding-window variant (from the OuteTTS notebook).
Applied once per process.
"""
if cls._repetition_penalty_patched:
return
@ -1743,10 +1659,10 @@ class InferenceBackend:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> str:
"""Render the chat prompt, peeling kwargs the template does not
understand. Delegates to the dependency-light helper module so
the fallback chain can be unit-tested without pulling unsloth /
torch into the test sandbox.
"""Render the chat prompt, peeling kwargs the template doesn't
understand. Delegates to the dependency-light helper module so the
fallback chain is unit-testable without pulling unsloth / torch into
the test sandbox.
"""
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
@ -1846,8 +1762,7 @@ class InferenceBackend:
return self._format_generic_template(chat_messages, {})
def _format_chat_manual(self, messages: list, template_type: str, special_tokens: dict) -> str:
"""
Manual chat formatting fallback for when tokenizer template fails
"""Manual chat-formatting fallback when the tokenizer template fails.
Args:
messages: List of message dictionaries
@ -1960,12 +1875,7 @@ class InferenceBackend:
return formatted
def check_vision_model_compatibility(self) -> bool:
"""
Check if current model supports vision.
Returns:
bool: True if current model supports vision, False otherwise
"""
"""Whether the current model supports vision."""
current_model = self.get_current_model()
if current_model and current_model in self.models:
return self.models[current_model].get("is_vision", False)
@ -1981,7 +1891,7 @@ class InferenceBackend:
return
try:
# This is a common pattern for Unsloth/Hugging Face models
# Common pattern for Unsloth/Hugging Face models
if hasattr(model, "past_key_values"):
model.past_key_values = None
if hasattr(model, "generation_config"):
@ -1995,7 +1905,7 @@ class InferenceBackend:
def reset_generation_state(self):
"""Reset any cached generation state to prevent hanging after errors"""
try:
# Clear cached states for ALL loaded models
# Clear cached state for ALL loaded models
for model_name in self.models.keys():
self._reset_model_generation_state(model_name)
@ -2029,9 +1939,9 @@ class InferenceBackend:
def _clean_generated_text(self, text: str) -> str:
"""Strip leaked special tokens using the tokenizer's own token list."""
if self._is_gpt_oss_model():
# HarmonyTextStreamer produces clean <think>...</think> output.
# Strip harmony protocol tokens and other gpt-oss added tokens
# (e.g. <|return|>) that may leak past the streamer.
# HarmonyTextStreamer emits clean <think>...</think>. Strip any
# harmony protocol tokens and other gpt-oss tokens (e.g.
# <|return|>) that leak past the streamer.
import re
text = re.sub(r"<\|[a-z_]+\|>", "", text)
return text.strip()
@ -2059,7 +1969,7 @@ class InferenceBackend:
try:
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER
# Try exact match first
# Exact match first
model_name_lower = model_name.lower()
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
@ -2067,7 +1977,7 @@ class InferenceBackend:
f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper"
)
else:
# Try partial match (for variants like model_name-bnb-4bit)
# Partial match (for variants like model_name-bnb-4bit)
for key in MODEL_TO_TEMPLATE_MAPPER:
if key in model_name_lower or model_name_lower in key:
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[key]
@ -2127,15 +2037,15 @@ class InferenceBackend:
logger.info(f"No built-in chat template for {model_name}, will use generic formatting")
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"""
"""Whether any model is currently 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 load_model_simple(
@ -2145,9 +2055,8 @@ class InferenceBackend:
max_seq_length: int = 2048,
load_in_4bit: bool = True,
) -> bool:
"""
Simple model loading wrapper for chat interface.
Accepts model path as string and handles ModelConfig creation internally.
"""Simple model-loading wrapper for the chat interface. Takes a string
path and builds the ModelConfig internally.
Args:
model_path: Model name or path (e.g., "unsloth/llama-3-8b")
@ -2159,14 +2068,12 @@ class InferenceBackend:
bool: True if successful, False otherwise
"""
try:
# Create config from string path
config = ModelConfig.from_ui_selection(
model_path,
lora_path = None, # No LoRA for chat
is_lora = False,
)
# Call existing load_model with config
return self.load_model(
config = config,
max_seq_length = max_seq_length,

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

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

File diff suppressed because it is too large Load diff

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

@ -3,11 +3,10 @@
"""Boundary validator for user-supplied llama-server pass-through args.
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.
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.
Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
"""
@ -19,11 +18,11 @@ from typing import Iterable, Optional
# 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], ...] = (
# Parallel slots: owned by typer --parallel; a pass-through would
# desync app.state.llama_parallel_slots from llama-server.
# 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.
# 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"}),
@ -40,15 +39,15 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"--path"}),
frozenset({"--api-prefix"}),
frozenset({"--reuse-port"}),
# Auth / TLS: Studio terminates auth; upstream --api-key / TLS
# shadows Studio's key and breaks the 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"}),
# Built-in web UI. --webui/--no-webui is the legacy spelling;
# upstream renamed to --ui/--no-ui + --ui-*. Keep both so prebuilt
# and system llama.cpp binaries both match.
# 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"}),
@ -62,8 +61,8 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# 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.
# llama-server's own built-in tools flag would silently stack on top of
# Studio's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
)
@ -73,11 +72,9 @@ _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
def _flag_name(token: str) -> Optional[str]:
"""Flag name for ``token``, or None if it isn't a flag.
Peels `--key=value` to `--key`, treats `-1` / `-0.5` as values
(llama-server shorts always start with a letter), strips
whitespace, and normalises attached `-np8` / signed `-np-1` /
digit-prefix-junk `-np8x` to `-np`. Mirrors the CLI's
`_expand_attached_np_short`.
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 {"-", "--"}:
@ -95,9 +92,9 @@ def _flag_name(token: str) -> Optional[str]:
def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
"""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."""
"""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] = []
@ -116,15 +113,15 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
def is_managed_flag(flag: str) -> bool:
"""True if ``flag`` is Studio-managed. Normalises via ``_flag_name``
so `-np8` / `--parallel=8` classify like the canonical tokens."""
"""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;
# stripped from inherited extras so they can't last-wins-override an
# Apply that re-sets the same 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"})
_SPEC_FLAGS: frozenset[str] = frozenset(
@ -157,16 +154,15 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset(
_SHADOWING_FLAGS: frozenset[str] = _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
# Shadowing flags that take no value -- strip the flag only, never the
# following token.
# 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 flag parsing for the one pass-through
numeric knob Studio's load-time fit logic needs to see.
Mirrors llama.cpp's last-wins parsing for the one numeric knob Studio's
load-time fit logic needs.
"""
if not args:
return None
@ -204,10 +200,8 @@ def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) -> int:
"""Return the context size load_model should treat as requested.
Single source of truth for the two-line ``ctx_override = parse_ctx_override(...);
requested_ctx = ctx_override if ctx_override is not None else n_ctx`` pattern
used by ``load_model`` so tests don't have to reimplement the conditional
locally and then assert against their own reimplementation.
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
@ -216,10 +210,8 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) ->
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.
Recognises -ctk (key) and -ctv (value); treats both as one setting,
since Studio's KV estimate has a single cache_type_kv knob.
"""
if not args:
return None
@ -256,8 +248,7 @@ def resolve_cache_type_kv(
) -> Optional[str]:
"""Return the cache type load_model should treat as requested.
Single source of truth for the cache override conditional used by
``load_model``.
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
@ -274,10 +265,8 @@ def strip_shadowing_flags(
"""Strip flags that shadow first-class Studio settings.
Used when inheriting a previous load's ``llama_extra_args`` so an
inherited `-c 4096` can't override the current `max_seq_length`
(same for cache / spec / template). Each ``strip_*`` toggle
controls one group; the route only strips groups whose first-class
field the caller actually supplied.
inherited `-c 4096` can't override the current `max_seq_length` (same for
cache / spec / template). Each ``strip_*`` toggle controls one group.
"""
shadowing: set[str] = set()
if strip_context:
@ -299,8 +288,8 @@ def strip_shadowing_flags(
out.append(tok)
i += 1
continue
# Drop the flag; consume the next token too unless it's
# boolean, already inline (`-c=4096`), or another flag.
# 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:

52
studio/backend/core/inference/mcp_client.py
View file

@ -31,20 +31,17 @@ def parse_stdio_command(address: str) -> list[str]:
posix = sys.platform != "win32"
parts = shlex.split(address, posix = posix)
if not posix:
# posix=False keeps backslash paths intact but also keeps the surrounding
# quotes on a token. Strip a matched pair so the argv reaches the
# subprocess clean ('"C:\\Program Files\\node"' -> C:\\Program Files\\node).
# posix=False keeps backslash paths but also keeps surrounding quotes;
# strip a matched pair so argv reaches the subprocess clean.
parts = [p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p for p in parts]
return parts
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (and bypass
the python/terminal sandbox), so they are only allowed when the backend
host is the user's own machine. The Tauri desktop app sets
UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see main.py); advanced localhost /
self-hosted users can opt in with the same variable. It stays off for
Colab and any network (0.0.0.0) bind."""
"""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. The Tauri
app sets UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1; localhost/self-hosted users can opt
in with the same var. Off for Colab and any network (0.0.0.0) bind."""
return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1"
@ -63,8 +60,8 @@ def probe_timeout(address: str, use_oauth: bool) -> float:
def parse_server_headers(server: dict) -> Optional[dict]:
"""Parsed headers_json. For stdio servers this dict is the process
environment instead of HTTP headers (see _client)."""
"""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
@ -82,8 +79,8 @@ def _oauth_store():
from key_value.aio.stores.filetree import FileTreeStore
from utils.paths.storage_roots import ensure_dir, studio_root
# Hash keys/collections — fastmcp uses raw URLs like https://x.com as
# keys and FileTreeStore would treat the "://" as nested directories.
# 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(),
@ -93,12 +90,10 @@ def _oauth_store():
async def clear_oauth_tokens_async(url: str) -> None:
"""Drop any persisted OAuth tokens for ``url``. fastmcp keys tokens by
MCP URL, so on server delete / URL change / OAuth disable we have to
clear the old credentials explicitly. Otherwise re-registering the
same URL would silently reuse the old account's token. The entire
body runs inside the protected block -- store / OAuth construction
failing must not make the delete / update route 500."""
"""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())
@ -124,9 +119,8 @@ def _client(
parts = parse_stdio_command(url)
if not parts:
raise ValueError(f"Empty stdio command: {url!r}")
# env vars ride the headers field (merged over the SDK's safe default env).
# keep_alive=False tears the subprocess down on exit, so a one-shot
# probe/tool call never leaves an orphan process.
# 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],
@ -192,10 +186,9 @@ def call_tool_sync(
) -> str:
"""Synchronously call an MCP tool.
``cancel_event``: optional ``threading.Event``. When set, the in-flight
HTTP call is cancelled and the function returns a cancellation Error.
Polled in parallel with the tool call via ``asyncio.wait`` so a /cancel
POST from the UI interrupts even mid-network-read.
``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:
@ -204,14 +197,13 @@ def call_tool_sync(
async def _watch_cancel() -> None:
# 50 ms cadence keeps cancellation responsive without busy-looping;
# matches the cadence routes/inference.py uses for cancel watchers.
# 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 (reviewer-reproduced race).
# 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())

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

@ -13,8 +13,7 @@ logger = get_logger(__name__)
def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
"""Map mlx_lm / mlx_vlm stream stats onto the usage/timings shape
llama-server emits so the chat speed popover renders the same."""
"""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)
@ -52,7 +51,6 @@ class MLXInferenceBackend:
# 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
@ -65,10 +63,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
@ -109,18 +106,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}': "
@ -187,9 +176,8 @@ class MLXInferenceBackend:
"audio_type": None,
"has_audio_input": False,
}
# 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)
@ -198,10 +186,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
@ -276,10 +262,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,
@ -308,7 +292,7 @@ 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)
)
@ -389,8 +373,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,
@ -425,7 +408,7 @@ class MLXInferenceBackend:
):
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,
@ -471,10 +454,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
@ -492,8 +474,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 = ""
@ -503,11 +484,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,

212
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.
"""
@ -51,9 +50,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,16 +59,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()
@ -86,13 +83,11 @@ 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()
# ------------------------------------------------------------------
@ -101,14 +96,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:
@ -133,8 +127,7 @@ 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.
# 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
]
@ -192,16 +185,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
@ -243,7 +236,7 @@ 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()
# ------------------------------------------------------------------
@ -277,14 +270,12 @@ class InferenceOrchestrator:
) -> 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
@ -345,8 +336,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:
@ -367,10 +358,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
@ -422,9 +412,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,13 +442,9 @@ class InferenceOrchestrator:
) -> 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"
@ -532,10 +517,9 @@ class InferenceOrchestrator:
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", "")
@ -574,8 +558,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
@ -588,9 +572,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:
@ -618,9 +602,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
@ -649,16 +632,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 (
@ -680,7 +661,7 @@ 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",
@ -689,14 +670,13 @@ class InferenceOrchestrator:
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", {})
@ -710,8 +690,8 @@ class InferenceOrchestrator:
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
}
# 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
@ -745,7 +725,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
@ -796,14 +776,12 @@ class InferenceOrchestrator:
) -> 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
receives the worker's usage/timings. Request-scoped by design so
concurrent streams cannot read each other's stats.
``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,
@ -847,22 +825,22 @@ class InferenceOrchestrator:
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 = "",
@ -874,11 +852,11 @@ 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, same as the GGUF tool loop's metadata
# last turn wins, like the GGUF tool loop
stats_holder = stats_holder,
)
if use_adapter is not None:
@ -914,9 +892,9 @@ class InferenceOrchestrator:
) -> 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,
@ -946,9 +924,8 @@ class InferenceOrchestrator:
) -> 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"
@ -958,14 +935,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,
@ -1029,8 +1003,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:
@ -1046,8 +1019,7 @@ 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)
@ -1071,12 +1043,11 @@ 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", "")
@ -1117,7 +1088,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")
@ -1242,7 +1213,7 @@ 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)
)
@ -1310,9 +1281,7 @@ class InferenceOrchestrator:
img,
max_size: int = 800,
):
"""Resize image while maintaining aspect ratio.
No ML imports needed — runs locally in parent process.
"""
"""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:
@ -1331,26 +1300,25 @@ 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 "")
@ -1360,7 +1328,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()

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

@ -1,8 +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
"""Static per-MTok pricing tables and ``calculate_cost`` helper for
turning an upstream ``usage`` block into a USD figure.
"""Per-MTok pricing tables and ``calculate_cost`` (usage block -> USD).
Sources: Anthropic prompt-caching docs (5m write 1.25x, 1h write 2x,
read 0.1x), web search ($10/1000), code execution; OpenAI pricing page.
@ -12,13 +11,13 @@ from __future__ import annotations
from typing import Any, Optional
# Per-MTok base pricing in USD. Cache multipliers are applied ON
# `input_per_mtok` (not absolute prices), matching Anthropic's docs.
# 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},
# Alias both the bare id and dated id: backend defaults reference
# the bare form, which won't prefix-match the dated key.
# 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},
@ -34,8 +33,8 @@ ANTHROPIC_PRICING: dict[str, dict[str, float]] = {
OPENAI_PRICING: dict[str, dict[str, float]] = {
# Verified against developers.openai.com/api/docs/pricing.
# `long_context_*` keys apply once input exceeds the threshold
# (gpt-5.5/5.4: 272k); families without these keys ship a single rate.
# `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,
@ -58,28 +57,28 @@ OPENAI_PRICING: dict[str, dict[str, float]] = {
# 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 are no longer on the pricing page; omit them
# so calculate_cost returns priced=False rather than silently $0.
# 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 standard on input + output.
# 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 0.1x; cache writes pay normal input price.
# OpenAI: cache reads 0.1x; cache writes pay input price.
OPENAI_CACHE_READ_MULT = 0.1
# Server-tool surcharges. Anthropic code_exec is $0.05/hr marginal
# (50 free hours/day per org, not visible here).
# 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 container bills per memory tier; we report the 1g default
# ($0.09/hour) since the tier isn't surfaced to the cost ledger.
# 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
@ -96,10 +95,8 @@ def _lookup(provider: str, model: str) -> Optional[dict[str, float]]:
return None
if model in table:
return table[model]
# Longest-prefix match on a dash boundary: lets dated snapshots
# inherit canonical prices while preventing "claude-opus-4-15"
# from matching "claude-opus-4-1" or "gpt-5.5-prod" from matching
# "gpt-5.5-pro". Sort longest-first to pick the most specific row.
# 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]
@ -107,11 +104,9 @@ def _lookup(provider: str, model: str) -> Optional[dict[str, float]]:
def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str, float]:
"""Return a per-turn USD cost breakdown with per-bucket + total
fields so the frontend can render either a single number or a
tooltip without re-doing the math. When the model isn't in the
static table, ``priced`` is False and USD fields are 0.0 (token
counts still report).
"""Return a per-turn USD cost breakdown (per-bucket + total).
Unknown model -> ``priced`` False and USD fields 0.0 (token counts still report).
"""
prices = _lookup(provider, model)
out: dict[str, float] = {
@ -128,22 +123,20 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
}
# Accept raw (input_tokens/output_tokens) and Studio chat-style
# (prompt_tokens/completion_tokens) envelopes. Cache buckets
# behave differently per envelope:
# (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 tokens >=0 so corrupted payloads can't produce a negative bill.
# 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))
# Fallback to mirrored prompt_tokens_details only when the native
# cache_read_input_tokens key is absent. An explicit native 0 is
# authoritative, so a stale mirrored block from a proxy can never
# inflate cache_read past the native count.
# 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):
@ -152,22 +145,22 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
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
# recover the raw uncached prompt count.
# 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` fallback would
# silently pick a stale completion_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 either input_tokens_details (raw
# Responses) or prompt_tokens_details (Studio chat-style).
# 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):
@ -200,8 +193,8 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
out_per = prices["output_per_mtok"]
# Anthropic fast-mode: 6x on input + output. Cache multipliers stack
# on top of fast-mode, so applying once to (base, out_per) propagates
# into the cache_*_usd buckets computed below.
# 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
@ -235,15 +228,15 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
+ code_exec_hours * ANTHROPIC_CODE_EXEC_USD_PER_HOUR
)
else:
# OpenAI: cache writes pay base input; only cache reads get
# 0.1x. Subtract cached from already-counted input_usd to
# avoid double-billing (OpenAI folds cache into input_tokens).
# 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
# OpenAI server-tool surcharges arrive under `openai_tool_use`
# (normalised by the SSE finaliser from output array items).
# (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)

127
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,23 +58,15 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
},
"gemini": {
"display_name": "Google Gemini",
# Native Gemini REST endpoint -- the Gemini API does NOT speak
# OpenAI Chat Completions on this base. Requests/responses are
# translated in `_stream_gemini` in external_provider.py.
# API reference: https://ai.google.dev/gemini-api/docs
# 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 -- the live ListModels response returns dozens
# of historical / experimental / embedding ids. Cap to the
# current chat-capable Gemini families (3.5 / 3.1 / 3 Flash /
# 2.5) plus the Nano Banana image trio and the rolling
# `*-latest` aliases. Excluded on purpose:
# - `gemini-2.0-flash*` (Google retired 2026-06-01; 404 on use)
# - `gemini-3-pro-preview` (shut down 2026-03-09; auto-redirects
# to `gemini-3.1-pro-preview` per Google's deprecation notice,
# so we surface 3.1 directly and skip the redirect).
# The allowlist below blocks the retired ids from re-appearing
# via the live ListModels fetch. Verified against the live
# `/v1beta/models` catalog 2026-05-24.
# 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",
@ -100,8 +85,7 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
# The native API takes the API key on the `x-goog-api-key`
# header. An empty `auth_prefix` ensures we send the bare key.
# Native API takes the bare key on `x-goog-api-key`.
"auth_header": "x-goog-api-key",
"auth_prefix": "",
"openai_compatible": False,
@ -110,20 +94,13 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"API key from https://aistudio.google.com/apikey. "
"See https://ai.google.dev/gemini-api/docs for endpoint shapes."
),
# Even after the regex match, drop ids that Google still
# returns from ListModels but routes via implicit redirect.
# gemini-3-pro-preview was shut down 2026-03-09 and is
# auto-aliased to gemini-3.1-pro-preview; we surface the
# canonical id only so users do not see two cards for the
# same underlying model.
# 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",),
# Matches the chat-capable 3.5 / 3.1 / 3 / 2.5 families plus the
# rolling *-latest aliases (which Google rolls forward as new
# generations ship). Image-tier ids (`-image`, `-image-preview`,
# `nano-banana-pro-preview`) flow through the Nano Banana
# `responseModalities` path in `_stream_gemini`. Retired 2.0
# ids ARE NOT in this regex on purpose -- Google's ListModels
# would otherwise re-surface them and they 404 on use.
# 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"^("
r"gemini-3\.5-(?:flash|pro)(?:-preview)?|"
@ -184,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",
@ -203,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": {
@ -228,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",
@ -248,28 +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|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,
@ -277,15 +242,11 @@ 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,
},
"ollama": {
@ -321,7 +282,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",
@ -371,10 +332,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():

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

@ -4,75 +4,72 @@
"""
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,
RENDER_HTML_REPEAT_NUDGE,
TOOL_ERROR_NUDGE,
TOOL_ERROR_PREFIXES,
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,
)
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}"
_CANONICAL_HEAL_ARG = {
"python": "code",
"terminal": "command",
"render_html": "code",
}
return status_for_tool(tool_name, arguments)
_FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
@ -96,34 +93,43 @@ def _detect_render_html_tool_start(content: str) -> bool:
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:
"""Normalise tool ``arguments`` to a dict.
arguments, _ = _coerce_arguments_with_provenance(
raw_args,
heal = heal,
tool_name = tool_name,
)
return arguments
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 {}
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(
@ -140,42 +146,42 @@ def run_safetensors_tool_loop(
) -> 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]] = []
render_html_succeeded = False
unrestricted_tools = not tools
tool_controller = ToolLoopController(
tools = None if unrestricted_tools else tools,
auto_heal_tool_calls = auto_heal_tool_calls,
)
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": ""}
@ -189,6 +195,17 @@ 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 = ""
@ -197,7 +214,7 @@ def run_safetensors_tool_loop(
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:
@ -205,7 +222,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
@ -215,7 +232,11 @@ def run_safetensors_tool_loop(
if detect_state == _state_draining:
if (
not render_html_succeeded
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)
):
@ -225,26 +246,35 @@ def run_safetensors_tool_loop(
"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 render_html_succeeded
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)
):
@ -254,10 +284,15 @@ def run_safetensors_tool_loop(
"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}
@ -271,7 +306,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
@ -280,9 +315,24 @@ 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 render_html_succeeded
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)
):
@ -292,13 +342,18 @@ def run_safetensors_tool_loop(
"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}
@ -308,16 +363,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
@ -325,19 +388,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),
@ -347,22 +421,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.
@ -372,100 +463,69 @@ 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)
repeat_render_html = tool_name == "render_html" and render_html_succeeded
if not repeat_render_html:
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 repeat_render_html:
result = RENDER_HTML_REPEAT_NUDGE
elif 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
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}"
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
if not repeat_render_html:
yield {
"type": "tool_end",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"result": result,
}
yield {"type": "status", "text": decision.status_text}
yield decision.tool_start_event()
is_error = isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
if tool_name == "render_html" and not is_error:
render_html_succeeded = True
tool_call_history.append((tc_key, is_error))
eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
try:
result = execute_tool(
decision.tool_name,
decision.arguments,
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", decision.tool_name, exc)
result = f"Error: tool raised an exception: {exc}"
# 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": ""}

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

@ -11,11 +11,9 @@ import json
import re
# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing
# unclosed runs so truncated tails don't leak markup.
# Function-name char set tracks OpenAI's ^[a-zA-Z0-9_-]{1,64}$ so MCP
# tool names that contain a hyphen (e.g. mcp__srv__list-issues) parse
# the same as the built-in web_search/python/terminal names.
# _TOOL_CLOSED_PATS: 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),
@ -72,8 +70,7 @@ _TC_JSON_START_RE = re.compile(r"<tool_call>\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*$")
# Parameter names can carry hyphens too (e.g. MCP tool schemas with
# `issue-number`, `repo-name`); using `\w+` here dropped those keys.
# [\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>"
@ -105,7 +102,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"}}``
@ -119,15 +121,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):
@ -147,27 +151,31 @@ 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 = [
fm
@ -185,18 +193,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()
@ -206,8 +233,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)}",

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

@ -0,0 +1,410 @@
# 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 image sentinels before feeding the model."""
if "__IMAGES__:" in result:
return result.split("__IMAGES__:", 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,
)

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

@ -1,11 +1,8 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Tool definitions and executors for LLM tool calling.
Supports web search (DuckDuckGo), Python code execution, and terminal commands.
"""
"""Tool definitions and executors for LLM tool calling: web search
(DuckDuckGo), Python code execution, and terminal commands."""
import ast
import http.client
@ -42,9 +39,8 @@ logger = get_logger(__name__)
_EXEC_TIMEOUT = 300 # 5 minutes
# Pre-import modules used in _sandbox_preexec at module level so that
# the preexec_fn closure does not trigger the import machinery in the
# forked child (which can deadlock in multi-threaded servers).
# Import these at module level so the preexec_fn closure triggers no imports in
# the forked child (which can deadlock multi-threaded servers).
_libc = None
if sys.platform == "linux":
try:
@ -64,8 +60,8 @@ if sys.platform != "win32":
except ImportError:
pass
# Strict raster-image allowlist for sandbox file serving.
# No .svg (XSS risk via embedded scripts), no .html, no .pdf.
# Raster-image allowlist for sandbox file serving.
# No .svg (XSS via embedded scripts), no .html, no .pdf.
_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
_MAX_OUTPUT_CHARS = 8000 # truncate long output
_BLOCKED_COMMANDS_COMMON = frozenset(
@ -122,9 +118,9 @@ _BLOCKED_COMMANDS = (
_SHELL_SEPARATORS = frozenset({";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"})
# Bash keywords that introduce a new command position (then $cmd, do $cmd, etc.).
# Bash keywords starting a new command position (then $cmd, do $cmd, etc.).
_SHELL_KEYWORDS_AS_SEP = frozenset({"then", "do", "else", "elif"})
# Wrappers whose next non-flag argument is itself the command Bash will exec.
# Wrappers whose next non-flag argument is the command Bash will exec.
_COMMAND_PREFIXES = frozenset(
{
"env",
@ -152,21 +148,18 @@ _FIND_EXEC_FLAGS = frozenset({"-exec", "-execdir", "-ok", "-okdir"})
def _find_blocked_commands(command: str) -> set[str]:
"""Detect blocked commands at shell command position only.
A token is at command position if it is the first token, or if the
preceding token is a shell separator / brace-group opener / keyword
that starts a new command (`then`, `do`, etc.), or a command-prefix
wrapper like `env` / `time` / `xargs` (the next token is the real
command). Tokens in argument position (`grep -r curl .`,
`echo source the data`, `ls /usr/bin/curl`) are passed through.
Also scans `find ... -exec CMD` and recurses into bash -c / cmd /c.
A token is at command position if it is the first token, or follows a
shell separator / brace-group opener / new-command keyword (`then`, `do`,
etc.), or a command-prefix wrapper like `env` / `time` / `xargs` (next
token is the real command). Tokens in argument position (`grep -r curl .`,
`echo source the data`, `ls /usr/bin/curl`) pass through. Also scans
`find ... -exec CMD` and recurses into bash -c / cmd /c.
"""
blocked: set[str] = set()
# shlex with punctuation_chars splits `;`, `&&`, `||`, `|`, `(`, `)`, `` ` ``
# off as their own tokens so we can detect command position even when a
# caller writes `echo done; rm -rf x` (no whitespace) or quote-splits the
# command name itself (`r''m` collapses to a single token `rm` at command
# position after the `;` separator).
# punctuation_chars splits separators into their own tokens, so command
# position is detected even in `echo done; rm -rf x` (no whitespace) or
# quote-split names (`r''m` collapses to `rm` after `;`).
try:
if sys.platform == "win32":
tokens = shlex.split(command, posix = False)
@ -178,9 +171,7 @@ def _find_blocked_commands(command: str) -> set[str]:
tokens = command.split()
def _token_basename(tok: str) -> str:
# shlex may glue trailing meta-chars onto a token (`rm;`); strip them
# so the basename match still hits `rm`. Leading shell-state chars
# likewise.
# Strip glued-on meta-chars (`rm;`) so the basename still matches `rm`.
tok = tok.strip(";&|()`{}")
base = os.path.basename(tok).lower()
stem, ext = os.path.splitext(base)
@ -189,33 +180,31 @@ def _find_blocked_commands(command: str) -> set[str]:
return base
expect_command = True # start of string is a command position
prefix_pending = False # last command-position token was env/time/timeout/xargs/...
prefix_pending = False # last cmd-position token was a wrapper (env/time/xargs/...)
for token in tokens:
if token in _SHELL_SEPARATORS or token in _SHELL_KEYWORDS_AS_SEP:
expect_command = True
prefix_pending = False
continue
if token.startswith("-"):
# Flags belong to the active command. While a wrapper prefix is
# waiting for its command (`stdbuf -oL cmd`, `xargs -- cmd`),
# keep expect_command intact.
# Flags belong to the active command, but keep expect_command while a
# wrapper prefix awaits its command (`stdbuf -oL cmd`, `xargs -- cmd`).
if not prefix_pending:
expect_command = False
continue
if not expect_command:
continue
# FOO=bar prefix: assignment list, next non-assignment token is the command.
# FOO=bar assignment prefix; next non-assignment token is the command.
if _ASSIGNMENT_RE.match(token):
continue
# `timeout 1 cmd` / `nice -n 5 cmd` style numeric wrapper arg.
# Numeric wrapper arg: `timeout 1 cmd` / `nice -n 5 cmd`.
if prefix_pending and token.lstrip("-").isdigit():
continue
base = _token_basename(token)
if base in _BLOCKED_COMMANDS:
blocked.add(base)
# Wrappers (`env` / `time` / `xargs` / `sudo`) consume one command; the
# next non-flag, non-numeric token is the real command. `sudo` is
# already in _BLOCKED_COMMANDS, so it's flagged AND we keep walking.
# Wrappers (env/time/xargs/sudo) consume one command; the next non-flag,
# non-numeric token is the real command. sudo is also in _BLOCKED_COMMANDS.
if base in _COMMAND_PREFIXES:
prefix_pending = True
continue
@ -229,10 +218,9 @@ def _find_blocked_commands(command: str) -> set[str]:
if base in _BLOCKED_COMMANDS:
blocked.add(base)
# Regex: blocked words at shell command boundaries that shlex won't see,
# e.g. inside an unquoted $(rm -rf), <(rm), backtick chain, or appended to
# a separator with no whitespace ("foo;rm"). Anchored to command-position
# delimiters; does not match in argument position.
# Regex catches blocked words at command boundaries shlex misses: inside
# $(rm -rf), <(rm), backtick chains, or "foo;rm". Anchored to command-position
# delimiters, so it doesn't match in argument position.
lowered = command.lower()
if _BLOCKED_COMMANDS:
words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS))
@ -243,11 +231,9 @@ def _find_blocked_commands(command: str) -> set[str]:
)
blocked.update(re.findall(pattern, lowered))
# Nested shell invocations (bash -c 'sudo whoami',
# bash -lc '...', bash --login -c '...', cmd /c '...').
# When a -c or /c flag is found, look backwards for a shell name
# (skipping intermediate flags like --login, -l, -x) and recursively
# scan the nested command string.
# Nested shell invocations (bash -c '...', bash -lc '...', cmd /c '...'):
# on a -c/-/c flag, look back for a shell name (skipping flags) and
# recursively scan the nested command string.
_SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"}
_SHELLS_WIN = {"cmd", "cmd.exe"}
for i, token in enumerate(tokens):
@ -259,10 +245,8 @@ def _find_blocked_commands(command: str) -> set[str]:
is_win_c = tok_lower == "/c"
if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens):
continue
# Look backwards past any flags to find the shell binary.
# On Unix, flags start with - (skip those). On Windows, flags
# start with / but so do absolute paths, so only skip short
# single-char /X flags (not /bin/bash style paths).
# Look back past flags for the shell binary. Windows flags and absolute
# paths both start with /, so only skip short /X flags (not /bin/bash).
for j in range(i - 1, -1, -1):
prev = tokens[j]
if prev.startswith("-"):
@ -282,17 +266,13 @@ def _find_blocked_commands(command: str) -> set[str]:
def _build_safe_env(workdir: str) -> dict[str, str]:
"""Build a minimal, credential-free environment for sandboxed subprocesses.
Whitelist-built from scratch -- the parent process env is NOT inherited.
Only PATH / HOME / TMPDIR / LANG / TERM / PYTHONIOENCODING (+ VIRTUAL_ENV
or Windows SystemRoot when applicable) reach the child. HF_TOKEN,
WANDB_API_KEY, AWS_*, GH_TOKEN, OPENAI_API_KEY, LD_PRELOAD, DYLD_*, and
every other parent var are absent by construction. HOME points at the
sandbox workdir so HF / wandb / aws SDKs cannot read cached credentials
from the operator's real ~/.
Whitelist-built from scratch (parent env NOT inherited): only PATH/HOME/
TMPDIR/LANG/TERM/PYTHONIOENCODING (+VIRTUAL_ENV or Windows SystemRoot) reach
the child; all credential vars (HF_TOKEN, AWS_*, etc.) are absent. HOME
points at the sandbox workdir so SDKs can't read the operator's cached creds.
"""
# Start with the directory containing the running Python interpreter
# so that subprocess calls to 'python', 'pip', etc. resolve to the
# same environment the Studio server is running in.
# Start from the running interpreter's dir so 'python'/'pip' resolve to the
# same environment the Studio server runs in.
exe_dir = os.path.dirname(sys.executable)
path_entries = [exe_dir] if exe_dir else []
@ -309,7 +289,7 @@ def _build_safe_env(workdir: str) -> dict[str, str]:
else:
path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"])
# Deduplicate while preserving order
# Deduplicate, preserving order.
deduped = list(dict.fromkeys(p for p in path_entries if p))
env = {
@ -322,17 +302,15 @@ def _build_safe_env(workdir: str) -> dict[str, str]:
}
if venv:
env["VIRTUAL_ENV"] = venv
# Windows needs SystemRoot for Python/subprocess to work
# Windows needs SystemRoot for Python/subprocess to work.
if sys.platform == "win32":
env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows")
return env
def _sandbox_preexec():
"""Best-effort sandbox setup for sandboxed subprocesses.
Modules are resolved at import time so the forked child runs no imports.
"""
"""Best-effort sandbox setup for sandboxed subprocesses (modules are
resolved at import time so the forked child runs no imports)."""
try:
os.setsid()
except OSError:
@ -354,9 +332,9 @@ def _sandbox_preexec():
except (OSError, AttributeError):
pass
# CLONE_NEWNET intentionally not applied: where userns is enabled it
# blocks all egress, including allowlisted hosts. Network policy is
# enforced by the AST host check and the bash blocklist.
# CLONE_NEWNET not applied: with userns enabled it blocks all egress,
# including allowlisted hosts. Network policy is enforced by the AST
# host check and the bash blocklist.
if _resource is not None:
# RLIMIT_NPROC is per-real-UID, so the cap is well above normal usage.
@ -380,11 +358,9 @@ def _sandbox_preexec():
except (ValueError, OSError, AttributeError):
pass
try:
# Default high enough for multi-shard safetensors mmaps + Python's
# own handle count; tunable via env for installs that hit the cap.
# High enough for multi-shard safetensors mmaps; tunable via env.
# Clamp to the inherited hard limit so setrlimit doesn't ValueError
# on machines where the parent's hard cap is below the requested
# value (would otherwise leave NOFILE at the parent's default).
# when the parent's hard cap is below the request.
nofile = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NOFILE", "16384"))
_soft_cur, hard_cur = _resource.getrlimit(_resource.RLIMIT_NOFILE)
target = nofile if hard_cur == _resource.RLIM_INFINITY else min(nofile, hard_cur)
@ -401,8 +377,7 @@ def _get_shell_cmd(command: str) -> list[str]:
# Per-session working directories so each chat thread gets its own sandbox.
# Falls back to a shared ~/studio_sandbox/_default for API callers without a
# session_id.
# Falls back to ~/studio_sandbox/_default for callers without a session_id.
_workdirs: dict[str, str] = {}
@ -567,10 +542,9 @@ RENDER_HTML_TOOL = {
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, RENDER_HTML_TOOL]
# OpenAI's function.name regex: ^[a-zA-Z0-9_-]{1,64}$ -- enforced before
# streaming starts. MCP servers can return tool names containing '.', '/',
# spaces, etc., which the prefix scheme would forward to OpenAI verbatim
# and 400 the whole request. Validate up front and skip with a warning.
# OpenAI's function.name regex ^[a-zA-Z0-9_-]{1,64}$, enforced before streaming.
# MCP tool names with '.', '/', spaces, etc. would 400 the whole request, so we
# validate up front and skip with a warning.
_OPENAI_FN_NAME_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")
@ -585,9 +559,8 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
logger.warning("Skipping MCP tool on '%s': empty name.", display)
continue
name = f"{MCP_TOOL_PREFIX}{server['id']}__{raw_name}"
# OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; bad chars
# (., /, spaces, etc.) or oversized names would 400 the whole
# request. Skip + warn so the rest of the tools still ship.
# Bad chars or oversized names would 400 the whole request; skip + warn
# so the rest of the tools still ship.
if not _OPENAI_FN_NAME_RE.fullmatch(name):
logger.warning(
"Skipping MCP tool '%s' on '%s': composed name '%s' is not "
@ -597,8 +570,7 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
name,
)
continue
# Same MCP server returning duplicate tool names would also 400
# OpenAI ("tools[N].function.name duplicates ..."). Drop dupes.
# Duplicate tool names would also 400 OpenAI; drop dupes.
if name in seen_names:
logger.warning("Skipping duplicate MCP tool '%s' on '%s'.", raw_name, display)
continue
@ -619,7 +591,7 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
async def get_enabled_mcp_tools() -> list[dict]:
servers = [s for s in mcp_servers_db.list_servers() if s.get("is_enabled")]
# Never spawn stdio servers when stdio is disabled on this host (e.g. a DB
# carried over from a desktop install onto a Colab / network deployment).
# carried from a desktop install onto a Colab/network deployment).
if not stdio_mcp_enabled():
servers = [s for s in servers if not is_stdio(s["url"])]
if not servers:
@ -681,11 +653,10 @@ def execute_tool(
timeout: int | None = _TIMEOUT_UNSET,
session_id: str | None = None,
) -> str:
"""Execute a tool by name with the given arguments. Returns result as a string.
"""Execute a tool by name with the given arguments; returns a string.
``timeout``: int sets per-call limit in seconds, ``None`` means no limit,
unset (default) uses ``_EXEC_TIMEOUT`` (300 s).
``session_id``: optional thread/session ID for per-conversation sandbox isolation.
``timeout``: int seconds, ``None`` = no limit, unset = ``_EXEC_TIMEOUT``.
``session_id``: optional ID for per-conversation sandbox isolation.
"""
logger.info(f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}")
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
@ -725,11 +696,10 @@ def execute_tool(
return f"Unknown tool: {name}"
_MAX_PAGE_CHARS = 16000 # limit fetched page text (after HTML-to-MD conversion)
# Raw download cap. Must be larger than _MAX_PAGE_CHARS because SSR pages
# embed large <head> sections (CSS, JS, SVGs) that are stripped during
# HTML-to-Markdown conversion. 512 KB is enough to reach article content
# on GitBook / Next.js / Docusaurus pages whose <head> alone can be 200 KB.
_MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion)
# Raw download cap > _MAX_PAGE_CHARS because SSR pages embed large <head>
# sections stripped during conversion; 512 KB reaches article content even
# where <head> alone is ~200 KB.
_MAX_FETCH_BYTES = 512 * 1024
_USER_AGENTS = (
@ -750,14 +720,12 @@ class _NoRedirect(urllib.request.HTTPRedirectHandler):
class _PinnedHTTPSConnection(http.client.HTTPSConnection):
"""HTTPS connection that connects to a pinned IP but uses a different
hostname for SNI and certificate verification.
"""HTTPS connection to a pinned IP, using a different hostname for SNI and
cert verification.
The SSRF IP-pinning rewrites URLs to raw IPs. A normal HTTPSConnection
would then send no SNI and verify the cert against the IP, both of which
fail. This subclass splits the two concerns: TCP connects to the pinned
IP (``host`` parameter) while TLS uses ``sni_hostname`` for the
ClientHello and cert check.
SSRF IP-pinning rewrites URLs to raw IPs; a normal HTTPSConnection would then
send no SNI and verify the cert against the IP (both fail). This splits the
concerns: TCP connects to the pinned IP (``host``), TLS uses ``sni_hostname``.
"""
def __init__(self, host: str, *, sni_hostname: str, **kwargs):
@ -765,8 +733,7 @@ class _PinnedHTTPSConnection(http.client.HTTPSConnection):
self._sni_hostname = sni_hostname
def connect(self):
# TCP connect to the pinned IP stored in self.host (+ tunnel if
# a proxy is configured via set_tunnel, though we do not use one).
# TCP connect to the pinned IP in self.host.
http.client.HTTPConnection.connect(self)
# TLS handshake with the real hostname for SNI + cert verification.
self.sock = self._context.wrap_socket(
@ -776,11 +743,11 @@ class _PinnedHTTPSConnection(http.client.HTTPSConnection):
class _SNIHTTPSHandler(urllib.request.HTTPSHandler):
"""HTTPS handler that sends the correct SNI hostname during TLS handshake.
"""HTTPS handler sending the correct SNI hostname during TLS handshake.
The SSRF IP-pinning rewrites URLs to raw IPs, which breaks SNI and cert
verification. This handler returns a ``_PinnedHTTPSConnection`` that
connects to the pinned IP but verifies TLS against the original hostname.
SSRF IP-pinning breaks SNI and cert verification; this returns a
``_PinnedHTTPSConnection`` that connects to the pinned IP but verifies TLS
against the original hostname.
"""
def __init__(self, hostname: str):
@ -798,9 +765,9 @@ class _SNIHTTPSHandler(urllib.request.HTTPSHandler):
def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]:
"""Resolve *hostname*, reject non-public IPs, return a pinned IP string.
Returns ``(ok, reason_or_empty, resolved_ip)``. The caller should
connect to *resolved_ip* (with a ``Host`` header) to prevent DNS
rebinding between validation and the actual fetch.
Returns ``(ok, reason_or_empty, resolved_ip)``. The caller should connect
to *resolved_ip* (with a ``Host`` header) to prevent DNS rebinding between
validation and the actual fetch.
"""
import ipaddress
import socket
@ -815,14 +782,10 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str
for *_, sockaddr in infos:
ip = ipaddress.ip_address(sockaddr[0])
# `not ip.is_global` rejects every category the denylist below
# also rejects PLUS shared address space (100.64.0.0/10 carrier-
# grade NAT) and benchmarking/documentation/exchange ranges that
# Python classifies with `is_private=False` and `is_global=False`
# (see https://docs.python.org/3/library/ipaddress.html#ipaddress.IPv4Address.is_global).
# The explicit predicates after it give human-readable categories
# in the error message, but a single non-global check is the
# source of truth and prevents future ranges from leaking.
# `not ip.is_global` is the source of truth: it rejects every category
# below PLUS shared/CGNAT (100.64.0.0/10) and benchmarking/doc ranges
# Python marks is_private=False and is_global=False. The explicit
# predicates only give human-readable categories in the error message.
if (
not ip.is_global
or ip.is_private
@ -834,7 +797,7 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str
):
return False, f"Blocked: refusing to fetch non-public address {ip}.", ""
# Return the first resolved address for pinning
# Return the first resolved address for pinning.
first_ip = infos[0][4][0]
return True, "", first_ip
@ -872,8 +835,8 @@ def _fetch_page_text(
ua = random.choice(_USER_AGENTS)
for _hop in range(5):
# Pin to the validated IP to prevent DNS rebinding.
# Rewrite the URL to use the IP and set the Host header.
# Pin to the validated IP (prevents DNS rebinding): rewrite URL to
# the IP, set the Host header.
cp = urlparse(current_url)
# Bracket IPv6 addresses so the netloc is valid in a URL.
ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip
@ -913,7 +876,7 @@ def _fetch_page_text(
return reason2
current_host = rp.hostname
continue
# Success -- read capped body
# Success: read capped body.
raw_bytes = resp.read(max_bytes)
break
else:
@ -926,7 +889,7 @@ def _fetch_page_text(
except Exception as e:
return f"Failed to fetch URL: {e}"
# Convert HTML to Markdown using the builtin converter (no external deps)
# Convert HTML to Markdown with the builtin converter (no external deps).
from ._html_to_md import html_to_markdown
text = html_to_markdown(raw_html)
@ -948,7 +911,7 @@ def _web_search(
If ``url`` is provided, fetches that page directly instead of searching.
"""
# Direct URL fetch mode
# Direct URL fetch mode.
if url and url.strip():
fetch_timeout = 60 if timeout is None else min(timeout, 60)
return _fetch_page_text(url.strip(), timeout = fetch_timeout)
@ -980,14 +943,9 @@ def _web_search(
def _check_signal_escape_patterns(code: str):
"""
Check if code contains patterns that could escape signal-based timeouts.
Vendored from unsloth_zoo.rl_environments to avoid importing unsloth_zoo
(which requires GPU drivers and fails on Mac/Apple Silicon).
Returns (safe: bool, details: dict)
"""
"""Check for patterns that could escape signal-based timeouts. Returns
(safe: bool, details: dict). Vendored from unsloth_zoo.rl_environments to
avoid importing unsloth_zoo (needs GPU drivers; fails on Apple Silicon)."""
try:
tree = ast.parse(code)
except SyntaxError as e:
@ -1018,7 +976,7 @@ def _check_signal_escape_patterns(code: str):
return full_name in names
return False
# Dangerous os/subprocess functions that can execute shell commands
# Dangerous os/subprocess functions that can execute shell commands.
_SHELL_EXEC_FUNCS = frozenset(
{
"os.system",
@ -1071,8 +1029,8 @@ def _check_signal_escape_patterns(code: str):
return parts
return []
# Keyword argument names that carry command content (as opposed to
# control flags like check=True, text=True, capture_output=True).
# Kwarg names that carry command content (not control flags like
# check=True, text=True, capture_output=True).
_CMD_KWARGS = frozenset({"args", "command", "executable", "path", "file"})
def _check_args_for_blocked(args_nodes):
@ -1093,8 +1051,8 @@ def _check_signal_escape_patterns(code: str):
self.signal_aliases = {"signal"}
self.os_aliases = {"os"}
self.subprocess_aliases = {"subprocess"}
# Maps bare function names to their fully-qualified form
# for from-import tracking (e.g. "system" -> "os.system")
# Bare name -> fully-qualified form for from-import tracking
# (e.g. "system" -> "os.system").
self.shell_exec_aliases: dict[str, str] = {}
self.loop_depth = 0
@ -1130,7 +1088,7 @@ def _check_signal_escape_patterns(code: str):
self.os_aliases.add("os")
else:
self.subprocess_aliases.add("subprocess")
# Track from-imports of dangerous functions
# Track from-imports of dangerous functions.
for alias in node.names:
fq = f"{node.module}.{alias.name}"
if fq in _SHELL_EXEC_FUNCS:
@ -1197,7 +1155,7 @@ def _check_signal_escape_patterns(code: str):
)
# --- Shell escape detection ---
# Resolve the fully qualified function name for os.*/subprocess.*
# Resolve the FQ function name for os.*/subprocess.*
shell_func = None
if isinstance(func, ast.Attribute):
if isinstance(func.value, ast.Name):
@ -1206,11 +1164,11 @@ def _check_signal_escape_patterns(code: str):
elif func.value.id in self.subprocess_aliases:
shell_func = f"subprocess.{func.attr}"
elif isinstance(func, ast.Name):
# Check from-import aliases: from os import system; system(...)
# from-import aliases: from os import system; system(...)
shell_func = self.shell_exec_aliases.get(func.id)
if shell_func and shell_func in _SHELL_EXEC_FUNCS:
# Expand **kwargs dicts to inspect their keys
# Expand **kwargs dicts to inspect their keys.
expanded_kwargs: dict[str, ast.AST] = {}
has_opaque_kwargs = False
for kw in node.keywords:
@ -1229,7 +1187,7 @@ def _check_signal_escape_patterns(code: str):
blocked_in_args = _check_args_for_blocked(all_call_args)
if has_opaque_kwargs:
# Can't inspect dynamic **kwargs -- flag as unsafe
# Can't inspect dynamic **kwargs; flag as unsafe.
shell_escapes.append(
{
"type": "shell_escape_dynamic",
@ -1249,10 +1207,9 @@ def _check_signal_escape_patterns(code: str):
}
)
else:
# Only flag dynamic args for functions that interpret
# strings as shell commands, or when shell= might be
# enabled. Treat any non-literal-False shell= value
# as potentially True (conservative).
# Only flag dynamic args for funcs that interpret strings as
# shell commands, or when shell= might be on. Any non-literal-
# False shell= is treated as potentially True (conservative).
_STRING_SHELL_FUNCS = frozenset(
{
"os.system",
@ -1310,11 +1267,9 @@ def _check_signal_escape_patterns(code: str):
}
)
elif isinstance(node.type, ast.Name):
# Only flag BaseException and TimeoutError, NOT Exception.
# except Exception does not catch SystemExit or
# KeyboardInterrupt, so it cannot suppress timeout
# enforcement. Flagging Exception causes false positives
# on normal error-handling patterns.
# Flag BaseException/TimeoutError but NOT Exception: `except
# Exception` can't catch SystemExit/KeyboardInterrupt, so it
# can't suppress timeout enforcement.
if node.type.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
@ -1342,9 +1297,9 @@ def _check_signal_escape_patterns(code: str):
if visitor.imports_signal and not signal_tampering:
warnings.append("Code imports 'signal' module - review manually for safety")
# Static host policy: block metadata hosts and any literal host outside
# the trusted allowlist; uploads blocked regardless of host. Dynamic hosts
# are caught by the bash blocklist instead.
# Static host policy: block metadata hosts and any literal host outside the
# trusted allowlist; uploads blocked regardless of host. Dynamic hosts are
# caught by the bash blocklist.
network_calls: list[dict] = []
sensitive_file_reads: list[dict] = []
_NETWORK_FQ_PREFIXES = (
@ -1590,10 +1545,9 @@ def _check_signal_escape_patterns(code: str):
return True
return False
# Bare method-name fallback (`x.upload_file(...)`) is intentionally fuzzy,
# but should only fire when huggingface_hub / hf_api is actually imported
# somewhere in the snippet -- otherwise paramiko.upload_file, boto3
# create_commit, etc. hit a false positive. We pre-scan for the imports.
# Bare method-name fallback (`x.upload_file(...)`) is fuzzy, so it fires only
# when huggingface_hub/hf_api is imported; else paramiko.upload_file,
# boto3.create_commit, etc. would false-positive. Pre-scan for the imports.
_HF_IMPORT_MODULES = (
"huggingface_hub",
"hf_api",
@ -1611,8 +1565,8 @@ def _check_signal_escape_patterns(code: str):
if root in _HF_IMPORT_MODULES:
return True
elif isinstance(n, ast.Call) and n.args:
# __import__('huggingface_hub'), importlib.import_module('huggingface_hub'),
# and bare import_module('huggingface_hub') (via `from importlib import ...`).
# __import__('huggingface_hub'), importlib.import_module(...),
# and bare import_module(...) (via `from importlib import ...`).
arg0 = n.args[0]
if not (isinstance(arg0, ast.Constant) and isinstance(arg0.value, str)):
continue
@ -1631,13 +1585,9 @@ def _check_signal_escape_patterns(code: str):
_hf_in_scope = _module_has_hf_import(tree)
def _method_call_hf_upload_name(node: ast.Call) -> str | None:
"""Return the HF upload method name (`upload_file`, ...) or None.
Catches `HfApi().upload_file(...)` (Attribute) and
`from huggingface_hub import upload_file; upload_file(...)` (Name).
The bare-name branch fires only when an HF import is in scope, mirroring
the Attribute branch's gating so paramiko/boto3 do not false-positive.
"""
"""Return the HF upload method name (`upload_file`, ...) or None. Covers
the Attribute and bare-Name forms; the bare-name branch fires only when
an HF import is in scope so paramiko/boto3 don't false-positive."""
if not _hf_in_scope:
return None
f = node.func
@ -1647,9 +1597,8 @@ def _check_signal_escape_patterns(code: str):
return f.id
return None
# Kwargs that ship a credential over the wire. Sandbox env strips HF_TOKEN
# / WANDB_API_KEY / AWS_* up front, so any value here is hard-coded or
# lifted from the parent process.
# Kwargs that ship a credential over the wire. The sandbox env strips
# credentials up front, so any value here is hard-coded or lifted from parent.
_HF_SENSITIVE_KWARGS = frozenset(
{
"token",
@ -1672,16 +1621,11 @@ def _check_signal_escape_patterns(code: str):
)
def _reads_env_or_secret(node: ast.AST | None) -> bool:
"""True if any node in the subtree resolves to an env / process read.
"""True if any node in the subtree resolves to an env/process read.
Walking the subtree (not just the root) means wrapper calls like
`str(os.environ)`, `json.dumps(os.environ)`, or
`'-'.join(os.environ.values())` are caught too.
Covers: `os.environ`, `os.environ[K]`, `os.environ.get(K)`, `os.getenv(K)`,
bare `getenv(K)` (after `from os import getenv`), and
`subprocess.{run,check_output,Popen,getoutput,getstatusoutput}` which
the LLM could use to lift parent env via `printenv` / `env` / `set`.
Walks the whole subtree (not just the root) to catch wrappers like
`str(os.environ)`. Covers os.environ[/.get]/os.getenv, bare getenv, and
subprocess.{run,check_output,...} that could lift parent env via printenv.
"""
if node is None:
return False
@ -1751,10 +1695,10 @@ def _check_signal_escape_patterns(code: str):
Policy: HF uploads are allowed only when (a) no sensitive kwarg is set,
(b) no positional / keyword value reads `os.environ` or related env
readers, and (c) the path argument is a sandbox-local literal -- a
relative string with no `..`, an `open(<literal>)`, or inline bytes.
Dynamic / variable paths are rejected; the policy cannot prove safety
statically and the cost of a wrong-allow is a credential exfiltration.
readers, and (c) the path arg is a sandbox-local literal: a relative
string with no `..`, an `open(<literal>)`, or inline bytes. Dynamic /
variable paths are rejected since safety can't be proven statically and
a wrong-allow means credential exfiltration.
"""
for kw in node.keywords or []:
if kw.arg in _HF_SENSITIVE_KWARGS:
@ -1813,7 +1757,7 @@ def _check_signal_escape_patterns(code: str):
}
)
# Direct sock.connect((host, port)) bypasses the FQ-prefix branch below.
# Direct sock.connect((host, port)) bypasses the FQ-prefix branch.
if isinstance(node.func, ast.Attribute) and node.func.attr == "connect" and node.args:
a0 = node.args[0]
host_lit = None
@ -1947,9 +1891,8 @@ def _check_code_safety(code: str) -> str | None:
"""
safe, info = _check_signal_escape_patterns(code)
if not safe:
# SyntaxError from ast.parse -- let these through so the subprocess
# produces a normal Python traceback instead of a misleading
# "unsafe code detected" message.
# Let SyntaxError from ast.parse through so the subprocess produces a
# normal Python traceback instead of a misleading "unsafe code" message.
if info.get("error"):
return None
@ -2032,7 +1975,7 @@ def _python_exec(
tmp_path = None
workdir = _get_workdir(session_id)
# Snapshot image mtimes so we detect both new and overwritten files.
# Snapshot image mtimes to detect new and overwritten files.
_before: dict[str, int] = {}
if os.path.isdir(workdir):
for _name in os.listdir(workdir):
@ -2088,7 +2031,7 @@ def _python_exec(
result = f"Exit code {proc.returncode}:\n{result}"
result = _truncate(result) if result.strip() else "(no output)"
# Detect new or overwritten image files and append sentinel for frontend
# Detect new/overwritten images and append sentinel for the frontend
if session_id and os.path.isdir(workdir):
new_images = []
for _name in os.listdir(workdir):

123
studio/backend/core/inference/worker.py
View file

@ -4,14 +4,10 @@
"""
Inference subprocess entry point.
Each inference 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
(generate, load, unload) via mp.Queue. It exits on shutdown or unload.
Pattern follows core/training/worker.py.
Each session runs in a persistent spawn subprocess, giving a clean interpreter
with no stale module state (solves transformers version-switching). It stays
alive while a model is loaded, taking commands (generate, load, unload) via
mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker.py.
"""
from __future__ import annotations
@ -35,7 +31,7 @@ from utils.hardware import apply_gpu_ids
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 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)
@ -96,18 +92,10 @@ def _build_model_config(config: dict):
def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, bool] | None:
"""Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
When *model_names* is provided, only those models' ``blobs/``
directories are checked instead of scanning every cached model --
much faster on systems with many models. Accepts multiple names so
that LoRA loads can watch both the adapter repo and the base model
repo simultaneously.
*has_incomplete* is True when any ``*.incomplete`` files exist in the
watched blobs directories, indicating that ``huggingface_hub`` is
actively downloading.
Returns None if the state cannot be determined (import error,
permission error, etc.) so callers can skip stall logic.
With *model_names*, only those models' ``blobs/`` dirs are checked (faster);
accepts multiple names so LoRA loads can watch adapter + base repos at once.
*has_incomplete* is True when any ``*.incomplete`` files exist (download
active). None means state could not be determined, so callers skip stall logic.
"""
try:
from huggingface_hub.constants import HF_HUB_CACHE
@ -125,14 +113,12 @@ def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, b
for name in model_names:
if not name:
continue
# Skip local filesystem paths -- HF model IDs use forward
# slashes (org/model) but never start with / . ~ or contain
# backslashes. This distinguishes them from absolute paths,
# relative paths, and Windows paths.
# Skip local filesystem paths -- HF IDs (org/model) never start
# with / . ~ or contain backslashes.
if name.startswith(("/", ".", "~")) or "\\" in name:
continue
name = resolve_cached_repo_id_case(name)
# HF cache dir format: models--org--name (slashes -> --)
# HF cache dir format: models--org--name (slashes -> --).
cache_dir_name = "models--" + name.replace("/", "--")
blobs_dir = cache / cache_dir_name / "blobs"
if blobs_dir.exists():
@ -165,16 +151,10 @@ def _start_heartbeat(
) -> threading.Event:
"""Start a daemon thread that sends periodic status heartbeats.
Monitors the HF Hub cache directory for download activity. A stall
is only reported when ``*.incomplete`` files are present (indicating
``huggingface_hub`` is actively downloading) **and** the total cache
size has not changed for *stall_timeout* seconds.
Once the download finishes (no more ``.incomplete`` files), the stall
timer resets, so post-download initialization (quantization, GPU
weight loading) is never misclassified as a stalled download.
Returns a stop event -- set it to terminate the heartbeat thread.
A stall is reported only when ``*.incomplete`` files are present (download
active) AND cache size hasn't changed for *stall_timeout* seconds. When the
download finishes the timer resets, so post-download init (quantization, GPU
weight load) isn't misclassified as a stall. Returns a stop event.
"""
stop = threading.Event()
transport = "https" if xet_disabled else "xet"
@ -188,7 +168,7 @@ def _start_heartbeat(
state = _get_hf_download_state(model_names)
now = time.monotonic()
# Skip stall logic if we cannot measure the cache
# Skip stall logic if we cannot measure the cache.
if state is None:
_send_response(
resp_queue,
@ -206,10 +186,8 @@ def _start_heartbeat(
last_size = current_size
last_change = now
# Only fire stall when .incomplete files are present,
# confirming a download is actively in progress.
# Once downloads finish (no .incomplete), reset the timer
# so model init time is not counted as a stall.
# Only fire stall while .incomplete files confirm an active download;
# reset the timer otherwise so model init isn't counted as a stall.
if not has_incomplete:
last_change = now
elif now - last_change >= stall_timeout:
@ -224,7 +202,7 @@ def _start_heartbeat(
"ts": time.time(),
},
)
# Only fire once -- the orchestrator will kill us
# fire once -- the orchestrator will kill us
return
_send_response(
@ -249,7 +227,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
hf_token = config.get("hf_token")
hf_token = hf_token if hf_token and hf_token.strip() else None
# Auto-detect quantization for LoRA adapters
# Auto-detect quantization for LoRA adapters.
load_in_4bit = config.get("load_in_4bit", True)
if mc.is_lora and mc.path:
import json
@ -280,10 +258,8 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
except Exception as e:
logger.warning("Could not read adapter_config.json: %s", e)
# Auto-enable trust_remote_code for NemotronH/Nano models only.
# NemotronH has config parsing bugs requiring trust_remote_code=True.
# Other transformers 5.x models are native and do NOT need it.
# NOTE: Must NOT match Llama-Nemotron (standard Llama architecture).
# Auto-enable trust_remote_code only for NemotronH/Nano (config parsing
# bugs require it). Must NOT match Llama-Nemotron (standard Llama arch).
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code:
@ -298,12 +274,10 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
model_name,
)
# Send heartbeats every 30s so the orchestrator knows we're still alive
# (download / weight loading can take a long time on slow connections)
# Heartbeat every 30s so the orchestrator knows we're alive during slow loads.
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
# Watch both the model repo and base model repo (for LoRA loads
# where the base model download is the actual bottleneck)
# Watch model + base repos (base download is the LoRA bottleneck).
watch_repos = [mc.identifier]
base = getattr(mc, "base_model", None)
if base and str(base) != mc.identifier:
@ -328,7 +302,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
heartbeat_stop.set()
if success:
# Build model_info for the parent to mirror
# Build model_info for the parent to mirror.
model_info = {
"identifier": mc.identifier,
"display_name": mc.display_name,
@ -341,8 +315,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
"audio_type": getattr(mc, "audio_type", None),
"has_audio_input": getattr(mc, "has_audio_input", False),
}
# Forward chat_template_info so the parent can classify
# capabilities without re-entering the subprocess.
# Forward chat_template_info so the parent can classify capabilities.
try:
_bm = getattr(backend, "models", {}) or {}
_entry = (
@ -397,22 +370,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"""Handle a generate command: stream tokens back via resp_queue.
cancel_event is an mp.Event shared with the parent process.
The parent can set it at any time (e.g. user stops generation,
or user loads a new model while generating) and generation
stops within 1-2 tokens.
cancel_event is an mp.Event the parent can set anytime (user stop, or new
model load mid-generate); generation stops within 1-2 tokens.
"""
request_id = cmd.get("request_id", "")
try:
# Decode image if provided
image = None
image_b64 = cmd.get("image_base64")
if image_b64:
image = _decode_image(image_b64)
image = _resize_image(image)
# Build generation kwargs
gen_kwargs = {
"messages": cmd["messages"],
"system_prompt": cmd.get("system_prompt", ""),
@ -426,9 +395,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"cancel_event": cancel_event,
}
# Optional template/tool plumbing: only forward keys that are
# actually present so the backend signature can evolve without
# breaking older command payloads.
# Forward only present optional keys so the backend signature can evolve.
for opt_key in (
"tools",
"enable_thinking",
@ -438,7 +405,6 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
if opt_key in cmd:
gen_kwargs[opt_key] = cmd[opt_key]
# Choose generation path
use_adapter = cmd.get("use_adapter")
if use_adapter is not None:
generator = backend.generate_with_adapter_control(
@ -451,7 +417,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
logger.info("Starting text generation for request_id=%s", request_id)
for cumulative_text in generator:
# cancel_event is an mp.Event — checked instantly, no queue polling
# cancel_event is an mp.Event — checked instantly, no queue polling.
if cancel_event.is_set():
logger.info("Generation cancelled for request %s", request_id)
break
@ -508,7 +474,7 @@ def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
use_adapter = cmd.get("use_adapter"),
)
# Send WAV bytes as base64 (bytes can't go through mp.Queue directly)
# Send WAV bytes as base64 (bytes can't go through mp.Queue directly).
_send_response(
resp_queue,
{
@ -542,7 +508,7 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
try:
import numpy as np
# Decode audio array from list (numpy arrays can't go through mp.Queue)
# numpy arrays can't go through mp.Queue, so decode from list.
audio_array = np.array(cmd["audio_data"], dtype = np.float32)
audio_type = cmd.get("audio_type")
@ -638,12 +604,12 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None:
"""Subprocess entrypoint. Persistent — runs command loop until shutdown.
"""Subprocess entrypoint. Persistent — runs the command loop until shutdown.
Args:
cmd_queue: mp.Queue for receiving commands from parent.
resp_queue: mp.Queue for sending responses to parent.
cancel_event: mp.Event shared with parent — set by parent to cancel generation.
cancel_event: mp.Event the parent sets to cancel generation.
config: Initial configuration dict with model info.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@ -668,7 +634,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
model_name = config["model_name"]
# ── 0. MLX fast-path — skip torch/transformers entirely ──
# ── 0. MLX fast-path — skip torch/transformers ──
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
@ -702,7 +668,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
)
return
# Enter same command loop as GPU path
# Enter the same command loop as the GPU path.
logger.info("MLX inference subprocess ready, entering command loop")
while True:
try:
@ -760,7 +726,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
)
return
# ── 1. Activate correct transformers version BEFORE any ML imports ──
# ── 1. Activate transformers version BEFORE any ML imports ──
try:
_activate_transformers_version(model_name)
except Exception as exc:
@ -775,7 +741,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
)
return
# ── 1b. On Windows, check Triton availability (must be before import torch) ──
# ── 1b. Windows: check Triton availability (must precede import torch) ──
if sys.platform == "win32":
try:
import triton # noqa: F401
@ -848,8 +814,8 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
return
# ── 4. Command loop — process commands until shutdown ──
# cancel_event is an mp.Event shared with parent — parent can set it
# at any time to cancel generation instantly (no queue polling needed).
# cancel_event is an mp.Event the parent can set anytime to cancel
# generation instantly (no queue polling needed).
logger.info("Inference subprocess ready, entering command loop")
while True:
@ -873,8 +839,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "load":
# Load a new model (reusing this subprocess)
# First unload current model
# Unload the current model before loading the new one.
if backend.active_model_name:
backend.unload_model(backend.active_model_name)
_handle_load(backend, cmd, resp_queue)
@ -891,7 +856,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
_handle_unload(backend, cmd, resp_queue)
elif cmd_type == "cancel":
# Redundant with mp.Event but handle gracefully
# Redundant with mp.Event but handle gracefully.
cancel_event.set()
logger.info("Cancel command received")
@ -907,7 +872,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
)
elif cmd_type == "status":
# Return current status
_send_response(
resp_queue,
{
@ -927,7 +891,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
elif cmd_type == "shutdown":
logger.info("Shutdown command received, exiting")
# Unload all models
for model_name in list(backend.models.keys()):
try:
backend.unload_model(model_name)

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

@ -3,24 +3,20 @@
"""Tool-call XML parsing and stripping helpers.
Extracted verbatim from studio/backend/core/inference/llama_cpp.py so that
external inference servers (llama-server wrappers, llama-swap, custom
shims) can reuse the same logic without importing the inference
Extracted verbatim from studio/backend/core/inference/llama_cpp.py so external
inference servers can reuse the logic without importing the inference
orchestrator, structlog, httpx, or the rest of the studio backend.
The regexes and function bodies are byte-for-byte identical to the
original inline implementation in llama_cpp.py. Any change made here must
preserve that equivalence; tests/python/test_tool_healing_extraction_is_exact.py
verifies it with AST comparison.
Regexes and bodies are byte-for-byte identical to the original; any change must
preserve that. test_tool_healing_extraction_is_exact.py verifies via AST.
"""
import json
import re
# Pre-compiled patterns for tool XML stripping. Hyphen in the
# function/parameter name char-class tracks OpenAI's allowed set so
# MCP tool names with dashes (mcp__srv__list-issues) and parameter
# names with dashes (`issue-number`) parse alongside the built-ins.
# Pre-compiled patterns for tool XML stripping. The hyphen in the name
# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues,
# issue-number) parse alongside the built-ins.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
@ -46,13 +42,13 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
Handles formats like:
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
Closing tags (</tool_call>, </function>, </parameter>) are all optional
since models frequently omit them.
Closing tags (</tool_call>, </function>, </parameter>) are all
optional since models frequently omit them.
"""
tool_calls = []
# Pattern 1: JSON inside <tool_call> tags.
# Use balanced-brace extraction that skips braces inside JSON strings.
# Pattern 1: JSON inside <tool_call> tags. Balanced-brace extraction that
# skips braces inside JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # position of the opening {
depth, i = 0, brace_start
@ -93,19 +89,16 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
pass
# Pattern 2: XML-style <function=name><parameter=key>value</parameter></function>
# All closing tags optional -- models frequently omit </parameter>,
# </function>, and/or </tool_call>.
# All closing tags optional; models frequently omit them.
if not tool_calls:
# Step 1: Find all <function=name> positions and extract their bodies.
# Body boundary: use only </tool_call> or next <function= as hard
# boundaries. We avoid using </function> as a boundary because
# code parameter values can contain that literal string.
# After extracting, we trim a trailing </function> if present.
# Step 1: Find <function=name> positions and extract bodies. Use only
# </tool_call> or the next <function= as hard boundaries (</function>
# can appear in code values); trim a trailing </function> afterwards.
func_starts = list(_TC_FUNC_START_RE.finditer(content))
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
body_start = fm.end()
# Hard boundaries: next <function= tag or </tool_call>
# Boundaries: next <function= tag or </tool_call>
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:
@ -114,18 +107,16 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
# Trim trailing </function> if present (it's the real closing tag)
body = _TC_FUNC_CLOSE_RE.sub("", body)
body = _TC_FUNC_CLOSE_RE.sub("", body) # trim closing </function>
# Step 2: Extract parameters from body.
# For single-parameter functions (the common case: code, command,
# query), use body end as the only boundary to avoid false matches
# on </parameter> inside code strings.
# Step 2: Extract parameters from body. For single-parameter
# functions, use body end as the only boundary to avoid matching
# </parameter> inside code strings.
arguments = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
# Single parameter: value is everything from after the tag
# to end of body, trimming any trailing </parameter>.
# Value is everything after the tag to end of body, less a
# trailing </parameter>.
pm = param_starts[0]
val = body[pm.end() :]
val = _TC_PARAM_CLOSE_RE.sub("", val)
@ -141,8 +132,7 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
else len(body)
)
val = body[val_start:next_param]
# Trim trailing </parameter> if present
val = _TC_PARAM_CLOSE_RE.sub("", val)
val = _TC_PARAM_CLOSE_RE.sub("", val) # trim trailing </parameter>
arguments[param_name] = val.strip()
tc = {

4
studio/backend/core/training/__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
"""
Training submodule - Training backends and trainer classes
"""
"""Training backends and trainer classes."""
from .training import TrainingBackend, TrainingProgress, get_training_backend

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

File diff suppressed because it is too large Load diff

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

@ -4,14 +4,10 @@
"""
Training backend — subprocess orchestrator.
Each training job runs in a fresh subprocess (mp.get_context("spawn")),
solving the transformers version-switching problem. The old in-process
UnslothTrainer singleton is only used inside the subprocess (worker.py).
This file orchestrates the subprocess lifecycle, pumps events from the
worker's mp.Queue, and exposes the same API surface to routes/training.py.
Pattern follows core/data_recipe/jobs/manager.py.
Each job runs in a fresh spawn subprocess (solving transformers version-switching);
the in-process UnslothTrainer singleton is only used inside the worker. This file
orchestrates the subprocess lifecycle, pumps events from the worker's mp.Queue, and
exposes the same API to routes/training.py. Pattern follows data_recipe/jobs/manager.py.
"""
import json as _json
@ -47,9 +43,8 @@ _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$")
def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
"""Remove only HF Trainer ``tmp-checkpoint-<step>/`` partials after a cancel.
Completed ``checkpoint-<int>/`` dirs and any non-numeric-suffix tmp dir
are user-owned and survive. Symlinked output_dir / children are skipped
so containment cannot be bypassed.
Completed ``checkpoint-<int>/`` dirs survive. Symlinked output_dir / children
are skipped so containment can't be bypassed.
"""
out = Path(output_dir)
if not out.exists() or not out.is_dir() or out.is_symlink():
@ -95,8 +90,7 @@ PLOT_HEIGHT = 3.5
@dataclass
class TrainingProgress:
"""Mirror of trainer.TrainingProgress — kept here so the parent process
never needs to import the heavy ML modules."""
"""Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules."""
epoch: float = 0
step: int = 0
@ -118,7 +112,7 @@ class TrainingProgress:
class TrainingBackend:
"""
Training orchestration backend — subprocess-based.
Launches a fresh subprocess per training job, communicates via mp.Queue.
Launches a fresh subprocess per job, communicates via mp.Queue.
"""
FLUSH_THRESHOLD: int = 10
@ -136,7 +130,7 @@ class TrainingBackend:
self._should_stop = False
self._cancel_requested = False # True only for stop(save=False)
# Training Metrics (consumed by routes for SSE and /metrics)
# Training metrics (consumed by routes for SSE and /metrics)
self.loss_history: list = []
self.lr_history: list = []
self.step_history: list = []
@ -169,7 +163,7 @@ class TrainingBackend:
"""Spawn a subprocess to run the full training pipeline.
All kwargs are serialized into a config dict and sent to the worker.
Returns True if the subprocess was started successfully.
Returns True if the subprocess started successfully.
"""
with self._lock:
if self._proc is not None and self._proc.is_alive():
@ -244,12 +238,11 @@ class TrainingBackend:
"gpu_ids": kwargs.get("gpu_ids"),
}
# Full finetuning always runs in 16-bit. LoRA/QLoRA and CPT preserve the
# explicit request so 4-bit adapter/raw-text runs remain possible.
# Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request.
if config["training_type"] == "Full Finetuning":
config["load_in_4bit"] = False
# Spawn subprocess — use locals so state is untouched on failure
# Spawn into locals so state is untouched on failure.
from utils.hardware import hardware as _hw
if _hw.DEVICE == _hw.DeviceType.MLX:
@ -296,8 +289,7 @@ class TrainingBackend:
logger.info("Training subprocess started (pid=%s)", proc.pid)
# Reset state — safe because old pump thread is confirmed dead
# and proc.start() succeeded
# Reset state (old pump thread dead, proc.start() succeeded).
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
@ -320,15 +312,14 @@ class TrainingBackend:
self._db_config = {k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}}
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Assign subprocess handles after state reset
# Assign subprocess handles after state reset.
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
# Eagerly create DB run row so the run appears in history during model loading
# Eagerly create DB run row so it appears in history during model loading.
self._ensure_db_run_created()
# Start event pump thread
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
@ -345,7 +336,7 @@ class TrainingBackend:
self._stop_queue.put({"type": "stop", "save": save})
except (OSError, ValueError):
pass
# Update progress immediately for responsive UI
# Update progress immediately for responsive UI.
self._progress.status_message = (
"Stopping training and saving checkpoint..." if save else "Cancelling training..."
)
@ -367,8 +358,7 @@ class TrainingBackend:
proc.kill()
proc.join(timeout = 2.0)
# Wait for pump thread to finish DB finalization before returning
# (8s covers SQLite's default 5s lock timeout plus execution overhead)
# Wait for pump thread to finish DB finalization (8s covers SQLite's 5s lock timeout).
if self._pump_thread is not None and self._pump_thread.is_alive():
self._pump_thread.join(timeout = 8.0)
@ -384,22 +374,19 @@ class TrainingBackend:
def is_training_active(self) -> bool:
"""Check if training is currently active."""
with self._lock:
# Subprocess alive = active
if self._proc is not None and self._proc.is_alive():
return True
# Stop was requested and process exited → inactive
if self._should_stop:
return False
# Check progress state
p = self._progress
if p.is_training:
return True
if p.is_completed or p.error:
return False
# Check status message for activity indicators
# Infer activity from the status message.
status_lower = (p.status_message or "").lower()
if any(
k in status_lower
@ -492,21 +479,19 @@ class TrainingBackend:
if self._proc is None or self._event_queue is None:
return
# Try to read an event
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
if event is not None:
self._handle_event(event)
continue
# No event — check if process is still alive
if self._proc.is_alive():
continue
# Process exited — drain remaining events
# Process exited — drain remaining events.
for e in self._drain_queue(self._event_queue):
self._handle_event(e)
# Mark as done if no explicit complete/error was received
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
if self._should_stop:
@ -530,9 +515,8 @@ class TrainingBackend:
def _handle_event(self, event: dict) -> None:
"""Apply a subprocess event to local state.
State updates happen inside self._lock; DB I/O happens after
releasing it so status-polling API endpoints are never blocked
by slow SQLite writes.
State updates happen inside self._lock; DB I/O happens after releasing
it so status-polling endpoints aren't blocked by slow SQLite writes.
"""
etype = event.get("type")
db_action: Optional[str] = None
@ -542,7 +526,7 @@ class TrainingBackend:
if etype == "progress":
self._progress.step = event.get("step", self._progress.step)
self._progress.epoch = event.get("epoch", self._progress.epoch)
# loss/lr are sanitized below; update progress after coercion
# loss/lr sanitized below.
_raw_loss = event.get("loss")
_raw_lr = event.get("learning_rate")
try:
@ -580,7 +564,7 @@ class TrainingBackend:
if status:
self._progress.status_message = status
# Update metric histories — reuse sanitized values from above
# Update metric histories using sanitized values.
step = event.get("step", 0)
loss = _safe_loss
lr = _safe_lr
@ -616,7 +600,7 @@ class TrainingBackend:
else:
eval_loss = None
# Buffer metric for DB flush (loss/lr already sanitized above)
# Buffer metric for DB flush.
self._metric_buffer.append(
{
"step": step,
@ -630,7 +614,7 @@ class TrainingBackend:
}
)
# Decide which DB action to take after releasing the lock
# Pick the DB action to run after releasing the lock.
if not self._db_run_created and self.current_job_id and self._db_config:
db_action = "create_run"
db_action_kwargs = {
@ -784,14 +768,14 @@ class TrainingBackend:
"""Flush buffered metrics to the database and update live progress."""
if not self._metric_buffer or not self.current_job_id or not self._db_run_created:
return
# Cap buffer to prevent unbounded memory growth
# Cap buffer to bound memory growth.
if len(self._metric_buffer) > 500:
logger.warning(
"Metric buffer exceeded 500 entries (%d) — trimming oldest",
len(self._metric_buffer),
)
self._metric_buffer = self._metric_buffer[-500:]
# Snapshot before insert so metrics arriving during the write are preserved
# Snapshot before insert so metrics arriving during the write survive.
batch = list(self._metric_buffer)
try:
from storage.studio_db import insert_metrics_batch, update_run_progress
@ -831,7 +815,7 @@ class TrainingBackend:
return events
# ------------------------------------------------------------------
# Plot generation (unchanged from original)
# Plot generation
# ------------------------------------------------------------------
def _create_loss_plot(
@ -954,9 +938,8 @@ class TrainingBackend:
def _transfer_to_inference_backend(self) -> bool:
"""Transfer model to inference backend.
With subprocess-based training, the model lives in the subprocess
and is freed when it exits. Inference must load from the saved
checkpoint on disk. This is a no-op placeholder.
No-op: with subprocess training the model is freed on exit, so inference
must load from the saved checkpoint on disk.
"""
logger.info(
"_transfer_to_inference_backend: subprocess training — "

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

@ -4,11 +4,9 @@
"""
Training subprocess entry point.
Each training job runs in a fresh subprocess (mp.get_context("spawn")).
This gives us a clean Python interpreter with no stale module state —
solving the transformers version-switching problem completely.
Pattern follows core/data_recipe/jobs/worker.py.
Each job runs in a fresh subprocess (mp.get_context("spawn")): a clean
interpreter with no stale module state, which solves transformers
version-switching. Pattern follows core/data_recipe/jobs/worker.py.
"""
from __future__ import annotations
@ -57,7 +55,7 @@ _FLASH_ATTN_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL"
_TILELANG_PACKAGE_VERSION = "0.1.8"
_APACHE_TVM_FFI_PACKAGE_VERSION = "0.1.9"
_TILELANG_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL"
# Pin both so plain pip cannot silently upgrade torch under the worker (fla-core needs torch>=2.7).
# Pin both so plain pip can't silently upgrade torch under the worker (fla-core needs torch>=2.7).
_FLA_PACKAGE_VERSION = "0.5.0"
_FLA_CORE_PACKAGE_VERSION = "0.5.0"
_FLA_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLA_INSTALL"
@ -72,14 +70,12 @@ _TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11")
_FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS"
# Module-level handle so the torch.library.Library registration survives past
# run_training_process() and is not garbage collected mid-run.
# run_training_process() and isn't GC'd mid-run.
_WINDOWS_ROCM_GROUPED_MM_LIB = None
# Worker subprocesses inherit the parent env but not the parent's
# os.add_dll_directory registrations. Replicate main.py's Windows ROCm DLL
# setup at module load so the first `import torch` can find amdhip64.dll even
# when HIP_PATH\bin is not on the system PATH. Handles retained at module
# scope so they are not garbage collected.
# Subprocesses don't inherit os.add_dll_directory registrations. Replicate
# main.py's Windows ROCm DLL setup so the first `import torch` finds
# amdhip64.dll. Handles retained at module scope so they aren't GC'd.
_ROCM_DLL_HANDLES: list = []
if sys.platform == "win32":
@ -94,7 +90,7 @@ if sys.platform == "win32":
)
def _ver_key(name: str) -> tuple:
# Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string.
# Numeric tuple key so "10.0" sorts after "7.0".
parts = []
for chunk in name.split("."):
try:
@ -146,22 +142,13 @@ def _model_wants_causal_conv1d(model_name: str) -> bool:
def _hipcc_gcc_install_dir() -> str | None:
"""Return the highest-numbered ``/usr/lib/gcc/x86_64-linux-gnu/<N>`` that has
BOTH the gcc runtime dir AND the corresponding ``/usr/include/c++/<N>`` C++
headers, or ``None`` if no match (or non-Linux / non-x86_64).
"""Highest-numbered ``/usr/lib/gcc/x86_64-linux-gnu/<N>`` that has BOTH the
gcc runtime dir AND ``/usr/include/c++/<N>`` headers, or None.
Ubuntu 24.04 ships ``/usr/lib/gcc/x86_64-linux-gnu/14/`` (gcc-14 runtime
objects) but does NOT ship ``/usr/include/c++/14`` in its default apt set;
libstdc++ headers come from ``libstdc++-13-dev``. ROCm clang-20 picks the
highest-numbered runtime dir by default, finds no ``<cstdlib>``, and the
HIP source build fails with::
/opt/rocm-X.Y/lib/llvm/lib/clang/20/include/__clang_hip_runtime_wrapper.h:112:10:
fatal error: 'cstdlib' file not found
Returning a path lets the caller pass ``--gcc-install-dir=<path>`` to clang
via ``HIPCC_COMPILE_FLAGS_APPEND``. Mirrors the same loop ``bbf004c`` added
to ``studio/setup.sh`` for the llama.cpp HIP build branch (PR #5301).
Ubuntu 24.04 ships gcc-14 runtime but not ``/usr/include/c++/14``; ROCm
clang-20 picks the highest runtime dir, finds no ``<cstdlib>``, and the HIP
build fails. The returned path is passed to clang via
``--gcc-install-dir``. Mirrors bbf004c in studio/setup.sh (PR #5301).
"""
if not sys.platform.startswith("linux"):
return None
@ -300,11 +287,10 @@ def _install_package_wheel_first(
pypi_spec,
]
# Source compilation on ROCm can take 10-30 minutes; use a generous
# timeout. Non-HIP installs preserve the pre-existing "no timeout"
# behaviour so unrelated slow installs (e.g. causal-conv1d source
# build on Linux aarch64 or unsupported torch/CUDA combinations)
# are not aborted at 5 minutes by this PR.
# ROCm source compilation can take 10-30 min; use a generous timeout.
# Non-HIP installs keep the pre-existing "no timeout" behaviour so unrelated
# slow installs (e.g. causal-conv1d source build on Linux aarch64, or
# unsupported torch/CUDA combos) aren't aborted at 5 minutes.
_run_kwargs: dict[str, Any] = {
"stdout": _sp.PIPE,
"stderr": _sp.STDOUT,
@ -312,16 +298,10 @@ def _install_package_wheel_first(
}
if is_hip:
_run_kwargs["timeout"] = 1800
# On Ubuntu 24.04 + ROCm clang-20, the HIP source build (causal-conv1d,
# mamba-ssm source fallback, flash-attn source fallback) defaults to
# /usr/lib/gcc/x86_64-linux-gnu/14/ which has the runtime dir but no
# /usr/include/c++/14 headers, and dies at:
# __clang_hip_runtime_wrapper.h:112:10:
# fatal error: 'cstdlib' file not found
# Inject --gcc-install-dir for a gcc whose C++ headers actually exist.
# Respect any pre-existing --gcc-install-dir in HIPCC_COMPILE_FLAGS_APPEND
# (user knows best); otherwise append. Mirrors the same fix bbf004c
# added to studio/setup.sh for the llama.cpp HIP build (PR #5301).
# On Ubuntu 24.04 + ROCm clang-20 the HIP source build dies on a missing
# <cstdlib> (gcc-14 runtime dir lacks C++ headers). Inject
# --gcc-install-dir for a gcc whose headers exist, respecting any
# pre-existing one. Mirrors bbf004c in studio/setup.sh (PR #5301).
_existing_flags = os.environ.get("HIPCC_COMPILE_FLAGS_APPEND", "")
if "--gcc-install-dir" not in _existing_flags:
_gcc_dir = _hipcc_gcc_install_dir()
@ -369,9 +349,9 @@ def _install_package_wheel_first(
)
else:
if sys.platform == "win32":
# No prebuilt wheel and no source build toolchain on Windows --
# this is expected for packages like causal-conv1d. Log at
# info so users aren't alarmed by what looks like an error.
# No prebuilt wheel and no source toolchain on Windows --
# expected for packages like causal-conv1d. Log at info so
# users aren't alarmed by what looks like an error.
logger.info(
"%s is not available on Windows (no prebuilt wheel); skipping",
display_name,
@ -487,7 +467,7 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
)
return False
# Probe once; reuse result so the --force-reinstall decision and the short-circuit
# Probe once; reuse so the --force-reinstall decision and the short-circuit
# share the same call count (stable for tests).
already_importable = _flash_linear_attention_importable()
if already_importable and _flash_linear_attention_current(already_importable = True):
@ -499,7 +479,7 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
f"Installing flash-linear-attention=={_FLA_PACKAGE_VERSION} for faster training...",
)
# `--no-deps` blocks the silent torch upgrade; we bring the non-torch runtime deps in by hand.
# `--no-deps` blocks the silent torch upgrade; bring non-torch runtime deps in by hand.
specs = [
*_FLA_RUNTIME_DEPS,
f"fla-core=={_FLA_CORE_PACKAGE_VERSION}",
@ -616,7 +596,7 @@ _MODEL_NAME_SEP_CHARS = ("-", ".", "/", " ")
def _discover_fla_model_types() -> frozenset[str]:
"""Model_types in the installed transformers whose modeling file imports `from fla.*`."""
"""Installed-transformers model_types whose modeling file imports `from fla.*`."""
global _TRANSFORMERS_FLA_MODEL_TYPES_CACHE
if _TRANSFORMERS_FLA_MODEL_TYPES_CACHE is not None:
return _TRANSFORMERS_FLA_MODEL_TYPES_CACHE
@ -690,17 +670,17 @@ def _torch_has_hip() -> bool:
def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
"""Classify a ROCm device as unified-memory (APU) or discrete.
Returns ``(gcn_arch, is_unified)`` where:
- ``gcn_arch`` is the canonical arch string (e.g. ``"gfx1151"``) when a
known attribute is present, or ``""`` when all arch attrs are absent.
- ``is_unified`` is ``True`` for AMD APUs with a shared GPU/system-RAM pool
Returns ``(gcn_arch, is_unified)``:
- ``gcn_arch``: canonical arch string (e.g. ``"gfx1151"``) when a known
attribute is present, else ``""``.
- ``is_unified``: ``True`` for AMD APUs with a shared GPU/system-RAM pool
(gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower
``set_per_process_memory_fraction`` cap to leave headroom for the OS.
``set_per_process_memory_fraction`` cap to leave OS headroom.
Classification priority:
1. ``gcnArchName`` / variant spellings (stable, naming-independent).
2. Device-name substring match as a last-resort fallback when all arch
attrs are absent (AMD SDK / Radeon wheels may not populate them):
2. Device-name substring match (last resort when all arch attrs absent;
AMD SDK / Radeon wheels may not populate them):
- gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M``
- gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395),
``Radeon 8050S`` (cut-down SKU)
@ -726,7 +706,7 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
def _tilelang_platform_supported() -> bool:
"""True iff a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch.
HIP excluded because tilelang 0.1.8 has no HIP GEMM instruction and crashes mid-backward.
HIP excluded: tilelang 0.1.8 has no HIP GEMM and crashes mid-backward.
"""
import platform as _platform
@ -770,9 +750,10 @@ def _run_pip(cmd: list[str], event_queue: Any, label: str) -> bool:
def _ensure_tilelang_backend_unconditional(event_queue: Any) -> bool:
"""Install pinned tilelang + apache-tvm-ffi; two-step repair if a broken tvm-ffi is present.
Returns True iff both import post-call. Step 1 surgically downgrades a broken tvm-ffi
with --force-reinstall --no-deps so torch / CUDA stay untouched; step 2 is a regular
install for missing transitive deps. Bypass via UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL=1.
Returns True iff both import post-call. Step 1 downgrades a broken tvm-ffi
with --force-reinstall --no-deps so torch / CUDA stay untouched; step 2 is a
regular install for missing transitive deps. Bypass via
UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL=1.
"""
if os.getenv(_TILELANG_SKIP_ENV) == "1":
return False
@ -800,7 +781,7 @@ def _ensure_tilelang_backend_unconditional(event_queue: Any) -> bool:
logger.info("tilelang + apache-tvm-ffi already installed")
return True
# Step 1: --no-deps keeps --force-reinstall from touching torch/CUDA via the dep graph.
# Step 1: --no-deps keeps --force-reinstall off torch/CUDA via the dep graph.
if needs_repair:
logger.info(
"Forcing apache-tvm-ffi downgrade: %s is on the broken list",
@ -822,7 +803,7 @@ def _ensure_tilelang_backend_unconditional(event_queue: Any) -> bool:
if not _run_pip(repair_cmd, event_queue, "TileLang backend repair"):
return False
# Step 2: regular install pulls in transitive deps (z3-solver, ml-dtypes) without touching torch.
# Step 2: regular install pulls transitive deps (z3-solver, ml-dtypes) without touching torch.
_send_status(
event_queue,
f"Installing TileLang=={_TILELANG_PACKAGE_VERSION} for faster training...",
@ -855,17 +836,17 @@ def _ensure_tilelang_backend(event_queue: Any, model_name: str) -> None:
# ── Fast-path hooks ──
# Wrap transformers' is_{flash_linear_attention,causal_conv1d}_available so the first call
# (at modeling import time) drives the install. Any model that queries the gate gets the
# install; models that never query it (Llama, Gemma, dense Qwen) pay nothing.
# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the legacy substring path.
# Wrap transformers' is_{flash_linear_attention,causal_conv1d}_available so the
# first call (at modeling import) drives the install. Models that never query
# the gate (Llama, Gemma, dense Qwen) pay nothing.
# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the substring path.
def _rebind_in_already_imported_modules(*, attr_name: str, old_obj: Any, new_obj: Any) -> int:
"""Rebind `attr_name -> new_obj` in every module that already imported `old_obj`.
"""Rebind `attr_name -> new_obj` in every module that imported `old_obj`.
`from X import Y` creates a local binding that reassigning X.Y won't reach.
Uses `__dict__.get` (not `getattr`) to skip lazy `__getattr__` aliases.
Uses `__dict__.get` to skip lazy `__getattr__` aliases.
"""
count = 0
missing = object()
@ -894,8 +875,8 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
logger.info("Fast-path hooks disabled via env; using substring fallback")
return
# On HIP torch, even already-installed tilelang crashes FLA's TileLang dispatch.
# User can override with FLA_TILELANG=1.
# On HIP torch, even installed tilelang crashes FLA's TileLang dispatch.
# Override with FLA_TILELANG=1.
if _torch_has_hip() and os.environ.get("FLA_TILELANG") is None:
os.environ["FLA_TILELANG"] = "0"
logger.info(
@ -937,8 +918,8 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
logger.warning("%s install raised: %s; falling back to torch", gate_name, exc)
ok = False
logger.info("%s hook done; available=%s", gate_name, ok)
# post_available_fn handles "gate already True but ancillary kernel broken" (e.g. tilelang
# missing while FLA imports fine); skip when install_fn already chained the follow-up.
# post_available_fn handles "gate already True but ancillary kernel broken"
# (e.g. tilelang missing while FLA imports); skip when install_fn already chained it.
if ok and not ran_install and post_available_fn is not None:
try:
post_available_fn(event_queue)
@ -966,7 +947,7 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
return True
def _fla_post_available(eq: Any) -> None:
# FLA already imports; repair tilelang if missing or on the broken tvm-ffi list.
# FLA imports; repair tilelang if missing or on the broken tvm-ffi list.
if not _model_wants_tilelang(model_name):
return
if _installed_tvm_ffi_version() not in _TVM_FFI_BROKEN_VERSIONS and _tilelang_importable():
@ -1057,8 +1038,8 @@ def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int
largest_side = max(width, height)
if largest_side <= target:
return width, height
# Integer formula matches unsloth_zoo's collator (Python round() differs
# by 1px on half-pixel cases). max(1, _) avoids zero-side degenerate output.
# Integer formula matches unsloth_zoo's collator (Python round() differs by
# 1px on half-pixel cases). max(1, _) avoids a zero-side degenerate output.
new_w = max(1, (width * target + largest_side // 2) // largest_side)
new_h = max(1, (height * target + largest_side // 2) // largest_side)
return new_w, new_h
@ -1079,9 +1060,9 @@ def _resize_mlx_vlm_image(image, resize):
if new_size != image.size:
resampling = getattr(Image, "Resampling", Image).LANCZOS
image = image.resize(new_size, resampling)
# When a resize is requested, hand mlx-vlm a writable RGB ndarray so its
# PIL-path square-resize is skipped and HF processors don't warn on
# non-writable views. resize=None (Default) above keeps the original PIL.
# On resize, hand mlx-vlm a writable RGB ndarray so its PIL-path
# square-resize is skipped and HF processors don't warn on non-writable
# views. resize=None above keeps the original PIL.
return np.array(image, copy = True)
@ -1092,12 +1073,12 @@ def _resize_mlx_vlm_images(value, resize):
def _adapt_for_mlx_vlm(items, resize = None):
"""Adapt GPU-path VLM dataset output for mlx-vlm consumption.
"""Adapt GPU-path VLM dataset output for mlx-vlm.
The GPU path embeds PIL images inside messages content as
{"type": "image", "image": PIL_Image}. mlx-vlm's prepare_inputs
needs images at top-level to produce pixel_values — regardless of
model type. Extract them and leave bare {"type": "image"} placeholders.
The GPU path embeds PIL images in message content as
{"type": "image", "image": PIL_Image}, but mlx-vlm's prepare_inputs needs
images at top-level to produce pixel_values (any model type). Extract them
and leave bare {"type": "image"} placeholders.
"""
adapted = []
for item in items:
@ -1218,8 +1199,8 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str:
def _run_mlx_training(event_queue, stop_queue, config):
"""Self-contained MLX training path for Apple Silicon.
Uses MLXTrainer from unsloth_zoo directly -- no torch/SFTTrainer needed.
Mirrors the event_queue protocol so the parent process pump works unchanged.
Uses unsloth_zoo's MLXTrainer directly (no torch/SFTTrainer). Mirrors the
event_queue protocol so the parent process pump works unchanged.
"""
import time
import math
@ -1276,8 +1257,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
lr_scheduler_type = _normalize_mlx_studio_scheduler(config.get("lr_scheduler_type", "linear"))
# ── 1. Load model ──
# Force text-only if the dataset is not an image dataset, even if the model
# has vision capabilities (e.g. Qwen3.5-VL trained on plain alpaca text).
# Force text-only for non-image datasets even on vision-capable models
# (e.g. Qwen3.5-VL trained on plain alpaca text).
_send("status", status_message = f"Loading {model_name}...")
is_dataset_image = bool(config.get("is_dataset_image", False))
training_type = config.get("training_type", "LoRA/QLoRA")
@ -1314,8 +1295,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
)
# ── 2. Apply LoRA / full FT ──
# Pass gradient_checkpointing as string ("mlx"/"unsloth"/"none"/etc.)
# get_peft_model and MLXTrainer both accept strings and handle them.
# gradient_checkpointing stays a string ("mlx"/"unsloth"/"none"/etc.);
# get_peft_model and MLXTrainer both accept and handle strings.
gc_setting = config.get("gradient_checkpointing", "mlx")
if isinstance(gc_setting, str):
use_grad_checkpoint = (
@ -1418,8 +1399,8 @@ def _run_mlx_training(event_queue, stop_queue, config):
eval_dataset = _load_local(config["local_eval_datasets"])
# ── 3b. Format dataset (VLM or text) ──
# Reuse the GPU path's format pipeline for both VLM (auto-detects OCR/caption/
# llava/sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column).
# Reuse the GPU format pipeline for VLM (auto-detects OCR/caption/llava/
# sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column).
format_type = config.get("format_type", "")
try:
from utils.datasets import format_and_template_dataset
@ -1511,7 +1492,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
output_dir = config.get("output_dir", "")
if not output_dir:
output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}"
# Resolve to ~/.unsloth/studio/outputs/ so the export page can find it
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import resolve_output_dir, ensure_dir
output_dir = str(resolve_output_dir(output_dir))
@ -1525,9 +1506,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
else:
eval_steps_val = int(eval_steps_val)
# MLX: per-element clip to [-1, 1]; norm clip disabled (it needs a
# global reduction that breaks MLX's eager pipeline). 1.0 (not 5.0):
# |g_i| > 5 rarely fires, so the historical 5.0 was effectively no-op.
# MLX: per-element clip to [-1, 1]; norm clip disabled (its global reduction
# breaks MLX's eager pipeline). 1.0 not 5.0: |g_i| > 5 rarely fires, so the
# historical 5.0 was effectively a no-op.
max_grad_norm = 0.0
max_grad_value = 1.0 # TODO: expose MLX grad-clip in Studio UI for power users
@ -1561,7 +1542,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
),
)
# Tell the parent that eval is configured so the frontend shows the eval chart
# Tell the parent eval is configured so the frontend shows the eval chart
if eval_dataset is not None and eval_steps_val > 0:
_send("eval_configured")
@ -1716,7 +1697,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
except _queue.Empty:
continue
except (EOFError, OSError):
# why safe: pipe permanently broken, no further messages can arrive
# Safe: pipe permanently broken, no more messages can arrive.
return
stop_thread = threading.Thread(target = _poll_stop, daemon = True)
@ -1753,15 +1734,14 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
"""Subprocess entrypoint. Fresh Python — no stale module state.
Args:
event_queue: mp.Queue for sending progress/status/error events to parent.
stop_queue: mp.Queue for receiving stop commands from parent.
config: Training configuration dict with all parameters.
event_queue: mp.Queue for progress/status/error events to the parent.
stop_queue: mp.Queue for stop commands from the parent.
config: Training config dict with all parameters.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
os.environ["PYTHONWARNINGS"] = "ignore" # before imports
# Offline auto-detect: skip ~25s of HF retries per call when DNS is
# dead. Scoped to this subprocess (orchestrator spawns a fresh one).
# Offline auto-detect: skip ~25s of HF retries per call when DNS is dead.
if "HF_HUB_OFFLINE" not in os.environ:
import socket as _socket
import threading as _threading
@ -1806,8 +1786,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
model_name = config["model_name"]
# ── 0. MLX FAST-PATH (must run before any torch/transformers imports) ──
# Apple Silicon uses MLXTrainer directly -- skip transformers version
# activation, causal-conv1d install, and torch imports entirely.
# Apple Silicon uses MLXTrainer directly -- skip torch imports / installs.
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
@ -1827,7 +1806,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
return
# Activate correct transformers version (Gemma-4 needs 5.5.0, etc.)
# Must happen before any transformers/mlx-lm imports in _run_mlx_training.
# before any transformers/mlx-lm imports in _run_mlx_training.
try:
_activate_transformers_version(model_name)
except Exception:
@ -1860,11 +1839,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
return
# ── 1a. Auto-enable trust_remote_code for NemotronH/Nano models ──
# NemotronH has config parsing bugs in transformers that require
# trust_remote_code=True as a workaround. Other transformers 5.x models
# (Qwen3.5, Gemma 4, etc.) are native and do NOT need it — enabling it
# bypasses the compiler (disabling fused CE).
# NOTE: Must NOT match Llama-Nemotron (standard Llama architecture).
# NemotronH needs trust_remote_code=True to work around config-parsing bugs.
# Other 5.x models are native and don't need it (it bypasses the compiler,
# disabling fused CE). Must NOT match Llama-Nemotron (standard Llama arch).
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_lowered = model_name.lower()
if (
@ -1879,20 +1856,12 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
# ── 1b. Install fast-path kernel libraries for the chosen model.
#
# 1) causal-conv1d ALWAYS runs eagerly via the substring path.
# Some SSM modeling files (nemotron_h, falcon_h1, granitemoehybrid)
# use `lazy_load_kernel("causal-conv1d")` directly and never call
# transformers' `is_causal_conv1d_available()`, so the runtime
# hook on that gate would not fire for them.
# 2) FLA + tilelang: primary gate is the runtime hook on transformers'
# `is_flash_linear_attention_available`. Models whose architecture
# queries that gate auto-trigger the install; others never pay.
# `_install_fast_path_hooks` also wraps `is_causal_conv1d_available`
# as a defence in depth for newer modeling files that do use it.
# 3) mamba-ssm + flash-attn keep their existing substring / size gates.
# 4) `UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1` falls back to the
# substring path for FLA / tilelang.
# 1) causal-conv1d ALWAYS runs eagerly via the substring path: some SSM
# modeling files lazy_load it without calling is_causal_conv1d_available.
# 2) FLA + tilelang: gated by the runtime hook on
# is_flash_linear_attention_available (hooks also wrap causal-conv1d).
# 3) mamba-ssm + flash-attn keep their substring / size gates.
# 4) UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the substring path.
try:
_ensure_causal_conv1d_fast_path(event_queue, model_name)
if os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1":
@ -1923,12 +1892,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
return
# ── 1c. Set fork start method so dataset.map() can multiprocess ──
# The parent launched us via spawn (clean process), but the compiled
# SFTTrainer checks get_start_method() and disables num_proc if not "fork".
# Linux only: fork is the default start method and is safe here (no CUDA
# context exists yet). macOS defaults to spawn since Python 3.8 because
# fork is unsafe with macOS frameworks (Metal/MPS, CoreFoundation) --
# do NOT override on macOS. Windows has no fork at all.
# The compiled SFTTrainer disables num_proc if start method isn't "fork".
# Linux only and safe here (no CUDA context yet); macOS/Windows excluded.
if sys.platform == "linux":
import multiprocessing as _mp
try:
@ -1949,17 +1914,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
# ── 1d. Stub torchao on Windows ROCm ──
# Shared with the export worker; see core/_torchao_stub.py for the full
# rationale (torchao -> torch.distributed._functional_collectives crashes on
# Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm.
# Must run before any import of transformers / unsloth_zoo.
# See core/_torchao_stub.py for the rationale (no RCCL backend on Windows
# ROCm). No-op elsewhere. Must run before importing transformers/unsloth_zoo.
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
# ── 1e. Ensure torch.distributed helper attrs are present ──
# Single-GPU training never initialises the process group, so these helpers
# are never called — but transformers/trl import them unconditionally.
# Single-GPU never inits the process group, but transformers/trl import
# these unconditionally.
_td_stubs = {
"is_initialized": lambda: False,
"is_available": lambda: False,
@ -1987,36 +1950,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# ── 1f. Windows ROCm runtime patches ──
# torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm ≤ 7.12 Windows),
# causing 0xC0000005 (access violation) during training.
#
# Root cause: the JitDecomp autograd decomposition system (NOT torch.compile)
# dispatches _grouped_mm → _fused_adagrad_ → _grouped_mm HIP → null crash.
# TORCHDYNAMO_DISABLE=1 stops the compiler frontend but does NOT stop
# JitDecomp, so we must also override the CUDA dispatch key for _grouped_mm
# with a safe Python fallback.
#
# Fixed in AMD's wheel: torch==2.11.0+rocm7.13.0 — the 3-D batch and grouped
# (with offs) variants of _grouped_mm now have working HIP kernels on gfx1200.
# We gate the dispatch override on HIP < 7.13 so users on the fixed wheel get
# the real GPU kernel rather than our Python fallback.
#
# Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0.
#
# Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None,
# Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor
# offs: optional group-split offsets (MoE-style variable-size batches)
#
# torch is already in sys.modules from section 1e's `import torch.distributed`.
# Module-level _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past
# function return / mid-run GC.
# causing 0xC0000005 during training. Root cause: JitDecomp (not
# torch.compile) dispatches _grouped_mm → null crash; TORCHDYNAMO_DISABLE
# doesn't cover JitDecomp, so we also override the CUDA dispatch key with a
# Python fallback. Fixed in torch==2.11.0+rocm7.13.0, so gate on HIP < 7.13.
# Schema: _grouped_mm(self, mat2, offs=None, bias=None, out_dtype=None);
# offs: optional group-split offsets (MoE-style variable-size batches).
# _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past return/GC.
global _WINDOWS_ROCM_GROUPED_MM_LIB
if sys.platform == "win32":
_torch_for_rocm = sys.modules.get("torch")
# Broad check: torch.version.hip OR "rocm" in torch.__version__.
# AMD SDK / Radeon Windows wheels do not always populate
# torch.version.hip; without the broad check the BNB version pin,
# dynamo-disable, and _grouped_mm fallback below silently skip
# (matches the torchao stub gate above and main.py).
# Broad check (torch.version.hip OR "rocm" in __version__): AMD SDK /
# Radeon wheels don't always set torch.version.hip, and without it the
# BNB pin, dynamo-disable, and _grouped_mm fallback would silently skip.
_build_version_for_rocm = (
getattr(_torch_for_rocm, "__version__", "").lower()
if _torch_for_rocm is not None
@ -2030,20 +1976,16 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
)
if _is_win_rocm_torch:
# Disable dynamo (belt-and-suspenders; JitDecomp patch below is the
# real fix, but keeping dynamo off avoids any other compile paths).
# Disable dynamo (belt-and-suspenders; the JitDecomp patch is the
# real fix, but this avoids other compile paths).
if "TORCHDYNAMO_DISABLE" not in os.environ:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.info("Windows ROCm: torch.compile (dynamo) disabled")
# BNB auto-detects the HIP version from torch.version.hip and uses
# it to choose which DLL to load (e.g. "7.13" → rocm713.dll).
# AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its
# version suffix does not always match the torch HIP version (e.g.
# torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still
# ships rocm72.dll). We detect the actual DLL name from the installed
# package and override BNB's auto-detection. "72" is a safe fallback
# if detection fails. Callers may override by pre-setting the var.
# BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB
# wheel may ship a DLL whose suffix doesn't match. Detect the actual
# DLL name and override; "72" is a safe fallback. Callers may
# pre-set the var to override.
if "BNB_ROCM_VERSION" not in os.environ:
_bnb_rocm_ver = None
try:
@ -2064,10 +2006,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
if _m:
_all_vers.append(_m.group(1))
# Pick the highest numeric suffix so that e.g. "713"
# wins over "72" when both variants are present.
# Filesystem glob order is not guaranteed, so always
# sort rather than stopping at the first match.
# Highest numeric suffix wins (glob order isn't sorted).
if _all_vers:
_bnb_rocm_ver = max(_all_vers, key = lambda v: int(v))
except Exception:
@ -2081,31 +2020,21 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
_bnb_rocm_ver,
)
# Parse HIP version for the kernel-fix gate below.
# torch.version.hip can be "7.13.99004", "7.2.0", etc.
# AMD SDK / Radeon wheels may leave torch.version.hip unset and
# encode the ROCm version in torch.__version__ instead
# (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back
# to that string when version.hip is missing.
# Parse HIP version for the kernel-fix gate below, falling back to
# the rocm version embedded in torch.__version__ when version.hip is
# unset (AMD SDK / Radeon wheels).
def _hip_ver_at_least(major: int, minor: int) -> bool:
_hip_str = getattr(getattr(_torch_for_rocm, "version", None), "hip", None)
if not _hip_str:
# Try the standard "+rocmX.Y.Z" embedded version first
# (e.g. "2.11.0+rocm7.13.0").
# Try the standard "+rocmX.Y.Z" embedded version first.
_ver_match = re.search(r"rocm(\d+)\.(\d+)", _build_version_for_rocm)
if _ver_match:
return (
int(_ver_match.group(1)),
int(_ver_match.group(2)),
) >= (major, minor)
# AMD SDK / Radeon Windows wheels encode the build as
# "+rocmsdk<date>" (e.g. "2.9.0+rocmsdk20251116") with no
# explicit rocmX.Y component. The rocmsdk format was
# introduced after the gfx120X null-kernel fix landed in
# ROCm 7.13, so any wheel with this suffix is new enough to
# have working HIP kernels. Treat as >= 7.13 rather than
# falling back to False and installing the Python workaround
# on a wheel that doesn't need it.
# "+rocmsdk<date>" wheels postdate the gfx120X null-kernel
# fix (ROCm 7.13), so treat them as >= 7.13 (no workaround).
if "rocmsdk" in _build_version_for_rocm:
logger.debug(
"Windows ROCm: AMD SDK wheel detected (%r); "
@ -2139,10 +2068,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
return False
# _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12,
# causing 0xC0000005. AMD fixed it in ROCm 7.13 (torch 2.11+).
# Only install the Python fallback on the affected versions so users
# on 7.13+ get the real GPU kernel for MoE workloads.
# Install the Python fallback only on affected versions (ROCm ≤ 7.12)
# so 7.13+ uses the real GPU kernel.
if not _hip_ver_at_least(7, 13):
try:
import warnings as _warnings
@ -2159,24 +2086,20 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
"""Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
_t = _torch_for_rocm
if offs is None:
# No offsets: behave like the real op, which
# accepts either (M, K) x (K, N) -> mm, or 3-D
# batched inputs -> bmm. Picking torch.mm
# unconditionally previously raised "self must be
# a matrix" on 3-D MoE workloads.
# No offsets: 2-D -> mm, 3-D batched -> bmm
# (unconditional mm broke 3-D MoE).
if self.dim() == 3 and mat2.dim() == 3:
result = _t.bmm(self.contiguous(), mat2.contiguous())
elif self.dim() == 3 and mat2.dim() == 2:
# Broadcast 2-D mat2 across the batch dim.
result = _t.matmul(self.contiguous(), mat2.contiguous())
elif self.dim() == 2 and mat2.dim() == 3:
# Broadcast 2-D self across batch via matmul semantics.
# Broadcast 2-D self across batch via matmul.
result = _t.matmul(self.contiguous(), mat2.contiguous())
else:
result = _t.mm(self.contiguous(), mat2.contiguous())
else:
# Grouped case: offs[i] is the exclusive end-row of
# group i in `self`; mat2 may be 3-D or 2-D.
# Grouped: offs[i] is the exclusive end-row of group i.
offs_list = offs.tolist()
pieces = []
prev = 0
@ -2189,7 +2112,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
b_part = mat2.contiguous()
pieces.append(_t.mm(a_part, b_part))
prev = end
# Include any trailing rows not covered by offs
# Include trailing rows not covered by offs.
if prev < self.shape[0]:
a_tail = self[prev:].contiguous()
b_tail = (
@ -2237,26 +2160,18 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
# ── 1g. ROCm OOM guard ──
# On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can
# cause a HIP driver hang that freezes the entire system rather than
# raising a Python exception. set_per_process_memory_fraction caps the
# HIP allocator so PyTorch raises OutOfMemoryError before hitting the
# hardware limit, giving the UI a clean error instead of a system freeze.
# Only applied on ROCm -- NVIDIA CUDA has a graceful OOM path and does
# not need this cap.
# Unified-memory APUs (gfx1150 Strix Point / gfx1151 Strix Halo) share GPU
# and system RAM in one pool: 0.90 of 128 GB starves the OS. Use 0.80 there.
# Primary classifier: gcnArchName from device properties — stable within a
# product family and naming-independent. AMD SDK / Radeon wheels may omit
# gcnArchName or expose it under a variant spelling, so we try several attr
# names then fall back to known device-name markers as a last resort.
# Non-fatal: silently skipped if torch is not importable.
# On ROCm, exhausting VRAM can hang the HIP driver instead of raising.
# set_per_process_memory_fraction caps the allocator so PyTorch raises
# OutOfMemoryError first (NVIDIA already has a graceful OOM path).
# Unified-memory APUs (gfx1150/gfx1151) share GPU+system RAM, so use 0.80
# vs 0.90 for discrete. Classify via gcnArchName, else device-name markers.
# Non-fatal: skipped if torch is not importable.
if _hw.IS_ROCM:
try:
import torch as _torch_mem
if _torch_mem.cuda.is_available():
# Classify unified vs discrete via _rocm_classify_unified_memory.
# See that function's docstring for classification priority.
# Classify unified vs discrete via _rocm_classify_unified_memory
# (see its docstring for classification priority).
_props = _torch_mem.cuda.get_device_properties(0)
_dev_name = _props.name
_gcn_arch, _is_unified = _rocm_classify_unified_memory(_props)
@ -2310,9 +2225,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
return
# ── 2b. EMBEDDING MODEL FAST-PATH ──
# Embedding models use a completely different pipeline (FastSentenceTransformer
# + SentenceTransformerTrainer + MultipleNegativesRankingLoss) so we branch
# early and handle the entire flow in a self-contained function.
# Embedding models use a different pipeline (FastSentenceTransformer +
# SentenceTransformerTrainer + MultipleNegativesRankingLoss), so branch early
# and handle the whole flow in a self-contained function.
if config.get("is_embedding", False):
try:
_run_embedding_training(event_queue, stop_queue, config)
@ -2380,9 +2295,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
stop_thread.start()
# ── 4. Execute the training pipeline ──
# Order: detect → dataset → model → prepare → train
# Dataset processing (including LLM-assisted detection) runs BEFORE model
# loading so both never occupy VRAM at the same time.
# Order: detect → dataset → model → prepare → train. Dataset processing runs
# BEFORE model loading so both never occupy VRAM at once.
try:
hf_token = config.get("hf_token", "")
hf_token = hf_token if hf_token and hf_token.strip() else None
@ -2427,32 +2341,13 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
dataset = dataset_result
eval_dataset = None
# [DEBUG] Print first sample before model is loaded
# dataset is a dict {"dataset": <Dataset>, "detected_format": ..., ...}
# or a raw Dataset for audio paths
# try:
# ds = dataset["dataset"] if isinstance(dataset, dict) else dataset
# print(
# f"\n[DEBUG] Dataset loaded BEFORE model. type={type(ds).__name__}, len={len(ds)}",
# flush = True,
# )
# print(f"[DEBUG] Columns: {ds.column_names}", flush = True)
# sample = ds[0]
# preview = {k: str(v)[:300] for k, v in sample.items()}
# print(f"[DEBUG] First sample: {preview}\n", flush = True)
# except Exception as e:
# print(
# f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}",
# flush = True,
# )
# Disable eval if eval_steps <= 0
eval_steps = config.get("eval_steps", 0.00)
if eval_steps is not None and float(eval_steps) <= 0:
eval_dataset = None
# Tell the parent process that eval is configured so the frontend
# shows "Waiting for first evaluation step..." instead of "not configured"
# Tell the parent eval is configured so the frontend shows
# "Waiting for first evaluation step..." instead of "not configured".
if eval_dataset is not None:
event_queue.put(
{
@ -2475,7 +2370,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
return
# ── Start tqdm monitor early so it captures download + tokenization bars ──
# ── Start tqdm monitor early to capture download + tokenization bars ──
import threading as _th
_tqdm_stop = _th.Event()
@ -2533,9 +2428,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# ── 4d. Prepare model (LoRA, full finetuning, or CPT) ──
if is_cpt:
_send_status(event_queue, "Configuring LoRA for continued pretraining...")
# embed_tokens (if the user included it) goes to modules_to_save —
# trained full-precision at embedding_learning_rate. lm_head stays as
# a LoRA target for merge compatibility (see unsloth PR #4106).
# embed_tokens (if included) goes to modules_to_save — trained
# full-precision at embedding_learning_rate. lm_head stays a LoRA
# target for merge compatibility (see unsloth PR #4106).
_user_modules = config.get("target_modules") or []
wants_embed = "embed_tokens" in _user_modules
cpt_trains_embeddings = wants_embed
@ -2610,13 +2505,13 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
return
# embedding_learning_rate is validated by the Pydantic model (Optional[float],
# gt=0, lt=1.0); if present it is already a finite float in range.
# embedding_learning_rate is validated by Pydantic (Optional[float],
# gt=0, lt=1.0); if present it's already a finite float in range.
embedding_lr_value = config.get("embedding_learning_rate")
if is_cpt:
if cpt_trains_embeddings:
if embedding_lr_value is None:
# Default embedding_learning_rate = lr/10 per Unsloth's CPT notebook.
# Default embedding_learning_rate = lr/10 (Unsloth CPT notebook).
embedding_lr_value = lr_value / 10.0
logger.info(
f"CPT: using default embedding_learning_rate={embedding_lr_value:.1e} "
@ -2644,7 +2539,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
tensorboard_dir = str(resolve_tensorboard_dir(tensorboard_dir))
ensure_dir(Path(tensorboard_dir))
# Start training (directly — no inner thread, we ARE the subprocess)
# Start training directly — no inner thread, we ARE the subprocess.
dataset_display = config.get("hf_dataset", "") or config.get("uploaded_file", "") or ""
_send_status(
event_queue,
@ -2761,10 +2656,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
"""Self-contained embedding model training pipeline.
Uses FastSentenceTransformer + SentenceTransformerTrainer +
MultipleNegativesRankingLoss — completely separate from the
LLM/VLM/audio paths in UnslothTrainer.
Mirrors the pattern from the reference embedding notebooks:
MultipleNegativesRankingLoss — separate from UnslothTrainer's LLM/VLM/audio
paths. Mirrors the reference embedding notebooks:
All_MiniLM_L6_v2.py, BGE_M3.py, EmbeddingGemma_300M.py,
ModernBert.py, Qwen3_Embedding_0_6B.py
"""
@ -2860,7 +2753,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
_send_status(event_queue, "Configuring LoRA adapters (FEATURE_EXTRACTION)...")
try:
gradient_checkpointing = config.get("gradient_checkpointing", False)
# Normalize: "none" or empty → False
# Normalize "none"/empty → False.
if gradient_checkpointing in ("none", "", None):
gradient_checkpointing = False
@ -2913,8 +2806,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
token = hf_token,
)
elif local_datasets:
# Load from local file(s) — mirrors the non-embedding pipeline's
# directory handling so recipe outputs (parquet-files/) work.
# Load local file(s) — mirrors the non-embedding pipeline's directory
# handling so recipe outputs (parquet-files/) work.
all_files: list[str] = []
for dataset_file in local_datasets:
file_path = (
@ -3050,7 +2943,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
else:
training_args_kwargs["num_train_epochs"] = num_epochs if num_epochs > 0 else 2
# warmup: prefer warmup_ratio (standard for embedding scripts), fallback to steps
# warmup: prefer warmup_ratio (standard for embedding scripts), else steps
if warmup_ratio is not None and warmup_ratio > 0:
training_args_kwargs["warmup_ratio"] = warmup_ratio
elif warmup_steps_val is not None and warmup_steps_val > 0:
@ -3074,7 +2967,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
# ── 8. Create progress callback ──
class _EmbeddingProgressCallback(TrainerCallback):
"""Sends training progress events to the parent process via event_queue."""
"""Send training progress events to the parent via event_queue."""
def on_log(
self,

8
studio/backend/hub/__init__.py Normal file
View file

@ -0,0 +1,8 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hub + Download Manager feature module.
Self-contained routes, schemas, utilities, workers, and storage for the model
inventory layer and the HuggingFace download manager. Wired into the FastAPI
app via two routers plus startup/shutdown hooks in main.py."""

24
studio/backend/hub/dependencies.py Normal file
View file

@ -0,0 +1,24 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared FastAPI dependencies for Hub routes."""
from __future__ import annotations
from typing import Optional
from fastapi import Header
HUB_HF_TOKEN_HEADER = "X-Unsloth-HF-Token"
HUB_HF_TOKEN_MAX_LENGTH = 512
def get_hf_token(
hf_token: Optional[str] = Header(
None,
alias = HUB_HF_TOKEN_HEADER,
max_length = HUB_HF_TOKEN_MAX_LENGTH,
),
) -> Optional[str]:
token = (hf_token or "").strip()
return token or None

12
studio/backend/hub/routes/__init__.py Normal file
View file

@ -0,0 +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
"""Hub routers exposed at /api/hub/* and /api/hub/datasets/*."""
from hub.routes.inventory import router as inventory_router
from hub.routes.datasets import router as datasets_router
__all__ = [
"inventory_router",
"datasets_router",
]

138
studio/backend/hub/routes/datasets.py Normal file
View file

@ -0,0 +1,138 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Endpoints mounted at /api/hub/datasets/*."""
from __future__ import annotations
from typing import Optional
from fastapi import APIRouter, Body, Depends, Query, UploadFile
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
from hub.schemas.datasets import (
AiAssistMappingRequest,
AiAssistMappingResponse,
CachedDatasetsResponse,
CheckFormatRequest,
CheckFormatResponse,
DeleteCachedDatasetResponse,
LocalDatasetsResponse,
UploadDatasetResponse,
)
from hub.schemas.downloads import (
ActiveDownloadsResponse,
CancelDatasetDownloadRequest,
CancelDatasetDownloadResponse,
DatasetDownloadJobStatus,
DatasetDownloadStartResponse,
DownloadProgressResponse,
DownloadDatasetRequest,
TransportStatusResponse,
)
from hub.services.datasets import cache_inventory, downloads, formatting, local
router = APIRouter()
@router.post("/upload", response_model = UploadDatasetResponse)
async def upload_dataset(
file: UploadFile, current_subject: str = Depends(get_current_subject)
) -> UploadDatasetResponse:
return await local.upload_dataset_response(file)
@router.get("/local", response_model = LocalDatasetsResponse)
def list_local_datasets(
current_subject: str = Depends(get_current_subject),
) -> LocalDatasetsResponse:
return local.list_local_datasets_response()
@router.get(
"/cached",
response_model = CachedDatasetsResponse,
response_model_exclude_unset = True,
)
async def list_cached_datasets(current_subject: str = Depends(get_current_subject)):
return await cache_inventory.list_cached_datasets_response()
@router.delete("/cached", response_model = DeleteCachedDatasetResponse)
async def delete_cached_dataset(
repo_id: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject)
):
return await cache_inventory.delete_cached_dataset_response(repo_id)
@router.get("/download-progress", response_model = DownloadProgressResponse)
async def get_dataset_download_progress(
repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"),
expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_dataset_download_progress_response(
repo_id,
expected_bytes = expected_bytes,
hf_token = hf_token,
)
@router.post("/download", response_model = DatasetDownloadStartResponse, status_code = 202)
async def download_dataset(
body: DownloadDatasetRequest,
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await downloads.download_dataset_response(body, hf_token)
@router.post("/download/cancel", response_model = CancelDatasetDownloadResponse, status_code = 202)
async def cancel_dataset_download(
body: CancelDatasetDownloadRequest, current_subject: str = Depends(get_current_subject)
):
return await downloads.cancel_dataset_download_response(body)
@router.get("/download-status", response_model = DatasetDownloadJobStatus)
async def get_dataset_download_status(
repo_id: str = Query(..., description = "HuggingFace dataset repo ID"),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_dataset_download_status_response(repo_id)
@router.get("/active-downloads", response_model = ActiveDownloadsResponse)
async def get_active_dataset_downloads(
repo_id: str = Query("", description = "HuggingFace dataset repo ID"),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_active_dataset_downloads_response(repo_id)
@router.get("/transport-status", response_model = TransportStatusResponse)
async def get_dataset_transport_status(
repo_id: str = Query(..., description = "HuggingFace dataset repo ID"),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_dataset_transport_status_response(repo_id)
@router.post("/check-format", response_model = CheckFormatResponse)
def check_format(
request: CheckFormatRequest,
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return formatting.check_format_response(request, hf_token)
@router.post("/ai-assist-mapping", response_model = AiAssistMappingResponse)
def ai_assist_mapping(
request: AiAssistMappingRequest,
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return formatting.ai_assist_mapping_response(request, hf_token)

222
studio/backend/hub/routes/inventory.py Normal file
View file

@ -0,0 +1,222 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Endpoints mounted at /api/hub/* for the model inventory."""
from __future__ import annotations
from typing import Optional
from fastapi import APIRouter, Body, Depends, Query
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
from hub.schemas.downloads import (
ActiveDownloadsResponse,
CancelDownloadResponse,
CancelDownloadRequest,
DownloadProgressResponse,
DownloadJobStatus,
DownloadModelRequest,
DownloadStartResponse,
TransportStatusResponse,
)
from hub.schemas.inventory import (
AddScanFolderRequest,
BrowseFoldersResponse,
CachedGgufResponse,
CachedModelsResponse,
DeleteCachedModelResponse,
GgufVariantsResponse,
LocalModelListResponse,
RecommendedFoldersResponse,
RemoveScanFolderResponse,
ScanFolderInfo,
ScanFoldersResponse,
)
from hub.services.models import (
cache_inventory,
deletion,
downloads,
folder_browser,
gguf_variants,
local_inventory,
)
router = APIRouter()
@router.get("/local", response_model = LocalModelListResponse)
async def list_local_models(
models_dir: str = Query(
default = "./models", description = "Directory to scan for local model folders"
),
current_subject: str = Depends(get_current_subject),
):
return await local_inventory.list_local_models_response(models_dir)
# Plain `def` (not async): synchronous SQLite + filesystem work runs in
# FastAPI's thread-pool instead of blocking the event loop.
@router.get("/scan-folders", response_model = ScanFoldersResponse)
def get_scan_folders(current_subject: str = Depends(get_current_subject)):
return local_inventory.get_scan_folders_response()
@router.post("/scan-folders", response_model = ScanFolderInfo, status_code = 201)
def add_scan_folder_endpoint(
body: AddScanFolderRequest, current_subject: str = Depends(get_current_subject)
):
return local_inventory.add_scan_folder_response(body.path)
@router.delete("/scan-folders/{folder_id}", response_model = RemoveScanFolderResponse)
def remove_scan_folder_endpoint(
folder_id: int, current_subject: str = Depends(get_current_subject)
):
return local_inventory.remove_scan_folder_response(folder_id)
@router.get("/recommended-folders", response_model = RecommendedFoldersResponse)
def get_recommended_folders(current_subject: str = Depends(get_current_subject)):
return folder_browser.get_recommended_folders_response()
@router.get("/browse-folders", response_model = BrowseFoldersResponse)
def browse_folders(
path: Optional[str] = Query(None),
show_hidden: bool = Query(False),
current_subject: str = Depends(get_current_subject),
):
return folder_browser.browse_folders_response(path, show_hidden)
@router.get("/gguf-variants", response_model = GgufVariantsResponse)
async def get_gguf_variants(
repo_id: str = Query(
..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
),
prefer_local_cache: bool = Query(False),
offline: bool = Query(False),
local_path: Optional[str] = Query(None),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await gguf_variants.get_gguf_variants_response(
repo_id,
prefer_local_cache = prefer_local_cache,
offline = offline,
local_path = local_path,
hf_token = hf_token,
)
@router.post("/download", response_model = DownloadStartResponse, status_code = 202)
async def download_model(
body: DownloadModelRequest,
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await downloads.download_model_response(body, hf_token)
@router.post("/download/cancel", response_model = CancelDownloadResponse, status_code = 202)
async def cancel_download_model(
body: CancelDownloadRequest, current_subject: str = Depends(get_current_subject)
):
return await downloads.cancel_download_model_response(body)
@router.get("/download-status", response_model = DownloadJobStatus)
async def get_download_status(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_download_status_response(repo_id, gguf_variant)
@router.get("/active-downloads", response_model = ActiveDownloadsResponse)
async def get_active_downloads(
repo_id: str = Query("", description = "HuggingFace repo ID"),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_active_downloads_response(repo_id)
@router.get("/transport-status", response_model = TransportStatusResponse)
async def get_model_transport_status(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_model_transport_status_response(
repo_id,
gguf_variant,
hf_token,
)
@router.get(
"/gguf-download-progress",
response_model = DownloadProgressResponse,
response_model_exclude_none = True,
)
async def get_gguf_download_progress(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
variant: str = Query("", description = "Quantization variant (e.g. UD-TQ1_0)"),
expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_gguf_download_progress_response(
repo_id,
variant = variant,
expected_bytes = expected_bytes,
hf_token = hf_token,
)
@router.get("/download-progress", response_model = DownloadProgressResponse)
async def get_download_progress(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_download_progress_response(
repo_id,
expected_bytes = expected_bytes,
hf_token = hf_token,
)
@router.get("/cached-gguf", response_model = CachedGgufResponse)
async def list_cached_gguf(
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await cache_inventory.list_cached_gguf_response(hf_token)
@router.get("/cached-models", response_model = CachedModelsResponse)
async def list_cached_models(
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await cache_inventory.list_cached_models_response(hf_token)
@router.delete(
"/delete-cached",
response_model = DeleteCachedModelResponse,
response_model_exclude_none = True,
)
async def delete_cached_model(
repo_id: str = Body(...),
variant: Optional[str] = Body(None),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
return await deletion.delete_cached_model_response(repo_id, variant, hf_token)

2
studio/backend/hub/schemas/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

108
studio/backend/hub/schemas/datasets.py Normal file
View file

@ -0,0 +1,108 @@
# 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
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, Field, model_validator
class CheckFormatRequest(BaseModel):
dataset_name: str
is_vlm: bool = False
subset: Optional[str] = None
train_split: Optional[str] = "train"
prefer_local_cache: bool = False
local_path: Optional[str] = None
@model_validator(mode = "before")
@classmethod
def _compat_split(cls, values: Any) -> Any:
if isinstance(values, dict) and "split" in values:
merged = {**values}
merged.setdefault("train_split", merged.pop("split"))
return merged
return values
class CheckFormatResponse(BaseModel):
requires_manual_mapping: bool
detected_format: str
columns: List[str]
is_image: bool = False
is_audio: bool = False
multimodal_columns: Optional[List[str]] = None
suggested_mapping: Optional[Dict[str, str]] = None
detected_image_column: Optional[str] = None
detected_audio_column: Optional[str] = None
detected_text_column: Optional[str] = None
detected_speaker_column: Optional[str] = None
preview_samples: Optional[List[Dict]] = None
total_rows: Optional[int] = None
warning: Optional[str] = None
class AiAssistMappingRequest(BaseModel):
columns: List[str]
samples: List[Dict[str, Any]]
dataset_name: Optional[str] = None
model_name: Optional[str] = None
model_type: Optional[str] = None
class AiAssistMappingResponse(BaseModel):
success: bool
suggested_mapping: Optional[Dict[str, str]] = None
warning: Optional[str] = None
system_prompt: Optional[str] = None
user_template: Optional[str] = None
assistant_template: Optional[str] = None
label_mapping: Optional[Dict[str, Dict[str, str]]] = None
dataset_type: Optional[str] = None
is_conversational: Optional[bool] = None
user_notification: Optional[str] = None
class UploadDatasetResponse(BaseModel):
filename: str = Field(..., description = "Original filename")
stored_path: str = Field(..., description = "Absolute path stored on backend")
class LocalDatasetItem(BaseModel):
class Metadata(BaseModel):
actual_num_records: Optional[int] = None
target_num_records: Optional[int] = None
total_num_batches: Optional[int] = None
num_completed_batches: Optional[int] = None
columns: Optional[List[str]] = None
id: str
label: str
path: str
source: Literal["recipe", "upload"]
rows: Optional[int] = None
updated_at: Optional[float] = None
metadata: Optional[Metadata] = None
class LocalDatasetsResponse(BaseModel):
datasets: List[LocalDatasetItem] = Field(default_factory = list)
class CachedDatasetItem(BaseModel):
repo_id: str
size_bytes: int = 0
cache_path: Optional[str] = None
processed_cache: bool = False
partial: bool = False
partial_transport: Optional[str] = None
class CachedDatasetsResponse(BaseModel):
cached: List[CachedDatasetItem] = Field(default_factory = list)
class DeleteCachedDatasetResponse(BaseModel):
status: str
repo_id: str

161
studio/backend/hub/schemas/downloads.py Normal file
View file

@ -0,0 +1,161 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Pydantic schemas for the Hub download manager (/api/hub/downloads/*)."""
from pydantic import BaseModel, Field
from typing import List, Literal, Optional
DownloadJobState = Literal["idle", "running", "cancelling", "cancelled", "complete", "error"]
class DownloadModelRequest(BaseModel):
"""Body for POST /api/hub/download.
The HuggingFace token travels in the internal Hub token header.
"""
repo_id: str = Field(
...,
description = "HuggingFace repo ID, e.g. 'unsloth/Qwen3-4B-GGUF'",
)
gguf_variant: Optional[str] = Field(
None,
description = "Quantization label (e.g. 'Q4_K_M'). Required for GGUF repos.",
)
use_xet: bool = Field(
False,
description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.",
)
class CancelDownloadRequest(BaseModel):
repo_id: str = Field(..., description = "HuggingFace repo ID")
gguf_variant: Optional[str] = Field(
None,
description = "GGUF variant label; omit for safetensors snapshots",
)
generation: Optional[int] = Field(
None,
description = "Download generation tag from a prior start; passing it scopes the cancel to that exact run.",
)
class DownloadJobStatus(BaseModel):
"""Live state of a background download job."""
state: DownloadJobState = Field(
...,
description = "Current download job state.",
)
error: Optional[str] = Field(None, description = "Error message if state == 'error'")
generation: int = Field(
0,
description = "Current run generation; an adopting client stores it so a later cancel is scoped to this exact run.",
)
class DownloadStartResponse(BaseModel):
job_key: str
state: str
accepted: bool
generation: int
class CancelDownloadResponse(BaseModel):
job_key: str
state: str
class ActiveDownload(BaseModel):
"""One in-flight download for a repo. ``variant`` is null for safetensors."""
repo_id: Optional[str] = None
variant: Optional[str] = None
transport: Optional[str] = None
state: str
generation: int = Field(
0,
description = "Current run generation; an adopting client stores it so a later cancel is scoped to this exact run.",
)
class ActiveDownloadsResponse(BaseModel):
downloads: List[ActiveDownload]
class TransportCapability(BaseModel):
available: bool
reason: Optional[str] = None
class TransportCapabilities(BaseModel):
http: TransportCapability
xet: TransportCapability
class TransportStatusResponse(BaseModel):
has_partial: bool
last_transport: Optional[str] = None
resumable: bool
class DownloadProgressResponse(BaseModel):
downloaded_bytes: int
# Finalized-blob bytes only (no ``.incomplete``). Registry-loss completion
# fallbacks key off this so a partial isn't mistaken for a finished download.
completed_bytes: int = 0
complete_on_disk: bool = Field(
False,
description = (
"True only when the backend verified a usable completed snapshot/variant on disk."
),
)
expected_bytes: int
progress: float
cache_path: Optional[str] = None
class DownloadDatasetRequest(BaseModel):
"""Body for POST /api/hub/datasets/download.
The HuggingFace token travels in the internal Hub token header.
"""
repo_id: str = Field(..., description = "HuggingFace dataset repo ID")
use_xet: bool = Field(
False,
description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.",
)
class CancelDatasetDownloadRequest(BaseModel):
repo_id: str = Field(..., description = "HuggingFace dataset repo ID")
generation: Optional[int] = Field(None, description = "Download generation")
class DatasetDownloadJobStatus(BaseModel):
"""Live state of a background dataset download job."""
state: DownloadJobState = Field(
...,
description = "Current dataset download job state.",
)
error: Optional[str] = Field(None, description = "Error message if state == 'error'")
generation: int = Field(
0,
description = "Current run generation; an adopting client stores it so a later cancel is scoped to this exact run.",
)
class DatasetDownloadStartResponse(BaseModel):
repo_id: str
state: str
accepted: bool
generation: int
class CancelDatasetDownloadResponse(BaseModel):
repo_id: str
state: str

286
studio/backend/hub/schemas/inventory.py Normal file
View file

@ -0,0 +1,286 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Pydantic schemas for the Hub inventory layer (/api/hub/*).
Kept independent from upstream models/models.py so the Hub module can ship
without modifying any upstream schema."""
from pydantic import BaseModel, Field
from typing import List, Literal, Optional
ModelFormat = Literal["gguf", "safetensors", "adapter", "checkpoint", "unknown"]
ModelRuntime = Literal["llama_cpp", "transformers", "adapter", "unknown"]
class GgufVariantDetail(BaseModel):
"""A single GGUF quantization variant in a HuggingFace repo."""
filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')")
quant: str = Field(..., description = "Quantization label or internal GGUF variant key")
display_label: Optional[str] = Field(
None, description = "Optional user-facing label when quant is an internal key"
)
size_bytes: int = Field(0, description = "File size in bytes")
download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant")
downloaded: bool = Field(
False, description = "Whether this variant is already in the local HF cache"
)
partial: bool = Field(
False,
description = "Whether this variant has an in-progress (.incomplete) blob in cache",
)
partial_transport: Optional[str] = Field(
None,
description = (
'Transport recorded for the partial state ("http" or '
'"xet"), or null if not partial / unknown. Frontend uses '
"this to pick Resume (http) vs Redownload (xet) labels."
),
)
class GgufVariantsResponse(BaseModel):
"""Response for listing GGUF quantization variants in a HuggingFace repo."""
repo_id: str = Field(..., description = "HuggingFace repo ID")
variants: List[GgufVariantDetail] = Field(
default_factory = list, description = "Available GGUF variants"
)
has_vision: bool = Field(
False, description = "Whether the model has vision support (mmproj files)"
)
default_variant: Optional[str] = Field(
None, description = "Recommended default quantization variant"
)
class LocalModelCapabilities(BaseModel):
can_train: bool = False
can_chat: bool = False
can_delete: bool = False
can_download: bool = False
requires_variant: bool = False
supports_lora: bool = False
supports_vision: bool = False
class LocalModelInfo(BaseModel):
"""Discovered local model candidate."""
id: str = Field(..., description = "Identifier to use for loading/training")
inventory_id: Optional[str] = Field(
None, description = "Stable semantic inventory row identifier"
)
load_id: Optional[str] = Field(
None, description = "Identifier/path to pass to load or train APIs"
)
display_name: str = Field(..., description = "Display label")
path: str = Field(..., description = "Local path where model data was discovered")
size_bytes: int = Field(0, description = "Observed model artifact size in bytes")
model_format: ModelFormat = Field("unknown", description = "Model file format")
runtime: ModelRuntime = Field("unknown", description = "Expected runtime backend")
format_variant: Optional[str] = Field(
None, description = "Format variant label, for example a GGUF quant"
)
capabilities: LocalModelCapabilities = Field(
default_factory = LocalModelCapabilities,
description = "Declared capabilities for this inventory row",
)
source: Literal["models_dir", "hf_cache", "lmstudio", "ollama", "custom"] = Field(
...,
description = "Discovery source",
)
model_id: Optional[str] = Field(
None,
description = "HF repo id for cached models, e.g. org/model",
)
base_model: Optional[str] = Field(
None,
description = "Base model from adapter_config.json when this is an adapter",
)
base_model_source: Optional[Literal["huggingface", "local", "unknown"]] = Field(
None,
description = "Whether the adapter base model is a HF repo id or local path",
)
adapter_type: Optional[str] = Field(
None,
description = "Adapter type from adapter_config.json, e.g. LORA",
)
training_method: Optional[str] = Field(
None,
description = "Training method hint from adapter_config.json",
)
updated_at: Optional[float] = Field(
None,
description = "Unix timestamp of latest observed update",
)
partial: bool = Field(
False,
description = "True when this hf_cache entry has incomplete blobs",
)
partial_transport: Optional[str] = Field(
None,
description = (
'Transport recorded for the partial state ("http" or '
'"xet"), or null if not partial / unknown.'
),
)
class LocalModelListResponse(BaseModel):
"""Response schema for listing local/cached models."""
models_dir: str = Field(..., description = "Directory scanned for custom local models")
hf_cache_dir: Optional[str] = Field(
None,
description = "HF cache root that was scanned",
)
lmstudio_dirs: List[str] = Field(
default_factory = list,
description = "LM Studio model directories that were scanned",
)
ollama_dirs: List[str] = Field(
default_factory = list,
description = "Ollama model directories that were scanned",
)
models: List[LocalModelInfo] = Field(
default_factory = list,
description = "Discovered local/cached models",
)
class CachedRepoBase(BaseModel):
"""Shared shape for a cached HF repo row surfaced under On Device."""
repo_id: str
size_bytes: int = 0
cache_path: Optional[str] = None
partial: bool = False
partial_transport: Optional[str] = None
inventory_id: Optional[str] = None
load_id: Optional[str] = None
model_format: ModelFormat = "unknown"
runtime: ModelRuntime = "unknown"
format_variant: Optional[str] = None
capabilities: LocalModelCapabilities = Field(default_factory = LocalModelCapabilities)
class CachedGgufRepo(CachedRepoBase):
model_format: ModelFormat = "gguf"
class CachedGgufResponse(BaseModel):
cached: List[CachedGgufRepo] = Field(default_factory = list)
class CachedModelRepo(CachedRepoBase):
quant_method: Optional[str] = None
pipeline_tag: Optional[str] = None
library_name: Optional[str] = None
tags: Optional[List[str]] = None
class CachedModelsResponse(BaseModel):
cached: List[CachedModelRepo] = Field(default_factory = list)
class AddScanFolderRequest(BaseModel):
"""Request body for adding a custom scan folder."""
path: str = Field(
...,
description = "Absolute or relative folder path, or a model weight file path",
)
class ScanFolderInfo(BaseModel):
"""A registered custom model scan folder."""
id: int = Field(..., description = "Database row ID")
path: str = Field(..., description = "Normalized absolute path")
created_at: str = Field(..., description = "ISO 8601 creation timestamp")
class ScanFoldersResponse(BaseModel):
folders: List[ScanFolderInfo] = Field(default_factory = list)
class RemoveScanFolderResponse(BaseModel):
ok: bool
class RecommendedFoldersResponse(BaseModel):
folders: List[str] = Field(default_factory = list)
class DeleteCachedModelResponse(BaseModel):
status: str
repo_id: str
variant: Optional[str] = None
class BrowseEntry(BaseModel):
"""A directory entry surfaced by the folder browser."""
name: str = Field(..., description = "Entry name (basename, not full path)")
has_models: bool = Field(
False,
description = (
"Hint that the directory likely contains models "
"(*.gguf, *.safetensors, config.json, or HF-style "
"`models--*` subfolders). Used by the UI to highlight "
"promising candidates; the scanner itself is authoritative."
),
)
hidden: bool = Field(
False,
description = "Name starts with a dot (e.g. `.cache`)",
)
class BrowseFoldersResponse(BaseModel):
"""Response schema for the folder browser endpoint."""
current: str = Field(..., description = "Absolute path of the directory just listed")
parent: Optional[str] = Field(
None,
description = (
"Parent directory of `current`, or null if `current` is the "
"filesystem root. The frontend uses this to render an `Up` row."
),
)
entries: List[BrowseEntry] = Field(
default_factory = list,
description = (
"Subdirectories of `current`. Sorted with model-bearing "
"directories first, then alphabetically case-insensitive; "
"hidden entries come last within each group."
),
)
suggestions: List[str] = Field(
default_factory = list,
description = (
"Handy starting points (home, HF cache, already-registered "
"scan folders). Rendered as quick-pick chips above the list."
),
)
truncated: bool = Field(
False,
description = (
"True when the listing was capped because the directory had "
"more subfolders than the server is willing to enumerate in "
"one request. The UI should show a hint telling the user to "
"narrow their path."
),
)
model_files_here: int = Field(
0,
description = (
"Count of GGUF/safetensors files immediately inside "
"``current``. Used by the UI to surface a hint on leaf "
"model directories (which otherwise look `empty` because "
"they contain only files, no subdirectories)."
),
)

29
studio/backend/hub/services/__init__.py Normal file
View file

@ -0,0 +1,29 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared helpers for the Hub service layer."""
from __future__ import annotations
from typing import Iterable
from fastapi import HTTPException
from hub.utils.hf_cache_state import resolve_destructive_case_matches
def resolve_destructive_repo_ids(repo_id: str, candidates: Iterable[str], *, noun: str) -> set[str]:
"""Cache-dir repo ids a destructive op on *repo_id* may target.
Refuses with 409 on ambiguous case-only matches so a delete never removes
the wrong casing. *noun* is the plural shown to the user."""
resolved = resolve_destructive_case_matches(repo_id, candidates)
if resolved is None:
raise HTTPException(
status_code = 409,
detail = (
f"Multiple cached {noun} differ only by case. "
"Delete the exact repo casing from On Device."
),
)
return resolved

4
studio/backend/hub/services/datasets/__init__.py Normal file
View file

@ -0,0 +1,4 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Dataset services for Hub routes."""

474
studio/backend/hub/services/datasets/cache_inventory.py Normal file
View file

@ -0,0 +1,474 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Cached dataset inventory and deletion services."""
from __future__ import annotations
import asyncio
import os
from pathlib import Path
from typing import Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.services import resolve_destructive_repo_ids
from hub.services.datasets import downloads
from hub.utils import download_manifest
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import (
purge_partial_repo,
purge_repo_cache_dirs,
resolve_destructive_case_matches,
)
from hub.utils.paths import (
hf_default_cache_dir,
is_valid_repo_id as _is_valid_repo_id,
legacy_hf_cache_dir,
resolve_cached_repo_id_case,
)
logger = get_logger(__name__)
def _collect_hf_cache_scans() -> tuple[list, set[str]]:
scans = hf_cache_scan.all_hf_cache_scans()
seen_roots = {
str(cache_dir)
for cache_dir in (getattr(scan, "cache_dir", None) for scan in scans)
if cache_dir is not None
}
return scans, seen_roots
def _hf_hub_cache_roots() -> list[Path]:
roots: list[Path] = []
seen: set[str] = set()
def _add(path: Optional[Path]) -> None:
if path is None or not path.is_dir():
return
try:
resolved = str(path.resolve())
except OSError:
return
if resolved in seen:
return
seen.add(resolved)
roots.append(path)
try:
from huggingface_hub.constants import HF_HUB_CACHE
_add(Path(HF_HUB_CACHE))
except Exception:
pass
hf_hub_cache = os.environ.get("HF_HUB_CACHE")
if hf_hub_cache:
_add(Path(hf_hub_cache).expanduser())
hf_home = os.environ.get("HF_HOME")
if hf_home:
_add(Path(hf_home).expanduser() / "hub")
_add(legacy_hf_cache_dir())
_add(hf_default_cache_dir())
return roots
def _repo_id_from_hub_dataset_dir(name: str) -> str | None:
if not name.startswith("datasets--"):
return None
encoded = name.removeprefix("datasets--")
owner, sep, repo = encoded.partition("--")
if not sep or not owner or not repo:
return None
repo_id = f"{owner}/{repo}"
return repo_id if _is_valid_repo_id(repo_id) else None
def _directory_size(path: Path) -> int:
total = 0
try:
for entry in path.rglob("*"):
try:
if entry.is_file() and not entry.is_symlink():
total += entry.stat().st_size
except OSError:
continue
except OSError:
return 0
return total
def _prefer_dataset_cache_row(candidate: dict, existing: Optional[dict]) -> bool:
if existing is None:
return True
candidate_partial = bool(candidate.get("partial"))
existing_partial = bool(existing.get("partial"))
if candidate_partial != existing_partial:
return not candidate_partial
return int(candidate.get("size_bytes") or 0) > int(existing.get("size_bytes") or 0)
def _hub_dataset_snapshot_count(path: Path) -> int:
snapshots = path / "snapshots"
try:
return sum(1 for entry in snapshots.iterdir() if entry.is_dir())
except OSError:
return 0
def _scan_hub_dataset_cache_dirs() -> list[dict]:
"""Fallback scanner: ``scan_cache_dir()`` skips repos when one cache entry is partially corrupt, so this keeps On Device matching disk."""
seen_lower: dict[str, dict] = {}
for root in _hf_hub_cache_roots():
try:
entries = [entry for entry in root.iterdir() if entry.is_dir()]
except OSError:
continue
for entry in entries:
repo_id = _repo_id_from_hub_dataset_dir(entry.name)
if repo_id is None:
continue
size_bytes = _directory_size(entry / "blobs")
if size_bytes <= 0:
size_bytes = _directory_size(entry)
if size_bytes <= 0:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
snapshot_partial = _hub_dataset_snapshot_count(
entry
) == 0 or hf_cache_scan.is_snapshot_partial("dataset", repo_id, entry)
row = {
"repo_id": repo_id,
"size_bytes": size_bytes,
"cache_path": str(entry.resolve()),
# snapshot_count == 0 catches blobs-but-no-snapshot;
# is_snapshot_partial adds active-row state checks.
"partial": snapshot_partial,
"partial_transport": (
hf_cache_scan.partial_transport_for(
"dataset",
repo_id,
repo_cache_dir = entry,
)
if snapshot_partial
else None
),
}
if _prefer_dataset_cache_row(row, existing):
seen_lower[key] = row
return sorted(seen_lower.values(), key = lambda c: c["repo_id"])
def _hf_datasets_cache_roots() -> list[Path]:
roots: list[Path] = []
seen: set[str] = set()
def _add(path: Optional[Path]) -> None:
if path is None or not path.is_dir():
return
try:
resolved = str(path.resolve())
except OSError:
return
if resolved in seen:
return
seen.add(resolved)
roots.append(path)
env_cache = os.environ.get("HF_DATASETS_CACHE")
if env_cache:
_add(Path(env_cache).expanduser())
try:
from datasets import config as datasets_config
_add(Path(datasets_config.HF_DATASETS_CACHE))
except Exception:
pass
hf_home = os.environ.get("HF_HOME")
if hf_home:
_add(Path(hf_home).expanduser() / "datasets")
xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser()
_add(xdg_cache / "huggingface" / "datasets")
return roots
def _repo_id_from_datasets_cache_dir(name: str) -> str | None:
if "___" not in name:
return None
owner, repo = name.split("___", 1)
repo_id = f"{owner}/{repo}"
return repo_id if _is_valid_repo_id(repo_id) else None
def _processed_dataset_cache_size(path: Path) -> int:
total = 0
try:
for entry in path.rglob("*"):
try:
if entry.is_file():
total += entry.stat().st_size
except OSError:
continue
except OSError:
return 0
return total
def _looks_like_processed_dataset_cache(path: Path) -> bool:
try:
for entry in path.rglob("*"):
if not entry.is_file():
continue
if entry.name in {"dataset_info.json", "state.json"}:
return True
if entry.suffix == ".arrow":
return True
except OSError:
return False
return False
def _scan_processed_dataset_caches() -> list[dict]:
"""`load_dataset()` stores processed Arrow caches separately from the Hub snapshot cache, so they're usable on-device but invisible to `scan_cache_dir()`."""
seen_lower: dict[str, dict] = {}
for root in _hf_datasets_cache_roots():
try:
entries = [entry for entry in root.iterdir() if entry.is_dir()]
except OSError:
continue
for entry in entries:
repo_id = _repo_id_from_datasets_cache_dir(entry.name)
if repo_id is None:
continue
if not _looks_like_processed_dataset_cache(entry):
continue
size_bytes = _processed_dataset_cache_size(entry)
if size_bytes <= 0:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
if existing is None or size_bytes > existing["size_bytes"]:
seen_lower[key] = {
"repo_id": repo_id,
"size_bytes": size_bytes,
"cache_path": str(entry.resolve()),
"processed_cache": True,
"partial": False,
}
return sorted(seen_lower.values(), key = lambda c: c["repo_id"])
def _scan_hf_dataset_caches() -> list[dict]:
scans, seen_roots = _collect_hf_cache_scans()
seen_lower: dict[str, dict] = {}
inspected = 0
for hf_cache in scans:
for repo_info in hf_cache.repos:
inspected += 1
try:
# str(...) guards against the library switching repo_type to an Enum.
if str(repo_info.repo_type) != "dataset":
continue
total_size = int(getattr(repo_info, "size_on_disk", 0) or 0)
if total_size == 0:
unique_blobs: dict[str, int] = {}
for rev in repo_info.revisions:
rev_id = getattr(rev, "commit_hash", None) or str(id(rev))
for f in rev.files:
blob_path = getattr(f, "blob_path", None)
key = str(blob_path) if blob_path else f"{rev_id}:{f.file_name}"
unique_blobs[key] = int(f.size_on_disk or 0)
total_size = sum(unique_blobs.values())
key = repo_info.repo_id.lower()
existing = seen_lower.get(key)
cache_dir = Path(repo_info.repo_path)
snapshot_partial = hf_cache_scan.is_snapshot_partial(
"dataset",
repo_info.repo_id,
cache_dir,
)
row = {
"repo_id": repo_info.repo_id,
"size_bytes": total_size,
"cache_path": str(repo_info.repo_path),
"partial": snapshot_partial,
"partial_transport": (
hf_cache_scan.partial_transport_for(
"dataset",
repo_info.repo_id,
repo_cache_dir = cache_dir,
)
if snapshot_partial
else None
),
}
if _prefer_dataset_cache_row(row, existing):
seen_lower[key] = row
except Exception as exc:
label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning("Skipping cached dataset repo %s: %s", label, exc)
for row in _scan_hub_dataset_cache_dirs():
key = row["repo_id"].lower()
existing = seen_lower.get(key)
if _prefer_dataset_cache_row(row, existing):
seen_lower[key] = row
elif existing is not None and bool(existing.get("partial")) == bool(row.get("partial")):
existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"])
existing["cache_path"] = existing.get("cache_path") or row.get("cache_path")
if (
existing.get("partial")
and not existing.get("partial_transport")
and row.get("partial_transport")
):
existing["partial_transport"] = row["partial_transport"]
for row in _scan_processed_dataset_caches():
key = row["repo_id"].lower()
existing = seen_lower.get(key)
if existing is None or (bool(existing.get("partial")) and not bool(row.get("partial"))):
seen_lower[key] = row
else:
existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"])
# Keep the processed-cache marker when a repo is both snapshot and
# processed Arrow cache; merging by size alone dropped it.
if row.get("processed_cache"):
existing["processed_cache"] = True
logger.info(
"Cached dataset scan: roots=%d inspected=%d returned=%d",
len(seen_roots) or len(scans),
inspected,
len(seen_lower),
)
return sorted(seen_lower.values(), key = lambda c: c["repo_id"])
async def list_cached_datasets_response() -> dict:
"""List dataset repos already downloaded into the HF cache."""
try:
return {"cached": await asyncio.to_thread(_scan_hf_dataset_caches)}
except Exception as exc:
logger.error("Error listing cached datasets: %s", exc, exc_info = True)
raise HTTPException(
status_code = 500,
detail = "Failed to read the local dataset cache.",
) from exc
async def delete_cached_dataset_response(repo_id: str) -> dict:
"""Remove a cached dataset repo from the HF cache."""
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
if not downloads.registry.begin_delete(repo_key):
raise HTTPException(
status_code = 400,
detail = "Cancel the active download before deleting.",
)
try:
return await asyncio.to_thread(_delete_cached_dataset_blocking, repo_key)
finally:
downloads.registry.end_delete(repo_key)
hf_cache_scan.invalidate_hf_cache_scans()
def _delete_cached_dataset_blocking(repo_id: str) -> dict:
scans, _seen_roots = _collect_hf_cache_scans()
candidate_entries = []
for hf_cache in scans:
for repo_info in hf_cache.repos:
if str(repo_info.repo_type) != "dataset":
continue
if repo_info.repo_id.lower() == repo_id.lower():
candidate_entries.append((hf_cache, repo_info))
matched_repo_ids = resolve_destructive_repo_ids(
repo_id,
[str(repo_info.repo_id) for _hf_cache, repo_info in candidate_entries],
noun = "datasets",
)
deleted = False
failures: list[str] = []
for hf_cache, repo_info in candidate_entries:
if str(repo_info.repo_id) not in matched_repo_ids:
continue
try:
strategy = hf_cache.delete_revisions(*(rev.commit_hash for rev in repo_info.revisions))
strategy.execute()
deleted = True
except Exception as exc:
failures.append(str(exc))
logger.error(
"Failed deleting cached dataset %s from %s: %s",
repo_id,
getattr(hf_cache, "cache_dir", "<unknown>"),
exc,
exc_info = True,
)
processed_deleted, processed_failures = _delete_processed_dataset_cache(repo_id)
failures.extend(processed_failures)
if failures:
raise HTTPException(
status_code = 500,
detail = (
f"Failed to delete dataset from {len(failures)} cache "
"location(s). Some files may remain."
),
)
# ``scan_cache_dir()`` skips blob-only/corrupt repos the revision delete
# can't touch, yet the fallback scanner shows them; purge the whole dir.
cache_purged = purge_repo_cache_dirs("dataset", repo_id)
partial_purged = purge_partial_repo("dataset", repo_id)
state_purged = download_manifest.purge_all_state_for_repo("dataset", repo_id) > 0
if not (deleted or processed_deleted or cache_purged or partial_purged or state_purged):
raise HTTPException(status_code = 404, detail = "Dataset not found in cache")
return {"status": "deleted", "repo_id": repo_id}
def _delete_processed_dataset_cache(repo_id: str) -> tuple[bool, list[str]]:
import shutil
target = repo_id.replace("/", "___")
folded_target = target.lower()
deleted = False
failures: list[str] = []
for root in _hf_datasets_cache_roots():
try:
entries = [
entry
for entry in root.iterdir()
if entry.is_dir() and entry.name.lower() == folded_target
]
except OSError:
continue
matched_names = resolve_destructive_case_matches(
target,
(entry.name for entry in entries),
)
if not matched_names:
continue
for entry in entries:
if entry.name not in matched_names:
continue
try:
shutil.rmtree(entry)
deleted = True
except Exception as exc:
failures.append(str(exc))
logger.error(
"Failed deleting processed dataset cache %s: %s",
repo_id,
exc,
exc_info = True,
)
return deleted, failures

268
studio/backend/hub/services/datasets/downloads.py Normal file
View file

@ -0,0 +1,268 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Start, cancel, and report progress for dataset downloads."""
from __future__ import annotations
import asyncio
import threading
import time
from collections import OrderedDict
from typing import Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.schemas.downloads import (
ActiveDownloadsResponse,
CancelDatasetDownloadRequest,
DatasetDownloadJobStatus,
DownloadDatasetRequest,
)
from hub.services import snapshot_progress
from hub.services import download_lifecycle
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import has_active_incomplete_blobs
from hub.utils.paths import (
is_valid_repo_id as _is_valid_repo_id,
resolve_cached_repo_id_case,
)
from hub.utils.snapshot_filters import (
blob_hashes_for_siblings,
total_size_for_siblings,
)
logger = get_logger(__name__)
_dataset_size_cache: "OrderedDict[str, tuple[int, frozenset[str], bool, str, float]]" = (
OrderedDict()
)
_dataset_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict()
_DATASET_SIZE_CACHE_MAX = 256
_DATASET_SIZE_POS_TTL = 60.0
_DATASET_SIZE_NEG_TTL = 60.0
_DATASET_SIZE_TIMEOUT_SECONDS = 5.0
_dataset_size_cache_lock = threading.Lock()
_registry = download_registry.get_datasets_registry()
def _download_job_key(repo_id: str) -> str:
return download_registry.normalize_repo_key(repo_id)
def get_dataset_snapshot_metadata_cached(
repo_id: str, hf_token: Optional[str] = None
) -> tuple[int, frozenset[str]]:
"""Raw snapshot size + expected blob hashes for a dataset repo.
The dataset worker downloads every sibling, so the denominator is the full
sibling-size sum and the hashes cover every file. Consumed by the shared
``snapshot_progress`` accounting."""
token_fp = hf_cache_scan.token_fingerprint(hf_token)
cache_key = (repo_id, token_fp)
with _dataset_size_cache_lock:
cached = _dataset_size_cache.get(repo_id)
if cached is not None:
size, hashes, restricted, cached_fp, ts = cached
if (time.monotonic() - ts) >= _DATASET_SIZE_POS_TTL:
del _dataset_size_cache[repo_id]
# A gated/private repo's metadata is only served back to the token
# that fetched it; another token may have no access at all.
elif not restricted or cached_fp == token_fp:
_dataset_size_cache.move_to_end(repo_id)
return size, hashes
neg_ts = _dataset_size_neg_cache.get(cache_key)
if neg_ts is not None and (time.monotonic() - neg_ts) < _DATASET_SIZE_NEG_TTL:
return 0, frozenset()
try:
from huggingface_hub import HfApi
info = HfApi(token = hf_token).dataset_info(
repo_id,
files_metadata = True,
timeout = _DATASET_SIZE_TIMEOUT_SECONDS,
)
total = total_size_for_siblings(info.siblings)
hashes = blob_hashes_for_siblings(info.siblings)
restricted = bool(getattr(info, "private", False) or getattr(info, "gated", False))
except Exception:
with _dataset_size_cache_lock:
_dataset_size_neg_cache[cache_key] = time.monotonic()
_dataset_size_neg_cache.move_to_end(cache_key)
while len(_dataset_size_neg_cache) > _DATASET_SIZE_CACHE_MAX:
_dataset_size_neg_cache.popitem(last = False)
return 0, frozenset()
with _dataset_size_cache_lock:
_dataset_size_cache[repo_id] = (
total,
hashes,
restricted,
token_fp,
time.monotonic(),
)
_dataset_size_cache.move_to_end(repo_id)
_dataset_size_neg_cache.pop(cache_key, None)
while len(_dataset_size_cache) > _DATASET_SIZE_CACHE_MAX:
_dataset_size_cache.popitem(last = False)
return total, hashes
async def get_dataset_download_progress_response(
repo_id: str,
expected_bytes: int = 0,
hf_token: Optional[str] = None,
) -> dict:
"""Return download progress for a HuggingFace dataset repo.
Scans the ``datasets--owner--name`` cache dir and shares the blob accounting
with the model path via ``snapshot_progress``. Returns ``cache_path`` for the
UI."""
return await snapshot_progress.snapshot_progress_response(
repo_type = "dataset",
repo_id = repo_id,
job_key = _download_job_key(repo_id),
expected_bytes = expected_bytes,
hf_token = hf_token,
registry = _registry,
metadata_resolver = get_dataset_snapshot_metadata_cached,
)
def _dataset_status(key: str, *, repo_id: Optional[str] = None) -> DatasetDownloadJobStatus:
state, error, generation = download_lifecycle.idle_status(
_registry,
key,
repo_type = "dataset",
repo_id = repo_id,
variant = None,
)
return DatasetDownloadJobStatus(state = state, error = error, generation = generation)
async def download_dataset_response(
body: DownloadDatasetRequest, hf_token: Optional[str] = None
) -> dict:
"""Start a background download for a HuggingFace dataset."""
repo_id = body.repo_id.strip()
if not _is_valid_repo_id(repo_id):
raise HTTPException(
status_code = 400,
detail = f"Invalid repo_id: {repo_id!r}",
)
# Canonicalize so two different-cased paste-ins share one job + cache dir.
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
key = _download_job_key(repo_id)
transport = download_lifecycle.resolve_transport(body.use_xet)
claimed, claim_state = _registry.claim(
key,
transport,
repo_type = "dataset",
repo_id = repo_id,
)
generation = _registry.current_generation(key)
if not claimed:
# Pollable when rejected by this repo's own in-flight job; an
# in-progress delete leaves no job, so flag it via ``adoptable``.
return {
"repo_id": repo_id,
"state": claim_state,
"accepted": _registry.adoptable(key),
"generation": generation,
}
download_manifest.clear_cancel_marker("dataset", repo_id, None)
state = download_lifecycle.launch_worker(
_registry,
key,
spawn = lambda: download_lifecycle.spawn_worker(
["--repo-id", repo_id, "--dataset"],
hf_token,
use_xet = body.use_xet,
),
hf_token = hf_token,
label = repo_id,
log_prefix = "Dataset download",
logger = logger,
repo_type = "dataset",
repo_id = repo_id,
transport = transport,
watch_name = f"hf-dataset-download-watch-{repo_id}",
)
return {
"repo_id": repo_id,
"state": state,
"accepted": True,
"generation": generation,
}
async def cancel_dataset_download_response(body: CancelDatasetDownloadRequest) -> dict:
"""Cancel an in-flight dataset download (SIGKILL; HF cache resumes on next download)."""
repo_id = body.repo_id.strip()
if not _is_valid_repo_id(repo_id):
raise HTTPException(
status_code = 400,
detail = f"Invalid repo_id: {repo_id!r}",
)
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
key = _download_job_key(repo_id)
state = download_lifecycle.cancel_worker(
_registry,
key,
generation = body.generation,
label = f"dataset {repo_id}",
logger = logger,
)
return {"repo_id": repo_id, "state": state}
async def get_dataset_download_status_response(repo_id: str) -> DatasetDownloadJobStatus:
"""Return the latest state of a background dataset download job."""
repo_id = repo_id.strip()
if not _is_valid_repo_id(repo_id):
return DatasetDownloadJobStatus(state = "idle")
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
return _dataset_status(_download_job_key(repo_id), repo_id = repo_id)
async def get_active_dataset_downloads_response(repo_id: str = "") -> ActiveDownloadsResponse:
repo_id = repo_id.strip()
if repo_id and not _is_valid_repo_id(repo_id):
return ActiveDownloadsResponse(downloads = [])
canonical_repo_id = (
await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
if repo_id
else None
)
return ActiveDownloadsResponse(
downloads = download_lifecycle.active_download_refs(
_registry,
canonical_repo_id,
with_variant = False,
)
)
async def get_dataset_transport_status_response(repo_id: str) -> dict:
"""Last transport used, whether partial blobs exist, and whether they
support byte-level resume. XET partials show via ``has_partial`` but are not
byte-level resumable (see ``models.get_model_transport_status``)."""
repo_id = repo_id.strip()
if not _is_valid_repo_id(repo_id):
return {"has_partial": False, "last_transport": None, "resumable": False}
return {
"has_partial": has_active_incomplete_blobs("dataset", repo_id),
"last_transport": download_registry.read_active_transport_marker("dataset", repo_id),
"resumable": download_registry.is_resumable_partial("dataset", repo_id),
}
registry = _registry

527
studio/backend/hub/services/datasets/formatting.py Normal file
View file

@ -0,0 +1,527 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Dataset preview, format-check, and mapping-assist services."""
from __future__ import annotations
import base64
import errno
import io
import re
from pathlib import Path
from typing import Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.schemas.datasets import (
AiAssistMappingRequest,
AiAssistMappingResponse,
CheckFormatRequest,
CheckFormatResponse,
)
from hub.services.datasets.local import (
DATA_EXTS,
_TABULAR_EXTS,
_load_local_preview_slice,
_stream_file_preview_slice,
)
from hub.utils.dataset_cache import (
cached_dataset_candidates as _shared_cached_dataset_candidates,
latest_cached_dataset_snapshot as _shared_latest_cached_dataset_snapshot,
split_label_matches as _split_label_matches,
)
from hub.utils import download_registry
from hub.utils.dataset_format import check_dataset_format, format_dataset_preview
from hub.utils.hf_errors import hf_error_status
from hub.utils.paths import (
is_valid_repo_id as _is_valid_repo_id,
resolve_dataset_path,
)
logger = get_logger(__name__)
_BINARY_IMAGE_PREVIEW_MAX_BYTES = 10 * 1024 * 1024
_IMAGE_PREVIEW_MAX_PIXELS = 16_000_000
_IMAGE_PREVIEW_THUMBNAIL_SIZE = (512, 512)
def _image_pixel_count(image) -> int:
width = max(int(getattr(image, "width", 0) or 0), 0)
height = max(int(getattr(image, "height", 0) or 0), 0)
return width * height
def _pil_image_has_transparency(image) -> bool:
if "A" in image.getbands():
extrema = image.getchannel("A").getextrema()
return bool(extrema and extrema[0] < 255)
if image.mode == "P":
transparency = image.info.get("transparency")
if transparency is None:
return False
if isinstance(transparency, bytes):
return any(alpha < 255 for alpha in transparency)
return True
return False
def _serialize_pil_image(image):
pixel_count = _image_pixel_count(image)
if pixel_count > _IMAGE_PREVIEW_MAX_PIXELS:
return (
f"<image preview omitted, {image.width}x{image.height} pixels "
f"exceeds {_IMAGE_PREVIEW_MAX_PIXELS:,} pixel limit>"
)
preview = image.copy()
preview.thumbnail(_IMAGE_PREVIEW_THUMBNAIL_SIZE)
buffer = io.BytesIO()
if _pil_image_has_transparency(preview):
preview.save(buffer, format = "PNG")
mime = "image/png"
else:
preview.convert("RGB").save(buffer, format = "JPEG", quality = 85)
mime = "image/jpeg"
return {
"type": "image",
"mime": mime,
"width": preview.width,
"height": preview.height,
"data": base64.b64encode(buffer.getvalue()).decode("ascii"),
}
def _serialize_binary_value(data):
if len(data) > _BINARY_IMAGE_PREVIEW_MAX_BYTES:
return (
f"<binary data omitted, {len(data)} bytes exceeds "
f"{_BINARY_IMAGE_PREVIEW_MAX_BYTES:,} byte preview limit>"
)
try:
from PIL import Image as PILImageModule
with PILImageModule.open(io.BytesIO(data)) as image:
return _serialize_pil_image(image)
except Exception:
return f"<binary data, {len(data)} bytes>"
def _serialize_preview_value(value):
if value is None or isinstance(value, (str, int, float, bool)):
return value
if isinstance(value, (bytes, bytearray, memoryview)):
return _serialize_binary_value(value)
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
return _serialize_pil_image(value)
except Exception:
pass
if isinstance(value, dict):
# Undecoded HF Image/Audio cells are {"bytes": b"...", "path": ...}.
raw = value.get("bytes")
if isinstance(raw, (bytes, bytearray, memoryview)) and not (
value.keys() - {"bytes", "path"}
):
return _serialize_binary_value(raw)
return {str(key): _serialize_preview_value(item) for key, item in value.items()}
if isinstance(value, (list, tuple)):
return [_serialize_preview_value(item) for item in value]
return str(value)
def _serialize_preview_rows(rows):
return [
{str(key): _serialize_preview_value(value) for key, value in dict(row).items()}
for row in rows
]
def _latest_cached_dataset_snapshot(
repo_id: str, local_path: Optional[str] = None
) -> Optional[Path]:
return _shared_latest_cached_dataset_snapshot(repo_id, local_path)
def _cached_dataset_candidates(
snapshot: Path, *, subset: Optional[str], train_split: str
) -> list[Path]:
return _shared_cached_dataset_candidates(
snapshot,
subset = subset,
train_split = train_split,
extensions = DATA_EXTS,
preferred_extensions = _TABULAR_EXTS,
)
def _repo_file_label_tokens(path: str) -> set[str]:
return {token for token in re.split(r"[^a-z0-9]+", path.lower()) if token}
def _repo_file_matches_label(path: str, label: str) -> bool:
return label.strip().lower() in _repo_file_label_tokens(path)
def _repo_file_matches_split(path: str, split: str) -> bool:
return _split_label_matches(path, split)
def _select_tier1_repo_file(
files: list[str], *, subset: Optional[str], train_split: str
) -> Optional[str]:
data_files = sorted(f for f in files if any(f.lower().endswith(ext) for ext in DATA_EXTS))
if not data_files:
return None
tabular_files = [f for f in data_files if any(f.lower().endswith(ext) for ext in _TABULAR_EXTS)]
candidates = tabular_files or data_files
if subset:
candidates = [f for f in candidates if _repo_file_matches_label(f, subset)]
if not candidates:
return None
candidates = [f for f in candidates if _repo_file_matches_split(f, train_split)]
return candidates[0] if candidates else None
def _load_cached_hf_preview_slice(request: CheckFormatRequest, preview_size: int):
if not _is_valid_repo_id(request.dataset_name):
return None
snapshot = _latest_cached_dataset_snapshot(
request.dataset_name,
request.local_path,
)
if snapshot is None:
return None
train_split = request.train_split or "train"
for candidate in _cached_dataset_candidates(
snapshot,
subset = request.subset,
train_split = train_split,
):
try:
preview = _stream_file_preview_slice(candidate, preview_size)
except Exception as exc:
logger.debug("Cached dataset preview failed for %s: %s", candidate, exc)
continue
if preview is not None:
return preview
return None
def _load_processed_hf_preview_slice(
request: CheckFormatRequest,
preview_size: int,
hf_token: Optional[str] = None,
):
if not _is_valid_repo_id(request.dataset_name):
return None
try:
from datasets import DownloadConfig, load_dataset
except Exception:
return None
load_kwargs = {
"path": request.dataset_name,
"split": request.train_split or "train",
"download_config": DownloadConfig(local_files_only = True),
}
if request.subset:
load_kwargs["name"] = request.subset
if hf_token:
load_kwargs["token"] = hf_token
dataset = load_dataset(**load_kwargs)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
return preview_slice, total_rows
def _load_any_cached_hf_preview_slice(
request: CheckFormatRequest,
preview_size: int,
hf_token: Optional[str] = None,
):
cached_preview = _load_cached_hf_preview_slice(request, preview_size)
if cached_preview is not None:
return cached_preview
try:
return _load_processed_hf_preview_slice(request, preview_size, hf_token)
except Exception as exc:
logger.debug(
"Processed dataset cache preview failed for %s: %s",
request.dataset_name,
exc,
)
return None
def check_format_response(
request: CheckFormatRequest, hf_token: Optional[str] = None
) -> CheckFormatResponse:
"""
Check if a dataset requires manual column mapping.
HF datasets: tier 1 loads a single requested split/subset file (avoids
resolving thousands of files); tier 2 falls back to full streaming. Local
files load directly. Plain `def` so FastAPI runs the blocking IO in a
thread-pool.
"""
try:
from itertools import islice
PREVIEW_SIZE = 10
logger.info(f"Checking format for dataset: {request.dataset_name}")
try:
dataset_path = resolve_dataset_path(request.dataset_name)
except ValueError as e:
# Malformed path (null bytes, '..', outside roots) is a client error:
# surface 400 rather than the generic 500 below.
raise HTTPException(status_code = 400, detail = str(e)) from e
total_rows = None
if dataset_path.exists():
train_split = request.train_split or "train"
preview_slice, total_rows = _load_local_preview_slice(
dataset_path = dataset_path,
train_split = train_split,
preview_size = PREVIEW_SIZE,
)
else:
from datasets import Dataset, load_dataset
# Tier 1: list_repo_files → load only the first data file
cached_preview = (
_load_any_cached_hf_preview_slice(request, PREVIEW_SIZE, hf_token)
if request.prefer_local_cache
else None
)
if cached_preview is not None:
preview_slice, total_rows = cached_preview
elif request.prefer_local_cache:
raise HTTPException(
status_code = 404,
detail = "Dataset is not available in the local cache.",
)
else:
preview_slice = None
try:
from huggingface_hub import HfApi
api = HfApi()
repo_files = api.list_repo_files(
request.dataset_name,
repo_type = "dataset",
token = hf_token or None,
)
train_split = request.train_split or "train"
first_file = _select_tier1_repo_file(
repo_files,
subset = request.subset,
train_split = train_split,
)
if first_file:
logger.info(f"Tier 1: loading single file {first_file}")
load_kwargs = {
"path": request.dataset_name,
"data_files": {train_split: [first_file]},
"split": train_split,
"streaming": True,
}
if hf_token:
load_kwargs["token"] = hf_token
streamed_ds = load_dataset(**load_kwargs)
rows = list(islice(streamed_ds, PREVIEW_SIZE))
if rows:
preview_slice = Dataset.from_list(rows)
except Exception as e:
logger.warning(
"Tier 1 (single-file) failed: %s",
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
if preview_slice is None:
# Tier 2: full streaming (resolves all files — slow for large repos)
logger.info("Tier 2: falling back to full streaming load_dataset")
try:
load_kwargs = {
"path": request.dataset_name,
"split": request.train_split or "train",
"streaming": True,
}
if request.subset:
load_kwargs["name"] = request.subset
if hf_token:
load_kwargs["token"] = hf_token
streamed_ds = load_dataset(**load_kwargs)
rows = list(islice(streamed_ds, PREVIEW_SIZE))
if not rows:
raise HTTPException(
status_code = 400,
detail = "Dataset appears to be empty or could not be streamed",
)
preview_slice = Dataset.from_list(rows)
total_rows = None
except Exception:
cached_preview = _load_any_cached_hf_preview_slice(
request,
PREVIEW_SIZE,
hf_token,
)
if cached_preview is None:
raise
preview_slice, total_rows = cached_preview
result = check_dataset_format(preview_slice, is_vlm = request.is_vlm)
logger.info(
f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}"
)
preview_samples = None
if not result["requires_manual_mapping"]:
if result.get("suggested_mapping"):
# Heuristic-detected: show raw data so columns match the response;
# column stripping happens at training time, not preview.
preview_samples = _serialize_preview_rows(preview_slice)
else:
try:
processed = format_dataset_preview(preview_slice)
preview_samples = _serialize_preview_rows(processed)
except Exception as e:
logger.warning(f"Processed preview generation failed (non-fatal): {e}")
preview_samples = _serialize_preview_rows(preview_slice)
else:
preview_samples = _serialize_preview_rows(preview_slice)
# Collect warnings: from check_dataset_format + URL-based image detection
warning = result.get("warning")
image_col = result.get("detected_image_column")
if image_col and image_col in (result.get("columns") or []):
try:
sample_val = preview_slice[0][image_col]
if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")):
url_warning = (
"This dataset contains image URLs instead of embedded images. "
"Images will be downloaded during training, which may be slow for large datasets."
)
logger.info(f"URL-based image column detected: {image_col}")
warning = f"{warning} {url_warning}" if warning else url_warning
except Exception:
pass
return CheckFormatResponse(
requires_manual_mapping = result["requires_manual_mapping"],
detected_format = result["detected_format"],
columns = result["columns"],
is_image = result.get("is_image", False),
is_audio = result.get("is_audio", False),
multimodal_columns = result.get("multimodal_columns"),
suggested_mapping = result.get("suggested_mapping"),
detected_image_column = result.get("detected_image_column"),
detected_audio_column = result.get("detected_audio_column"),
detected_text_column = result.get("detected_text_column"),
detected_speaker_column = result.get("detected_speaker_column"),
preview_samples = preview_samples,
total_rows = total_rows,
warning = warning,
)
except HTTPException:
raise
except Exception as e:
scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token)
# Missing/gated/bad-token and malformed names are client errors, not 500s.
status = hf_error_status(e)
if (
status is None
and isinstance(e, OSError)
and getattr(e, "errno", None) == errno.ENAMETOOLONG
):
status, scrubbed = 400, "Invalid dataset name"
elif status is None and isinstance(e, FileNotFoundError):
# datasets raises DatasetNotFoundError (FileNotFoundError) for missing/gated.
status = 404
elif status is None and isinstance(e, ValueError):
status = 400
if status is not None:
raise HTTPException(status_code = status, detail = scrubbed)
logger.error("Error checking dataset format: %s", scrubbed)
raise HTTPException(
status_code = 500,
detail = "Failed to check dataset format: " + scrubbed,
)
def ai_assist_mapping_response(
request: AiAssistMappingRequest, hf_token: Optional[str] = None
) -> AiAssistMappingResponse:
"""
Run the LLM-assisted dataset conversion advisor (user-triggered).
Multi-pass analysis with a 7B helper model: classify dataset type, generate
a conversion strategy, then validate it. Falls back to simple column
classification if the advisor fails.
"""
try:
from hub.utils.llm_assist import llm_conversion_advisor
truncated = [
{col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5]
]
result = llm_conversion_advisor(
column_names = request.columns,
samples = truncated,
dataset_name = request.dataset_name,
hf_token = hf_token,
model_name = request.model_name,
model_type = request.model_type,
)
if result and result.get("success"):
return AiAssistMappingResponse(
success = True,
suggested_mapping = result.get("suggested_mapping"),
system_prompt = result.get("system_prompt"),
user_template = result.get("user_template"),
assistant_template = result.get("assistant_template"),
label_mapping = result.get("label_mapping"),
dataset_type = result.get("dataset_type"),
is_conversational = result.get("is_conversational"),
user_notification = result.get("user_notification"),
warning = result.get("warning"),
)
return AiAssistMappingResponse(
success = False,
warning = "AI could not determine column roles. Please assign them manually.",
)
except Exception as e:
scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token)
status = hf_error_status(e)
if status is None and isinstance(e, FileNotFoundError):
status = 404
elif status is None and isinstance(e, ValueError):
status = 400
if status is not None:
raise HTTPException(status_code = status, detail = scrubbed)
logger.error("AI assist mapping failed: %s", scrubbed)
raise HTTPException(
status_code = 500,
detail = "AI assist failed: " + scrubbed,
)

330
studio/backend/hub/services/datasets/local.py Normal file
View file

@ -0,0 +1,330 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Local dataset upload and listing services."""
from __future__ import annotations
import asyncio
import json
import uuid
from pathlib import Path
from fastapi import HTTPException, UploadFile
from hub.schemas.datasets import (
LocalDatasetItem,
LocalDatasetsResponse,
UploadDatasetResponse,
)
from hub.utils.paths import dataset_uploads_root, ensure_dir, recipe_datasets_root
# Tabular formats are preferred over archives for Tier 1 preview: archives
# (e.g. images.zip) load as ImageFolder with synthetic columns that don't
# match the real schema.
_TABULAR_EXTS = (".parquet", ".json", ".jsonl", ".csv", ".tsv", ".arrow")
_ARCHIVE_EXTS = (".tar", ".tar.gz", ".tgz", ".gz", ".zst", ".zip", ".txt")
DATA_EXTS = _TABULAR_EXTS + _ARCHIVE_EXTS
LOCAL_FILE_EXTS = (".json", ".jsonl", ".csv", ".parquet")
LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"}
LOCAL_UPLOAD_CHUNK_BYTES = 1024 * 1024
LOCAL_UPLOAD_MAX_BYTES = 500 * 1024 * 1024
LOCAL_DATASETS_ROOT = recipe_datasets_root()
DATASET_UPLOAD_DIR = dataset_uploads_root()
def _safe_read_metadata(path: Path) -> dict | None:
try:
payload = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError, TypeError):
return None
if not isinstance(payload, dict):
return None
return payload
def _safe_read_rows_from_metadata(payload: dict | None) -> int | None:
if not payload:
return None
for key in ("actual_num_records", "target_num_records"):
value = payload.get(key)
if isinstance(value, int):
return value
return None
def _safe_read_metadata_summary(payload: dict | None) -> dict | None:
if not payload:
return None
actual_num_records = (
payload.get("actual_num_records")
if isinstance(payload.get("actual_num_records"), int)
else None
)
target_num_records = (
payload.get("target_num_records")
if isinstance(payload.get("target_num_records"), int)
else actual_num_records
)
columns: list[str] | None = None
schema = payload.get("schema")
if isinstance(schema, dict):
columns = [str(key) for key in schema.keys()]
if not columns:
stats = payload.get("column_statistics")
if isinstance(stats, list):
derived = [
str(item.get("column_name"))
for item in stats
if isinstance(item, dict) and item.get("column_name")
]
columns = derived or None
parquet_files_count = None
file_paths = payload.get("file_paths")
if isinstance(file_paths, dict):
parquet_files = file_paths.get("parquet-files")
if isinstance(parquet_files, list):
parquet_files_count = len(parquet_files)
total_num_batches = (
payload.get("total_num_batches")
if isinstance(payload.get("total_num_batches"), int)
else parquet_files_count
)
num_completed_batches = (
payload.get("num_completed_batches")
if isinstance(payload.get("num_completed_batches"), int)
else total_num_batches
)
return {
"actual_num_records": actual_num_records,
"target_num_records": target_num_records,
"total_num_batches": total_num_batches,
"num_completed_batches": num_completed_batches,
"columns": columns,
}
def _safe_mtime(path: Path) -> float | None:
try:
return path.stat().st_mtime
except OSError:
return None
def _display_uploaded_dataset_name(path: Path) -> str:
stem = path.stem
prefix, sep, rest = stem.partition("_")
if sep and len(prefix) == 32 and all(c in "0123456789abcdef" for c in prefix):
return f"{rest}{path.suffix}"
return path.name
def _build_recipe_dataset_items() -> list[LocalDatasetItem]:
if not LOCAL_DATASETS_ROOT.exists():
return []
items: list[LocalDatasetItem] = []
for entry in LOCAL_DATASETS_ROOT.iterdir():
if not entry.is_dir() or not entry.name.startswith("recipe_"):
continue
parquet_dir = entry / "parquet-files"
if not parquet_dir.exists() or not any(parquet_dir.glob("*.parquet")):
continue
rows = None
metadata_summary = None
metadata_path = entry / "metadata.json"
if metadata_path.exists():
metadata_payload = _safe_read_metadata(metadata_path)
rows = _safe_read_rows_from_metadata(metadata_payload)
metadata_summary = _safe_read_metadata_summary(metadata_payload)
items.append(
LocalDatasetItem(
id = entry.name,
label = entry.name,
path = str(parquet_dir.resolve()),
source = "recipe",
rows = rows,
updated_at = _safe_mtime(entry),
metadata = metadata_summary,
)
)
return items
def _build_uploaded_dataset_items() -> list[LocalDatasetItem]:
if not DATASET_UPLOAD_DIR.exists():
return []
items: list[LocalDatasetItem] = []
for path in DATASET_UPLOAD_DIR.iterdir():
if not path.is_file() or path.suffix.lower() not in LOCAL_UPLOAD_EXTS:
continue
try:
if path.stat().st_size == 0:
continue
except OSError:
continue
label = _display_uploaded_dataset_name(path)
items.append(
LocalDatasetItem(
id = path.name,
label = label,
path = str(path.resolve()),
source = "upload",
updated_at = _safe_mtime(path),
)
)
return items
def _build_local_dataset_items() -> list[LocalDatasetItem]:
items = _build_recipe_dataset_items() + _build_uploaded_dataset_items()
items.sort(key = lambda item: item.updated_at or 0, reverse = True)
return items
def _stream_file_preview_slice(path: Path, preview_size: int):
"""Stream the first ``preview_size`` rows so a large file is never fully parsed into Arrow; returns ``(Dataset, None)`` or ``None`` if empty/unsupported."""
from itertools import islice
from datasets import Dataset, load_dataset
name = path.name.lower()
if name.endswith((".json", ".jsonl")):
loader = "json"
elif name.endswith((".csv", ".tsv")):
loader = "csv"
elif name.endswith(".parquet"):
loader = "parquet"
elif name.endswith(".arrow"):
loader = "arrow"
elif name.endswith(".txt"):
loader = "text"
else:
return None
streamed = load_dataset(
loader,
data_files = str(path),
split = "train",
streaming = True,
)
rows = list(islice(streamed, preview_size))
if not rows:
return None
return Dataset.from_list(rows), None
def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
from datasets import load_dataset
if dataset_path.is_dir():
parquet_dir = (
dataset_path / "parquet-files"
if (dataset_path / "parquet-files").exists()
else dataset_path
)
parquet_files = sorted(parquet_dir.glob("*.parquet"))
if parquet_files:
dataset = load_dataset(
"parquet",
data_files = [str(path) for path in parquet_files],
split = train_split,
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
return preview_slice, total_rows
candidate_files: list[Path] = []
for ext in LOCAL_FILE_EXTS:
candidate_files.extend(sorted(dataset_path.glob(f"*{ext}")))
if not candidate_files:
raise HTTPException(
status_code = 400,
detail = "Unsupported local dataset directory (expected parquet/json/jsonl/csv files)",
)
dataset_path = candidate_files[0]
suffix = dataset_path.suffix.lower()
# Parquet/Arrow give a cheap exact total_rows via len()+select; JSON/CSV
# carry no such metadata, so stream them and report total_rows=None.
if suffix == ".parquet":
dataset = load_dataset("parquet", data_files = str(dataset_path), split = train_split)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
return preview_slice, total_rows
if suffix in (".json", ".jsonl", ".csv"):
preview = _stream_file_preview_slice(dataset_path, preview_size)
if preview is None:
raise HTTPException(
status_code = 400,
detail = "Dataset appears to be empty or could not be read",
)
return preview
raise HTTPException(status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}")
def _sanitize_filename(filename: str) -> str:
name = Path(filename).name.strip().replace("\x00", "")
if not name:
return "dataset_upload"
return name
def _upload_too_large(size_bytes: int) -> HTTPException:
return HTTPException(
status_code = 413,
detail = (f"Upload is too large " f"({size_bytes:,} bytes; max {LOCAL_UPLOAD_MAX_BYTES:,})."),
)
async def upload_dataset_response(file: UploadFile) -> UploadDatasetResponse:
filename = _sanitize_filename(file.filename or "dataset_upload")
ext = Path(filename).suffix.lower()
if ext not in LOCAL_UPLOAD_EXTS:
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
raise HTTPException(
status_code = 400,
detail = f"Unsupported file type: {ext}. Allowed: {allowed}",
)
declared_size = getattr(file, "size", None)
if isinstance(declared_size, int) and declared_size > LOCAL_UPLOAD_MAX_BYTES:
raise _upload_too_large(declared_size)
ensure_dir(DATASET_UPLOAD_DIR)
stem = Path(filename).stem
stored_name = f"{uuid.uuid4().hex}_{stem}{ext}"
stored_path = DATASET_UPLOAD_DIR / stored_name
written = 0
try:
with open(stored_path, "wb") as f:
while chunk := await file.read(LOCAL_UPLOAD_CHUNK_BYTES):
written += len(chunk)
if written > LOCAL_UPLOAD_MAX_BYTES:
raise _upload_too_large(written)
await asyncio.to_thread(f.write, chunk)
except Exception:
stored_path.unlink(missing_ok = True)
raise
if written == 0:
stored_path.unlink(missing_ok = True)
raise HTTPException(status_code = 400, detail = "Empty upload payload")
return UploadDatasetResponse(filename = filename, stored_path = str(stored_path))
def list_local_datasets_response() -> LocalDatasetsResponse:
return LocalDatasetsResponse(datasets = _build_local_dataset_items())

449
studio/backend/hub/services/download_lifecycle.py Normal file
View file

@ -0,0 +1,449 @@
# 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 os
import signal
import subprocess
import sys
import threading
from pathlib import Path
from typing import Callable, Optional
from fastapi import HTTPException
from hub.schemas.downloads import ActiveDownload, DownloadJobState
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import EXIT_CANCELLED
from hub.utils.state_dir import RepoType
def backend_dir() -> Path:
return Path(__file__).resolve().parent.parent.parent
def resolve_transport(use_xet: bool) -> str:
transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
unavailable_reason = download_registry.download_transport_unavailable_reason(transport)
if unavailable_reason is not None:
raise HTTPException(status_code = 400, detail = unavailable_reason)
return transport
def spawn_worker(
args: list[str],
hf_token: Optional[str],
*,
use_xet: bool,
protected_blob_hashes: Optional[frozenset[str]] = None,
) -> subprocess.Popen:
"""Spawn the download worker.
XET and ``hf_transfer`` write chunks out of order, so their partials can't
resume under a sequential writer; the HTTP path stays sequential so
SIGKILL -> resume is byte-identical. ``protected_blob_hashes`` are blobs a
concurrent same-repo peer is writing, excluded from the cache-prep purge so a
shared ``.incomplete`` (e.g. bundled mmproj) is never deleted.
"""
cwd = backend_dir()
mode = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
env = os.environ.copy()
if protected_blob_hashes:
env["UNSLOTH_PROTECTED_BLOB_HASHES"] = ",".join(sorted(protected_blob_hashes))
else:
env.pop("UNSLOTH_PROTECTED_BLOB_HASHES", None)
env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
env["HF_HUB_DISABLE_TELEMETRY"] = "1"
env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1"
env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0" if hf_token else "1"
# hf_transfer's parallel Range chunks can leave sparse partials even in
# "http" mode; disable so the worker's writer is always sequential.
env["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
for token_key in (
"HF_TOKEN",
"HF_HUB_TOKEN",
"HUGGING_FACE_HUB_TOKEN",
"HUGGINGFACE_HUB_TOKEN",
"HUGGINGFACEHUB_API_TOKEN",
):
env.pop(token_key, None)
if hf_token:
env["HF_TOKEN"] = hf_token
existing_path = env.get("PYTHONPATH", "")
env["PYTHONPATH"] = f"{cwd}{os.pathsep}{existing_path}" if existing_path else str(cwd)
return subprocess.Popen(
[
sys.executable,
"-m",
"hub.workers.hf_download",
*args,
"--parent-pid",
str(os.getpid()),
"--transport",
mode,
],
env = env,
cwd = str(cwd),
stdout = subprocess.DEVNULL,
stderr = subprocess.PIPE,
start_new_session = sys.platform != "win32",
)
def drain_stderr_excerpt(stream, edge_bytes: int = 500) -> bytes:
"""Drain a worker's stderr to EOF, retaining the first and last bytes.
Incremental reads keep the pipe from filling while bounding memory; long
messages keep both ends since stderr prefixes often name the failing repo."""
if stream is None:
return b""
edge_bytes = max(1, edge_bytes)
max_bytes = edge_bytes * 2
full = bytearray()
head = bytearray()
tail = bytearray()
truncated = False
for chunk in iter(lambda: stream.read(4096), b""):
if not truncated:
full.extend(chunk)
if len(full) <= max_bytes:
continue
truncated = True
head.extend(full[:edge_bytes])
tail.extend(full[-edge_bytes:])
full.clear()
continue
tail.extend(chunk)
if len(tail) > edge_bytes:
del tail[:-edge_bytes]
if not truncated:
return bytes(full)
return bytes(head + b"\n...[stderr truncated]...\n" + tail)
def _cancellation_return_codes() -> frozenset[int]:
"""Returncodes for intentional cancellation only (SIGKILL/SIGTERM/SIGINT); crash signals stay errors, and ``getattr`` tolerates Windows where these signals are absent."""
codes: set[int] = set()
for name in ("SIGKILL", "SIGTERM", "SIGINT"):
sig = getattr(signal, name, None)
if sig is not None:
codes.add(-int(sig))
return frozenset(codes)
_CANCELLATION_RETURN_CODES = _cancellation_return_codes()
def _sigpipe_return_codes() -> frozenset[int]:
sig = getattr(signal, "SIGPIPE", None)
if sig is None:
return frozenset()
value = int(sig)
return frozenset({-value, 128 + value})
_SIGPIPE_RETURN_CODES = _sigpipe_return_codes()
def classify_exit(rc: int, *, cancel_requested: bool = False) -> str:
"""Map a worker process exit code to a job state.
- rc == 0: clean completion.
- rc == EXIT_CANCELLED (130): the worker trapped a stop signal and exited
cleanly with a resumable partial. In-app cancel uses untrappable SIGKILL
and the OOM killer never produces 130, so 130 is always a resumable cancel.
- rc killed by SIGKILL/SIGTERM/SIGINT: a cancel only when *we* asked for it.
The OOM killer also sends SIGKILL, so an unrequested kill surfaces as error.
- rc killed by SIGPIPE (or 128+SIGPIPE): parent pipe is gone; treated as
cancelled.
- any other non-zero rc (incl. crash signals): worker errored out.
Windows has no POSIX signal exit encoding, so a user cancel can't be told from
an error by code alone; there ``cancel_requested`` decides.
"""
if rc == 0:
return "complete"
if rc == EXIT_CANCELLED:
return "cancelled"
if rc in _SIGPIPE_RETURN_CODES:
return "cancelled"
if rc in _CANCELLATION_RETURN_CODES:
return "cancelled" if cancel_requested else "error"
if cancel_requested and sys.platform == "win32":
return "cancelled"
return "error"
def finalize_worker_exit(
registry: download_registry.DownloadRegistry,
key: str,
proc: subprocess.Popen,
*,
hf_token: Optional[str],
label: str,
log_prefix: str,
logger,
repo_type: Optional[RepoType] = None,
repo_id: Optional[str] = None,
transport: Optional[str] = None,
) -> None:
"""Block until *proc* exits, then record the job's terminal state in
*registry*. Drains and scrubs stderr first, then classifies the exit code.
A no-op when the process was already dropped (e.g. superseded).
No stall watchdog: huggingface_hub already times out chunk reads and raises
a resumable error on a dead connection, so the worker's exit code is the
single source of truth."""
stderr_data = drain_stderr_excerpt(proc.stderr)
rc = proc.wait()
cancel_requested = registry.cancel_requested(key)
if not registry.drop_process(key, proc):
return
stderr_text = download_registry.scrub_secrets(
(stderr_data or b"").decode("utf-8", "replace").strip(),
hf_token = hf_token,
)
state = classify_exit(rc, cancel_requested = cancel_requested)
if state == "complete":
registry.set_job(key, "complete")
if stderr_text:
if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text:
logger.warning(
f"{log_prefix} complete with degraded diagnostics for "
f"{label}: {stderr_text}"
)
else:
logger.info(f"{log_prefix} worker diagnostics for {label}: {stderr_text}")
logger.info(f"{log_prefix} complete: {label}")
# Defensive cleanup: the canonical clear is at download-start; this
# catches the rare case where that failed but the download succeeded.
if repo_type and repo_id:
try:
download_manifest.clear_cancel_marker(
repo_type,
repo_id,
download_registry.variant_from_key(key),
)
except Exception as exc:
logger.debug(f"clear_cancel_marker failed for {repo_id} (rc=0): {exc}")
elif state == "cancelled":
# Read metadata before the terminal set_job so a concurrent eviction
# can't drop it; the job key is the fallback variant label.
metadata = registry.get_job_metadata(key)
registry.set_job(key, "cancelled")
logger.info(f"{log_prefix} cancelled: {label} (rc={rc})")
download_registry.persist_cancel_marker(
repo_type,
repo_id,
metadata.variant
if metadata is not None and metadata.variant
else download_registry.variant_from_key(key),
transport,
logger = logger,
)
else:
registry.set_job(
key,
"error",
stderr_text or f"worker exited with code {rc}",
)
logger.error(
f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}",
)
def kill_and_reap_process(
proc: subprocess.Popen,
*,
label: str,
logger,
timeout: float = 10.0,
) -> None:
try:
proc.kill()
except ProcessLookupError:
pass
except Exception as exc:
logger.warning(f"Cancel SIGKILL for {label} failed: {exc}")
try:
proc.wait(timeout = timeout)
except subprocess.TimeoutExpired:
logger.warning(f"Cancelled worker for {label} did not exit after SIGKILL")
except Exception:
pass
def register_worker(
registry: download_registry.DownloadRegistry,
key: str,
proc: subprocess.Popen,
*,
hf_token: Optional[str],
label: str,
log_prefix: str,
logger,
repo_type: RepoType,
repo_id: str,
transport: str,
watch_name: str,
) -> bool:
if not registry.register_process(key, proc):
kill_and_reap_process(proc, label = label, logger = logger)
return False
worker_token = hf_token
def _watch() -> None:
finalize_worker_exit(
registry,
key,
proc,
hf_token = worker_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
)
if registry.get_job(key).state in ("error", "cancelled"):
download_registry.purge_empty_marker_dir(
repo_type,
repo_id,
download_registry.variant_from_key(key),
)
hf_cache_scan.invalidate_hf_cache_scans()
threading.Thread(target = _watch, name = watch_name, daemon = True).start()
return True
def launch_worker(
registry: download_registry.DownloadRegistry,
key: str,
*,
spawn: Callable[[], subprocess.Popen],
hf_token: Optional[str],
label: str,
log_prefix: str,
logger,
repo_type: RepoType,
repo_id: str,
transport: str,
watch_name: str,
) -> str:
try:
proc = spawn()
except Exception as e:
scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token)
logger.error(
f"Failed to spawn {log_prefix.lower()} worker for {label}: {scrubbed}",
exc_info = True,
)
registry.set_job(key, "error", scrubbed)
raise HTTPException(
status_code = 500,
detail = f"Failed to start {log_prefix.lower()}: {scrubbed}",
) from e
register_worker(
registry,
key,
proc,
hf_token = hf_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
watch_name = watch_name,
)
return registry.get_job(key).state
def cancel_worker(
registry: download_registry.DownloadRegistry,
key: str,
*,
generation: Optional[int],
label: str,
logger,
) -> str:
proc = registry.get_process(key)
# No worker process yet: arm a pending cancel so register_process kills it on
# arrival during the claim-to-register window.
if proc is None:
if registry.mark_pending_cancel(key, generation):
return "cancelling"
return registry.get_job(key).state
# Worker already exited; let its watcher classify the real return code.
# Arming a pending cancel here could mislabel a genuine failure as a cancel.
if proc.poll() is not None:
return registry.get_job(key).state
if not registry.request_cancel(key, proc, generation):
return registry.get_job(key).state
# No eager marker: finalize_worker_exit writes it on a "cancelled" exit.
# Persisting before the kill races a clean completion and strands a stale marker.
try:
proc.kill()
except ProcessLookupError:
pass
except Exception as e:
logger.warning(f"Cancel SIGKILL for {label} failed: {e}")
return "cancelling"
def idle_status(
registry: download_registry.DownloadRegistry,
key: str,
*,
repo_type: RepoType,
repo_id: Optional[str],
variant: Optional[str],
) -> tuple[DownloadJobState, Optional[str], int]:
state = registry.get_job(key)
generation = registry.current_generation(key)
if (
state.state == "idle"
and repo_id
and download_manifest.has_cancel_marker(
repo_type,
repo_id,
variant,
)
):
return ("cancelled", None, generation)
return (state.state, state.error, generation)
def active_download_refs(
registry: download_registry.DownloadRegistry, repo_id: Optional[str], *, with_variant: bool
) -> list[ActiveDownload]:
downloads: list[ActiveDownload] = []
for ref in registry.active_job_refs(repo_id):
metadata = ref.metadata
if with_variant:
ref_repo_id = metadata.repo_id if metadata is not None else ref.key.split("::", 1)[0]
if metadata is not None:
variant = metadata.variant
else:
_repo, sep, raw_variant = ref.key.partition("::")
variant = raw_variant if sep and raw_variant else None
else:
ref_repo_id = metadata.repo_id if metadata is not None else ref.key
variant = None
downloads.append(
ActiveDownload(
repo_id = ref_repo_id,
variant = variant,
transport = metadata.transport if metadata is not None else None,
state = ref.state,
generation = ref.generation,
)
)
return downloads

4
studio/backend/hub/services/models/__init__.py Normal file
View file

@ -0,0 +1,4 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Model service layer."""

461
studio/backend/hub/services/models/cache_inventory.py Normal file
View file

@ -0,0 +1,461 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Cached model inventory."""
from __future__ import annotations
import json
import asyncio
import threading
import time
from collections import OrderedDict
from pathlib import Path
from typing import NamedTuple, Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.schemas.inventory import ModelFormat
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils import download_registry
from hub.utils.snapshot_filters import (
snapshot_download_blob_hashes,
snapshot_download_size,
)
from hub.services.models.common import (
_capabilities_for_format,
_classify_non_gguf_model_format,
_gguf_variant_state_summary,
_is_adapter_weight_name,
_is_checkpoint_weight_name,
_is_gguf_filename,
_is_main_gguf_filename,
_is_transformers_safetensors_weight_name,
_local_inventory_id,
_prefer_complete_larger,
_runtime_for_format,
)
logger = get_logger(__name__)
_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = OrderedDict()
_repo_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict()
_REPO_SIZE_CACHE_MAX = 256
_REPO_SIZE_POS_TTL = 60.0
_REPO_SIZE_NEG_TTL = 60.0
_MODEL_METADATA_TIMEOUT_SECONDS = 5.0
_repo_size_cache_lock = threading.Lock()
def get_repo_snapshot_metadata_cached(
repo_id: str, hf_token: Optional[str] = None
) -> tuple[int, frozenset[str]]:
token_fp = hf_cache_scan.token_fingerprint(hf_token)
cache_key = (repo_id, token_fp)
with _repo_size_cache_lock:
cached = _repo_size_cache.get(cache_key)
if cached is not None:
total, blob_hashes, ts = cached
if (time.monotonic() - ts) < _REPO_SIZE_POS_TTL:
_repo_size_cache.move_to_end(cache_key)
return total, blob_hashes
del _repo_size_cache[cache_key]
neg_ts = _repo_size_neg_cache.get(cache_key)
if neg_ts is not None and (time.monotonic() - neg_ts) < _REPO_SIZE_NEG_TTL:
return 0, frozenset()
try:
from huggingface_hub import HfApi
info = HfApi(token = hf_token).model_info(
repo_id,
files_metadata = True,
timeout = _MODEL_METADATA_TIMEOUT_SECONDS,
)
total = snapshot_download_size(info.siblings)
blob_hashes = snapshot_download_blob_hashes(info.siblings)
except Exception as e:
logger.warning(
"Failed to get repo size for %s: %s",
repo_id,
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
with _repo_size_cache_lock:
_repo_size_neg_cache[cache_key] = time.monotonic()
_repo_size_neg_cache.move_to_end(cache_key)
while len(_repo_size_neg_cache) > _REPO_SIZE_CACHE_MAX:
_repo_size_neg_cache.popitem(last = False)
return 0, frozenset()
with _repo_size_cache_lock:
_repo_size_cache[cache_key] = (total, blob_hashes, time.monotonic())
_repo_size_cache.move_to_end(cache_key)
_repo_size_neg_cache.pop(cache_key, None)
while len(_repo_size_cache) > _REPO_SIZE_CACHE_MAX:
_repo_size_cache.popitem(last = False)
return total, blob_hashes
def all_hf_cache_scans():
return hf_cache_scan.all_hf_cache_scans()
def _repo_gguf_size_bytes(repo_info) -> int:
"""Sum primary GGUF blob sizes across revisions, deduped by blob path (HF hardlinks shared blobs); mmproj is excluded so a vision-adapter-only repo isn't classed as GGUF."""
unique_blobs: dict[str, int] = {}
for revision in repo_info.revisions:
rev_id = getattr(revision, "commit_hash", None) or str(id(revision))
for f in revision.files:
if _is_main_gguf_filename(f.file_name):
blob_path = getattr(f, "blob_path", None)
size = f.size_on_disk or 0
if blob_path:
unique_blobs[str(blob_path)] = size
else:
unique_blobs[f"{rev_id}:{f.file_name}"] = size
return sum(unique_blobs.values())
def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool:
if existing is None:
return True
return _prefer_complete_larger(
bool(candidate.get("partial")),
int(candidate.get("size_bytes") or 0),
bool(existing.get("partial")),
int(existing.get("size_bytes") or 0),
)
def _cache_inventory_fields(
repo_id: str,
model_format: ModelFormat,
*,
partial: bool = False,
requires_variant: bool = False,
) -> dict:
return {
"inventory_id": _local_inventory_id("cache", model_format, repo_id),
"load_id": repo_id,
"model_format": model_format,
"runtime": _runtime_for_format(model_format),
"format_variant": None,
"capabilities": _capabilities_for_format(
model_format,
"hf_cache",
partial = partial,
requires_variant = requires_variant,
).model_dump(),
}
def invalidate_hf_cache_scans() -> None:
hf_cache_scan.invalidate_hf_cache_scans()
def _scan_cached_gguf() -> list[dict]:
"""Synchronous HF-cache disk walk for GGUF repos; runs in a worker thread."""
cache_scans = all_hf_cache_scans()
seen_lower: dict[str, dict] = {}
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
try:
if str(repo_info.repo_type) != "model":
continue
repo_id = repo_info.repo_id
total_size = _repo_gguf_size_bytes(repo_info)
has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id)
if total_size == 0 and not has_variant_state:
continue
partial = hf_cache_scan.is_gguf_repo_partial(
repo_id,
Path(repo_info.repo_path),
)
if total_size == 0 and not partial:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
row = {
"repo_id": repo_id,
"size_bytes": max(total_size, variant_state_size),
"cache_path": str(repo_info.repo_path),
"partial": partial,
# GGUF row-level transport is ambiguous (variants may differ);
# per-variant detail lives on GgufVariantDetail.
"partial_transport": None,
}
row.update(
_cache_inventory_fields(
repo_id,
"gguf",
partial = bool(row["partial"]),
requires_variant = True,
)
)
if _prefer_cache_row(row, existing):
seen_lower[key] = row
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}")
continue
return sorted(seen_lower.values(), key = lambda c: c["repo_id"])
async def list_cached_gguf_response(hf_token: Optional[str] = None):
"""List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
try:
cached = await asyncio.to_thread(_scan_cached_gguf)
return {"cached": cached}
except Exception as e:
logger.error(
"Error listing cached GGUF repos: %s",
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
raise HTTPException(
status_code = 500,
detail = "Failed to read the local model cache.",
) from e
class _CachedNonGgufPayload(NamedTuple):
size_bytes: int
has_runnable_weights: bool
model_format: ModelFormat
def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
all_weight_blobs: dict[str, int] = {}
adapter_blobs: dict[str, int] = {}
safetensors_blobs: dict[str, int] = {}
checkpoint_blobs: dict[str, int] = {}
has_config = False
has_adapter_config = False
has_adapter_weights = False
has_safetensors = False
has_transformers_safetensors = False
has_checkpoint = False
def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None:
blob_path = getattr(file_obj, "blob_path", None)
size = int(file_obj.size_on_disk or 0)
key = str(blob_path) if blob_path else f"{rev_id}:{file_name}"
target[key] = size
all_weight_blobs[key] = size
for revision in repo_info.revisions:
rev_id = getattr(revision, "commit_hash", None) or str(id(revision))
for f in revision.files:
file_name = str(f.file_name)
lower = file_name.lower()
name = lower.replace("\\", "/").rsplit("/", 1)[-1]
if _is_gguf_filename(lower):
continue
if name == "config.json":
has_config = True
continue
if name == "adapter_config.json":
has_adapter_config = True
continue
is_adapter = _is_adapter_weight_name(name)
is_safetensors = name.endswith(".safetensors") and not is_adapter
is_checkpoint = _is_checkpoint_weight_name(name)
if is_adapter:
has_adapter_weights = True
_record_blob(adapter_blobs, f, rev_id, file_name)
if is_safetensors:
has_safetensors = True
if _is_transformers_safetensors_weight_name(name):
has_transformers_safetensors = True
_record_blob(safetensors_blobs, f, rev_id, file_name)
if is_checkpoint:
has_checkpoint = True
_record_blob(checkpoint_blobs, f, rev_id, file_name)
model_format = (
_classify_non_gguf_model_format(
has_config = has_config,
has_adapter_config = has_adapter_config,
has_adapter_weights = has_adapter_weights,
has_safetensors = has_safetensors,
has_transformers_safetensors = has_transformers_safetensors,
has_checkpoint_weights = has_checkpoint,
trusted_hf_cache_repo = True,
)
or "unknown"
)
if model_format == "adapter":
size_bytes = sum(adapter_blobs.values())
elif model_format == "safetensors":
size_bytes = sum(safetensors_blobs.values())
elif model_format == "checkpoint":
size_bytes = sum(checkpoint_blobs.values())
else:
size_bytes = sum(all_weight_blobs.values())
return _CachedNonGgufPayload(
size_bytes = size_bytes,
has_runnable_weights = model_format != "unknown",
model_format = model_format,
)
def _cached_model_snapshot_path(repo_path: Path) -> Optional[Path]:
resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_path)
if not resolved:
return None
path = Path(resolved)
return path if path.is_dir() else None
def _read_json_object(path: Path) -> dict:
try:
with open(path, "r", encoding = "utf-8") as f:
data = json.load(f)
return data if isinstance(data, dict) else {}
except Exception:
return {}
def _read_model_card_frontmatter(path: Path) -> dict:
try:
text = path.read_text(encoding = "utf-8")
except Exception:
return {}
lines = text.splitlines()
if not lines or lines[0].strip() != "---":
return {}
body: list[str] = []
for line in lines[1:]:
if line.strip() == "---":
break
body.append(line)
if not body:
return {}
try:
import yaml
data = yaml.safe_load("\n".join(body)) or {}
return data if isinstance(data, dict) else {}
except Exception:
return {}
def _cached_model_local_metadata(repo_path: Path) -> dict:
snapshot = _cached_model_snapshot_path(repo_path)
if snapshot is None:
return {}
result: dict = {}
config = _read_json_object(snapshot / "config.json")
quant_method = (
config.get("quantization_config", {}).get("quant_method")
if isinstance(config.get("quantization_config"), dict)
else None
)
if isinstance(quant_method, str) and quant_method.strip():
result["quant_method"] = quant_method.strip()
card = _read_model_card_frontmatter(snapshot / "README.md")
pipeline_tag = card.get("pipeline_tag")
if isinstance(pipeline_tag, str) and pipeline_tag.strip():
result["pipeline_tag"] = pipeline_tag.strip()
library_name = card.get("library_name")
if isinstance(library_name, str) and library_name.strip():
result["library_name"] = library_name.strip()
tags = card.get("tags")
if isinstance(tags, list):
clean_tags = [tag.strip() for tag in tags if isinstance(tag, str) and tag.strip()]
if clean_tags:
result["tags"] = clean_tags
return result
def _scan_cached_models() -> list[dict]:
"""Synchronous HF-cache disk walk for non-GGUF model repos; runs in a worker thread."""
cache_scans = all_hf_cache_scans()
seen_lower: dict[str, dict] = {}
inspected = 0
skipped_gguf = 0
skipped_no_weights = 0
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
inspected += 1
try:
if str(repo_info.repo_type) != "model":
continue
repo_id = repo_info.repo_id
has_main_gguf = _repo_has_gguf_files(repo_info)
payload = _repo_non_gguf_model_payload(repo_info)
if payload.size_bytes == 0:
if has_main_gguf:
skipped_gguf += 1
continue
if not payload.has_runnable_weights:
skipped_no_weights += 1
continue
key = repo_id.lower()
existing = seen_lower.get(key)
repo_path = Path(repo_info.repo_path)
snapshot_partial = hf_cache_scan.is_snapshot_partial(
"model",
repo_id,
repo_path,
)
row = {
"repo_id": repo_id,
"size_bytes": payload.size_bytes,
"cache_path": str(repo_info.repo_path),
"partial": snapshot_partial,
"partial_transport": (
hf_cache_scan.partial_transport_for(
"model",
repo_id,
repo_cache_dir = repo_path,
)
if snapshot_partial
else None
),
**_cached_model_local_metadata(repo_path),
}
row.update(
_cache_inventory_fields(
repo_id,
payload.model_format,
partial = bool(row["partial"]),
)
)
if _prefer_cache_row(row, existing):
seen_lower[key] = row
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")
continue
cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"])
logger.info(
"Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d returned=%d",
inspected,
skipped_gguf,
skipped_no_weights,
len(cached),
)
return cached
async def list_cached_models_response(hf_token: Optional[str] = None):
"""List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
try:
cached = await asyncio.to_thread(_scan_cached_models)
return {"cached": cached}
except Exception as e:
logger.error(
"Error listing cached models: %s",
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
raise HTTPException(
status_code = 500,
detail = "Failed to read the local model cache.",
) from e

610
studio/backend/hub/services/models/common.py Normal file
View file

@ -0,0 +1,610 @@
# 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 model inventory helpers for the Hub service layer."""
from __future__ import annotations
import json
import os
from pathlib import Path
from typing import List, Literal, Optional
from urllib.parse import quote
from hub.schemas.inventory import (
LocalModelCapabilities,
LocalModelInfo,
ModelFormat,
ModelRuntime,
)
from hub.utils.gguf import (
extract_quant_label,
is_gguf_filename as _is_gguf_filename,
is_mmproj_filename as _is_mmproj_filename,
)
from hub.utils.paths import is_valid_repo_id as _is_valid_repo_id
ModelType = Literal["text", "vision", "audio", "embeddings"]
LocalModelSource = Literal["models_dir", "hf_cache", "lmstudio", "ollama", "custom"]
def _safe_is_dir(path) -> bool:
# Py >= 3.12 propagates PermissionError (EACCES) from is_dir(); folder scans
# probe root-owned system dirs, so treat un-stat-able paths as not-a-dir.
try:
return Path(path).is_dir()
except OSError:
return False
_LOCAL_CHECKPOINT_EXTENSIONS = (
".bin",
".pt",
".pth",
".ckpt",
".h5",
".msgpack",
".npz",
)
_LOCAL_BASE_MODEL_PREFIXES = {
"checkpoint",
"checkpoints",
"export",
"exports",
"model",
"models",
"output",
"outputs",
"run",
"runs",
"train",
}
_HF_CACHE_MODEL_FILE_PROBE_LIMIT = 2000
def _is_model_directory(d: Path) -> bool:
"""True when *d* has a config plus real weights; excludes mmproj GGUFs and non-weight ``.bin`` files (``tokenizer.bin``) to avoid false positives."""
def _is_weight_file(f: Path) -> bool:
suffix = f.suffix.lower()
if suffix == ".safetensors":
return True
if suffix == ".gguf":
return "mmproj" not in f.name.lower()
if suffix == ".bin":
name = f.name.lower()
return (
name.startswith("pytorch_model")
or name.startswith("model")
or name.startswith("adapter_model")
or name.startswith("consolidated")
)
return False
try:
has_config = (d / "config.json").exists() or (d / "adapter_config.json").exists()
if not has_config:
return False
return any(_is_weight_file(f) for f in d.iterdir() if f.is_file())
except OSError:
return False
def _local_inventory_id(
source: str,
model_format: ModelFormat,
semantic_id: str,
variant: Optional[str] = None,
) -> str:
parts = [
source,
model_format,
quote(semantic_id, safe = ""),
]
if variant:
parts.append(quote(variant, safe = ""))
return ":".join(parts)
def _runtime_for_format(model_format: ModelFormat) -> ModelRuntime:
if model_format == "gguf":
return "llama_cpp"
if model_format == "adapter":
return "adapter"
if model_format in {"safetensors", "checkpoint"}:
return "transformers"
return "unknown"
def _capabilities_for_format(
model_format: ModelFormat,
source: str,
*,
partial: bool = False,
requires_variant: bool = False,
) -> LocalModelCapabilities:
is_complete = not partial
can_chat = model_format in {"gguf", "safetensors", "adapter", "checkpoint"}
can_train = model_format in {"safetensors", "checkpoint"} and is_complete
return LocalModelCapabilities(
can_train = can_train,
can_chat = can_chat and is_complete,
can_delete = source == "hf_cache",
can_download = False,
requires_variant = requires_variant,
supports_lora = model_format in {"safetensors", "checkpoint"} and is_complete,
supports_vision = False,
)
def _prefer_complete_larger(
candidate_partial: bool,
candidate_size_bytes: int,
existing_partial: bool,
existing_size_bytes: int,
) -> bool:
if candidate_partial != existing_partial:
return not candidate_partial
return candidate_size_bytes > existing_size_bytes
def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]:
"""Whether GGUF variant-scoped state exists and its expected size; a cancelled/in-progress variant may have only manifests/markers/`.incomplete` blobs, which inventory needs to avoid a generic fallback row."""
from hub.utils import download_manifest
variant_keys: set[str] = set()
size_by_variant: dict[str, int] = {}
for variant, _path in download_manifest.iter_variant_manifests(
"model",
repo_id,
):
key = variant.lower()
variant_keys.add(key)
manifest = download_manifest.read_manifest("model", repo_id, variant)
if manifest is None:
continue
size_by_variant[key] = max(
size_by_variant.get(key, 0),
sum(max(0, int(file.size or 0)) for file in manifest.expected_files),
)
for variant, _path in download_manifest.iter_variant_markers(
"model",
repo_id,
):
variant_keys.add(variant.lower())
return bool(variant_keys), sum(size_by_variant.values())
def _apply_format_aware_partial(
rows: List[LocalModelInfo],
*,
snapshot_partial: bool,
gguf_partial: bool,
snapshot_partial_transport: Optional[str] = None,
) -> List[LocalModelInfo]:
"""Rewrite each row's partial flag with format-aware predicates so a hybrid (gguf + safetensors) repo's broken format doesn't taint the clean one; capabilities are recomputed from the new flag."""
rewritten: List[LocalModelInfo] = []
for row in rows:
target = gguf_partial if row.model_format == "gguf" else snapshot_partial
if not target:
rewritten.append(row)
continue
# GGUF row-level transport is ambiguous (variants may differ); per-variant
# detail lives on GgufVariantDetail.partial_transport via the variants endpoint.
partial_transport = None if row.model_format == "gguf" else snapshot_partial_transport
rewritten.append(
row.model_copy(
update = {
"partial": True,
"partial_transport": partial_transport,
"capabilities": _capabilities_for_format(
row.model_format,
row.source,
partial = True,
requires_variant = row.capabilities.requires_variant,
),
}
)
)
return rewritten
def _weight_basename(name: str) -> str:
return name.replace("\\", "/").rsplit("/", 1)[-1].lower()
def _is_adapter_weight_name(name: str) -> bool:
lower = _weight_basename(name)
return lower.startswith("adapter_model") and lower.endswith((".safetensors", ".bin"))
def _is_transformers_safetensors_weight_name(name: str) -> bool:
lower = _weight_basename(name)
return lower.endswith(".safetensors") and lower.startswith(
("model", "pytorch_model", "consolidated")
)
def _is_transformers_bin_weight_name(name: str) -> bool:
lower = _weight_basename(name)
if not lower.endswith(".bin"):
return False
return lower.startswith(("pytorch_model", "model", "consolidated", "adapter_model"))
def _is_checkpoint_weight_name(name: str) -> bool:
lower = _weight_basename(name)
if lower.endswith(".bin"):
return _is_transformers_bin_weight_name(lower)
return lower.endswith(_LOCAL_CHECKPOINT_EXTENSIONS)
def _is_adapter_weight_file(path: Path) -> bool:
return _is_adapter_weight_name(path.name)
def _is_transformers_safetensors_weight_file(path: Path) -> bool:
return _is_transformers_safetensors_weight_name(path.name)
def _is_transformers_bin_weight_file(path: Path) -> bool:
return _is_transformers_bin_weight_name(path.name)
def _is_checkpoint_weight_file(path: Path) -> bool:
return _is_checkpoint_weight_name(path.name)
def _classify_non_gguf_model_format(
*,
has_config: bool,
has_adapter_config: bool,
has_adapter_weights: bool,
has_safetensors: bool,
has_transformers_safetensors: bool,
has_checkpoint_weights: bool,
trusted_hf_cache_repo: bool = False,
) -> Optional[ModelFormat]:
if has_safetensors and (has_config or (trusted_hf_cache_repo and has_transformers_safetensors)):
return "safetensors"
if has_adapter_config and has_adapter_weights:
return "adapter"
if has_config and has_checkpoint_weights:
return "checkpoint"
return None
def _is_main_gguf_filename(name: str) -> bool:
return _is_gguf_filename(name) and not _is_mmproj_filename(name)
def _iter_gguf_paths(root: Path):
stack = [root]
while stack:
current = stack.pop()
try:
entries = list(current.iterdir())
except OSError:
continue
for path in entries:
try:
if path.is_dir() and not path.is_symlink():
stack.append(path)
elif path.is_file() and _is_gguf_filename(path.name):
yield path
except OSError:
continue
def _iter_immediate_files(path: Path, *, include_symlinks: bool = False) -> list[Path]:
if path.is_file():
return [path]
if not path.is_dir():
return []
try:
return [
entry
for entry in path.iterdir()
if entry.is_file() or (include_symlinks and entry.is_symlink())
]
except OSError:
return []
def _iter_hf_cache_model_files(path: Path) -> list[Path]:
files = _iter_immediate_files(path, include_symlinks = True)
if not path.is_dir():
return files
if any(
_is_main_gguf_filename(entry.name)
or _is_transformers_safetensors_weight_file(entry)
or _is_checkpoint_weight_file(entry)
for entry in files
):
return files
try:
bounded: list[Path] = []
for index, entry in enumerate(path.rglob("*"), start = 1):
if index > _HF_CACHE_MODEL_FILE_PROBE_LIMIT:
break
if entry.is_file() or entry.is_symlink():
bounded.append(entry)
return bounded
except OSError:
return []
def _file_size_bytes(path: Path) -> int:
try:
if path.is_file() or path.is_symlink():
return path.stat().st_size
except OSError:
return 0
return 0
def _sum_file_sizes(paths) -> int:
return sum(_file_size_bytes(path) for path in paths)
def _main_gguf_files(path: Path, *, include_symlinks: bool = False) -> list[Path]:
return [
entry
for entry in _iter_immediate_files(path, include_symlinks = include_symlinks)
if _is_main_gguf_filename(entry.name)
]
def _format_label(model_format: ModelFormat) -> str:
if model_format == "gguf":
return "GGUF"
if model_format == "safetensors":
return "Safetensors"
if model_format == "adapter":
return "Adapter"
if model_format == "checkpoint":
return "Checkpoint"
return "Unknown"
def _read_adapter_config(path: Path) -> dict:
if not path.is_dir():
return {}
try:
with (path / "adapter_config.json").open("r", encoding = "utf-8") as f:
data = json.load(f)
except Exception:
return {}
return data if isinstance(data, dict) else {}
def _clean_optional_string(value: object) -> Optional[str]:
return value.strip() if isinstance(value, str) and value.strip() else None
def _base_model_looks_local(value: str) -> bool:
raw = value.strip()
normalized = raw.replace("\\", "/")
if raw.startswith(("/", "./", "../", "~", "\\\\")) or (
len(raw) >= 3 and raw[1] == ":" and raw[0].isalpha()
):
return True
first = normalized.split("/", 1)[0].lower()
return "/" in normalized and first in _LOCAL_BASE_MODEL_PREFIXES
def _base_model_source(value: Optional[str], adapter_dir: Path) -> Optional[str]:
if not value:
return None
candidates = [value, value.replace("\\", "/")]
for candidate in candidates:
try:
expanded = Path(os.path.expanduser(candidate))
if expanded.exists() or (adapter_dir / candidate).exists():
return "local"
except (OSError, ValueError):
return "unknown"
if _base_model_looks_local(value):
return "local"
if _is_valid_repo_id(value):
return "huggingface"
return "unknown"
def _local_model_info(
*,
scan_path: Path,
load_path: Path,
source: LocalModelSource,
model_format: ModelFormat,
display_name: Optional[str] = None,
model_id: Optional[str] = None,
updated_at: Optional[float] = None,
partial: bool = False,
requires_variant: bool = False,
format_variant: Optional[str] = None,
size_bytes: int = 0,
base_model: Optional[str] = None,
base_model_source: Optional[str] = None,
adapter_type: Optional[str] = None,
training_method: Optional[str] = None,
) -> LocalModelInfo:
load_id = model_id if source == "hf_cache" and model_id else str(load_path)
semantic_id = model_id or str(load_path)
return LocalModelInfo(
id = load_id,
inventory_id = _local_inventory_id(
source,
model_format,
semantic_id,
format_variant,
),
load_id = load_id,
model_id = model_id,
display_name = display_name or (scan_path.stem if scan_path.is_file() else scan_path.name),
path = str(load_path),
size_bytes = max(0, int(size_bytes or 0)),
source = source,
base_model = base_model,
base_model_source = base_model_source,
adapter_type = adapter_type,
training_method = training_method,
updated_at = updated_at,
partial = partial,
model_format = model_format,
runtime = _runtime_for_format(model_format),
format_variant = format_variant,
capabilities = _capabilities_for_format(
model_format,
source,
partial = partial,
requires_variant = requires_variant,
),
)
def _classify_local_path(
scan_path: Path,
source: LocalModelSource,
*,
load_path: Optional[Path] = None,
display_name: Optional[str] = None,
model_id: Optional[str] = None,
updated_at: Optional[float] = None,
partial: bool = False,
) -> list[LocalModelInfo]:
load_path = load_path or scan_path
files = (
_iter_hf_cache_model_files(scan_path)
if source == "hf_cache"
else _iter_immediate_files(scan_path)
)
if not files:
return []
rows: list[LocalModelInfo] = []
include_broken_snapshot_symlinks = source == "hf_cache"
gguf_files = _main_gguf_files(
scan_path,
include_symlinks = include_broken_snapshot_symlinks,
)
if gguf_files:
gguf_size_bytes = _sum_file_sizes(gguf_files)
variant = (
extract_quant_label(gguf_files[0].name)
if scan_path.is_file() and len(gguf_files) == 1
else None
)
rows.append(
_local_model_info(
scan_path = scan_path,
load_path = load_path,
source = source,
model_format = "gguf",
display_name = display_name,
model_id = model_id,
updated_at = updated_at,
partial = partial,
requires_variant = scan_path.is_dir(),
format_variant = variant,
size_bytes = gguf_size_bytes,
)
)
has_config = (scan_path / "config.json").is_file() if scan_path.is_dir() else False
has_adapter_config = (
(scan_path / "adapter_config.json").is_file() if scan_path.is_dir() else False
)
adapter_config = _read_adapter_config(scan_path) if has_adapter_config else {}
adapter_base_model = _clean_optional_string(adapter_config.get("base_model_name_or_path"))
adapter_type = _clean_optional_string(adapter_config.get("peft_type"))
training_method = _clean_optional_string(adapter_config.get("unsloth_training_method"))
has_adapter_weights = any(_is_adapter_weight_file(f) for f in files)
has_safetensors = any(
f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f) for f in files
)
has_transformers_safetensors = any(
_is_transformers_safetensors_weight_file(f) and not _is_adapter_weight_file(f)
for f in files
)
has_checkpoint_weights = any(_is_checkpoint_weight_file(f) for f in files)
trusted_hf_cache_repo = source == "hf_cache" and bool(model_id)
model_format = _classify_non_gguf_model_format(
has_config = has_config,
has_adapter_config = has_adapter_config,
has_adapter_weights = has_adapter_weights,
has_safetensors = has_safetensors,
has_transformers_safetensors = has_transformers_safetensors,
has_checkpoint_weights = has_checkpoint_weights,
trusted_hf_cache_repo = trusted_hf_cache_repo,
)
if model_format is not None:
if model_format == "adapter":
size_bytes = _sum_file_sizes(f for f in files if _is_adapter_weight_file(f))
elif model_format == "safetensors":
size_bytes = _sum_file_sizes(
f
for f in files
if f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f)
)
else:
size_bytes = _sum_file_sizes(f for f in files if _is_checkpoint_weight_file(f))
rows.append(
_local_model_info(
scan_path = scan_path,
load_path = load_path,
source = source,
model_format = model_format,
display_name = display_name,
model_id = model_id,
updated_at = updated_at,
partial = partial,
size_bytes = size_bytes,
base_model = adapter_base_model if model_format == "adapter" else None,
base_model_source = (
_base_model_source(adapter_base_model, scan_path)
if model_format == "adapter"
else None
),
adapter_type = adapter_type if model_format == "adapter" else None,
training_method = training_method if model_format == "adapter" else None,
)
)
elif not rows:
fallback_format: ModelFormat = (
"safetensors" if trusted_hf_cache_repo and has_config else "unknown"
)
size_bytes = _sum_file_sizes(files)
rows.append(
_local_model_info(
scan_path = scan_path,
load_path = load_path,
source = source,
model_format = fallback_format,
display_name = display_name,
model_id = model_id,
updated_at = updated_at,
partial = partial or trusted_hf_cache_repo,
size_bytes = size_bytes,
)
)
if len(rows) > 1:
rows = [
row.model_copy(
update = {
"display_name": f"{row.display_name} ({_format_label(row.model_format)})",
"inventory_id": _local_inventory_id(
row.source,
row.model_format,
row.model_id or row.path,
row.format_variant,
),
}
)
for row in rows
]
return rows

455
studio/backend/hub/services/models/deletion.py Normal file
View file

@ -0,0 +1,455 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Cached model deletion."""
from __future__ import annotations
import asyncio
from pathlib import Path
from typing import Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.gguf import extract_quant_label
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
purge_partial_repo,
purge_repo_cache_dirs,
)
from hub.utils.paths import (
is_valid_gguf_variant as _is_valid_gguf_variant,
is_valid_repo_id as _is_valid_repo_id,
resolve_cached_repo_id_case,
)
from hub.services import resolve_destructive_repo_ids
from hub.services.models import cache_inventory, downloads, gguf_variants
from hub.services.models.common import (
_is_gguf_filename,
_is_main_gguf_filename,
_is_mmproj_filename,
)
logger = get_logger(__name__)
def _snapshot_blob_reference_counts(repo_dir: Optional[Path]) -> dict[Path, int]:
"""Map each blob's realpath to its live snapshot symlink count, so per-variant deletion never unlinks a blob another revision still references (call after the target variant's own symlinks are removed)."""
counts: dict[Path, int] = {}
if repo_dir is None:
return counts
snapshots = repo_dir / "snapshots"
if not snapshots.is_dir():
return counts
try:
entries = list(snapshots.rglob("*"))
except OSError:
return counts
for link in entries:
try:
if not link.is_symlink():
continue
target = link.resolve()
except OSError:
continue
counts[target] = counts.get(target, 0) + 1
return counts
def _blob_hash_from_path(blob: Path) -> Optional[str]:
name = blob.name
if not name or name.endswith(INCOMPLETE_SUFFIX):
return None
return name
def _path_exists_or_symlink(path: Path) -> bool:
try:
return path.is_symlink() or path.exists()
except OSError:
return False
def _repo_file_matches(target_repo, predicate) -> list[tuple[Path, Optional[Path], str]]:
matches: list[tuple[Path, Optional[Path], str]] = []
for rev in getattr(target_repo, "revisions", ()):
for f in getattr(rev, "files", ()):
name = str(getattr(f, "file_name", ""))
if not predicate(name):
continue
file_path = getattr(f, "file_path", None)
if not file_path:
continue
blob_path = getattr(f, "blob_path", None)
matches.append(
(
Path(file_path),
Path(blob_path) if blob_path else None,
name,
)
)
return matches
def _has_remaining_main_gguf(target_repo) -> bool:
return any(
_path_exists_or_symlink(snap)
for snap, _blob, _name in _repo_file_matches(
target_repo,
_is_main_gguf_filename,
)
)
def _delete_gguf_variant_from_repos(
repo_id: str,
variant: str,
target_repos: list,
hf_token: Optional[str],
*,
sibling_active: bool = False,
) -> dict:
failures: list[str] = []
removed_snapshots = 0
deleted_bytes = 0
deleted_blobs = 0
completed_hashes: set[str] = set()
for target_repo in target_repos:
repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None
matched = _repo_file_matches(
target_repo,
lambda name: _is_main_gguf_filename(name)
and extract_quant_label(name).lower() == variant.lower(),
)
for snap, _blob, name in matched:
try:
if _path_exists_or_symlink(snap):
snap.unlink()
removed_snapshots += 1
except OSError as e:
failures.append(f"{name}: {e}")
companion_matches: list[tuple[Path, Optional[Path], str]] = []
if matched and not sibling_active and not _has_remaining_main_gguf(target_repo):
companion_matches = _repo_file_matches(
target_repo,
lambda name: _is_gguf_filename(name) and _is_mmproj_filename(name),
)
for snap, _blob, name in companion_matches:
try:
if _path_exists_or_symlink(snap):
snap.unlink()
removed_snapshots += 1
except OSError as e:
failures.append(f"{name}: {e}")
ref_counts = _snapshot_blob_reference_counts(repo_dir)
seen_blobs: set[Path] = set()
for _snap, blob, name in [*matched, *companion_matches]:
if blob is None:
continue
blob_hash = _blob_hash_from_path(blob)
if blob_hash:
completed_hashes.add(blob_hash)
try:
blob_key = blob.resolve()
except OSError:
blob_key = blob
if blob_key in seen_blobs:
continue
seen_blobs.add(blob_key)
if ref_counts.get(blob_key, 0) > 0:
continue
try:
if blob.exists():
deleted_bytes += blob.stat().st_size
blob.unlink()
deleted_blobs += 1
except OSError as e:
failures.append(f"{name}: {e}")
if failures:
raise HTTPException(
status_code = 409,
detail = (
f"Couldn't fully delete {variant} for {repo_id}: "
f"{len(failures)} file(s) are in use. "
"Unload the model and try again."
),
)
incomplete_result = gguf_variants.delete_variant_incomplete_blobs_result(
repo_id,
variant,
hf_token,
extra_hashes = frozenset(completed_hashes),
companions = not sibling_active,
)
if incomplete_result.unresolved:
raise HTTPException(
status_code = 409,
detail = (
f"Couldn't fully delete {variant} for {repo_id}: partial "
"download bytes exist but this variant's blob hashes are unavailable. "
"Reconnect or provide access to the repo, then try again."
),
)
state_purged = download_manifest.purge_state("model", repo_id, variant)
if (
removed_snapshots == 0
and deleted_blobs == 0
and incomplete_result.deleted == 0
and not state_purged
):
raise HTTPException(
status_code = 404,
detail = f"Variant {variant} not found in cache for {repo_id}",
)
freed_mb = deleted_bytes / (1024 * 1024)
logger.info(
f"Deleted {removed_snapshots} file(s) for {repo_id} variant {variant}: "
f"{freed_mb:.1f} MB freed"
)
return {"status": "deleted", "repo_id": repo_id, "variant": variant}
def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool:
"""True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``."""
rid = repo_id.lower()
lid = loaded_id.lower()
return lid == rid or lid.startswith(f"{rid}/")
def _loaded_repo_variant_blocks_delete(
loaded_id: str, repo_id: str, delete_variant: Optional[str], loaded_variant: Optional[str]
) -> bool:
if not _loaded_id_matches_repo(loaded_id, repo_id):
return False
if not delete_variant:
return True
if not loaded_variant:
return True
return loaded_variant.lower() == delete_variant.lower()
_LOAD_STATE_UNVERIFIABLE_DETAIL = (
"Couldn't verify whether this model is still loaded for inference. "
"Unload it if it is active, then try deleting again."
)
def _llama_cpp_blocks_delete(repo_id: str, variant: Optional[str]) -> bool:
"""Whether the llama.cpp backend holds *repo_id* (/variant). Acquiring fails open (import error means nothing loaded); reading load state is unguarded so a raise propagates and the caller fails closed rather than delete a live model."""
try:
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
except Exception as e:
logger.debug(f"llama.cpp backend unavailable during delete guard for {repo_id}: {e}")
return False
loaded_id = backend.model_identifier
loaded_variant = getattr(backend, "hf_variant", None)
if backend.is_active and not backend.is_loaded and loaded_id:
return _loaded_repo_variant_blocks_delete(
loaded_id,
repo_id,
variant,
loaded_variant,
)
if backend.is_loaded and loaded_id:
return _loaded_repo_variant_blocks_delete(
loaded_id,
repo_id,
variant,
loaded_variant,
)
return False
def _inference_backend_blocks_delete(repo_id: str) -> bool:
"""Whether the subprocess inference backend holds *repo_id*; same fail-open-on-acquire / surface-on-query contract as :func:`_llama_cpp_blocks_delete`."""
try:
from core.inference import get_inference_backend
backend = get_inference_backend()
except Exception as e:
logger.debug(f"Inference backend unavailable during delete guard for {repo_id}: {e}")
return False
active_name = backend.active_model_name
return bool(active_name) and _loaded_id_matches_repo(active_name, repo_id)
async def delete_cached_model_response(
repo_id: str,
variant: Optional[str] = None,
hf_token: Optional[str] = None,
):
"""Delete a cached model repo (or a specific GGUF variant) from the HF cache.
When *variant* is provided, only the GGUF files matching that quant label
are removed (e.g. ``UD-Q4_K_XL``). Otherwise the entire repo is deleted.
Refuses if the model is currently loaded for inference.
"""
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
variant = (variant or "").strip() or None
if variant is not None and not _is_valid_gguf_variant(variant):
raise HTTPException(
status_code = 400,
detail = f"Invalid gguf_variant: {variant!r}",
)
# Guard fails closed: if a live backend's load state can't be read, abort
# with 503 rather than risk unlinking weights under a running process.
try:
blocks_delete = _llama_cpp_blocks_delete(repo_id, variant) or (
_inference_backend_blocks_delete(repo_id)
)
except Exception as e:
logger.warning(f"Load-state verification failed for {repo_id}; refusing delete: {e}")
raise HTTPException(
status_code = 503,
detail = _LOAD_STATE_UNVERIFIABLE_DETAIL,
)
if blocks_delete:
raise HTTPException(
status_code = 400,
detail = "Unload the model before deleting",
)
repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
if not downloads.registry.begin_delete(repo_key, variant):
detail = (
f"Cancel the {variant} download before deleting it."
if variant is not None
else "Cancel the active downloads before deleting."
)
raise HTTPException(status_code = 400, detail = detail)
try:
return await asyncio.to_thread(_delete_cached_model_blocking, repo_id, variant, hf_token)
finally:
downloads.registry.end_delete(repo_key, variant)
cache_inventory.invalidate_hf_cache_scans()
def _delete_cached_model_blocking(
repo_id: str, variant: Optional[str], hf_token: Optional[str]
) -> dict:
try:
# If a sibling quant is downloading concurrently, restrict this delete to
# the variant's own files and leave the shared mmproj companion for it.
sibling_active = bool(
variant and downloads.registry.has_active_peer_variant(repo_id, variant)
)
cache_scans = cache_inventory.all_hf_cache_scans()
candidate_entries = []
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
if str(repo_info.repo_type) != "model":
continue
if repo_info.repo_id.lower() == repo_id.lower():
candidate_entries.append((hf_cache, repo_info))
matched_repo_ids = resolve_destructive_repo_ids(
repo_id,
[str(repo_info.repo_id) for _hf_cache, repo_info in candidate_entries],
noun = "models",
)
target_entries = [
(hf_cache, repo_info)
for hf_cache, repo_info in candidate_entries
if str(repo_info.repo_id) in matched_repo_ids
]
if not target_entries:
if variant is None:
cache_purged = purge_repo_cache_dirs("model", repo_id) or purge_partial_repo(
"model", repo_id
)
state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0
if cache_purged or state_purged:
return {"status": "deleted", "repo_id": repo_id}
if variant:
incomplete_result = gguf_variants.delete_variant_incomplete_blobs_result(
repo_id,
variant,
hf_token,
companions = not sibling_active,
)
if incomplete_result.unresolved:
raise HTTPException(
status_code = 409,
detail = (
f"Couldn't fully delete {variant} for {repo_id}: partial "
"download bytes exist but this variant's blob hashes are unavailable. "
"Reconnect or provide access to the repo, then try again."
),
)
state_purged = download_manifest.purge_state(
"model",
repo_id,
variant,
)
if incomplete_result.deleted > 0 or state_purged:
return {
"status": "deleted",
"repo_id": repo_id,
"variant": variant,
}
raise HTTPException(status_code = 404, detail = "Model not found in cache")
if variant:
return _delete_gguf_variant_from_repos(
repo_id,
variant,
[repo for _cache, repo in target_entries],
hf_token,
sibling_active = sibling_active,
)
deleted_revisions = False
for hf_cache, repo_info in target_entries:
revision_hashes = [
rev.commit_hash for rev in repo_info.revisions if getattr(rev, "commit_hash", None)
]
if not revision_hashes:
continue
delete_strategy = hf_cache.delete_revisions(*revision_hashes)
logger.info(
f"Deleting cached model {repo_id} from "
f"{getattr(hf_cache, 'cache_dir', '<unknown>')}: "
f"{delete_strategy.expected_freed_size_str} will be freed"
)
delete_strategy.execute()
deleted_revisions = True
cache_purged = purge_repo_cache_dirs("model", repo_id)
partial_purged = purge_partial_repo("model", repo_id)
state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0
if not (deleted_revisions or cache_purged or partial_purged or state_purged):
raise HTTPException(status_code = 404, detail = "No revisions found for model")
return {"status": "deleted", "repo_id": repo_id}
except HTTPException:
raise
except Exception as e:
logger.error(
"Error deleting cached model %s: %s",
repo_id,
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
raise HTTPException(
status_code = 500,
detail = "Failed to delete cached model: "
+ download_registry.scrub_secrets(str(e), hf_token = hf_token),
)

411
studio/backend/hub/services/models/downloads.py Normal file
View file

@ -0,0 +1,411 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Download orchestration."""
from __future__ import annotations
import asyncio
from typing import Optional, TYPE_CHECKING
from fastapi import HTTPException
from loggers import get_logger
from hub.schemas.downloads import (
ActiveDownloadsResponse,
CancelDownloadRequest,
DownloadJobStatus,
DownloadModelRequest,
)
from hub.utils import download_registry
from hub.utils import download_manifest
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import has_active_incomplete_blobs
from hub.utils.paths import (
is_valid_gguf_variant as _is_valid_gguf_variant,
is_valid_repo_id as _is_valid_repo_id,
resolve_cached_repo_id_case,
)
from hub.services import snapshot_progress
from hub.services import download_lifecycle
from hub.services.models import cache_inventory, gguf_variants
logger = get_logger(__name__)
if TYPE_CHECKING:
import subprocess
_registry = download_registry.get_models_registry()
def _download_job_key(repo_id: str, variant: Optional[str]) -> str:
return download_registry.normalize_job_key(
f"{download_registry.normalize_repo_key(repo_id)}::{variant or ''}"
)
def _job_status(
key: str,
*,
repo_id: Optional[str] = None,
variant: Optional[str] = None,
) -> DownloadJobStatus:
state, error, generation = download_lifecycle.idle_status(
_registry,
key,
repo_type = "model",
repo_id = repo_id,
variant = variant,
)
return DownloadJobStatus(state = state, error = error, generation = generation)
def _spawn_download_worker(
repo_id: str,
variant: Optional[str],
hf_token: Optional[str],
use_xet: bool = False,
protected_blob_hashes: Optional[frozenset[str]] = None,
) -> subprocess.Popen:
args = ["--repo-id", repo_id]
if variant:
args.extend(["--variant", variant])
return download_lifecycle.spawn_worker(
args,
hf_token,
use_xet = use_xet,
protected_blob_hashes = protected_blob_hashes,
)
async def download_model_response(body: DownloadModelRequest, hf_token: Optional[str] = None):
"""Start a background download for a HuggingFace model."""
repo_id = body.repo_id.strip()
if not _is_valid_repo_id(repo_id):
raise HTTPException(
status_code = 400,
detail = f"Invalid repo_id: {repo_id!r}",
)
# Canonicalize so two different-cased paste-ins share one job + cache dir.
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
variant = (body.gguf_variant or "").strip() or None
if variant is not None and not _is_valid_gguf_variant(variant):
raise HTTPException(
status_code = 400,
detail = f"Invalid gguf_variant: {variant!r}",
)
key = _download_job_key(repo_id, variant)
transport = download_lifecycle.resolve_transport(body.use_xet)
variant_blob_hashes = frozenset()
variant_progress_blob_hashes = frozenset()
completed_baseline_bytes = 0
if variant is not None:
try:
variant_blob_hashes = await asyncio.to_thread(
gguf_variants.gguf_variant_blob_hashes,
repo_id,
variant,
hf_token,
include_companions = False,
)
variant_progress_blob_hashes = await asyncio.to_thread(
gguf_variants.gguf_variant_blob_hashes,
repo_id,
variant,
hf_token,
include_companions = True,
)
except Exception as e:
logger.warning(
"GGUF hash pre-resolution failed for %s [%s]; continuing without "
"a completed-bytes baseline or peer-protection hashes (the worker "
"re-resolves its own blobs before purging): %s",
repo_id,
variant,
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
has_variant_resume_state = (
download_manifest.has_cancel_marker("model", repo_id, variant)
or download_manifest.read_manifest("model", repo_id, variant) is not None
)
if variant_progress_blob_hashes and not has_variant_resume_state:
completed_baseline_bytes = await asyncio.to_thread(
download_registry.completed_blob_bytes,
"model",
repo_id,
variant_progress_blob_hashes,
)
claimed, claim_state = _registry.claim(
key,
transport,
repo_type = "model",
repo_id = repo_id,
variant = variant,
blob_hashes = variant_blob_hashes,
progress_blob_hashes = variant_progress_blob_hashes,
completed_baseline_bytes = completed_baseline_bytes,
)
generation = _registry.current_generation(key)
if not claimed:
# claim_state is the blocking job's state. The client can attach only
# when the blocker is this key's own in-flight job (adoptable); a
# cross-variant conflict or in-progress delete is not accepted.
return {
"job_key": key,
"state": claim_state,
"accepted": _registry.adoptable(key),
"generation": generation,
}
download_manifest.clear_cancel_marker("model", repo_id, variant)
# Blobs a concurrent same-repo variant is already writing (e.g. a shared
# mmproj). The worker must not purge these during cache preparation.
protected_blob_hashes = _registry.peer_blob_hashes(key) if variant else frozenset()
label = f"{repo_id}{f' [{variant}]' if variant else ''}"
state = download_lifecycle.launch_worker(
_registry,
key,
spawn = lambda: _spawn_download_worker(
repo_id,
variant,
hf_token,
use_xet = body.use_xet,
protected_blob_hashes = protected_blob_hashes,
),
hf_token = hf_token,
label = label,
log_prefix = "Download",
logger = logger,
repo_type = "model",
repo_id = repo_id,
transport = transport,
watch_name = f"hf-download-watch-{repo_id}",
)
return {
"job_key": key,
"state": state,
"accepted": True,
"generation": generation,
}
async def cancel_download_model_response(body: CancelDownloadRequest):
"""Cancel an in-flight model download (SIGKILL; HF cache resumes on next download)."""
repo_id = body.repo_id.strip()
if not _is_valid_repo_id(repo_id):
raise HTTPException(
status_code = 400,
detail = f"Invalid repo_id: {repo_id!r}",
)
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
variant = (body.gguf_variant or "").strip() or None
if variant is not None and not _is_valid_gguf_variant(variant):
raise HTTPException(
status_code = 400,
detail = f"Invalid gguf_variant: {variant!r}",
)
key = _download_job_key(repo_id, variant)
state = download_lifecycle.cancel_worker(
_registry,
key,
generation = body.generation,
label = repo_id,
logger = logger,
)
return {"job_key": key, "state": state}
async def get_download_status_response(repo_id: str, gguf_variant: str = "") -> DownloadJobStatus:
"""Return the latest state of a background download job."""
repo_id = repo_id.strip()
if not _is_valid_repo_id(repo_id):
return DownloadJobStatus(state = "idle")
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
variant = (gguf_variant or "").strip() or None
key = _download_job_key(repo_id, variant)
return _job_status(key, repo_id = repo_id, variant = variant)
async def get_active_downloads_response(repo_id: str = "") -> ActiveDownloadsResponse:
"""Return every in-flight download for a repo in a single call."""
repo_id = repo_id.strip()
if repo_id and not _is_valid_repo_id(repo_id):
return ActiveDownloadsResponse(downloads = [])
canonical_repo_id = (
await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
if repo_id
else None
)
return ActiveDownloadsResponse(
downloads = download_lifecycle.active_download_refs(
_registry,
canonical_repo_id,
with_variant = True,
)
)
def _variant_transport_status(repo_id: str, variant: str, hf_token: Optional[str]) -> dict:
incomplete_hashes = download_registry.incomplete_blob_hashes(
"model",
repo_id,
active_only = True,
)
variant_hashes = gguf_variants.gguf_variant_blob_hashes(
repo_id,
variant,
hf_token,
allow_remote = False,
)
has_partial = hf_cache_scan.is_variant_partial(
repo_id,
variant,
incomplete_blob_hashes = incomplete_hashes,
variant_blob_hashes = variant_hashes,
)
last_transport = hf_cache_scan.partial_transport_for("model", repo_id, variant)
if (
last_transport is None
and has_partial
and incomplete_hashes
and variant_hashes
and incomplete_hashes.intersection(variant_hashes)
):
last_transport = download_registry.read_active_transport_marker(
"model",
repo_id,
variant,
)
has_matching_incomplete = bool(
incomplete_hashes and variant_hashes and incomplete_hashes.intersection(variant_hashes)
)
return {
"has_partial": has_partial,
"last_transport": last_transport,
"resumable": (
has_matching_incomplete and last_transport == download_registry.TRANSPORT_HTTP
),
}
async def get_model_transport_status_response(
repo_id: str,
gguf_variant: str = "",
hf_token: Optional[str] = None,
) -> dict:
"""Return last transport used for this repo + whether any partial blobs
exist + whether that partial supports byte-level resume.
``resumable`` is True only when an HTTP partial exists. XET partials
are reported via ``has_partial`` but always have ``resumable=False``
because ``hf_xet`` rewrites the destination from scratch on every
call (network resume happens transparently via its chunk cache).
"""
repo_id = repo_id.strip()
if not _is_valid_repo_id(repo_id):
return {"has_partial": False, "last_transport": None, "resumable": False}
variant = (gguf_variant or "").strip()
if variant:
if not _is_valid_gguf_variant(variant):
return {"has_partial": False, "last_transport": None, "resumable": False}
return _variant_transport_status(repo_id, variant, hf_token)
return {
"has_partial": has_active_incomplete_blobs("model", repo_id),
"last_transport": download_registry.read_active_transport_marker("model", repo_id),
"resumable": download_registry.is_resumable_partial("model", repo_id),
}
async def get_gguf_download_progress_response(
repo_id: str,
variant: str = "",
expected_bytes: int = 0,
hf_token: Optional[str] = None,
) -> dict:
"""Return download progress for a specific GGUF variant."""
expected_total = max(expected_bytes, 0)
progress_variant = variant.strip() or None
if progress_variant is not None and not _is_valid_gguf_variant(progress_variant):
return {
"downloaded_bytes": 0,
"completed_bytes": 0,
"complete_on_disk": False,
"expected_bytes": expected_total,
"progress": 0,
"cache_path": None,
}
def _metadata_resolver(
resolved_repo_id: str, token: Optional[str]
) -> tuple[int, frozenset[str]]:
if progress_variant is None:
return expected_total, frozenset()
requirement = gguf_variants.gguf_variant_requirements(
resolved_repo_id,
progress_variant,
token,
)
if requirement is not None:
return requirement.download_size_bytes, requirement.required_hashes
manifest = download_manifest.read_manifest(
"model",
resolved_repo_id,
progress_variant,
)
if manifest is not None:
return (
sum(max(0, int(file.size or 0)) for file in manifest.expected_files),
frozenset(file.sha256 for file in manifest.expected_files if file.sha256),
)
return (
expected_total,
gguf_variants.gguf_variant_blob_hashes(
resolved_repo_id,
progress_variant,
token,
allow_remote = False,
),
)
return await snapshot_progress.snapshot_progress_response(
repo_type = "model",
repo_id = repo_id,
job_key = _download_job_key(repo_id, progress_variant),
expected_bytes = expected_total,
hf_token = hf_token,
registry = _registry,
metadata_resolver = _metadata_resolver,
variant = progress_variant,
)
async def get_download_progress_response(
repo_id: str,
expected_bytes: int = 0,
hf_token: Optional[str] = None,
) -> dict:
"""Return download progress for any HuggingFace model repo.
Checks the local HF cache for completed blobs and in-progress
(.incomplete) downloads. Uses the caller-supplied expected total
when available; otherwise queries HF metadata and caches it.
Also returns ``cache_path``: the realpath of the snapshot directory
(or the cache repo root if no snapshot exists yet) so the UI can
show users where the weights actually live on disk.
"""
return await snapshot_progress.snapshot_progress_response(
repo_type = "model",
repo_id = repo_id,
job_key = _download_job_key(repo_id, None),
expected_bytes = expected_bytes,
hf_token = hf_token,
registry = _registry,
metadata_resolver = cache_inventory.get_repo_snapshot_metadata_cached,
)
registry = _registry

518
studio/backend/hub/services/models/folder_browser.py Normal file
View file

@ -0,0 +1,518 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Model folder recommendation and browsing services."""
from __future__ import annotations
import os
from pathlib import Path
from typing import Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.schemas.inventory import BrowseEntry, BrowseFoldersResponse
from hub.storage.scan_folders import (
contains_sensitive_path_component,
list_scan_folders,
)
from hub.utils.paths import (
exports_root,
hf_default_cache_dir,
legacy_hf_cache_dir,
lmstudio_model_dirs,
normalize_path,
outputs_root,
studio_root,
well_known_model_dirs,
)
from hub.services.models.common import _safe_is_dir
from hub.services.models.local_inventory import _resolve_hf_cache_dir
logger = get_logger(__name__)
def get_recommended_folders_response() -> dict:
"""Return well-known model directories that exist on this machine."""
folders: list[str] = []
seen: set[str] = set()
def _add(p: Optional[Path]) -> None:
if p is None:
return
try:
resolved = str(p.resolve())
except OSError:
return
if resolved in seen:
return
if _safe_is_dir(resolved) and os.access(resolved, os.R_OK | os.X_OK):
seen.add(resolved)
folders.append(resolved)
try:
for p in lmstudio_model_dirs():
_add(p)
except Exception as e:
logger.warning("Failed to scan for LM Studio model directories: %s", e)
ollama_env = os.environ.get("OLLAMA_MODELS")
if ollama_env:
_add(Path(normalize_path(ollama_env)).expanduser())
for candidate in (
Path.home() / ".ollama" / "models",
Path("/usr/share/ollama/.ollama/models"),
Path("/var/lib/ollama/.ollama/models"),
):
_add(candidate)
return {"folders": folders}
# Ceiling on children to stat when guessing if a directory holds models.
_BROWSE_MODEL_HINT_PROBE = 64
# Hard cap on returned subdirectory entries so pointing at ``/usr/lib`` or
# ``/proc`` can't stat-storm the process or flood the client.
_BROWSE_ENTRY_CAP = 2000
def _count_model_files(directory: Path, cap: int = 200) -> int:
"""Count GGUF/safetensors files immediately inside *directory*, bounded by visited entries (not matches) so the hint never costs more than ``cap`` stats."""
n = 0
visited = 0
try:
for f in directory.iterdir():
visited += 1
if visited > cap:
break
try:
if f.is_file():
low = f.name.lower()
if low.endswith((".gguf", ".safetensors")):
n += 1
except OSError:
continue
except PermissionError as e:
logger.debug("browse-folders: permission denied counting %s: %s", directory, e)
return 0
except OSError as e:
logger.debug("browse-folders: OS error counting %s: %s", directory, e)
return 0
return n
def _has_direct_model_signal(directory: Path) -> bool:
"""True if an immediate child signals a model (GGUF/safetensors/config file or ``models--*`` HF-cache subdir); bounded by the hint probe."""
try:
it = directory.iterdir()
except OSError:
return False
try:
for i, child in enumerate(it):
if i >= _BROWSE_MODEL_HINT_PROBE:
break
try:
name = child.name
if child.is_file():
low = name.lower()
if low.endswith((".gguf", ".safetensors")):
return True
if low in ("config.json", "adapter_config.json"):
return True
elif child.is_dir() and name.startswith("models--"):
return True
except OSError:
continue
except OSError:
return False
return False
def _looks_like_model_dir(directory: Path) -> bool:
"""Bounded heuristic flagging dirs worth exploring (false negatives are fine; the scanner is authoritative). Three signals, cheapest first: a ``models--*`` name, a direct child signal, or a grandchild signal (LM Studio / Ollama ``publisher/model/weights.gguf`` layout)."""
if directory.name.startswith("models--"):
return True
if _has_direct_model_signal(directory):
return True
try:
it = directory.iterdir()
except OSError:
return False
try:
for i, child in enumerate(it):
if i >= _BROWSE_MODEL_HINT_PROBE:
break
try:
if not child.is_dir():
continue
except OSError:
continue
if child.name.startswith("models--"):
return True
if _has_direct_model_signal(child):
return True
except OSError:
return False
return False
def _build_browse_allowlist() -> list[Path]:
"""Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary."""
from hub.storage.scan_folders import list_scan_folders
candidates: list[Path] = []
def _add(p: Optional[Path | str]) -> None:
if p is None:
return
try:
p = Path(normalize_path(str(p))).expanduser()
resolved = p.resolve()
except (OSError, RuntimeError, ValueError):
return
if _safe_is_dir(resolved):
candidates.append(resolved)
_add(Path.home())
_add(_resolve_hf_cache_dir())
try:
_add(hf_default_cache_dir())
except Exception: # noqa: BLE001 -- best-effort
pass
try:
_add(legacy_hf_cache_dir())
except Exception: # noqa: BLE001 -- best-effort
pass
try:
_add(studio_root())
_add(outputs_root())
_add(exports_root())
except Exception as exc: # noqa: BLE001 -- best-effort
logger.debug("browse-folders: studio roots unavailable: %s", exc)
try:
for folder in list_scan_folders():
p = folder.get("path")
if p:
_add(p)
except Exception as exc: # noqa: BLE001 -- best-effort
logger.debug("browse-folders: could not load scan folders: %s", exc)
try:
for p in well_known_model_dirs():
_add(p)
except Exception as exc: # noqa: BLE001 -- best-effort
logger.debug("browse-folders: well-known dirs unavailable: %s", exc)
seen: set[str] = set()
deduped: list[Path] = []
for p in candidates:
key = os.path.normcase(os.path.realpath(str(p)))
if key in seen:
continue
seen.add(key)
deduped.append(p)
return deduped
def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
"""True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox."""
try:
target_real = os.path.normcase(os.path.realpath(str(target)))
except OSError:
return False
for root in allowed_roots:
try:
root_real = os.path.normcase(os.path.realpath(str(root)))
except OSError:
continue
try:
if os.path.commonpath([target_real, root_real]) == root_real:
return True
except ValueError:
continue
if target_real == root_real:
return True
return False
def _normalize_browse_request_path(path: Optional[str], *, relative_root: Path) -> str:
"""Normalize the browse request path lexically, without touching the FS."""
if path is None or not path.strip():
return os.path.normpath(str(Path.home()))
expanded = os.path.expanduser(normalize_path(path.strip()))
if not os.path.isabs(expanded):
expanded = os.path.join(str(relative_root), expanded)
return os.path.normpath(expanded)
def _browse_relative_parts(requested_path: str, root: Path) -> Optional[list[str]]:
if "\x00" in requested_path:
raise HTTPException(
status_code = 400,
detail = "Path cannot contain null bytes",
)
root_text = os.path.normcase(os.path.normpath(str(root)))
requested_text = os.path.normcase(os.path.normpath(requested_path))
try:
rel_text = os.path.relpath(requested_text, root_text)
except ValueError:
return None
if rel_text == ".":
return []
if rel_text == ".." or rel_text.startswith(f"..{os.sep}"):
return None
parts = [part for part in rel_text.split(os.sep) if part not in ("", ".")]
altsep = os.altsep
for part in parts:
if part == ".." or "\x00" in part or os.sep in part or (altsep and altsep in part):
return None
return parts
def _match_browse_child(current: Path, name: str) -> Optional[Path]:
"""Immediate child named ``name`` under ``current``, or None. ``name`` is pre-validated as a safe single component, so the join is O(1); case resolution follows OS filesystem semantics."""
child = current / name
try:
child.stat()
except (FileNotFoundError, NotADirectoryError):
return None
except PermissionError:
raise HTTPException(
status_code = 403,
detail = f"Permission denied reading {current}",
) from None
except OSError as exc:
raise HTTPException(
status_code = 500,
detail = f"Could not read {current}: {exc}",
) from exc
return child
def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path:
"""Resolve a requested browse path by walking from trusted allowlist roots."""
requested_path = _normalize_browse_request_path(path, relative_root = Path.home())
resolved_roots: list[Path] = []
seen_roots: set[str] = set()
for root in sorted(allowed_roots, key = lambda p: len(str(p)), reverse = True):
try:
resolved = root.resolve()
except OSError:
continue
key = os.path.normcase(os.path.realpath(str(resolved)))
if key in seen_roots:
continue
seen_roots.add(key)
resolved_roots.append(resolved)
for root in resolved_roots:
parts = _browse_relative_parts(requested_path, root)
if parts is None:
continue
current = root
for part in parts:
child = _match_browse_child(current, part)
if child is None:
raise HTTPException(
status_code = 404,
detail = f"Path does not exist: {requested_path}",
)
try:
resolved_child = child.resolve()
except OSError as exc:
raise HTTPException(
status_code = 400,
detail = f"Invalid path: {exc}",
) from exc
if not _is_path_inside_allowlist(resolved_child, resolved_roots):
raise HTTPException(
status_code = 403,
detail = (
"Path is not in the browseable allowlist. Register it via "
"POST /api/hub/scan-folders first, or pick a directory "
"under your home folder."
),
)
# HOME is in the allowlist, so without this denylist (same one
# registration enforces) a user could browse into ~/.ssh, ~/.aws, etc.
if contains_sensitive_path_component(str(resolved_child)):
raise HTTPException(
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
current = resolved_child
if not current.is_dir():
raise HTTPException(
status_code = 400,
detail = f"Not a directory: {current}",
)
return current
raise HTTPException(
status_code = 403,
detail = (
"Path is not in the browseable allowlist. Register it via "
"POST /api/hub/scan-folders first, or pick a directory "
"under your home folder."
),
)
def browse_folders_response(
path: Optional[str] = None, show_hidden: bool = False
) -> BrowseFoldersResponse:
"""List immediate subdirectories of *path* for the Custom Folders picker.
Requests are bounded to the :func:`_build_browse_allowlist` roots; paths
outside it return 403 (symlinks resolved via realpath first, so traversal
can't escape). Sorting: model-bearing dirs first, then plain, then hidden.
"""
from hub.storage.scan_folders import list_scan_folders
# Build the allowlist once -- the sandbox check and suggestion chips share
# it so chips are always navigable.
allowed_roots = _build_browse_allowlist()
try:
target = _resolve_browse_target(path, allowed_roots)
except HTTPException:
requested_path = _normalize_browse_request_path(
path,
relative_root = Path.home(),
)
if path is not None and path.strip():
logger.warning(
"browse-folders: rejected path %r (normalized=%s)",
path,
requested_path,
)
raise
# Enumerate immediate subdirectories with a bounded cap so a stray
# query against ``/usr/lib`` or ``/proc`` can't stat-storm the process.
entries: list[BrowseEntry] = []
truncated = False
visited = 0
try:
it = target.iterdir()
except PermissionError:
raise HTTPException(
status_code = 403,
detail = f"Permission denied reading {target}",
)
except OSError as exc:
raise HTTPException(
status_code = 500,
detail = f"Could not read {target}: {exc}",
)
try:
for child in it:
# Bound by visited entries, not appended ones, so a directory full
# of files still caps work at ``_BROWSE_ENTRY_CAP`` stats.
visited += 1
if visited > _BROWSE_ENTRY_CAP:
truncated = True
break
try:
if not child.is_dir():
continue
except OSError:
continue
name = child.name
is_hidden = name.startswith(".")
if is_hidden and not show_hidden:
continue
# Don't surface credential/config dirs even with show_hidden:
# descending into them is refused and registration rejects them.
if contains_sensitive_path_component(name):
continue
entries.append(
BrowseEntry(
name = name,
has_models = _looks_like_model_dir(child),
hidden = is_hidden,
)
)
except PermissionError as exc:
logger.debug(
"browse-folders: permission denied during enumeration of %s: %s",
target,
exc,
)
except OSError as exc:
# Rare: iterdir succeeded but reading a specific entry failed.
logger.warning("browse-folders: partial enumeration of %s: %s", target, exc)
# Model-bearing dirs first, then plain, then hidden; case-insensitive
# alphabetical within each bucket.
def _sort_key(e: BrowseEntry) -> tuple[int, str]:
bucket = 0 if e.has_models else (2 if e.hidden else 1)
return (bucket, e.name.lower())
entries.sort(key = _sort_key)
# Parent is None at the FS root and when it would step outside the sandbox,
# so the up-row never 403s on click.
parent: Optional[str]
if target.parent == target or not _is_path_inside_allowlist(target.parent, allowed_roots):
parent = None
else:
parent = str(target.parent)
# Handy starting points for the quick-pick chips.
suggestions: list[str] = []
seen_sug: set[str] = set()
def _add_sug(p: Optional[Path | str]) -> None:
if p is None:
return
try:
p = Path(normalize_path(str(p))).expanduser()
resolved = str(p.resolve())
except (OSError, RuntimeError, ValueError):
return
if resolved in seen_sug:
return
if _safe_is_dir(resolved):
seen_sug.add(resolved)
suggestions.append(resolved)
# Home first as the safe fallback.
_add_sug(Path.home())
# The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default.
try:
_add_sug(_resolve_hf_cache_dir())
except Exception:
pass
try:
_add_sug(hf_default_cache_dir())
except Exception:
pass
# Already-registered scan folders (what the user has curated).
try:
for folder in list_scan_folders():
_add_sug(folder.get("path", ""))
except Exception as exc:
logger.debug("browse-folders: could not load scan folders: %s", exc)
# Well-known third-party dirs (LM Studio, Ollama, ~/models). Each helper
# only returns existing paths so we never show dead chips.
try:
for p in well_known_model_dirs():
_add_sug(p)
except Exception as exc:
logger.debug("browse-folders: could not load well-known dirs: %s", exc)
return BrowseFoldersResponse(
current = str(target),
parent = parent,
entries = entries,
suggestions = suggestions,
truncated = truncated,
model_files_here = _count_model_files(target),
)

643
studio/backend/hub/services/models/gguf_variants.py Normal file
View file

@ -0,0 +1,643 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GGUF variant resolution."""
from __future__ import annotations
import asyncio
import threading
import time
from collections import OrderedDict
from typing import NamedTuple, Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.schemas.inventory import GgufVariantDetail, GgufVariantsResponse
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_errors import hf_error_status
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
iter_destructive_repo_cache_dirs,
)
from hub.utils.gguf import (
extract_quant_label,
iter_hf_cache_snapshots,
list_gguf_variants,
list_gguf_variants_from_hf_cache,
list_local_gguf_variants,
list_partial_gguf_variants_from_state,
pick_best_gguf,
)
from hub.utils.paths import (
is_local_path,
is_valid_repo_id as _is_valid_repo_id,
)
from hub.services.models.common import (
_is_mmproj_filename,
_iter_gguf_paths,
)
from hub.utils.gguf_plan import (
GgufVariantPlan as _GgufVariantRequirement,
build_gguf_variant_plans,
is_main_gguf_variant_path,
)
logger = get_logger(__name__)
_VARIANT_HASH_CACHE: "OrderedDict[tuple[str, str, str, bool], tuple[frozenset[str], float]]" = (
OrderedDict()
)
_VARIANT_REQUIREMENT_CACHE: "OrderedDict[tuple[str, str, str], tuple[_GgufVariantRequirement, float]]" = OrderedDict()
_VARIANT_REQUIREMENT_NEG_CACHE: "OrderedDict[tuple[str, str], float]" = OrderedDict()
_VARIANT_HASH_MAX = 512
# Blob hashes are derived from the same mutable remote revision metadata as
# variant requirements, so they must not outlive that freshness window.
_VARIANT_HASH_POS_TTL = 60.0
# Refresh resolved variant requirements so a moved repo revision is picked up
# within the session instead of being pinned for the backend's lifetime.
_VARIANT_REQUIREMENT_POS_TTL = 60.0
# Suppress retries on a metadata-fetch failure so a slow/flaky link doesn't
# re-hammer the API on every page refresh.
_VARIANT_REQUIREMENT_NEG_TTL = 60.0
# Fail fast on a slow link so the variant render isn't blocked for seconds.
_GGUF_METADATA_TIMEOUT_SECONDS = 5.0
_VARIANT_HASH_LOCK = threading.Lock()
class VariantIncompleteDeleteResult(NamedTuple):
deleted: int
unresolved: bool
def _variant_hash_cache_key(
repo_id: str, variant: str, hf_token: Optional[str]
) -> tuple[str, str, str]:
return (
repo_id.lower(),
variant.lower(),
hf_cache_scan.token_fingerprint(hf_token),
)
def _variant_blob_hash_cache_key(
repo_id: str, variant: str, hf_token: Optional[str], include_companions: bool
) -> tuple[str, str, str, bool]:
base = _variant_hash_cache_key(repo_id, variant, hf_token)
return (*base, include_companions)
def _variant_repo_cache_key(repo_id: str, hf_token: Optional[str]) -> tuple[str, str]:
return (repo_id.lower(), hf_cache_scan.token_fingerprint(hf_token))
def _variant_requirement_neg_cache_active(key: tuple[str, str]) -> bool:
with _VARIANT_HASH_LOCK:
cached_at = _VARIANT_REQUIREMENT_NEG_CACHE.get(key)
if cached_at is None:
return False
if (time.monotonic() - cached_at) < _VARIANT_REQUIREMENT_NEG_TTL:
_VARIANT_REQUIREMENT_NEG_CACHE.move_to_end(key)
return True
_VARIANT_REQUIREMENT_NEG_CACHE.pop(key, None)
return False
def _variant_requirement_neg_cache_set(key: tuple[str, str]) -> None:
with _VARIANT_HASH_LOCK:
_VARIANT_REQUIREMENT_NEG_CACHE[key] = time.monotonic()
_VARIANT_REQUIREMENT_NEG_CACHE.move_to_end(key)
while len(_VARIANT_REQUIREMENT_NEG_CACHE) > _VARIANT_HASH_MAX:
_VARIANT_REQUIREMENT_NEG_CACHE.popitem(last = False)
def _variant_requirement_neg_cache_clear(key: tuple[str, str]) -> None:
with _VARIANT_HASH_LOCK:
_VARIANT_REQUIREMENT_NEG_CACHE.pop(key, None)
def _variant_hash_cache_get(key: tuple[str, str, str, bool]) -> Optional[frozenset[str]]:
with _VARIANT_HASH_LOCK:
cached = _VARIANT_HASH_CACHE.get(key)
if cached is None:
return None
hashes, ts = cached
if (time.monotonic() - ts) >= _VARIANT_HASH_POS_TTL:
_VARIANT_HASH_CACHE.pop(key, None)
return None
_VARIANT_HASH_CACHE.move_to_end(key)
return hashes
def _variant_hash_cache_set(key: tuple[str, str, str, bool], hashes: frozenset[str]) -> None:
with _VARIANT_HASH_LOCK:
_VARIANT_HASH_CACHE[key] = (hashes, time.monotonic())
_VARIANT_HASH_CACHE.move_to_end(key)
while len(_VARIANT_HASH_CACHE) > _VARIANT_HASH_MAX:
_VARIANT_HASH_CACHE.popitem(last = False)
def _variant_requirement_cache_get(key: tuple[str, str, str]) -> Optional[_GgufVariantRequirement]:
with _VARIANT_HASH_LOCK:
cached = _VARIANT_REQUIREMENT_CACHE.get(key)
if cached is None:
return None
requirement, ts = cached
if (time.monotonic() - ts) >= _VARIANT_REQUIREMENT_POS_TTL:
_VARIANT_REQUIREMENT_CACHE.pop(key, None)
return None
_VARIANT_REQUIREMENT_CACHE.move_to_end(key)
return requirement
def _variant_requirement_cache_set_many(
repo_id: str, hf_token: Optional[str], requirements: dict[str, _GgufVariantRequirement]
) -> None:
with _VARIANT_HASH_LOCK:
now = time.monotonic()
for quant, requirement in requirements.items():
key = _variant_hash_cache_key(repo_id, quant, hf_token)
_VARIANT_REQUIREMENT_CACHE[key] = (requirement, now)
_VARIANT_REQUIREMENT_CACHE.move_to_end(key)
while len(_VARIANT_REQUIREMENT_CACHE) > _VARIANT_HASH_MAX:
_VARIANT_REQUIREMENT_CACHE.popitem(last = False)
def _build_gguf_variant_requirements(siblings: list) -> dict[str, _GgufVariantRequirement]:
return build_gguf_variant_plans(siblings)
def gguf_variant_requirements(
repo_id: str,
variant: str,
hf_token: Optional[str] = None,
) -> Optional[_GgufVariantRequirement]:
key = _variant_hash_cache_key(repo_id, variant, hf_token)
cached = _variant_requirement_cache_get(key)
if cached is not None:
return cached
requirements = _fetch_gguf_variant_requirements(repo_id, hf_token)
return requirements.get(variant.lower())
def _fetch_gguf_variant_requirements(
repo_id: str,
hf_token: Optional[str] = None,
*,
siblings: Optional[list] = None,
) -> dict[str, _GgufVariantRequirement]:
repo_key = _variant_repo_cache_key(repo_id, hf_token)
if siblings is None:
if _variant_requirement_neg_cache_active(repo_key):
return {}
try:
from huggingface_hub import HfApi
info = HfApi(token = hf_token).model_info(
repo_id,
files_metadata = True,
timeout = _GGUF_METADATA_TIMEOUT_SECONDS,
)
except Exception as e:
logger.warning(
"model_info failed resolving GGUF files for %s: %s",
repo_id,
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
_variant_requirement_neg_cache_set(repo_key)
return {}
siblings = list(info.siblings)
requirements = _build_gguf_variant_requirements(siblings)
if requirements:
_variant_requirement_cache_set_many(repo_id, hf_token, requirements)
_variant_requirement_neg_cache_clear(repo_key)
return requirements
def _gguf_all_variant_requirements(
repo_id: str,
hf_token: Optional[str] = None,
*,
siblings: Optional[list] = None,
) -> dict[str, _GgufVariantRequirement]:
return _fetch_gguf_variant_requirements(repo_id, hf_token, siblings = siblings)
def _manifest_variant_blob_hashes(
repo_id: str,
variant: str,
*,
include_companions: bool = True,
) -> frozenset[str]:
manifest = download_manifest.read_manifest("model", repo_id, variant)
if manifest is None:
return frozenset()
variant_key = variant.lower()
hashes: set[str] = set()
for expected in manifest.expected_files:
if not expected.sha256:
continue
if include_companions:
hashes.add(expected.sha256)
continue
if is_main_gguf_variant_path(expected.path, variant_key):
hashes.add(expected.sha256)
return frozenset(hashes)
def gguf_variant_blob_hashes(
repo_id: str,
variant: str,
hf_token: Optional[str] = None,
*,
include_companions: bool = True,
allow_remote: bool = True,
) -> frozenset[str]:
key = _variant_blob_hash_cache_key(
repo_id,
variant,
hf_token,
include_companions,
)
cached = _variant_hash_cache_get(key)
if cached is not None:
return cached
hashes = _manifest_variant_blob_hashes(
repo_id,
variant,
include_companions = include_companions,
)
if hashes:
_variant_hash_cache_set(key, hashes)
return hashes
requirement_key = _variant_hash_cache_key(repo_id, variant, hf_token)
requirement = _variant_requirement_cache_get(requirement_key)
if requirement is None and allow_remote:
requirement = gguf_variant_requirements(repo_id, variant, hf_token)
if requirement is not None:
hashes = requirement.required_hashes if include_companions else requirement.main_hashes
if hashes:
_variant_hash_cache_set(key, hashes)
return hashes
return frozenset()
def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]:
return hf_cache_scan.partial_transport_for("model", repo_id, variant)
def delete_variant_incomplete_blobs_result(
repo_id: str,
variant: str,
hf_token: Optional[str],
*,
extra_hashes: frozenset[str] = frozenset(),
companions: bool = True,
) -> VariantIncompleteDeleteResult:
# With a sibling still downloading, ``companions=False`` keeps a shared mmproj
# from being unlinked out from under it; the repo's last delete reclaims it.
target_hashes = (
gguf_variant_blob_hashes(repo_id, variant, hf_token, include_companions = companions)
| extra_hashes
)
if not target_hashes:
has_variant_partial_state = hf_cache_scan.is_variant_partial(
repo_id,
variant,
incomplete_blob_hashes = set(),
variant_blob_hashes = frozenset(),
)
has_repo_partials = bool(download_registry.incomplete_blob_hashes("model", repo_id))
return VariantIncompleteDeleteResult(
deleted = 0,
unresolved = has_variant_partial_state and has_repo_partials,
)
deleted = 0
# Destructive iterator: only the exact-case match (or abort if ambiguous),
# so a case-variant sibling repo's partials are never unlinked.
for entry in iter_destructive_repo_cache_dirs("model", repo_id):
blobs_dir = entry / "blobs"
if not blobs_dir.is_dir():
continue
for h in target_hashes:
incomplete = blobs_dir / f"{h}{INCOMPLETE_SUFFIX}"
if incomplete.exists():
try:
incomplete.unlink()
deleted += 1
except OSError as e:
logger.warning(f"Failed to unlink {incomplete}: {e}")
return VariantIncompleteDeleteResult(deleted = deleted, unresolved = False)
async def get_gguf_variants_response(
repo_id: str,
prefer_local_cache: bool = False,
offline: bool = False,
local_path: Optional[str] = None,
hf_token: Optional[str] = None,
):
"""
List available GGUF quantization variants for a HuggingFace repo
or a local directory (e.g. LM Studio model folder).
Returns all available quantization variants (Q4_K_M, Q8_0, BF16, etc.)
with file sizes, whether the model supports vision, and the recommended
default variant.
"""
def _compute() -> GgufVariantsResponse:
def _local_response(
response_repo_id: str, variants, has_vision: bool
) -> GgufVariantsResponse:
filenames = [v.filename for v in variants]
best = pick_best_gguf(filenames)
default_variant = extract_quant_label(best) if best else None
return GgufVariantsResponse(
repo_id = response_repo_id,
variants = [
GgufVariantDetail(
filename = v.filename,
quant = v.quant,
display_label = v.display_label,
size_bytes = v.size_bytes,
download_size_bytes = v.size_bytes,
downloaded = True,
)
for v in variants
],
has_vision = has_vision,
default_variant = default_variant,
)
def _partial_local_response(
response_repo_id: str, variants, has_vision: bool
) -> GgufVariantsResponse:
filenames = [v.filename for v in variants]
best = pick_best_gguf(filenames)
default_variant = extract_quant_label(best) if best else None
return GgufVariantsResponse(
repo_id = response_repo_id,
variants = [
GgufVariantDetail(
filename = v.filename,
quant = v.quant,
display_label = v.display_label,
size_bytes = v.size_bytes,
download_size_bytes = v.download_size_bytes or v.size_bytes,
downloaded = False,
partial = True,
partial_transport = _partial_transport_for_variant(
response_repo_id,
v.quant,
),
)
for v in variants
],
has_vision = has_vision,
default_variant = default_variant,
)
# Local directory path (e.g. LM Studio models) — scan filesystem
if is_local_path(repo_id):
variants, has_vision = list_local_gguf_variants(repo_id)
return _local_response(repo_id, variants, has_vision)
# Reject invalid remote repo_ids up front (like download/delete) so a
# malformed id returns 400 instead of a 500 from the HF client.
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = f"Invalid repo_id: {repo_id!r}")
local_only = prefer_local_cache or offline
if local_only:
cached = list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
variants, has_vision = cached
return _local_response(repo_id, variants, has_vision)
if local_path and is_local_path(local_path):
variants, has_vision = list_local_gguf_variants(local_path)
if variants or has_vision:
return _local_response(repo_id, variants, has_vision)
partial = list_partial_gguf_variants_from_state(repo_id)
if partial is not None:
variants, has_vision = partial
return _partial_local_response(repo_id, variants, has_vision)
if local_path and offline:
return GgufVariantsResponse(
repo_id = repo_id,
variants = [],
has_vision = False,
default_variant = None,
)
if offline:
raise HTTPException(
status_code = 404,
detail = "No cached GGUF variants available while offline.",
)
try:
variants, has_vision, siblings = list_gguf_variants(repo_id, hf_token = hf_token)
except Exception:
cached = list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
variants, has_vision = cached
return _local_response(repo_id, variants, has_vision)
partial = list_partial_gguf_variants_from_state(repo_id)
if partial is not None:
variants, has_vision = partial
return _partial_local_response(repo_id, variants, has_vision)
raise
filenames = [v.filename for v in variants]
best = pick_best_gguf(filenames)
default_variant = extract_quant_label(best) if best else None
# Per-snapshot accounting: a variant counts as present only when one
# snapshot holds all its files (split GGUFs need every shard together),
# sizes are max across snapshots so shared blobs aren't double-counted,
# and keys are lowercased since cache dir casing can differ from repo_id.
cached_filenames_by_snapshot: list[dict[str, int]] = []
cached_quant_bytes_by_snapshot: list[dict[str, int]] = []
if _is_valid_repo_id(repo_id):
for snap in iter_hf_cache_snapshots(repo_id):
try:
gguf_paths = list(_iter_gguf_paths(snap))
except (OSError, RuntimeError, ValueError) as e:
logger.debug("Skipping GGUF cache snapshot %s: %s", snap, e)
continue
by_filename: dict[str, int] = {}
by_quant: dict[str, int] = {}
for f in gguf_paths:
try:
rel = f.relative_to(snap).as_posix()
size = f.stat().st_size
except (OSError, RuntimeError, ValueError) as e:
logger.debug("Skipping GGUF cache file %s: %s", f, e)
continue
key = rel.lower()
by_filename[key] = max(by_filename.get(key, 0), size)
if _is_mmproj_filename(f.name):
continue
q = extract_quant_label(rel).lower()
by_quant[q] = by_quant.get(q, 0) + size
if by_filename:
cached_filenames_by_snapshot.append(by_filename)
if by_quant:
cached_quant_bytes_by_snapshot.append(by_quant)
requirements_by_quant = {
v.quant.lower(): _variant_requirement_cache_get(
_variant_hash_cache_key(repo_id, v.quant, hf_token)
)
for v in variants
}
if any(req is None for req in requirements_by_quant.values()):
fetched_requirements = _gguf_all_variant_requirements(
repo_id, hf_token, siblings = siblings
)
for v in variants:
key = v.quant.lower()
if requirements_by_quant.get(key) is None:
requirements_by_quant[key] = fetched_requirements.get(key)
def _filenames_cached(filenames: frozenset[str], expected_size: int) -> bool:
if not filenames:
return False
wanted = [name.lower() for name in filenames]
# All files must live in a single snapshot, not spread across several.
for by_filename in cached_filenames_by_snapshot:
cached = 0
for name in wanted:
size = by_filename.get(name)
if size is None:
break
cached += size
else:
return expected_size <= 0 or cached >= expected_size * 0.99
return False
def _any_mmproj_cached(filenames: frozenset[str]) -> bool:
return any(
by_filename.get(name.lower()) is not None
for by_filename in cached_filenames_by_snapshot
for name in filenames
)
def _is_fully_downloaded(variant) -> bool:
requirement = requirements_by_quant.get(variant.quant.lower())
if requirement is None:
if variant.size_bytes == 0:
return False
quant = variant.quant.lower()
# Allow small rounding tolerance (symlinks vs real sizes).
return any(
by_quant.get(quant, 0) >= variant.size_bytes * 0.99
for by_quant in cached_quant_bytes_by_snapshot
)
if not _filenames_cached(
requirement.main_filenames,
requirement.main_size_bytes,
):
return False
# Vision repos ship an mmproj adapter per variant. Any mmproj
# precision on disk suffices (the loader picks whichever is present);
# requiring the API-preferred one would falsely demote variants.
if requirement.mmproj_filenames and not _any_mmproj_cached(
requirement.mmproj_filenames,
):
return False
return True
partial_quants: set[str] = set()
partial_quant_transports: dict[str, Optional[str]] = {}
try:
incomplete_hashes = download_registry.incomplete_blob_hashes("model", repo_id)
except Exception as e:
logger.warning(f"Failed to compute partial GGUF variants for {repo_id}: {e}")
incomplete_hashes = set()
scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan("model", repo_id)
# Manifest + marker + main incomplete-blob check: catches variants whose
# download was cancelled or whose expected shards are missing/undersized.
for variant in variants:
try:
requirement = requirements_by_quant.get(variant.quant.lower())
variant_hashes = requirement.main_hashes if requirement is not None else None
if variant_hashes is None and incomplete_hashes:
variant_hashes = gguf_variant_blob_hashes(
repo_id,
variant.quant,
hf_token,
include_companions = False,
)
if hf_cache_scan.is_variant_partial(
repo_id,
variant.quant,
scan_snapshot_dir,
incomplete_blob_hashes = incomplete_hashes,
variant_blob_hashes = variant_hashes,
):
partial_quants.add(variant.quant)
partial_quant_transports[variant.quant] = _partial_transport_for_variant(
repo_id,
variant.quant,
)
except Exception as e:
logger.warning(
f"Manifest-based partial check failed for " f"{repo_id}/{variant.quant}: {e}"
)
if incomplete_hashes:
for variant in variants:
requirement = requirements_by_quant.get(variant.quant.lower())
if requirement is None:
continue
if requirement.mmproj_hashes & incomplete_hashes and _filenames_cached(
requirement.main_filenames,
requirement.main_size_bytes,
):
partial_quants.add(variant.quant)
partial_quant_transports.setdefault(
variant.quant,
_partial_transport_for_variant(repo_id, variant.quant),
)
def _variant_detail(v) -> GgufVariantDetail:
is_partial = v.quant in partial_quants
requirement = requirements_by_quant.get(v.quant.lower())
return GgufVariantDetail(
filename = v.filename,
quant = v.quant,
display_label = v.display_label,
size_bytes = v.size_bytes,
download_size_bytes = (
requirement.download_size_bytes if requirement is not None else v.size_bytes
),
downloaded = _is_fully_downloaded(v) and not is_partial,
partial = is_partial,
partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None),
)
return GgufVariantsResponse(
repo_id = repo_id,
variants = [_variant_detail(v) for v in variants],
has_vision = has_vision,
default_variant = default_variant,
)
try:
return await asyncio.to_thread(_compute)
except HTTPException:
raise
except Exception as e:
scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token)
# Client-side HF error (missing repo, gated, bad token): pass the status through.
status = hf_error_status(e)
if status is not None:
raise HTTPException(status_code = status, detail = scrubbed)
logger.error("Error listing GGUF variants for %s: %s", repo_id, scrubbed)
raise HTTPException(
status_code = 500,
detail = "Failed to list GGUF variants: " + scrubbed,
)

679
studio/backend/hub/services/models/local_inventory.py Normal file
View file

@ -0,0 +1,679 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Local model, HF cache, LM Studio, and Ollama inventory services.
Ollama logic lives in :mod:`hub.services.models.ollama`; this module
orchestrates all on-device sources and exposes the route handlers.
"""
from __future__ import annotations
import asyncio
import os
from pathlib import Path
from typing import List, Optional
from fastapi import HTTPException
from loggers import get_logger
from hub.schemas.inventory import LocalModelInfo, LocalModelListResponse, ModelFormat
from hub.storage.scan_folders import (
add_scan_folder,
list_scan_folders,
remove_scan_folder,
)
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
lmstudio_model_dirs,
normalize_path,
ollama_model_dirs,
outputs_root,
path_is_same_or_child,
studio_root,
)
from hub.services.models import common as model_common
from hub.services.models.ollama import scan_ollama_dir
logger = get_logger(__name__)
_MAX_MODELS_PER_CUSTOM_FOLDER = 200
_MAX_CUSTOM_FOLDER_ENTRIES = 2000
_MODEL_SIGNAL_PROBE_LIMIT = 200
# Local aliases keep the extracted code close to the original implementation.
_is_model_directory = model_common._is_model_directory
_local_inventory_id = model_common._local_inventory_id
_local_model_info = model_common._local_model_info
_capabilities_for_format = model_common._capabilities_for_format
_apply_format_aware_partial = model_common._apply_format_aware_partial
_classify_local_path = model_common._classify_local_path
_is_main_gguf_filename = model_common._is_main_gguf_filename
_is_transformers_bin_weight_file = model_common._is_transformers_bin_weight_file
_prefer_complete_larger = model_common._prefer_complete_larger
_gguf_variant_state_summary = model_common._gguf_variant_state_summary
def _is_immediate_model_weight_file(path: Path) -> bool:
suffix = path.suffix.lower()
if suffix == ".safetensors":
return True
if suffix == ".gguf":
return _is_main_gguf_filename(path.name)
if suffix == ".bin":
return _is_transformers_bin_weight_file(path)
return False
def _has_immediate_model_weight(
path: Path, *, probe_limit: int = _MODEL_SIGNAL_PROBE_LIMIT
) -> bool:
try:
for index, entry in enumerate(path.iterdir(), start = 1):
if index > probe_limit:
break
try:
if entry.is_file() and _is_immediate_model_weight_file(entry):
return True
except OSError:
continue
except OSError:
return False
return False
def _has_immediate_model_signal(
path: Path, *, probe_limit: int = _MODEL_SIGNAL_PROBE_LIMIT
) -> bool:
try:
if (path / "config.json").exists() or (path / "adapter_config.json").exists():
return True
except OSError:
return False
return _has_immediate_model_weight(path, probe_limit = probe_limit)
def _is_model_directory_for_scan(path: Path, *, entry_limit: int | None) -> bool:
if entry_limit is None:
return _is_model_directory(path)
try:
has_config = (path / "config.json").exists() or (path / "adapter_config.json").exists()
except OSError:
return False
return has_config and _has_immediate_model_weight(path)
def _resolve_hf_cache_dir() -> Path:
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
def _scan_models_dir(
models_dir: Path,
*,
limit: int | None = None,
entry_limit: int | None = None,
) -> List[LocalModelInfo]:
if not models_dir.exists() or not models_dir.is_dir():
return []
_is_self_model = _is_model_directory_for_scan(
models_dir,
entry_limit = entry_limit,
)
if _is_self_model:
try:
updated_at = models_dir.stat().st_mtime
except OSError:
updated_at = None
return _classify_local_path(
models_dir,
"models_dir",
updated_at = updated_at,
)
found: List[LocalModelInfo] = []
visited = 0
try:
children = models_dir.iterdir()
except OSError:
return found
for child in children:
if limit is not None and len(found) >= limit:
break
visited += 1
if entry_limit is not None and visited > entry_limit:
break
try:
is_dir = child.is_dir()
is_gguf_file = not is_dir and child.suffix.lower() == ".gguf" and child.is_file()
if not is_dir and not is_gguf_file:
continue
has_model_files = is_gguf_file or _has_immediate_model_signal(child)
except OSError:
# Skip individual children that are unreadable (permissions, broken
# symlinks, etc.) rather than failing the entire scan.
continue
if not has_model_files:
continue
try:
updated_at = child.stat().st_mtime
except OSError:
updated_at = None
rows = _classify_local_path(
child,
"models_dir",
updated_at = updated_at,
)
if limit is not None:
rows = rows[: max(0, limit - len(found))]
found.extend(rows)
return found
def _hf_repo_dir_has_content(repo_dir: Path) -> bool:
blobs_dir = repo_dir / "blobs"
if not blobs_dir.is_dir():
return False
try:
for entry in blobs_dir.iterdir():
if entry.is_file() or entry.is_symlink():
return True
except OSError:
return False
return False
def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]:
if not cache_dir.exists() or not cache_dir.is_dir():
return []
discovered: List[tuple[Path, str, Optional[float]]] = []
visited = 0
try:
entries = cache_dir.iterdir()
except OSError:
return []
for repo_dir in entries:
visited += 1
if entry_limit is not None and visited > entry_limit:
break
if not repo_dir.name.startswith("models--"):
continue
if not repo_dir.is_dir():
continue
if not _hf_repo_dir_has_content(repo_dir):
continue
repo_name = repo_dir.name[len("models--") :]
if not repo_name:
continue
model_id = repo_name.replace("--", "/")
try:
updated_at = repo_dir.stat().st_mtime
except OSError:
updated_at = None
discovered.append((repo_dir, model_id, updated_at))
found: list[LocalModelInfo] = []
for repo_dir, model_id, updated_at in discovered:
snapshot_partial = hf_cache_scan.is_snapshot_partial(
"model",
model_id,
repo_dir,
)
gguf_partial = hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir)
has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(model_id)
snapshot_partial_transport = (
hf_cache_scan.partial_transport_for(
"model",
model_id,
repo_cache_dir = repo_dir,
)
if snapshot_partial
else None
)
resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_dir)
scan_path = Path(resolved) if resolved else repo_dir
# partial=False here; _apply_format_aware_partial below rewrites per-row
# so a hybrid repo's gguf row doesn't taint its safetensors row.
rows = _classify_local_path(
scan_path,
"hf_cache",
load_path = repo_dir,
display_name = model_id.split("/")[-1],
model_id = model_id,
updated_at = updated_at,
partial = False,
)
if not rows:
if has_gguf_variant_state and gguf_partial:
rows = [
_local_model_info(
scan_path = repo_dir,
load_path = repo_dir,
source = "hf_cache",
model_format = "gguf",
display_name = model_id.split("/")[-1],
model_id = model_id,
updated_at = updated_at,
partial = True,
requires_variant = True,
size_bytes = gguf_variant_state_size,
)
]
else:
# Fallback row's model_format is "unknown"; either signal
# applies because we can't dispatch to a specific predicate.
rows = [
_local_model_info(
scan_path = repo_dir,
load_path = repo_dir,
source = "hf_cache",
model_format = "unknown",
display_name = model_id.split("/")[-1],
model_id = model_id,
updated_at = updated_at,
partial = snapshot_partial or gguf_partial,
)
]
elif (
has_gguf_variant_state
and gguf_partial
and not any(row.model_format == "gguf" for row in rows)
):
rows.append(
_local_model_info(
scan_path = repo_dir,
load_path = repo_dir,
source = "hf_cache",
model_format = "gguf",
display_name = model_id.split("/")[-1],
model_id = model_id,
updated_at = updated_at,
partial = True,
requires_variant = True,
size_bytes = gguf_variant_state_size,
)
)
rows = _apply_format_aware_partial(
rows,
snapshot_partial = snapshot_partial,
gguf_partial = gguf_partial,
snapshot_partial_transport = snapshot_partial_transport,
)
found.extend(rows)
return found
def _scan_lmstudio_dir(lm_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]:
"""Scan an LM Studio models dir (``publisher/model-name`` folders of GGUFs, or top-level standalone GGUFs)."""
if not lm_dir.exists() or not lm_dir.is_dir():
return []
# If the dir is itself a model dir (config + weights), it's not an LM Studio
# publisher structure -- return it as a single entry rather than descend.
if _is_model_directory(lm_dir):
try:
updated_at = lm_dir.stat().st_mtime
except OSError:
updated_at = None
return _classify_local_path(
lm_dir,
"lmstudio",
updated_at = updated_at,
)
found: List[LocalModelInfo] = []
visited = 0
exhausted = False
def _consume_visit() -> bool:
nonlocal visited
visited += 1
return entry_limit is not None and visited > entry_limit
try:
children = lm_dir.iterdir()
except OSError:
return found
for child in children:
if _consume_visit():
break
try:
if not child.is_dir():
if child.suffix == ".gguf" and child.is_file():
try:
updated_at = child.stat().st_mtime
except OSError:
updated_at = None
found.extend(
_classify_local_path(
child,
"lmstudio",
updated_at = updated_at,
)
)
continue
# Child is itself a model dir: surface it directly, not as a publisher.
if _is_model_directory(child):
try:
updated_at = child.stat().st_mtime
except OSError:
updated_at = None
found.extend(
_classify_local_path(
child,
"lmstudio",
updated_at = updated_at,
)
)
continue
# child is a publisher directory -- scan its sub-directories
for model_dir in child.iterdir():
if _consume_visit():
exhausted = True
break
try:
if model_dir.is_dir():
has_model = _has_immediate_model_signal(model_dir)
if not has_model:
continue
model_id = f"{child.name}/{model_dir.name}"
try:
updated_at = model_dir.stat().st_mtime
except OSError:
updated_at = None
found.extend(
_classify_local_path(
model_dir,
"lmstudio",
display_name = model_dir.name,
model_id = model_id,
updated_at = updated_at,
)
)
elif model_dir.suffix == ".gguf" and model_dir.is_file():
try:
updated_at = model_dir.stat().st_mtime
except OSError:
updated_at = None
found.extend(
_classify_local_path(
model_dir,
"lmstudio",
model_id = f"{child.name}/{model_dir.stem}",
updated_at = updated_at,
)
)
except OSError:
continue
if exhausted:
break
except OSError:
continue
return found
def _resolve_allowed_models_dir(models_dir: str, allowed_roots: list[Path]) -> Path:
"""Resolve a requested model scan directory without widening subpaths."""
if not models_dir or not models_dir.strip():
raise ValueError("Directory not allowed")
requested = Path(os.path.realpath(os.path.expanduser(normalize_path(models_dir.strip()))))
if any(path_is_same_or_child(requested, root) for root in allowed_roots):
return requested
raise ValueError("Directory not allowed")
def _coerce_scan_folder_path(raw_path: str) -> str:
"""Normalize a scan registration target; the registry stores directories, so a pasted weight-file path is reduced to its parent folder."""
if not raw_path or not raw_path.strip():
raise ValueError("Path cannot be empty")
raw = raw_path.strip()
if "\x00" in raw:
raise ValueError("Path cannot contain null bytes")
def normalize(value: str) -> Path:
return Path(os.path.realpath(os.path.expanduser(normalize_path(value))))
try:
normalized = normalize(raw)
except (OSError, ValueError) as e:
raise ValueError(f"Path is not readable: {e}") from e
try:
exists = normalized.exists()
is_dir = normalized.is_dir()
is_file = normalized.is_file()
except (OSError, ValueError) as e:
raise ValueError(f"Path is not readable: {e}") from e
if not exists and "\\" in raw:
try:
slash_normalized = normalize(raw.replace("\\", "/"))
slash_exists = slash_normalized.exists()
except (OSError, ValueError) as e:
raise ValueError(f"Path is not readable: {e}") from e
if slash_exists:
normalized = slash_normalized
try:
is_dir = normalized.is_dir()
is_file = normalized.is_file()
except (OSError, ValueError) as e:
raise ValueError(f"Path is not readable: {e}") from e
exists = True
if not exists:
return str(normalized)
if is_dir:
return str(normalized)
if is_file:
suffix = normalized.suffix.lower()
if suffix not in {".gguf", ".safetensors", ".bin"}:
raise ValueError("Path must be a folder or model weight file")
return str(normalized.parent)
return str(normalized)
async def _scan_source(label: str, scanner, path: Path) -> List[LocalModelInfo]:
try:
return await asyncio.to_thread(scanner, path)
except Exception as e:
logger.warning("Skipping %s scan for %s: %s", label, path, e)
return []
async def _collect_models_from_default_sources(
models_root: Path,
hf_cache_dir: Path,
legacy_hf: Path,
hf_default: Path,
lm_dirs: list[Path],
ollama_dirs: list[Path],
) -> List[LocalModelInfo]:
local_models = await _scan_source("models directory", _scan_models_dir, models_root)
local_models += await _scan_source("HF cache", _scan_hf_cache, hf_cache_dir)
if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve():
local_models += await _scan_source("legacy HF cache", _scan_hf_cache, legacy_hf)
if (
hf_default.is_dir()
and hf_default.resolve() != hf_cache_dir.resolve()
and hf_default.resolve() != legacy_hf.resolve()
):
local_models += await _scan_source("default HF cache", _scan_hf_cache, hf_default)
for lm_dir in lm_dirs:
local_models += await _scan_source("LM Studio", _scan_lmstudio_dir, lm_dir)
for ollama_dir in ollama_dirs:
local_models += await _scan_source("Ollama", scan_ollama_dir, ollama_dir)
return local_models
def _scan_custom_folder(folder_path: Path) -> List[LocalModelInfo]:
supported_formats: set[ModelFormat] = {"gguf", "safetensors", "adapter"}
generic = [
m
for m in (
_scan_models_dir(
folder_path,
limit = _MAX_MODELS_PER_CUSTOM_FOLDER,
entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES,
)
+ _scan_hf_cache(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES)
+ _scan_lmstudio_dir(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES)
)
if m.model_format in supported_formats
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
return generic[:_MAX_MODELS_PER_CUSTOM_FOLDER]
def _promote_to_custom_source(model: LocalModelInfo) -> LocalModelInfo:
if model.source == "hf_cache":
return model
return model.model_copy(
update = {
"source": "custom",
"model_id": None,
"inventory_id": _local_inventory_id(
"custom",
model.model_format,
model.path,
model.format_variant,
),
"capabilities": _capabilities_for_format(
model.model_format,
"custom",
partial = model.partial,
requires_variant = model.capabilities.requires_variant,
),
}
)
async def _collect_models_from_custom_folders() -> List[LocalModelInfo]:
try:
custom_folders = await asyncio.to_thread(list_scan_folders)
except Exception as e:
logger.warning("Could not load custom scan folders: %s", e)
return []
local_models: List[LocalModelInfo] = []
for folder in custom_folders:
folder_path = Path(normalize_path(folder["path"])).expanduser()
try:
custom_models = await asyncio.to_thread(_scan_custom_folder, folder_path)
except Exception as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models.extend(_promote_to_custom_source(m) for m in custom_models)
return local_models
def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelInfo]:
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
if model.source == "hf_cache" and model.model_id:
key = "\x00".join(
(
"hf_cache",
model.model_id.strip().lower(),
model.model_format,
model.format_variant or "",
)
)
else:
row_key = model.inventory_id or model.id
key = f"{row_key}\x00custom" if model.source == "custom" else row_key
existing = deduped.get(key)
if existing is None or _prefer_complete_larger(
model.partial,
model.size_bytes,
existing.partial,
existing.size_bytes,
):
deduped[key] = model
return sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
async def list_local_models_response(models_dir: str = "./models") -> LocalModelListResponse:
"""List local model candidates from every supported on-device source."""
hf_cache_dir = _resolve_hf_cache_dir()
legacy_hf = legacy_hf_cache_dir()
hf_default = hf_default_cache_dir()
lm_dirs = lmstudio_model_dirs()
ollama_dirs = ollama_model_dirs()
allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir]
if legacy_hf.is_dir():
allowed_roots.append(legacy_hf)
if hf_default.is_dir():
allowed_roots.append(hf_default)
allowed_roots.extend([studio_root(), outputs_root()])
try:
models_root = _resolve_allowed_models_dir(models_dir, allowed_roots)
except ValueError:
raise HTTPException(status_code = 403, detail = "Directory not allowed")
try:
local_models = await _collect_models_from_default_sources(
models_root,
hf_cache_dir,
legacy_hf,
hf_default,
lm_dirs,
ollama_dirs,
)
local_models += await _collect_models_from_custom_folders()
models = _dedupe_local_models(local_models)
return LocalModelListResponse(
models_dir = str(models_root),
hf_cache_dir = str(hf_cache_dir),
lmstudio_dirs = [str(d) for d in lm_dirs],
ollama_dirs = [str(d) for d in ollama_dirs],
models = models,
)
except Exception as e:
logger.error(f"Error listing local models: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
detail = f"Failed to list local models: {str(e)}",
)
def get_scan_folders_response() -> dict:
return {"folders": list_scan_folders()}
def add_scan_folder_response(path: str) -> dict:
try:
folder = add_scan_folder(_coerce_scan_folder_path(path))
except ValueError as e:
logger.warning("Scan folder rejected: %s (path=%s)", e, path)
raise HTTPException(status_code = 400, detail = str(e))
logger.info("Scan folder added: %s", folder.get("path"))
return folder
def remove_scan_folder_response(folder_id: int) -> dict:
remove_scan_folder(folder_id)
logger.info("Scan folder removed: id=%s", folder_id)
return {"ok": True}

394
studio/backend/hub/services/models/ollama.py Normal file
View file

@ -0,0 +1,394 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Ollama model inventory: manifest parsing and writable-symlink materialization.
Ollama stores models content-addressed under ``<root>/manifests/`` and
``<root>/blobs/``. Inventory scans read the manifests directly (no writes),
returning rows whose ``id`` is an opaque ``ollama-manifest:`` reference. The
load path then calls :func:`materialize_ollama_model_ref`, which creates a
``.gguf``-named symlink (or hardlink) so that downstream loaders see a path
with the GGUF suffix without copying multi-GB blobs inside an API request.
"""
from __future__ import annotations
import hashlib
import json
import os
import uuid
from pathlib import Path
from typing import List, Optional
from urllib.parse import quote, unquote
from loggers import get_logger
from hub.schemas.inventory import LocalModelInfo
from hub.services.models.common import (
_capabilities_for_format,
_local_inventory_id,
)
from hub.utils.paths import (
cache_root,
ollama_model_dirs,
path_is_same_or_child,
tmp_root,
)
logger = get_logger(__name__)
_OLLAMA_MANIFEST_REF_PREFIX = "ollama-manifest:"
_OLLAMA_BLOB_NAME_CHARS = frozenset(
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._+-"
)
def _ollama_manifest_ref(tag_file: Path) -> str:
return f"{_OLLAMA_MANIFEST_REF_PREFIX}{quote(str(tag_file), safe = '')}"
def _safe_is_file(path: Path) -> bool:
try:
return path.is_file()
except OSError:
return False
def _ollama_blob_path(blobs_dir: Path, digest: object) -> Optional[Path]:
if not isinstance(digest, str):
return None
algorithm, separator, value = digest.partition(":")
if separator != ":" or not algorithm or not value:
return None
name = f"{algorithm}-{value}"
if (
not name
or name in (".", "..")
or any(char not in _OLLAMA_BLOB_NAME_CHARS for char in name)
or not name.isprintable()
):
return None
return blobs_dir / name
def _contained_link_path(link_dir: Path, link_name: str) -> Optional[Path]:
"""Resolve *link_name* to a direct child of *link_dir*, or ``None``. ``link_name`` derives from manifest fields, so requiring a direct child keeps a crafted value with separators, ``..``, or a drive prefix from escaping the links dir."""
if not link_name or link_name in (".", ".."):
return None
link_path = link_dir / link_name
try:
if link_path.parent.resolve() != link_dir.resolve():
return None
except OSError:
return None
return link_path
def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
"""Writable directory for Ollama ``.gguf`` symlinks. Prefers ``<ollama_dir>/.studio_links/`` next to the blobs; falls back to Studio's cache (read-only system installs), then the temp dir (sandboxed installs)."""
def _ensure_writable_dir(path: Path) -> Optional[Path]:
try:
path.mkdir(parents = True, exist_ok = True)
probe = path / f".write-test-{uuid.uuid4().hex[:8]}"
probe.mkdir()
probe.rmdir()
return path
except OSError as e:
logger.debug("Ollama link dir %s is not writable: %s", path, e)
return None
primary = ollama_dir / ".studio_links"
if _ensure_writable_dir(primary) is not None:
return primary
# Namespace by a hash of the ollama_dir so two different Ollama roots
# don't collide. This is a cache path, not a security boundary.
try:
digest = hashlib.sha256(str(ollama_dir.resolve()).encode()).hexdigest()[:12]
except (OSError, RuntimeError):
digest = "default"
fallback = cache_root() / "ollama_links" / digest
if _ensure_writable_dir(fallback) is not None:
return fallback
tmp_fallback = tmp_root() / "ollama_links" / digest
if _ensure_writable_dir(tmp_fallback) is not None:
return tmp_fallback
logger.warning(
"Could not create a writable Ollama link directory for %s",
ollama_dir,
)
return None
def _make_ollama_blob_link(link_dir: Path, link_name: str, target: Path) -> Optional[str]:
"""Create a .gguf-named link to an Ollama blob: tries symlink then hardlink, skips the model if neither works (a full multi-GB copy would block the API). Idempotent."""
try:
link_dir.mkdir(parents = True, exist_ok = True)
except OSError as e:
logger.warning(
"Could not create Ollama link directory %s: %s",
link_dir,
e,
)
return None
link_path = _contained_link_path(link_dir, link_name)
if link_path is None:
logger.warning("Refusing unsafe Ollama link name %r under %s", link_name, link_dir)
return None
try:
resolved = target.resolve()
except OSError as e:
logger.debug("Could not resolve Ollama blob %s: %s", target, e)
return None
# Skip if the link already points at the same blob. Use samefile, not size:
# `ollama pull` can swap a tag to a same-sized blob, leaving a stale link.
try:
if link_path.exists() and os.path.samefile(str(link_path), str(resolved)):
return str(link_path)
except OSError as e:
logger.debug("Error checking existing link %s: %s", link_path, e)
tmp_path = link_dir / f".{link_name}.tmp-{uuid.uuid4().hex[:8]}"
try:
if tmp_path.is_symlink() or tmp_path.exists():
tmp_path.unlink()
try:
tmp_path.symlink_to(resolved)
except OSError:
try:
os.link(str(resolved), str(tmp_path))
except OSError:
logger.warning(
"Could not create link for Ollama blob %s "
"(symlinks and hardlinks both failed). "
"Skipping model to avoid blocking the API.",
target,
)
return None
os.replace(str(tmp_path), str(link_path))
return str(link_path)
except OSError as e:
logger.debug("Could not create Ollama link %s: %s", link_path, e)
try:
if tmp_path.is_symlink() or tmp_path.exists():
tmp_path.unlink()
except OSError as cleanup_err:
logger.debug("Could not clean up tmp path %s: %s", tmp_path, cleanup_err)
return None
def _ollama_model_info_from_manifest(
ollama_dir: Path,
tag_file: Path,
*,
materialize_links: bool = False,
links_root: Optional[Path] = None,
) -> Optional[LocalModelInfo]:
manifests_root = ollama_dir / "manifests"
blobs_dir = ollama_dir / "blobs"
try:
rel = tag_file.relative_to(manifests_root)
except ValueError:
return None
parts = rel.parts
if len(parts) < 3:
return None
host = parts[0]
repo_parts = list(parts[1:-1])
tag = parts[-1]
if host == "registry.ollama.ai" and repo_parts and repo_parts[0] == "library":
repo_name = "/".join(repo_parts[1:])
elif host == "registry.ollama.ai":
repo_name = "/".join(repo_parts)
else:
repo_name = "/".join([host] + repo_parts)
if not repo_name:
return None
try:
manifest = json.loads(tag_file.read_text())
except (json.JSONDecodeError, OSError) as e:
logger.debug("Skipping unreadable/invalid Ollama manifest %s: %s", tag_file, e)
return None
config = manifest.get("config", {})
config_digest = config.get("digest", "") if isinstance(config, dict) else ""
model_type = ""
file_type = ""
if config_digest and blobs_dir.is_dir():
config_blob = _ollama_blob_path(blobs_dir, config_digest)
if config_blob is not None and _safe_is_file(config_blob):
try:
cfg = json.loads(config_blob.read_text())
model_type = cfg.get("model_type", "")
file_type = cfg.get("file_type", "")
except (json.JSONDecodeError, OSError) as e:
logger.debug("Could not parse Ollama config blob %s: %s", config_blob, e)
layers = manifest.get("layers") or []
if not isinstance(layers, list):
return None
model_blob: Optional[Path] = None
gguf_link_path: Optional[str] = None
stem_hash = hashlib.sha256(rel.as_posix().encode()).hexdigest()[:10]
model_link_dir = links_root / stem_hash if links_root is not None else None
safe_name = repo_name.replace("/", "-")
quant = f"-{file_type}" if file_type else ""
for layer in layers:
if not isinstance(layer, dict):
continue
media = layer.get("mediaType", "")
digest = layer.get("digest", "")
if not digest:
continue
if media == "application/vnd.ollama.image.model":
candidate = _ollama_blob_path(blobs_dir, digest)
if candidate is None or not _safe_is_file(candidate):
continue
model_blob = candidate
if materialize_links and model_link_dir is not None:
link_name = f"{safe_name}-{tag}{quant}.gguf"
gguf_link_path = _make_ollama_blob_link(model_link_dir, link_name, candidate)
elif materialize_links and media == "application/vnd.ollama.image.projector":
candidate = _ollama_blob_path(blobs_dir, digest)
if candidate is not None and _safe_is_file(candidate) and model_link_dir is not None:
mmproj_name = f"{safe_name}-{tag}-mmproj.gguf"
_make_ollama_blob_link(model_link_dir, mmproj_name, candidate)
if model_blob is None:
return None
if materialize_links and not gguf_link_path:
return None
suffix = ""
if model_type:
suffix += f" ({model_type}"
if file_type:
suffix += f" {file_type}"
suffix += ")"
try:
updated_at = tag_file.stat().st_mtime
except OSError:
updated_at = None
display = f"{repo_name}:{tag}"
model_id = f"ollama/{repo_name}:{tag}"
path = gguf_link_path if materialize_links and gguf_link_path else str(model_blob)
load_id = path if materialize_links else _ollama_manifest_ref(tag_file)
return LocalModelInfo(
id = load_id,
inventory_id = _local_inventory_id("ollama", "gguf", model_id),
load_id = load_id,
model_id = model_id,
display_name = display + suffix,
path = path,
source = "ollama",
updated_at = updated_at,
model_format = "gguf",
runtime = "llama_cpp",
capabilities = _capabilities_for_format("gguf", "ollama"),
)
def scan_ollama_dir(
ollama_dir: Path,
*,
limit: Optional[int] = None,
materialize_links: bool = False,
) -> List[LocalModelInfo]:
"""Scan an Ollama models directory for downloaded models.
Ollama uses a content-addressable layout
(``manifests/<host>/<namespace>/<model>/<tag>`` + ``blobs/sha256-...``),
iterated via ``rglob`` to find every depth. Each manifest's ``model`` layer
holds the GGUF weights (vision models add a projector layer).
Scans are read-only by default and return an opaque manifest reference;
the load route later calls :func:`materialize_ollama_model_ref` to create a
``.gguf`` symlink/hardlink, keeping GET /local free of filesystem writes.
"""
manifests_root = ollama_dir / "manifests"
if not manifests_root.is_dir():
return []
found: List[LocalModelInfo] = []
links_root = _ollama_links_dir(ollama_dir) if materialize_links else None
if materialize_links and links_root is None:
logger.warning(
"Skipping Ollama scan for %s: no writable location for .gguf links",
ollama_dir,
)
return []
try:
for tag_file in manifests_root.rglob("*"):
if not _safe_is_file(tag_file):
continue
info = _ollama_model_info_from_manifest(
ollama_dir,
tag_file,
materialize_links = materialize_links,
links_root = links_root,
)
if info is None:
continue
found.append(info)
if limit is not None and len(found) >= limit:
return found
except OSError as e:
logger.warning("Error scanning Ollama directory %s: %s", ollama_dir, e)
return found
def _ollama_dir_for_manifest(tag_file: Path) -> Optional[Path]:
"""Discovered Ollama root whose ``manifests/`` contains *tag_file*, or ``None``. Validating against known roots keeps a crafted reference from driving materialization to an arbitrary path."""
for ollama_dir in ollama_model_dirs():
if path_is_same_or_child(tag_file, ollama_dir / "manifests"):
return ollama_dir
return None
def materialize_ollama_model_ref(ref: str) -> str:
"""Resolve an ``ollama-manifest:`` reference to a loadable ``.gguf`` path,
creating the writable symlink/hardlink on demand.
Raises ``ValueError`` if the reference is malformed, points outside a
discovered Ollama models directory, or cannot be materialized.
"""
if not ref.startswith(_OLLAMA_MANIFEST_REF_PREFIX):
raise ValueError("Not an Ollama manifest reference")
tag_file = Path(unquote(ref[len(_OLLAMA_MANIFEST_REF_PREFIX) :]))
ollama_dir = _ollama_dir_for_manifest(tag_file)
if ollama_dir is None:
raise ValueError("Reference is outside any known Ollama models directory")
links_root = _ollama_links_dir(ollama_dir)
if links_root is None:
raise ValueError("No writable location for Ollama .gguf links")
info = _ollama_model_info_from_manifest(
ollama_dir,
tag_file,
materialize_links = True,
links_root = links_root,
)
if info is None or not info.path:
raise ValueError("Could not materialize Ollama model from manifest")
return info.path

260
studio/backend/hub/services/snapshot_progress.py Normal file
View file

@ -0,0 +1,260 @@
# 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 snapshot download-progress computation for models and datasets.
Both scan the cache's ``blobs/`` dir, split finalized vs ``.incomplete`` bytes,
filter to the target revision's expected hashes, and divide by its total size;
only the ``metadata_resolver`` differs. One copy keeps the two from drifting (a
prior hash-filter fix once landed only on the model copy, leaving datasets
summing stale blobs against the wrong total)."""
from __future__ import annotations
import asyncio
from pathlib import Path
from typing import Callable, Optional
from loggers import get_logger
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.state_dir import RepoType
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
blob_bytes_present,
latest_snapshot_dir,
preferred_repo_cache_dirs,
)
from hub.utils.paths import is_valid_repo_id as _is_valid_repo_id
logger = get_logger(__name__)
# (repo_id, hf_token) -> (expected_total_bytes, expected_blob_hashes)
SnapshotMetadataResolver = Callable[[str, Optional[str]], "tuple[int, frozenset[str]]"]
def _empty_progress(expected_bytes: int) -> dict:
return {
"downloaded_bytes": 0,
"completed_bytes": 0,
"complete_on_disk": False,
"expected_bytes": max(expected_bytes, 0),
"progress": 0,
"cache_path": None,
}
def _snapshot_complete_on_disk(
*,
repo_type: RepoType,
repo_id: str,
variant: Optional[str],
entry: Path,
expected_total: int,
completed_bytes: int,
in_progress_bytes: int,
) -> bool:
if expected_total <= 0 or completed_bytes < expected_total or in_progress_bytes > 0:
return False
snapshot_dir = latest_snapshot_dir(entry)
if snapshot_dir is None:
return False
if variant is None and hf_cache_scan.repo_cache_dir_has_incomplete_blobs(entry):
return False
if download_manifest.has_cancel_marker(repo_type, repo_id, variant):
return False
manifest = download_manifest.read_manifest(repo_type, repo_id, variant)
if manifest is None:
return False
return download_manifest.verify_against_disk(manifest, snapshot_dir).ok
def compute_snapshot_progress(
*,
repo_type: RepoType,
repo_id: str,
job_key: str,
expected_bytes: int,
hf_token: Optional[str],
registry,
metadata_resolver: SnapshotMetadataResolver,
variant: Optional[str] = None,
) -> dict:
"""Synchronous progress reading. Safe to run under ``asyncio.to_thread``."""
empty = _empty_progress(expected_bytes)
if not _is_valid_repo_id(repo_id):
return empty
job_state = registry.get_job(job_key).state
force_active = job_state in {"running", "cancelling"}
get_job_metadata = getattr(registry, "get_job_metadata", None)
metadata = get_job_metadata(job_key) if callable(get_job_metadata) else None
completed_baseline_bytes = max(
0,
int(getattr(metadata, "completed_baseline_bytes", 0) or 0),
)
expected_total = max(expected_bytes, 0)
# Always resolve the revision's blob hashes so stale blobs from a superseded
# revision can't inflate the count; hashes degrade to empty (count-all) only
# when metadata is unavailable (e.g. offline). Take the larger total so a low
# caller hint can't cap the bar below the revision's real size.
meta_total, expected_hashes = metadata_resolver(repo_id, hf_token)
expected_total = max(expected_total, meta_total)
# Without resolved hashes, a variant must not count unscoped blobs: sibling
# quants share one blobs/ dir, so a sibling's bytes (or .incomplete) would be
# misattributed and make the bar jump backward. A no-variant snapshot owns
# the whole dir, so it counts unscoped.
count_finalized_unscoped = variant is None
readings: list[tuple[int, int, Optional[str], bool]] = []
for entry in preferred_repo_cache_dirs(
repo_type,
repo_id,
force_active = force_active,
):
completed_bytes = 0
in_progress_bytes = 0
cache_path = hf_cache_scan.resolve_hf_cache_realpath(entry)
blobs_dir = entry / "blobs"
if blobs_dir.is_dir():
try:
blob_entries = list(blobs_dir.iterdir())
except OSError:
blob_entries = []
for f in blob_entries:
# Skip a blob that vanished mid-poll rather than zeroing the reading.
try:
if not f.is_file():
continue
if f.name.endswith(INCOMPLETE_SUFFIX):
blob_hash = f.name[: -len(INCOMPLETE_SUFFIX)]
if expected_hashes:
if blob_hash not in expected_hashes:
continue
elif not count_finalized_unscoped:
continue
in_progress_bytes += blob_bytes_present(f)
else:
if expected_hashes:
if f.name not in expected_hashes:
continue
elif not count_finalized_unscoped:
continue
completed_bytes += f.stat().st_size
except OSError:
continue
readings.append(
(
completed_bytes,
in_progress_bytes,
cache_path,
_snapshot_complete_on_disk(
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
entry = entry,
expected_total = expected_total,
completed_bytes = completed_bytes,
in_progress_bytes = in_progress_bytes,
),
)
)
selected = max(
readings,
key = lambda item: (item[0] + item[1], item[0]),
default = None,
)
if selected is None:
return empty
completed_bytes, in_progress_bytes, cache_path, complete_on_disk = selected
downloaded_bytes = completed_bytes + in_progress_bytes
# Subtract the companion baseline only while still counted in completed_bytes
# and the variant is not yet verified complete, else genuine progress reads as
# 0-byte.
effective_baseline_bytes = (
completed_baseline_bytes
if not complete_on_disk and completed_baseline_bytes <= completed_bytes
else 0
)
display_completed_bytes = max(0, completed_bytes - effective_baseline_bytes)
display_downloaded_bytes = max(0, downloaded_bytes - effective_baseline_bytes)
if expected_total <= 0:
# Cannot determine total; report bytes only, no percentage.
return {
"downloaded_bytes": display_downloaded_bytes,
"completed_bytes": display_completed_bytes,
"complete_on_disk": False,
"expected_bytes": 0,
"progress": 0,
"cache_path": cache_path,
}
display_expected_total = max(0, expected_total - effective_baseline_bytes)
if downloaded_bytes == 0:
return {
**empty,
"expected_bytes": display_expected_total,
"cache_path": cache_path,
}
# Cap at 0.99 until the manifest-backed disk check verifies completion: on
# resume, completed bytes can sit above the threshold while files still download.
progress = (
1.0
if complete_on_disk
else (
min(display_downloaded_bytes / display_expected_total, 0.99)
if display_expected_total > 0
else 0
)
)
return {
"downloaded_bytes": display_downloaded_bytes,
"completed_bytes": display_completed_bytes,
"complete_on_disk": complete_on_disk,
"expected_bytes": display_expected_total,
"progress": round(progress, 3),
"cache_path": cache_path,
}
async def snapshot_progress_response(
*,
repo_type: RepoType,
repo_id: str,
job_key: str,
expected_bytes: int,
hf_token: Optional[str],
registry,
metadata_resolver: SnapshotMetadataResolver,
variant: Optional[str] = None,
) -> dict:
"""Async wrapper: offloads the blocking cache walk and never raises."""
try:
return await asyncio.to_thread(
compute_snapshot_progress,
repo_type = repo_type,
repo_id = repo_id,
job_key = job_key,
expected_bytes = expected_bytes,
hf_token = hf_token,
registry = registry,
metadata_resolver = metadata_resolver,
variant = variant,
)
except Exception as e:
logger.warning(
"Error checking %s download progress for %s: %s: %s",
repo_type,
repo_id,
type(e).__name__,
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
return _empty_progress(expected_bytes)

2
studio/backend/hub/storage/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

194
studio/backend/hub/storage/scan_folders.py Normal file
View file

@ -0,0 +1,194 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Persistence for user-registered custom model scan folders.
Self-bootstrapping table inside the existing studio SQLite so the Hub module
doesn't have to modify upstream studio_db.py's schema init."""
from __future__ import annotations
import os
import platform
import sqlite3
import threading
from datetime import datetime, timezone
from storage.studio_db import get_connection
from hub.utils.paths import normalize_path
_schema_lock = threading.Lock()
_schema_ready = False
_SENSITIVE_PATH_COMPONENTS = {
".aws",
".azure",
".config",
".docker",
".gcloud",
".gnupg",
".huggingface",
".kaggle",
".kube",
".modelscope",
".ngc",
".local",
".mozilla",
".pki",
".thunderbird",
".ssh",
".1password",
".bitwarden",
".password-store",
"1password",
"bitwarden",
"keychains",
"keyrings",
"mozilla",
"thunderbird",
}
def _denied_path_prefixes() -> list[str]:
system = platform.system()
if system == "Linux":
return ["/proc", "/sys", "/dev", "/etc", "/boot", "/run"]
if system == "Darwin":
# realpath() resolves /etc -> /private/etc, /tmp -> /private/tmp on macOS,
# so include the /private variants to avoid bypasses.
return [
"/System",
"/Library",
"/dev",
"/etc",
"/private/etc",
"/tmp",
"/private/tmp",
"/var",
"/private/var",
]
if system == "Windows":
win = os.environ.get("SystemRoot", r"C:\Windows")
pf = os.environ.get("ProgramFiles", r"C:\Program Files")
pf86 = os.environ.get("ProgramFiles(x86)", r"C:\Program Files (x86)")
return [os.path.normcase(p) for p in [win, pf, pf86]]
return []
def _contains_sensitive_path_component(path: str) -> bool:
parts = os.path.normpath(path).split(os.sep)
return any(part.lower() in _SENSITIVE_PATH_COMPONENTS for part in parts)
def contains_sensitive_path_component(path: str) -> bool:
"""Public predicate for the credential/config denylist (.ssh, .aws, ...).
Shared with the folder browser so browse and register enforce one policy."""
return _contains_sensitive_path_component(path)
def _ensure_schema(conn: sqlite3.Connection) -> None:
global _schema_ready
if _schema_ready:
return
with _schema_lock:
if _schema_ready:
return
collation = "COLLATE NOCASE" if platform.system() == "Windows" else ""
conn.execute(
f"""
CREATE TABLE IF NOT EXISTS scan_folders (
id INTEGER PRIMARY KEY AUTOINCREMENT,
path TEXT NOT NULL UNIQUE {collation},
created_at TEXT NOT NULL
)
"""
)
conn.commit()
_schema_ready = True
def list_scan_folders() -> list[dict]:
conn = get_connection()
try:
_ensure_schema(conn)
rows = conn.execute(
"SELECT id, path, created_at FROM scan_folders ORDER BY created_at"
).fetchall()
return [dict(row) for row in rows]
finally:
conn.close()
def add_scan_folder(path: str) -> dict:
"""Add a readable directory for the local OS user; not a multi-user sandbox."""
if not path or not path.strip():
raise ValueError("Path cannot be empty")
normalized = os.path.realpath(os.path.expanduser(normalize_path(path.strip())))
if not os.path.exists(normalized):
raise ValueError("Path does not exist")
if not os.path.isdir(normalized):
raise ValueError("Path must be a directory, not a file")
if not os.access(normalized, os.R_OK | os.X_OK):
raise ValueError("Path is not readable")
if os.path.dirname(normalized) == normalized:
# Registering a filesystem root would expose denied system dirs via browse.
raise ValueError("The filesystem root cannot be registered")
if _contains_sensitive_path_component(normalized):
raise ValueError("Credential or configuration directories are not allowed")
is_win = platform.system() == "Windows"
check = os.path.normcase(normalized) if is_win else normalized
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
raise ValueError(f"Path under {prefix} is not allowed")
conn = get_connection()
try:
_ensure_schema(conn)
now = datetime.now(timezone.utc).isoformat()
if is_win:
existing = conn.execute(
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE",
(normalized,),
).fetchone()
else:
existing = conn.execute(
"SELECT id, path, created_at FROM scan_folders WHERE path = ?",
(normalized,),
).fetchone()
if existing is not None:
return dict(existing)
try:
conn.execute(
"INSERT INTO scan_folders (path, created_at) VALUES (?, ?)",
(normalized, now),
)
conn.commit()
except sqlite3.IntegrityError:
pass
fallback_sql = (
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE"
if is_win
else "SELECT id, path, created_at FROM scan_folders WHERE path = ?"
)
row = conn.execute(fallback_sql, (normalized,)).fetchone()
if row is None:
raise ValueError("Folder was concurrently removed")
return dict(row)
finally:
conn.close()
def remove_scan_folder(id: int) -> None:
# sqlite INTEGER is signed 64-bit; ids outside that range cannot exist.
if not -(2**63) <= id < 2**63:
return
conn = get_connection()
try:
_ensure_schema(conn)
conn.execute("DELETE FROM scan_folders WHERE id = ?", (id,))
conn.commit()
finally:
conn.close()

106
studio/backend/hub/tests/conftest.py Normal file
View file

@ -0,0 +1,106 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import sys
import types
class _BaseModel:
def __init__(self, **kwargs):
for name, value in self.__class__.__dict__.items():
if name.startswith("_") or callable(value):
continue
if name not in kwargs:
setattr(self, name, value)
for key, value in kwargs.items():
setattr(self, key, value)
def model_dump(self):
return dict(self.__dict__)
def model_copy(self, update = None):
data = self.model_dump()
if update:
data.update(update)
return self.__class__(**data)
def _field(default = ..., **kwargs):
if "default_factory" in kwargs:
return kwargs["default_factory"]()
return None if default is ... else default
def _model_validator(*args, **kwargs):
def decorator(fn):
return fn
return decorator
class _HTTPException(Exception):
def __init__(
self,
status_code: int,
detail = None,
):
super().__init__(detail)
self.status_code = status_code
self.detail = detail
class _APIRouter:
def get(self, *args, **kwargs):
return lambda fn: fn
def post(self, *args, **kwargs):
return lambda fn: fn
def delete(self, *args, **kwargs):
return lambda fn: fn
def _fastapi_marker(
default = None,
*args,
**kwargs,
):
return default
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
sys.modules.setdefault(
"pydantic",
types.SimpleNamespace(
BaseModel = _BaseModel,
Field = _field,
model_validator = _model_validator,
),
)
sys.modules.setdefault(
"fastapi",
types.SimpleNamespace(
APIRouter = _APIRouter,
Body = _fastapi_marker,
Depends = _fastapi_marker,
Header = _fastapi_marker,
HTTPException = _HTTPException,
Query = _fastapi_marker,
UploadFile = object,
),
)
sys.modules.setdefault(
"loggers",
types.SimpleNamespace(get_logger = lambda *args, **kwargs: _DummyLogger()),
)
sys.modules.setdefault(
"structlog",
types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
),
)

356
studio/backend/hub/tests/test_dataset_services.py Normal file
View file

@ -0,0 +1,356 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
from pathlib import Path
from types import SimpleNamespace
import pytest
from fastapi import HTTPException
from hub.schemas.datasets import CheckFormatRequest, LocalDatasetItem
from hub.services.datasets import cache_inventory, downloads, formatting, local
from hub.utils import download_manifest, download_registry, state_dir
class _Upload:
def __init__(self, filename: str, payload: bytes):
self.filename = filename
self._payload = payload
self._offset = 0
async def read(self, size: int) -> bytes:
if self._offset >= len(self._payload):
return b""
chunk = self._payload[self._offset : self._offset + size]
self._offset += len(chunk)
return chunk
def test_dataset_cache_scan_merges_raw_and_processed_rows(monkeypatch):
raw_repo = SimpleNamespace(
repo_id = "Org/Data",
repo_type = "dataset",
repo_path = "/cache/datasets--Org--Data",
size_on_disk = 100,
revisions = [SimpleNamespace(files = [], commit_hash = "abc")],
)
monkeypatch.setattr(
cache_inventory,
"_collect_hf_cache_scans",
lambda: ([SimpleNamespace(repos = [raw_repo])], {"/cache"}),
)
monkeypatch.setattr(
cache_inventory.hf_cache_scan,
"is_snapshot_partial",
lambda _repo_type, _repo_id, _cache_dir: False,
)
monkeypatch.setattr(
cache_inventory,
"_scan_hub_dataset_cache_dirs",
lambda: [],
)
monkeypatch.setattr(
cache_inventory,
"_scan_processed_dataset_caches",
lambda: [
{
"repo_id": "org/data",
"size_bytes": 250,
"cache_path": "/processed/org___data",
"processed_cache": True,
"partial": False,
}
],
)
rows = cache_inventory._scan_hf_dataset_caches()
assert len(rows) == 1
assert rows[0]["repo_id"] == "Org/Data"
assert rows[0]["size_bytes"] == 250
assert rows[0]["partial"] is False
def test_delete_cached_dataset_attempts_all_roots_before_raising(monkeypatch):
calls = []
purged_state = []
class _DeleteStrategy:
def __init__(self, label: str, fail: bool):
self.label = label
self.fail = fail
def execute(self):
calls.append(self.label)
if self.fail:
raise RuntimeError(f"{self.label} failed")
class _Cache:
def __init__(self, label: str, fail: bool):
self.cache_dir = label
self.repos = [
SimpleNamespace(
repo_type = "dataset",
repo_id = "Org/Data",
revisions = [SimpleNamespace(commit_hash = f"{label}-rev")],
)
]
self.fail = fail
def delete_revisions(self, *_revisions):
return _DeleteStrategy(self.cache_dir, self.fail)
monkeypatch.setattr(
cache_inventory,
"_collect_hf_cache_scans",
lambda: ([_Cache("first", True), _Cache("second", False)], set()),
)
monkeypatch.setattr(
cache_inventory,
"_delete_processed_dataset_cache",
lambda _repo_id: (True, []),
)
monkeypatch.setattr(
cache_inventory.download_manifest,
"purge_all_state_for_repo",
lambda *_args: purged_state.append(True) or 1,
)
with pytest.raises(HTTPException) as exc_info:
cache_inventory._delete_cached_dataset_blocking("Org/Data")
assert exc_info.value.status_code == 500
assert calls == ["first", "second"]
assert purged_state == []
def test_delete_cached_dataset_purges_blob_only_repo_dir(monkeypatch):
"""A blob-only ``datasets--owner--repo`` dir (no usable snapshot/refs) is
fully removable: purge_partial_repo alone clears only ``.incomplete`` files
and would leave the complete blobs and the row."""
purged_dirs: list[str] = []
monkeypatch.setattr(
cache_inventory,
"_collect_hf_cache_scans",
lambda: ([], set()),
)
monkeypatch.setattr(
cache_inventory,
"_delete_processed_dataset_cache",
lambda _repo_id: (False, []),
)
monkeypatch.setattr(
cache_inventory,
"purge_repo_cache_dirs",
lambda _repo_type, repo_id: purged_dirs.append(repo_id) or True,
)
monkeypatch.setattr(
cache_inventory,
"purge_partial_repo",
lambda *_args: False,
)
monkeypatch.setattr(
cache_inventory.download_manifest,
"purge_all_state_for_repo",
lambda *_args: 0,
)
result = cache_inventory._delete_cached_dataset_blocking("Org/Data")
assert result == {"status": "deleted", "repo_id": "Org/Data"}
assert purged_dirs == ["Org/Data"]
def test_delete_cached_dataset_absent_everywhere_raises_404(monkeypatch):
monkeypatch.setattr(
cache_inventory,
"_collect_hf_cache_scans",
lambda: ([], set()),
)
monkeypatch.setattr(
cache_inventory,
"_delete_processed_dataset_cache",
lambda _repo_id: (False, []),
)
monkeypatch.setattr(
cache_inventory,
"purge_repo_cache_dirs",
lambda *_args: False,
)
monkeypatch.setattr(
cache_inventory,
"purge_partial_repo",
lambda *_args: False,
)
monkeypatch.setattr(
cache_inventory.download_manifest,
"purge_all_state_for_repo",
lambda *_args: 0,
)
with pytest.raises(HTTPException) as exc_info:
cache_inventory._delete_cached_dataset_blocking("Org/Missing")
assert exc_info.value.status_code == 404
def test_check_format_rejects_invalid_path_as_400():
with pytest.raises(HTTPException) as exc_info:
formatting.check_format_response(CheckFormatRequest(dataset_name = "../../etc/passwd"))
assert exc_info.value.status_code == 400
def test_dataset_download_status_preserves_idle_shape():
status = downloads._dataset_status("Org/Data")
assert status.state == "idle"
assert status.error is None
def test_dataset_download_registry_key_is_case_insensitive():
registry = download_registry.DownloadRegistry()
claimed, state = registry.claim(
"Org/Data",
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "Org/Data",
)
duplicate_claimed, duplicate_state = registry.claim(
"org/data",
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "org/data",
)
assert claimed is True
assert state == "running"
assert duplicate_claimed is False
assert duplicate_state == "running"
assert registry.active_jobs("ORG/DATA") == {"org/data": "running"}
def test_dataset_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry())
assert download_manifest.write_cancel_marker("dataset", "Owner/Data", None, "http")
status = asyncio.run(downloads.get_dataset_download_status_response("owner/data"))
assert status.state == "cancelled"
assert status.error is None
def test_dataset_claim_register_cancel_uses_registry_marker_owner(monkeypatch):
killed = []
class _Registry:
def claim(self, *_args, **_kwargs):
return True, "running"
def current_generation(self, _key):
return 1
def register_process(self, _key, _proc):
return False
def persist_cancel_for_key(self, *_args, **_kwargs):
raise AssertionError("register_process owns pending-cancel markers")
def get_job(self, _key):
return SimpleNamespace(state = "cancelled", error = None)
monkeypatch.setattr(downloads, "_registry", _Registry())
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, **_kwargs: repo_id,
)
monkeypatch.setattr(
downloads.download_registry,
"download_transport_unavailable_reason",
lambda _transport: None,
)
monkeypatch.setattr(
downloads.download_lifecycle,
"spawn_worker",
lambda *_args, **_kwargs: object(),
)
monkeypatch.setattr(
downloads.download_lifecycle,
"kill_and_reap_process",
lambda proc, **_kwargs: killed.append(proc),
)
result = asyncio.run(
downloads.download_dataset_response(SimpleNamespace(repo_id = "Org/Data", use_xet = False))
)
assert result["state"] == "cancelled"
assert killed
def test_dataset_cancel_pending_spawn_arms_pending_cancel(monkeypatch):
events = []
class _Registry:
def get_process(self, _key):
return None
def mark_pending_cancel(self, key, generation):
events.append(("pending", key, generation))
return True
def get_job(self, _key):
return SimpleNamespace(state = "running")
monkeypatch.setattr(downloads, "_registry", _Registry())
monkeypatch.setattr(
downloads,
"resolve_cached_repo_id_case",
lambda repo_id, **_kwargs: repo_id,
)
result = asyncio.run(
downloads.cancel_dataset_download_response(
SimpleNamespace(repo_id = "Org/Data", generation = 5)
)
)
assert result == {"repo_id": "Org/Data", "state": "cancelling"}
assert events == [("pending", "org/data", 5)]
def test_upload_dataset_response_writes_non_empty_file(monkeypatch, tmp_path):
payload = b'{"text":"hello"}\n'
monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", tmp_path)
response = asyncio.run(local.upload_dataset_response(_Upload("../train.jsonl", payload)))
stored_path = Path(response.stored_path)
assert response.filename == "train.jsonl"
assert stored_path.parent == tmp_path
assert stored_path.name.endswith("_train.jsonl")
assert stored_path.read_bytes() == payload
def test_local_dataset_items_expose_recipe_and_upload_source(monkeypatch, tmp_path):
recipe_root = tmp_path / "recipes"
upload_root = tmp_path / "uploads"
parquet_dir = recipe_root / "recipe_alpha" / "parquet-files"
parquet_dir.mkdir(parents = True)
(parquet_dir / "part.parquet").write_bytes(b"parquet")
upload_root.mkdir()
(upload_root / "manual.jsonl").write_text('{"text":"hello"}\n', encoding = "utf-8")
monkeypatch.setattr(local, "LOCAL_DATASETS_ROOT", recipe_root)
monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", upload_root)
response = local.list_local_datasets_response()
assert "source" in LocalDatasetItem.__annotations__
by_id = {item.id: item for item in response.datasets}
assert by_id["recipe_alpha"].source == "recipe"
assert by_id["manual.jsonl"].source == "upload"

2962
studio/backend/hub/tests/test_model_services.py Normal file
View file

File diff suppressed because it is too large Load diff

2
studio/backend/hub/utils/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

145
studio/backend/hub/utils/dataset_cache.py Normal file
View file

@ -0,0 +1,145 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import re
from pathlib import Path
from typing import Optional
from hub.utils.hf_cache_state import iter_repo_cache_dirs
TRAINING_DATA_EXTS = (".parquet", ".json", ".jsonl", ".csv")
def _rel_lower(snapshot: Path, path: Path) -> str:
return path.relative_to(snapshot).as_posix().lower()
_SPLIT_ALIASES = {
"validation": frozenset({"validation", "valid", "val"}),
"valid": frozenset({"validation", "valid", "val"}),
"val": frozenset({"validation", "valid", "val"}),
"eval": frozenset({"eval", "validation", "valid", "val"}),
}
def _label_tokens(text: str) -> set[str]:
return {token for token in re.split(r"[^a-z0-9]+", text.lower()) if token}
def split_label_matches(text: str, split: str) -> bool:
"""Match a split name against a file path's tokens, expanding split aliases
(validation/valid/val, eval) so cached and remote selection agree."""
normalized = split.strip().lower()
if not normalized:
return False
labels = _SPLIT_ALIASES.get(normalized, frozenset({normalized}))
return bool(labels.intersection(_label_tokens(text)))
def _matches_label(snapshot: Path, path: Path, label: str) -> bool:
label = label.strip().lower()
if not label:
return False
rel = _rel_lower(snapshot, path)
tokens = [token for token in re.split(r"[^a-z0-9]+", rel) if token]
if label in tokens:
return True
if label in {"train", "test", "validation", "valid", "val", "eval"}:
return False
return label in rel
def dataset_snapshot_from_cache_path(local_path: Optional[str], repo_id: str) -> Optional[Path]:
if not local_path or not repo_id:
return None
try:
root = Path(local_path).expanduser()
if not root.exists():
return None
expected_repo_dir = f"datasets--{repo_id.replace('/', '--')}".lower()
if expected_repo_dir not in {part.lower() for part in root.parts}:
return None
if root.is_dir() and root.parent.name == "snapshots":
return root.resolve()
snapshots = root / "snapshots" if root.is_dir() else None
if snapshots is None or not snapshots.is_dir():
return None
candidates = [p for p in snapshots.iterdir() if p.is_dir()]
if not candidates:
return None
candidates.sort(
key = lambda path: path.stat().st_mtime if path.exists() else 0,
reverse = True,
)
return candidates[0].resolve()
except Exception:
return None
def latest_cached_dataset_snapshot(
repo_id: str, local_path: Optional[str] = None
) -> Optional[Path]:
local_snapshot = dataset_snapshot_from_cache_path(local_path, repo_id)
if local_snapshot is not None:
return local_snapshot
newest: Optional[Path] = None
newest_mtime = -1.0
for entry in iter_repo_cache_dirs("dataset", repo_id):
snapshots = entry / "snapshots"
if not snapshots.is_dir():
continue
try:
candidates = [s for s in snapshots.iterdir() if s.is_dir()]
except OSError:
continue
for snap in candidates:
try:
mtime = snap.stat().st_mtime
except OSError:
continue
if mtime > newest_mtime:
newest = snap
newest_mtime = mtime
return newest
def cached_dataset_candidates(
snapshot: Path,
*,
subset: Optional[str],
train_split: str,
extensions: tuple[str, ...],
preferred_extensions: tuple[str, ...] = TRAINING_DATA_EXTS,
) -> list[Path]:
try:
files = [
p for p in snapshot.rglob("*") if p.is_file() and p.name.lower().endswith(extensions)
]
except OSError:
return []
if not files:
return []
subset_lower = subset.lower() if subset else ""
split_lower = train_split.lower()
def score(path: Path) -> tuple[int, int, str]:
rel = _rel_lower(snapshot, path)
subset_match = bool(subset_lower and _matches_label(snapshot, path, subset_lower))
split_match = bool(split_lower and split_label_matches(rel, split_lower))
location_rank = 3
if split_match and (not subset_lower or subset_match):
location_rank = 0
elif split_match:
location_rank = 1
elif subset_match:
location_rank = 2
return (
0 if path.name.lower().endswith(preferred_extensions) else 1,
location_rank,
rel,
)
return sorted(files, key = score)

749
studio/backend/hub/utils/dataset_format.py Normal file
View file

@ -0,0 +1,749 @@
# 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 re
from typing import Any, Optional
def _first_row(dataset) -> Optional[dict]:
try:
row = next(iter(dataset))
except StopIteration:
return None
return row if isinstance(row, dict) else None
def _column_names(dataset, sample: Optional[dict] = None) -> list[str]:
names = getattr(dataset, "column_names", None)
if names is not None:
return list(names)
return list((sample or {}).keys())
def _keyword_in_column(keyword: str, col_name: str) -> bool:
return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None
def _unknown_dataset_format(
chat_column: Optional[str] = None, sample_keys: Optional[list[str]] = None
) -> dict:
return {
"format": "unknown",
"chat_column": chat_column,
"needs_standardization": None,
"sample_keys": sample_keys or [],
}
def detect_dataset_format(dataset) -> dict:
sample = _first_row(dataset)
if sample is None:
return _unknown_dataset_format()
column_names = set(sample.keys())
if {"instruction", "output"}.issubset(column_names):
return {
"format": "alpaca",
"chat_column": None,
"needs_standardization": False,
"sample_keys": [],
}
chat_column = None
if "messages" in column_names:
chat_column = "messages"
elif "conversations" in column_names:
chat_column = "conversations"
elif "texts" in column_names:
chat_column = "texts"
if not chat_column:
return _unknown_dataset_format()
chat_data = sample.get(chat_column)
if not isinstance(chat_data, (list, tuple)) or not chat_data:
return _unknown_dataset_format(chat_column)
first_msg = chat_data[0]
if not isinstance(first_msg, dict):
return _unknown_dataset_format(chat_column)
msg_keys = set(first_msg.keys())
sample_keys = [str(key) for key in msg_keys]
if "from" in msg_keys or "value" in msg_keys:
return {
"format": "sharegpt",
"chat_column": chat_column,
"needs_standardization": True,
"sample_keys": sample_keys,
}
if "role" in msg_keys and "content" in msg_keys:
return {
"format": "chatml",
"chat_column": chat_column,
"needs_standardization": False,
"sample_keys": sample_keys,
}
return _unknown_dataset_format(chat_column, sample_keys)
def detect_custom_format_heuristic(dataset):
sample = _first_row(dataset)
if sample is None:
return None
all_columns = list(sample.keys())
mapping = {}
assistant_words = [
"output",
"answer",
"response",
"assistant",
"completion",
"expected",
"recommendation",
"reply",
"result",
"target",
"solution",
"explanation",
"solve",
]
user_words_high_priority = [
"input",
"question",
"query",
"prompt",
"instruction",
"request",
"snippet",
"user",
"text",
"problem",
"exercise",
]
user_words_low_priority = ["task"]
user_words = user_words_high_priority + user_words_low_priority
system_words = [
"system",
"context",
"description",
"persona",
"role",
"template",
"task",
]
metadata_exact_match = {
"id",
"idx",
"index",
"key",
"timestamp",
"date",
"metadata",
"source",
"kind",
"type",
"category",
"score",
"label",
"tag",
"inference_mode",
}
metadata_prefix_patterns = [
"problem_type",
"problem_source",
"generation_model",
"pass_rate",
]
priority_patterns = {
"generated": 100,
"gen_": 90,
"model_": 80,
"predicted": 70,
"completion": 60,
}
def has_keyword(col_name, keywords):
col_lower = col_name.lower()
col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "")
return any(keyword in col_lower or keyword in col_normalized for keyword in keywords)
def is_metadata(col_name):
col_lower = col_name.lower()
if col_lower in metadata_exact_match or col_lower in metadata_prefix_patterns:
return True
for pattern in metadata_prefix_patterns:
if col_lower.startswith(pattern.split("_")[0] + "_") and col_lower != pattern:
if "_" in col_lower:
prefix = col_lower.split("_")[0]
if prefix in ["generation", "pass", "inference"]:
return True
return len(col_lower) <= 2 and col_lower not in ["qa", "q", "a"]
def get_priority_score(col_name):
col_lower = col_name.lower()
return sum(score for pattern, score in priority_patterns.items() if pattern in col_lower)
def get_content_length(col_name):
try:
return len(str(sample[col_name])) if sample.get(col_name) else 0
except Exception:
return 0
def score_column(col_name, keywords, role_type, num_candidates):
if not has_keyword(col_name, keywords):
return 0
score = 10
if role_type == "user":
col_lower = col_name.lower()
if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
score -= 15
score += get_priority_score(col_name)
if role_type in ["assistant", "user"]:
avg_length = get_content_length(col_name)
if num_candidates > 1:
if avg_length > 1000:
score += 50
elif avg_length > 200:
score += 30
elif avg_length > 50:
score += 10
elif avg_length < 50:
score -= 20
else:
if avg_length > 1000:
score += 50
elif avg_length > 200:
score += 30
elif avg_length > 50:
score += 10
return score
content_columns = [col for col in all_columns if not is_metadata(col)]
assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)]
user_potential = [col for col in content_columns if has_keyword(col, user_words)]
assistant_candidates = [
(col, score)
for col in assistant_potential
if (score := score_column(col, assistant_words, "assistant", len(assistant_potential))) > 0
]
if assistant_candidates:
assistant_candidates.sort(key = lambda item: item[1], reverse = True)
assistant_col = assistant_candidates[0][0]
mapping[assistant_col] = "assistant"
else:
assistant_col = None
user_candidates = []
for col in user_potential:
if col == assistant_col:
continue
score = score_column(col, user_words, "user", len(user_potential))
if score > 0:
user_candidates.append((col, score))
if user_candidates:
user_candidates.sort(key = lambda item: item[1], reverse = True)
user_col = user_candidates[0][0]
mapping[user_col] = "user"
else:
user_col = None
remaining_columns = [col for col in content_columns if col not in mapping]
system_col = None
for col in remaining_columns:
if has_keyword(col, system_words):
mapping[col] = "system"
system_col = col
break
if system_col:
remaining_columns = [col for col in remaining_columns if col != system_col]
if remaining_columns:
remaining_col = remaining_columns[0]
if not has_keyword(remaining_col, user_words + assistant_words):
mapping[remaining_col] = "system"
elif user_col is None:
mapping[remaining_col] = "user"
else:
mapping[remaining_col] = "system"
has_user = any(role == "user" for role in mapping.values())
has_assistant = any(role == "assistant" for role in mapping.values())
if not has_user:
for col in remaining_columns:
if col not in mapping:
mapping[col] = "user"
has_user = True
break
return mapping if has_user and has_assistant else None
_AUDIO_EXTENSIONS = (
".wav",
".mp3",
".flac",
".ogg",
".opus",
".m4a",
".aac",
".wma",
".webm",
)
def _is_audio_value(value) -> bool:
if value is None:
return False
if isinstance(value, dict):
if "array" in value and "sampling_rate" in value:
return True
if "bytes" in value or "path" in value:
path = value.get("path") or ""
return isinstance(path, str) and any(
path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
)
return False
def _has_image_header(data: bytes) -> bool:
if len(data) < 4:
return False
return (
data[:2] == b"\xff\xd8"
or data[:4] == b"\x89PNG"
or data[:3] == b"GIF"
or (data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP")
or data[:2] == b"BM"
)
def _is_image_value(value) -> bool:
if value is None:
return False
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
return True
except ImportError:
pass
if isinstance(value, dict):
if "array" in value and "sampling_rate" in value:
return False
if "bytes" in value and "path" in value:
path = value.get("path") or ""
if isinstance(path, str) and any(
path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
):
return False
return True
if isinstance(value, (bytes, bytearray)):
return _has_image_header(value)
image_exts = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg")
if isinstance(value, str) and len(value) < 1000:
lower = value.strip().lower()
if lower.startswith(("http://", "https://")):
return any(lower.split("?")[0].endswith(ext) for ext in image_exts)
return any(lower.endswith(ext) for ext in image_exts)
return False
def detect_multimodal_dataset(dataset):
sample = _first_row(dataset)
if sample is None:
return {
"is_image": False,
"multimodal_columns": [],
"modality_types": [],
"is_audio": False,
"audio_columns": [],
"detected_audio_column": None,
"detected_text_column": None,
"detected_speaker_column": None,
}
column_names = list(sample.keys())
image_keywords = [
"image",
"img",
"pixel",
"jpg",
"jpeg",
"png",
"webp",
"bmp",
"gif",
"tiff",
"svg",
"photo",
"pic",
"picture",
"visual",
"file_name",
"filename",
]
audio_keywords = ["audio", "speech", "wav", "waveform", "sound"]
multimodal_columns = []
audio_columns = []
modality_types = set()
for col_name in column_names:
if any(_keyword_in_column(keyword, col_name) for keyword in image_keywords):
multimodal_columns.append(col_name)
modality_types.add("image")
for col_name in column_names:
if col_name not in multimodal_columns and _is_image_value(sample[col_name]):
multimodal_columns.append(col_name)
modality_types.add("image")
for col_name in column_names:
if any(_keyword_in_column(keyword, col_name) for keyword in audio_keywords):
audio_columns.append(col_name)
modality_types.add("audio")
for col_name in column_names:
if col_name not in audio_columns and _is_audio_value(sample[col_name]):
audio_columns.append(col_name)
modality_types.add("audio")
if audio_columns:
multimodal_columns = [col for col in multimodal_columns if col not in set(audio_columns)]
detected_text_col = None
if audio_columns:
for col_name in column_names:
if col_name.lower() in [
"text",
"sentence",
"transcript",
"transcription",
"label",
]:
detected_text_col = col_name
break
detected_speaker_col = None
if audio_columns:
for col_name in column_names:
if col_name.lower() in ["source", "speaker", "speaker_id"]:
detected_speaker_col = col_name
break
return {
"is_image": len(multimodal_columns) > 0,
"multimodal_columns": multimodal_columns,
"modality_types": list(modality_types),
"is_audio": len(audio_columns) > 0,
"audio_columns": audio_columns,
"detected_audio_column": audio_columns[0] if audio_columns else None,
"detected_text_column": detected_text_col,
"detected_speaker_column": detected_speaker_col,
}
def detect_vlm_dataset_structure(dataset):
sample = _first_row(dataset)
if sample is None:
return {
"format": "unknown",
"needs_conversion": None,
"image_column": None,
"text_column": None,
"messages_column": None,
}
column_names = set(sample.keys())
if "messages" in column_names:
messages = sample["messages"]
if messages and len(messages) > 0:
first_msg = messages[0]
if "content" in first_msg:
content = first_msg["content"]
if (
isinstance(content, list)
and content
and isinstance(content[0], dict)
and "type" in content[0]
):
has_index = any("index" in item for item in content if isinstance(item, dict))
if has_index and "images" in column_names:
return {
"format": "vlm_messages_llava",
"needs_conversion": True,
"messages_column": "messages",
"image_column": "images",
"text_column": None,
}
has_image = any("image" in item for item in content if isinstance(item, dict))
if has_image:
return {
"format": "vlm_messages",
"needs_conversion": False,
"messages_column": "messages",
"image_column": None,
"text_column": None,
}
for chat_col in ("conversations", "messages"):
if chat_col not in column_names:
continue
chat_data = sample[chat_col]
if not isinstance(chat_data, list) or not chat_data:
continue
has_image_placeholder = any(
"<image>" in str(message.get("value", "") or message.get("content", ""))
for message in chat_data
if isinstance(message, dict)
)
if not has_image_placeholder:
continue
image_col = next(
(
col
for col in column_names
if col != chat_col
and (_keyword_in_column("image", col) or _keyword_in_column("img", col))
),
None,
)
if image_col:
return {
"format": "sharegpt_with_images",
"needs_conversion": True,
"image_column": image_col,
"text_column": None,
"messages_column": chat_col,
}
metadata_suffixes = (
"_id",
"_url",
"_name",
"_filename",
"_uri",
"_link",
"_key",
"_index",
)
metadata_prefixes = (
"id_",
"url_",
"name_",
"filename_",
"uri_",
"link_",
"key_",
"index_",
)
image_keywords = [
"image",
"img",
"photo",
"picture",
"pic",
"visual",
"scan",
"file_name",
"filename",
]
text_keywords = [
"text",
"caption",
"captions",
"description",
"answer",
"output",
"response",
"label",
]
def is_metadata_column(col_name):
lower = col_name.lower()
return any(lower.endswith(suffix) for suffix in metadata_suffixes) or any(
lower.startswith(prefix) for prefix in metadata_prefixes
)
image_candidates = []
for col in column_names:
value = sample[col]
if any(_keyword_in_column(keyword, col) for keyword in image_keywords) or _is_image_value(
value
):
if hasattr(value, "size") and hasattr(value, "mode"):
score = 100
elif isinstance(value, dict) and ("bytes" in value or "path" in value):
score = 75
elif isinstance(value, str):
score = (
55
if is_metadata_column(col)
else 70
if value.startswith(("http://", "https://"))
else 50
)
else:
score = 0
if score > 0:
image_candidates.append((col, score))
image_candidates.sort(key = lambda item: item[1], reverse = True)
text_candidates = []
for col in column_names:
if is_metadata_column(col) or not any(
_keyword_in_column(keyword, col) for keyword in text_keywords
):
continue
value = sample[col]
if isinstance(value, str) and value:
text_candidates.append((col, min(len(value), 1000)))
elif isinstance(value, list) and value and isinstance(value[0], str):
text_candidates.append((col, min(len(value[0]), 1000) // 2))
text_candidates.sort(key = lambda item: item[1], reverse = True)
found_image = image_candidates[0][0] if image_candidates else None
found_text = text_candidates[0][0] if text_candidates else None
if found_image and found_text:
return {
"format": "simple_image_text",
"needs_conversion": True,
"image_column": found_image,
"text_column": found_text,
"messages_column": None,
}
return {
"format": "unknown",
"needs_conversion": None,
"image_column": found_image,
"text_column": found_text,
"messages_column": None,
}
def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
sample = _first_row(dataset)
columns = _column_names(dataset, sample)
multimodal_info = detect_multimodal_dataset(dataset)
is_audio = multimodal_info.get("is_audio", False)
audio_fields = {
"is_audio": is_audio,
"detected_audio_column": multimodal_info.get("detected_audio_column"),
"detected_speaker_column": multimodal_info.get("detected_speaker_column"),
}
if is_vlm:
vlm_structure = detect_vlm_dataset_structure(dataset)
requires_mapping = vlm_structure["format"] == "unknown"
warning = None
if requires_mapping:
missing = []
if not vlm_structure.get("image_column"):
missing.append("image")
if not vlm_structure.get("text_column"):
missing.append("text")
if missing:
warning = (
f"Could not auto-detect {' or '.join(missing)} column. "
"Please assign image and text columns manually."
)
return {
"requires_manual_mapping": requires_mapping,
"detected_format": vlm_structure["format"],
"columns": columns,
"suggested_mapping": None,
"detected_image_column": vlm_structure.get("image_column"),
"detected_text_column": vlm_structure.get("text_column"),
"is_image": multimodal_info["is_image"],
"multimodal_columns": multimodal_info.get("multimodal_columns"),
"warning": warning,
**audio_fields,
}
if is_audio:
detected_audio = multimodal_info.get("detected_audio_column")
detected_text = multimodal_info.get("detected_text_column")
return {
"requires_manual_mapping": not detected_audio or not detected_text,
"detected_format": "audio",
"columns": columns,
"suggested_mapping": None,
"detected_image_column": None,
"detected_text_column": detected_text,
"is_image": False,
"multimodal_columns": multimodal_info.get("audio_columns"),
**audio_fields,
}
detected = detect_dataset_format(dataset)
if detected["format"] == "unknown":
heuristic_mapping = detect_custom_format_heuristic(dataset)
if heuristic_mapping:
return {
"requires_manual_mapping": False,
"detected_format": "custom_heuristic",
"columns": columns,
"suggested_mapping": heuristic_mapping,
"detected_image_column": None,
"detected_text_column": None,
"is_image": multimodal_info["is_image"],
"multimodal_columns": multimodal_info.get("multimodal_columns"),
**audio_fields,
}
return {
"requires_manual_mapping": True,
"detected_format": "unknown",
"columns": columns,
"suggested_mapping": None,
"detected_image_column": None,
"detected_text_column": None,
"is_image": multimodal_info["is_image"],
"multimodal_columns": multimodal_info.get("multimodal_columns"),
"warning": (
f"Could not auto-detect column roles for columns: {columns}. "
"Please assign roles manually, or use AI Assist."
),
**audio_fields,
}
return {
"requires_manual_mapping": False,
"detected_format": detected["format"],
"columns": columns,
"suggested_mapping": None,
"detected_image_column": None,
"detected_text_column": None,
"is_image": multimodal_info["is_image"],
"multimodal_columns": multimodal_info.get("multimodal_columns"),
**audio_fields,
}
_ROLE_MAP = {
"human": "user",
"user": "user",
"input": "user",
"gpt": "assistant",
"assistant": "assistant",
"output": "assistant",
"system": "system",
}
def _standardize_sharegpt_row(row: dict[str, Any], chat_column: str) -> dict[str, Any]:
chat_data = row.get(chat_column)
if not isinstance(chat_data, list):
return row
messages = []
for message in chat_data:
if not isinstance(message, dict):
continue
role = message.get("role") or message.get("from")
content = message.get("content") if "content" in message else message.get("value")
messages.append(
{
"role": _ROLE_MAP.get(str(role), str(role or "user")),
"content": "" if content is None else content,
}
)
return {chat_column: messages}
def format_dataset_preview(dataset):
detected = detect_dataset_format(dataset)
if detected.get("format") != "sharegpt":
return dataset
chat_column = detected.get("chat_column")
if not isinstance(chat_column, str):
return dataset
if hasattr(dataset, "map"):
return dataset.map(lambda row: _standardize_sharegpt_row(row, chat_column))
return dataset

487
studio/backend/hub/utils/download_manifest.py Normal file
View file

@ -0,0 +1,487 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hub download manifest + cancel-marker primitives.
Manifests record what a download was supposed to fetch (path + declared
size per expected file). Consumed by:
- the worker post-download, to verify on-disk sizes match what HF
declared, so a resume that no-ops doesn't get classified as success;
- the inventory scanner, to mark a row partial when expected files
are absent or undersized, so a half-finished GGUF/dataset doesn't
masquerade as a complete on-device row.
Cancel markers record that a user-initiated cancel landed for a
(repo_type, repo_id, variant) triple. *Existence* is the signal the
scanner reads; the body carries debuggability metadata. Markers are
cleared at the start of a new download attempt (supersedes prior cancel)
and on successful completion (defensive, in case the start clear failed).
I/O contracts:
- Writes are atomic via ``tmp + os.replace``: a SIGKILL mid-write
cannot leave a half-written file readable to the next reader.
- Manifest reads fail *open*: missing/corrupt/schema-mismatched
manifests return ``None`` and the scanner falls through to the
legacy on-disk-only check (matches HF-cache imports and pre-fix
downloads that never wrote a manifest).
- Cancel-marker reads fail *closed*: file existence is the signal
regardless of body parseability, so a corrupt marker still
suppresses the "on device" classification.
"""
from __future__ import annotations
import json
import os
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator, Optional, Sequence
from loggers import get_logger
from hub.utils.state_dir import (
RepoType,
cancelled_dir,
manifest_path,
manifests_dir,
marker_path,
variant_filename_prefix,
)
logger = get_logger(__name__)
_MANIFEST_VERSION = 1
_MARKER_VERSION = 2
_LEGACY_MARKER_VERSION = 1
# Verbatim phrase the worker emits on a degraded completion and the download
# lifecycle escalates to a warning log. Shared so the emit and match stay coupled.
MANIFEST_DEGRADED_MARKER = "completed without a manifest so partial detection is degraded"
@dataclass(frozen = True)
class ExpectedFile:
path: str
size: int
sha256: Optional[str] = None
@dataclass(frozen = True)
class Manifest:
repo_type: RepoType
repo_id: str
variant: Optional[str]
started_at: str
expected_files: tuple[ExpectedFile, ...]
transport: Optional[str] = None
@dataclass(frozen = True)
class VerifyResult:
ok: bool
missing: tuple[str, ...]
size_mismatched: tuple[str, ...]
def _atomic_write_json(path: Path, payload: dict) -> bool:
# Per-write uuid suffix so a concurrent caller or a stale tmp from a
# previous crash cannot collide with the in-flight write.
tmp = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex[:8]}")
try:
with tmp.open("w", encoding = "utf-8") as handle:
handle.write(json.dumps(payload, indent = 2))
handle.flush()
os.fsync(handle.fileno())
os.replace(tmp, path)
except OSError as exc:
logger.debug("Atomic write failed for %s: %s", path, exc)
try:
tmp.unlink(missing_ok = True)
except OSError:
pass
return False
if os.name != "nt":
try:
flags = os.O_RDONLY
if hasattr(os, "O_DIRECTORY"):
flags |= os.O_DIRECTORY
parent_fd = os.open(path.parent, flags)
try:
os.fsync(parent_fd)
finally:
os.close(parent_fd)
except OSError as exc:
logger.debug("Parent dir fsync failed for %s: %s", path, exc)
return True
def write_manifest(
repo_type: RepoType,
repo_id: str,
variant: Optional[str],
expected_files: Sequence[ExpectedFile],
transport: Optional[str] = None,
) -> bool:
"""Write/overwrite the manifest for this triple. Best-effort.
``False`` on write failure must not be treated as fatal: the
worst-case fallback is the pre-fix scanner behavior (one missed
partial detection), which is no regression.
"""
path = manifest_path(repo_type, repo_id, variant)
if path is None:
return False
payload = {
"version": _MANIFEST_VERSION,
"repo_type": repo_type,
"repo_id": repo_id,
"variant": variant,
"started_at": datetime.now(timezone.utc).isoformat(),
"expected_files": [
{
"path": f.path,
"size": int(f.size),
**({"sha256": f.sha256} if f.sha256 else {}),
}
for f in expected_files
],
"transport": transport,
}
return _atomic_write_json(path, payload)
def read_manifest(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
) -> Optional[Manifest]:
"""Return the manifest if present and parseable; ``None`` otherwise.
Treats missing-file, parse-error, and any schema mismatch all as
``None`` (fail-open). Scanner callers fall through to on-disk-only
behavior on ``None`` so this never regresses legacy/imported repos
that have no manifest.
Forward-compat: accepts only ``version == 1``; an unknown version is
treated as no manifest. A future v2 schema MUST either keep v1's
``expected_files`` shape on the same filename (bump
``_MANIFEST_VERSION`` and widen this check) or live under a different
filename, so an incompatible payload can never mis-classify rows.
"""
path = manifest_path(repo_type, repo_id, variant)
if path is None or not path.is_file():
return None
try:
data = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError) as exc:
logger.debug("Could not read manifest %s: %s", path, exc)
return None
if not isinstance(data, dict):
return None
if data.get("version") != _MANIFEST_VERSION:
logger.debug(
"Manifest %s has unknown version %r; ignoring.",
path,
data.get("version"),
)
return None
raw_files = data.get("expected_files")
if not isinstance(raw_files, list):
return None
expected: list[ExpectedFile] = []
for item in raw_files:
if not isinstance(item, dict):
return None
file_path = item.get("path")
size = item.get("size")
if not isinstance(file_path, str) or not isinstance(size, int):
return None
sha256 = item.get("sha256")
expected.append(
ExpectedFile(
path = file_path,
size = size,
sha256 = sha256 if isinstance(sha256, str) and sha256 else None,
)
)
raw_variant = data.get("variant")
transport = data.get("transport")
return Manifest(
repo_type = repo_type,
repo_id = str(data.get("repo_id", repo_id)),
variant = raw_variant if raw_variant else None,
started_at = str(data.get("started_at", "")),
expected_files = tuple(expected),
transport = transport if transport in ("http", "xet") else None,
)
def verify_against_disk(manifest: Manifest, snapshot_dir: Path) -> VerifyResult:
"""Check every expected file is present in *snapshot_dir* at its declared size.
Presence + size only, not content integrity: it converts a
no-op-on-cached ``snapshot_download`` into a clear error when shards are
missing or truncated, and marks a scanner row partial when expected bytes
aren't on disk. Byte-level integrity is already covered upstream by
``huggingface_hub`` (size check on HTTP, content-addressed chunk hashes on
XET), so re-hashing finalized multi-GB weights here would only duplicate
that at a large cost. ``Path.stat()`` follows symlinks, so HF's symlink and
Windows copy cache layouts both verify correctly.
"""
missing: list[str] = []
mismatched: list[str] = []
for expected in manifest.expected_files:
target = snapshot_dir / expected.path
try:
actual_size = target.stat().st_size
except OSError:
missing.append(expected.path)
continue
# expected.size == 0 means HF metadata had no declared size: verify
# existence only rather than flagging every such file as mismatched.
if expected.size > 0 and actual_size != expected.size:
mismatched.append(expected.path)
return VerifyResult(
ok = not missing and not mismatched,
missing = tuple(missing),
size_mismatched = tuple(mismatched),
)
def expected_files_from_snapshot_dir(snapshot_dir: Path) -> list[ExpectedFile]:
"""Derive expected-file entries from a completed snapshot directory.
Last-resort manifest source for when HF metadata was unreachable for the
whole download. ``snapshot_download`` has already exited cleanly, so every
regular file is a finished, correctly-sized blob; recording them keeps the
scanner's completion check in agreement with the worker's exit-0 success
instead of leaving a finished repo perpetually partial. ``stat()`` follows
HF's symlink layout and Windows copies, so the recorded sizes match what
``verify_against_disk`` later reads.
"""
out: list[ExpectedFile] = []
try:
entries = sorted(snapshot_dir.rglob("*"))
except OSError:
return out
for path in entries:
try:
if not path.is_file():
continue
relative = path.relative_to(snapshot_dir).as_posix()
out.append(
ExpectedFile(
path = relative,
size = path.stat().st_size,
sha256 = None,
)
)
except OSError:
continue
return out
def write_cancel_marker(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
transport: Optional[str] = None,
) -> bool:
"""Record that this triple was cancelled. Idempotent across repeated cancels.
``transport`` ("http"/"xet") is surfaced via partial_transport on
inventory rows so the UI labels HTTP retries as continuable and XET
retries as full redownloads. None is accepted for forward-compat.
"""
path = marker_path(repo_type, repo_id, variant)
if path is None:
return False
payload = {
"version": _MARKER_VERSION,
"repo_type": repo_type,
"repo_id": repo_id,
"variant": variant,
"transport": transport,
"cancelled_at": datetime.now(timezone.utc).isoformat(),
}
return _atomic_write_json(path, payload)
def read_cancel_marker_transport(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
) -> Optional[str]:
"""Return the transport recorded in the cancel marker, or ``None`` if no
marker exists or it is unreadable.
Cases:
* No marker on disk → ``None``.
* Legacy v1 marker → ``"http"``: v1 markers were only written by the
HTTP path, so the transport is unambiguous despite the absent field.
* v2 marker with a valid ``"http"`` / ``"xet"`` transport → that value.
* Corrupt, non-dict, or v2-with-missing-transport marker → ``None``.
Defaulting these to ``"http"`` misled the UI into showing a
byte-resume "Continue" label for what may have been an XET cancel;
``None`` keeps the neutral "Retry" label.
* Unknown future versions → ``None`` (unknown layout, unknown transport).
"""
path = marker_path(repo_type, repo_id, variant)
if path is None or not path.is_file():
return None
try:
data = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError) as exc:
logger.debug("Could not read cancel marker %s: %s", path, exc)
return None
if not isinstance(data, dict):
return None
version = data.get("version")
if version == _LEGACY_MARKER_VERSION:
return "http"
if version != _MARKER_VERSION:
return None
transport = data.get("transport")
if isinstance(transport, str) and transport in ("http", "xet"):
return transport
return None
def clear_cancel_marker(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
) -> None:
"""Remove the cancel marker for this triple if present.
Idempotent: a missing marker is not an error. Called at
download-start (a fresh attempt supersedes prior cancel state) and
again at successful completion (cleans up if the start clear failed).
"""
path = marker_path(repo_type, repo_id, variant)
if path is None:
return
try:
path.unlink(missing_ok = True)
except OSError as exc:
logger.debug("Could not clear cancel marker %s: %s", path, exc)
def has_cancel_marker(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
) -> bool:
"""File-existence check only. Body is never read.
Fail-closed: a corrupt marker still returns ``True`` because the
file's existence is the signal (the user once cancelled this
triple, even if the body is unreadable).
"""
path = marker_path(repo_type, repo_id, variant)
if path is None:
return False
try:
return path.is_file()
except OSError:
return False
def delete_manifest(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
) -> bool:
path = manifest_path(repo_type, repo_id, variant)
if path is None:
return False
try:
if not path.is_file():
return False
path.unlink()
return True
except OSError as exc:
logger.debug("Could not delete manifest %s: %s", path, exc)
return False
def purge_state(
repo_type: RepoType,
repo_id: str,
variant: Optional[str] = None,
) -> bool:
"""Remove manifest + cancel marker for this triple. Returns ``True``
when anything was present on disk before the call. Idempotent."""
marker_existed = has_cancel_marker(repo_type, repo_id, variant)
manifest_removed = delete_manifest(repo_type, repo_id, variant)
clear_cancel_marker(repo_type, repo_id, variant)
return marker_existed or manifest_removed
def purge_all_state_for_repo(repo_type: RepoType, repo_id: str) -> int:
"""Remove the snapshot-level manifest + marker AND every variant-keyed
manifest + marker for this repo. Used by the route delete handlers so
scanner state never outlives the cache it described. Returns the count
of (repo, variant) triples that had any state on disk."""
removed = 0
if purge_state(repo_type, repo_id, None):
removed += 1
variants: set[str] = set()
for variant, _ in iter_variant_manifests(repo_type, repo_id):
variants.add(variant)
for variant, _ in iter_variant_markers(repo_type, repo_id):
variants.add(variant)
for variant in variants:
if purge_state(repo_type, repo_id, variant):
removed += 1
return removed
def _variant_from_state_file(path: Path, fallback: str) -> str:
try:
data = json.loads(path.read_text(encoding = "utf-8"))
except (OSError, ValueError):
return fallback
if not isinstance(data, dict):
return fallback
variant = data.get("variant")
return variant if isinstance(variant, str) and variant else fallback
def _iter_variant_state_files(
parent: Optional[Path], repo_type: RepoType, repo_id: str
) -> Iterator[tuple[str, Path]]:
if parent is None:
return
prefix = variant_filename_prefix(repo_type, repo_id)
try:
entries = list(parent.iterdir())
except OSError:
return
for entry in entries:
if not entry.is_file() or not entry.name.endswith(".json"):
continue
stem = entry.name[: -len(".json")]
if not stem.lower().startswith(prefix):
continue
variant = stem[len(prefix) :]
if variant:
yield _variant_from_state_file(entry, variant), entry
def iter_variant_manifests(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]:
"""Yield (variant, manifest_path) for every variant-keyed manifest
written for this repo. Used by is_gguf_repo_partial to enumerate all
variants present on disk so the all-variants-broken gate can run."""
yield from _iter_variant_state_files(manifests_dir(), repo_type, repo_id)
def iter_variant_markers(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]:
"""Yield (variant, marker_path) for every variant-keyed cancel marker.
Companion to iter_variant_manifests: catches variants cancelled
before download-start ever wrote a manifest (very early failures)."""
yield from _iter_variant_state_files(cancelled_dir(), repo_type, repo_id)

1263
studio/backend/hub/utils/download_registry.py Normal file
View file

File diff suppressed because it is too large Load diff

406
studio/backend/hub/utils/gguf.py Normal file
View file

@ -0,0 +1,406 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GGUF filename helpers. Quantization variants are derived from filenames, not parsed from binary GGUF headers."""
from __future__ import annotations
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from loggers import get_logger
logger = get_logger(__name__)
_GGUF_MODEL_INFO_TIMEOUT_SECONDS = 5.0
@dataclass
class GgufVariantInfo:
filename: str
quant: str
size_bytes: int
display_label: Optional[str] = None
download_size_bytes: int = 0
GGUF_QUANT_PREFERENCE = [
"UD-Q4_K_XL",
"UD-Q4_K_L",
"UD-Q5_K_XL",
"UD-Q3_K_XL",
"UD-Q6_K_XL",
"UD-Q6_K_S",
"UD-Q8_K_XL",
"UD-Q2_K_XL",
"UD-IQ4_NL",
"UD-IQ4_XS",
"UD-IQ3_S",
"UD-IQ3_XXS",
"UD-IQ2_M",
"UD-IQ2_XXS",
"UD-IQ1_M",
"UD-IQ1_S",
"Q4_K_M",
"Q4_K_S",
"Q5_K_M",
"Q5_K_S",
"Q6_K",
"Q8_0",
"Q3_K_M",
"Q3_K_L",
"Q3_K_S",
"Q2_K",
"Q2_K_L",
"IQ4_NL",
"IQ4_XS",
"IQ3_M",
"IQ3_XXS",
"IQ2_M",
"IQ1_M",
"F16",
"BF16",
"F32",
]
_GGUF_SPLIT_SUFFIX_RE = re.compile(r"-\d{3,}-of-\d{3,}", re.IGNORECASE)
_GGUF_QUANT_RE = re.compile(
r"(UD-)?"
r"(MXFP[0-9]+(?:_[A-Z0-9]+)*"
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
r"|TQ[0-9]+_[0-9]+"
r"|Q[0-9]+_K_[A-Z]+"
r"|Q[0-9]+_[0-9]+"
r"|Q[0-9]+_K"
r"|BF16|F16|F32)",
re.IGNORECASE,
)
def is_mmproj_filename(filename: str) -> bool:
return "mmproj" in filename.lower()
def is_gguf_filename(filename: str) -> bool:
return filename.lower().endswith(".gguf")
# Cap recursive walks so a huge or system path cannot run unbounded.
_MAX_LOCAL_SCAN_ENTRIES = 100_000
def iter_gguf_files(directory: Path, recursive: bool = False):
if not directory.is_dir():
return
if recursive:
seen = 0
# os.walk skips unreadable subdirs instead of raising (e.g. /proc).
for dirpath, dirnames, filenames in os.walk(directory, onerror = lambda _e: None):
for name in filenames:
if is_gguf_filename(name):
yield Path(dirpath) / name
seen += len(dirnames) + len(filenames)
if seen > _MAX_LOCAL_SCAN_ENTRIES:
return
return
try:
entries = list(directory.iterdir())
except OSError:
return
for file in entries:
try:
if file.is_file() and is_gguf_filename(file.name):
yield file
except OSError:
continue
def pick_best_gguf(filenames: list[str]) -> Optional[str]:
gguf_files = [
name for name in filenames if is_gguf_filename(name) and not is_mmproj_filename(name)
]
if not gguf_files:
return None
by_quant: dict[str, str] = {}
for name in gguf_files:
by_quant.setdefault(extract_quant_label(name).upper(), name)
for quant in GGUF_QUANT_PREFERENCE:
filename = by_quant.get(quant.upper())
if filename is not None:
return filename
return gguf_files[0]
def _gguf_stem(filename: str) -> str:
basename = filename.rsplit("/", 1)[-1]
return _GGUF_SPLIT_SUFFIX_RE.sub("", basename.rsplit(".", 1)[0]).strip()
_FLOAT_PRECISION_QUANTS = frozenset({"BF16", "F16", "F32"})
def _select_quant_match(text: str) -> Optional[re.Match]:
fallback: Optional[re.Match] = None
for match in _GGUF_QUANT_RE.finditer(text):
if match.group(2).upper() in _FLOAT_PRECISION_QUANTS:
if fallback is None:
fallback = match
continue
return match
return fallback
def extract_quant_token(filename: str) -> Optional[str]:
stem = _gguf_stem(filename)
match = _select_quant_match(stem)
if not match and "/" in filename:
parents = filename.rsplit("/", 1)[0]
for segment in reversed(parents.split("/")):
parent_match = _select_quant_match(segment)
if parent_match:
match = parent_match
break
if match:
prefix = match.group(1) or ""
return f"{prefix}{match.group(2)}"
return None
def _unknown_gguf_variant_key(filename: str) -> str:
stem = _gguf_stem(filename)
if "/" not in filename:
return stem or "gguf"
parents = filename.rsplit("/", 1)[0].strip("/")
return f"{parents}/{stem}" if parents and stem else stem or "gguf"
def extract_quant_label(filename: str) -> str:
return extract_quant_token(filename) or _unknown_gguf_variant_key(filename)
def _apply_gguf_display_labels(variants: list[GgufVariantInfo]) -> None:
unknown_variants = [
variant for variant in variants if extract_quant_token(variant.filename) is None
]
if not unknown_variants:
return
ambiguous = len(unknown_variants) > 1
for variant in unknown_variants:
variant.display_label = f"GGUF · {variant.filename}" if ambiguous else "GGUF"
def _env_offline() -> bool:
return os.environ.get("HF_HUB_OFFLINE", "").lower() in (
"1",
"true",
"yes",
) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes")
def iter_hf_cache_snapshots(repo_id: str):
from hub.utils.hf_cache_state import iter_repo_cache_dirs
snapshots: list[Path] = []
for repo_dir in iter_repo_cache_dirs("model", repo_id):
snapshots_dir = repo_dir / "snapshots"
if not snapshots_dir.is_dir():
continue
try:
snapshots.extend(snap for snap in snapshots_dir.iterdir() if snap.is_dir())
except OSError:
continue
def _mtime(path: Path) -> float:
try:
return path.stat().st_mtime
except OSError:
return 0.0
snapshots.sort(key = _mtime, reverse = True)
yield from snapshots
def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]:
for snapshot in iter_hf_cache_snapshots(repo_id):
variants, has_vision = list_local_gguf_variants(str(snapshot))
if variants or has_vision:
return variants, has_vision
return None
def list_partial_gguf_variants_from_state(
repo_id: str,
) -> Optional[tuple[list[GgufVariantInfo], bool]]:
"""Reconstruct GGUF variants from download manifests/markers alone.
Used when no completed snapshot exists (download cancelled or interrupted)
and the HF API is unreachable (offline/gated/private). Each variant's
``quant`` is the stored variant key so a resume passes the matching
``--variant`` back to the worker.
"""
from hub.utils import download_manifest
# Variant identity on disk is case-insensitive (_entry_key lowercases it), so
# dedupe on the lowercased key. Manifests are read first to keep their
# original-casing label over a lowercased cancel marker for the same variant.
seen: set[str] = set()
ordered: list[str] = []
for source in (
download_manifest.iter_variant_manifests("model", repo_id),
download_manifest.iter_variant_markers("model", repo_id),
):
for variant, _path in source:
key = variant.lower()
if key not in seen:
seen.add(key)
ordered.append(variant)
if not ordered:
return None
variants: list[GgufVariantInfo] = []
has_vision = False
for variant in ordered:
manifest = download_manifest.read_manifest("model", repo_id, variant)
main_filename: Optional[str] = None
size_bytes = 0
companion_bytes = 0
if manifest is not None:
for expected in manifest.expected_files:
if not is_gguf_filename(expected.path):
continue
if is_mmproj_filename(expected.path):
has_vision = True
companion_bytes += max(0, int(expected.size or 0))
continue
if main_filename is None:
main_filename = expected.path
size_bytes += max(0, int(expected.size or 0))
if main_filename is None:
main_filename = f"{variant}.gguf"
variants.append(
GgufVariantInfo(
filename = main_filename,
quant = variant,
size_bytes = size_bytes,
download_size_bytes = size_bytes + companion_bytes,
)
)
variants.sort(key = lambda variant: -variant.size_bytes)
_apply_gguf_display_labels(variants)
return variants, has_vision
def list_gguf_variants(
repo_id: str, hf_token: Optional[str] = None
) -> tuple[list[GgufVariantInfo], bool, Optional[list]]:
from huggingface_hub import HfApi
if _env_offline():
cached = list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
return (*cached, None)
try:
info = HfApi(token = hf_token).model_info(
repo_id,
files_metadata = True,
timeout = _GGUF_MODEL_INFO_TIMEOUT_SECONDS,
)
except Exception as exc:
if type(exc).__name__ in (
"RepositoryNotFoundError",
"GatedRepoError",
"RevisionNotFoundError",
"EntryNotFoundError",
):
raise
cached = list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
logger.warning(
"HF API unreachable for %s (%s); using local cache snapshot.",
repo_id,
exc.__class__.__name__,
)
return (*cached, None)
raise
variants: list[GgufVariantInfo] = []
has_vision = False
quant_totals: dict[str, int] = {}
quant_first_file: dict[str, str] = {}
for sibling in info.siblings:
filename = getattr(sibling, "rfilename", None)
if not isinstance(filename, str) or not is_gguf_filename(filename):
continue
if is_mmproj_filename(filename):
has_vision = True
continue
quant = extract_quant_label(filename)
quant_totals[quant] = quant_totals.get(quant, 0) + int(getattr(sibling, "size", 0) or 0)
quant_first_file.setdefault(quant, filename)
for quant, total_size in quant_totals.items():
variants.append(
GgufVariantInfo(
filename = quant_first_file[quant],
quant = quant,
size_bytes = total_size,
)
)
variants.sort(key = lambda variant: -variant.size_bytes)
_apply_gguf_display_labels(variants)
return variants, has_vision, list(info.siblings)
def _resolve_gguf_dir(path: Path) -> Optional[Path]:
if path.is_dir():
return path
if path.is_file() and path.suffix.lower() == ".gguf":
parent = path.parent
if (
(parent / "config.json").exists()
or (parent / "adapter_config.json").exists()
or (parent / "export_metadata.json").exists()
):
return parent
return None
def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], bool]:
root = _resolve_gguf_dir(Path(directory))
if root is None:
return [], False
quant_totals: dict[str, int] = {}
quant_first_file: dict[str, str] = {}
has_vision = False
for file in sorted(iter_gguf_files(root, recursive = True)):
if is_mmproj_filename(file.name):
has_vision = True
continue
try:
size = file.stat().st_size
except OSError:
size = 0
rel = file.relative_to(root).as_posix()
quant = extract_quant_label(rel)
quant_totals[quant] = quant_totals.get(quant, 0) + size
quant_first_file.setdefault(quant, rel)
variants = [
GgufVariantInfo(
filename = quant_first_file[quant],
quant = quant,
size_bytes = size,
)
for quant, size in quant_totals.items()
]
variants.sort(key = lambda variant: -variant.size_bytes)
_apply_gguf_display_labels(variants)
return variants, has_vision

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