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

Studio: make code comments and docstrings more succinct (#6029)

Browse source 
Trim and tighten code comments and docstrings across studio/ Python. Comment-only: every changed file verified code-identical to main via AST/token comparison.
This commit is contained in:
Daniel Han 2026-06-08 23:07:28 -07:00 • committed by GitHub
commit 8292e699e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
205 changed files with 8065 additions and 8912 deletions
Download patch file Download diff file Expand all files Collapse all files

26
studio/backend/_platform_compat.py
View file

@ -4,17 +4,17 @@
"""
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 hardcoded
regex in platform._sys_version() can't parse this and raises ValueError;
CPython closed it as "not planned" (cpython#102396).
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 (no-op if the
cache is already seeded or there are no pipes).
"""
import platform
@ -29,12 +29,12 @@ def _seed_sys_version_cache() -> None:
# 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.
# Format B: "ver (build) | label | (build_dup) \n[compiler]" leaves two
# consecutive (...) groups after pipe-strip; 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:
@ -47,9 +47,9 @@ def _seed_sys_version_cache() -> None:
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,

14
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.
"""
@ -157,9 +157,9 @@ async def _get_current_subject(
credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool
) -> str:
"""
FastAPI dependency to validate the JWT and return the subject.
FastAPI dependency: validate the JWT and return the subject.
Use this as a dependency on routes that should be protected, e.g.:
Use as a dependency on protected routes, e.g.:
@router.get("/secure")
async def secure_endpoint(current_subject: str = Depends(get_current_subject)):

85
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,20 +15,19 @@ 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.
Written to ``_BOOTSTRAP_PW_PATH`` so it survives restarts (the DB only
stores the *hash*). Later calls / restarts return the persisted value.
"""
global _bootstrap_password
@ -51,7 +48,7 @@ def generate_bootstrap_password() -> str:
options = diceware.handle_options(args = ["-n", "4", "-d", "", "-c"])
)
# Persist so the *same* passphrase is used if the server restarts
# Persist so the *same* passphrase is reused if the server restarts
# before the user changes the password.
ensure_dir(_BOOTSTRAP_PW_PATH.parent)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password)
@ -88,21 +85,19 @@ 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.
Plain SHA-256 is intentional: refresh tokens are high-entropy random
strings from ``secrets.token_urlsafe(48)`` (384 bits), so a slow KDF
(Argon2 / bcrypt / PBKDF2) adds zero security (2^384 is unbruteforceable
regardless of hash speed) while costing tens of ms per refresh. See the
OWASP Password Storage Cheat Sheet on hashing 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.
API keys use the separate ``_pbkdf2_api_key`` helper (PBKDF2-HMAC-SHA256
with a persistent server-side salt) — not for crypto reasons (128-bit
random tokens don't need slow hashing) but because CodeQL's
``py/weak-sensitive-data-hashing`` query mislabels them as passwords and
demands a KDF.
"""
return hashlib.sha256(token.encode("utf-8")).hexdigest()
@ -176,22 +171,20 @@ 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).
Stored as a hex-encoded 32-byte random value in ``app_secrets`` under key
``"api_key_pbkdf2_salt"``. Regenerated only when the row is missing (fresh
install, or operator deleted it and accepts invalidating existing keys).
"""
global _api_key_pbkdf2_salt_cache
if _api_key_pbkdf2_salt_cache is not None:
@ -233,22 +226,18 @@ _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.
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 its credential-storage KDF.
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.
The slow KDF here is *only* a CodeQL appeasement, not a crypto
requirement: API keys are random 128-bit tokens (``secrets.token_hex``),
so brute force against 2^128 is infeasible regardless of hash speed.
CodeQL's ``py/weak-sensitive-data-hashing`` query mislabels them as
"password" data and demands a KDF from its allowlist (Argon2 / scrypt /
bcrypt / PBKDF2); we use PBKDF2 per its recommendation page. The salt is
still loaded from ``app_secrets`` so dumping the ``api_keys`` table alone
can't derive hashes for candidate tokens without the salt row.
"""
salt = _get_or_create_api_key_pbkdf2_salt()
dk = hashlib.pbkdf2_hmac(

63
studio/backend/colab.py
View file

@ -2,15 +2,14 @@
# 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.
# Anaconda/conda-forge Python: seed platform._sys_version_cache before any
# import that triggers attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
_backend_dir = str(Path(__file__).parent)
if _backend_dir not in sys.path:
@ -25,11 +24,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 +41,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 Colab proxy URL starts with https:// and embeds the port.
if url and isinstance(url, str) and url.startswith("https://") and str(port) in url:
return url.rstrip("/")
except Exception as e:
@ -61,16 +59,16 @@ 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 Colab proxy URL; when omitted,
``get_colab_url(port)`` is called. Pass it from ``_show_and_embed`` 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 unexpected URL shape never
# prevents the link from rendering.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
@ -79,8 +77,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 is visible even if HTML display fails.
logger.info(f"🌐 Unsloth Studio URL: {url}")
html = f"""
@ -123,12 +120,10 @@ 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 Colab proxy URL once (also registering the port with Colab's
reverse-proxy) then renders a header bar + full-height iframe as one HTML
block. Falls back to ``serve_kernel_port_as_iframe`` if
``IPython.display.HTML`` is unavailable.
"""
url = get_colab_url(port)
logger.info(f"🌐 Unsloth Studio URL: {url}")
@ -138,7 +133,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)
@ -188,8 +183,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.
# Re-launching would collide on the port or silently shift to a new one.
# Just re-show the link and iframe instead.
if _is_studio_healthy(port):
logger.info(f" Studio is already running on port {port} — reusing existing server.")
_show_and_embed(port)
@ -222,16 +217,16 @@ 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 the requested one is in use (e.g.
# Jupyter on 8888). Read back the bound port so the Colab 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 to confirm reachability before showing the link and
# registering the iframe — avoids the race where ready_event fires but the
# process hasn't finished binding.
import urllib.request
server_ready = False
@ -252,9 +247,9 @@ 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 stays running. Handle
# KeyboardInterrupt cleanly so interrupting the cell gives a readable
# message, not a raw traceback.
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").
# Put the backend dir on sys.path so bare "from utils.*" imports work when core
# is imported as a package (e.g. CLI: "from studio.backend.core import ModelConfig").
_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",

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

@ -6,15 +6,14 @@
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.
torch._C._distributed_c10d (the RCCL backend) is absent. `import
torch.distributed` alone is safe (guarded by `if is_available()`); the crash
only comes via torchao's import chain, so stubbing torchao short-circuits it.
_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
Previously duplicated inline in run_export_process() and run_training_process();
now both worker subprocesses call the single
`install_torchao_windows_rocm_stub()` entrypoint before importing
transformers / unsloth_zoo.
"""
@ -28,9 +27,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:
# Metaclass for stub types so 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.
@ -71,8 +70,8 @@ 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 a module) so isinstance(x, attr) returns
# False instead of raising TypeError.
child = _make_stub_type(f"{_n}.{attr}")
setattr(_m, attr, child)
return child
@ -114,14 +113,13 @@ 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 on every other platform (incl. Windows CUDA — there torchao is real
and shadowing it would break torchao quantization paths). Must run before
any import of 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
# ROCM_PATH stay set after installing the HIP SDK and reverting to a CUDA
# torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip
# but still encode "rocm" in torch.__version__, so accept either.
_is_win32_rocm = False
if sys.platform == "win32":
@ -135,8 +133,8 @@ 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 -- elsewhere no stub modules
# are seeded, so appending would be pure accumulation.
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/huggingface.py
View file

@ -95,8 +95,8 @@ 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 when pushing the dataset
# card, so push it ourselves to keep auth request-local.
card.push_to_hub(repo_id, token = hf_token, repo_type = "dataset")
client._upload_main_dataset_files(

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

@ -134,8 +134,8 @@ class JobManager:
``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.
JobManager revokes it 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 +202,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,9 +537,9 @@ 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: revocation failures are swallowed. The key expires on
its own after 24h, so a missed revoke is a latency concern, not a
correctness one.
"""
key_id = getattr(job, "internal_api_key_id", None)
if not key_id:

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 entering the next stage.
job._in_usage_summary = False
if update.usage_section_start is not None:

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

@ -91,8 +91,8 @@ class Job:
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.
# workflow. Revoked when the job ends so the key's window matches the
# run rather than its 24h TTL.
internal_api_key_id: int | None = None
_current_usage_model: str | None = None
_in_usage_summary: bool = False

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

@ -73,10 +73,7 @@ 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
@ -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 configures root logging in __init__. Attach the queue
# logger to `data_designer` directly so parser events survive root resets.
handler = _QueueLogHandler(event_queue)
handler.setLevel(logging.INFO)
for logger_name in (

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

@ -167,9 +167,10 @@ def _validate_recipe_runtime_support(recipe: dict[str, Any], model_providers: li
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.
# Same gate as the chat MCP path: stdio providers spawn a local
# subprocess, so build them only when this host allows it (desktop /
# explicit opt-in). Otherwise a recipe carried onto a hosted host
# cannot spawn.
from core.inference.mcp_client import stdio_mcp_enabled
stdio_allowed = stdio_mcp_enabled()
@ -258,7 +259,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
@ -283,8 +284,8 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None =
_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.
# when the pipeline has no LLM columns. Supply a lightweight stub so
# sampler/expression-only recipes can run without a real provider.
if not model_providers:
from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports]
model_providers = [

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

@ -5,8 +5,8 @@
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.
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

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

@ -2,9 +2,7 @@
# 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 +44,10 @@ 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() runs
`sudo apt-get update`, which hangs on a stdin password in a
non-interactive subprocess. setup.sh pre-installs the build deps on
WSL, so sudo isn't needed at runtime.
"""
if not _is_wsl():
return
@ -110,7 +108,7 @@ class ExportBackend:
try:
logger.info("Starting memory cleanup...")
# Unload all models from inference backend
# Unload all inference-backend models
model_names = list(self.inference_backend.models.keys())
for model_name in model_names:
self.inference_backend.unload_model(model_name)
@ -121,7 +119,7 @@ class ExportBackend:
self.current_checkpoint = None
self._audio_type = None
# Clear GPU memory cache (handles gc + backend-specific cleanup)
# Clear GPU cache (handles gc + backend-specific cleanup)
clear_gpu_cache()
logger.info("Memory cleanup completed successfully")
@ -137,8 +135,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 +156,12 @@ class ExportBackend:
try:
logger.info(f"Loading checkpoint: {checkpoint_path}")
# First, cleanup existing models
# Cleanup existing models first
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():
@ -246,7 +243,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 +255,13 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
)
# Check if PEFT / LoRA model
# Detect 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
@ -459,7 +455,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,
@ -509,12 +505,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 +519,7 @@ class ExportBackend:
)
username = repo_id.split("/")[0]
# Create and push model card
# Build and push model card
content = MODEL_CARD.format(
username = username,
base_model = base_model,
@ -535,7 +530,7 @@ class ExportBackend:
card = ModelCard(content)
card.push_to_hub(repo_id, token = hf_token, commit_message = "Unsloth Model Card")
# Upload model files
# Upload files
if save_directory:
hf_api.upload_folder(
folder_path = save_directory,
@ -585,7 +580,7 @@ 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
@ -613,27 +608,26 @@ class ExportBackend:
# 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
# Use 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.
# 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.
# Snapshot existing .gguf files in cwd before conversion;
# unsloth's convert_to_gguf writes output relative to cwd
# (repo root), so we diff afterwards and relocate.
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.
# 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,
@ -642,17 +636,17 @@ class ExportBackend:
)
# 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".
# convert_to_gguf writes .gguf to cwd (repo root) because
# --outfile is a relative path like "model.Q4_K_M.gguf".
new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
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.
# Flatten any .gguf from subdirs into abs_save_dir.
# save_pretrained_gguf may create subdirs (e.g. model_gguf/)
# with a name different from model_save_path.
# named differently from model_save_path.
for sub in list(Path(abs_save_dir).iterdir()):
if not sub.is_dir():
continue
@ -660,13 +654,13 @@ 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.)
# Clean up the subdir (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.
# checkpoint path, leaving a *_gguf dir in outputs/. Relocate
# any GGUFs from there and clean it up.
if self.current_checkpoint:
ckpt = Path(self.current_checkpoint)
gguf_dir = ckpt.parent / f"{ckpt.name}_gguf"
@ -686,8 +680,7 @@ 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.
# Log final file locations (post-relocation).
final_ggufs = sorted(glob.glob(os.path.join(abs_save_dir, "*.gguf")))
logger.info(
"GGUF export complete. Final files in %s:\n %s",

138
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,9 @@ 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 captured log lines kept in memory per export orchestrator;
# scrollback for the live export log panel. 4000 lines is ~1 MB
# worst-case at 256 chars/line.
_LOG_BUFFER_MAXLEN = 4000
@ -41,46 +39,44 @@ 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
# 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 (load_checkpoint, export_*, cleanup) so
# concurrent HTTP requests can't interleave commands on the
# subprocess queue.
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 log lines from the worker
# subprocess. Powers the GET /api/export/logs/stream SSE
# endpoint the export dialog consumes.
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 sequence number. Never reset across operations, so
# SSE clients can use it as a stable cursor even if clear_logs()
# runs mid-session.
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.
# Snapshot of _log_seq at the start of the current run (set by
# clear_logs()). The SSE endpoint defaults its cursor here so a
# client connecting AFTER the worker's first lines still sees the
# full run. Lines in the current run have seq > _run_start_seq;
# prior-run lines have seq <= it.
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 (load/export/cleanup) is running. The
# SSE endpoint ends the stream 1s after this flips False to drain
# trailing log lines.
self._export_active: bool = False
atexit.register(self._cleanup)
@ -91,12 +87,11 @@ class ExportOrchestrator:
# ------------------------------------------------------------------
def _append_log(self, entry: Dict[str, Any]) -> None:
"""Append a log line from the worker subprocess to the buffer.
"""Append a worker-subprocess log line 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.
"line": "...", "ts": ...}. Each gets a monotonic seq number so
SSE clients can cursor through new lines.
"""
line = entry.get("line")
if not line:
@ -113,17 +108,16 @@ class ExportOrchestrator:
)
def clear_logs(self) -> None:
"""Drop any buffered log lines from a previous operation.
"""Drop buffered log lines from a previous operation.
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).
current run. The seq counter is NOT reset, so an SSE client that
captured the cursor before clear_logs() still sees new lines
(with strictly greater seq).
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
Also snapshots the current seq into ``_run_start_seq`` so the SSE
endpoint can anchor its default cursor at this run's start.
Anything appended after has seq > the snapshot, reachable via
``get_logs_since(get_run_start_seq())``.
"""
with self._log_lock:
@ -144,11 +138,11 @@ 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.
"""Return the seq 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.
connecting AFTER the worker started emitting still sees every
line from the current run.
"""
with self._log_lock:
return self._run_start_seq
@ -193,22 +187,22 @@ class ExportOrchestrator:
self._proc = None
return
# 1. Drain stale responses
# 1. Drain stale responses.
self._drain_queue()
# 2. Send shutdown command
# 2. Send shutdown command.
try:
self._cmd_queue.put({"type": "shutdown"})
except (OSError, ValueError):
pass
# 3. Wait for graceful shutdown
# 3. Wait for graceful shutdown.
try:
self._proc.join(timeout = timeout)
except Exception:
pass
# 4. Force kill if still alive
# 4. 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 +262,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 +272,7 @@ class ExportOrchestrator:
resp = self._read_resp(timeout = min(remaining, 2.0))
if resp is None:
# Check subprocess health
# Check subprocess health.
if not self._ensure_subprocess_alive():
raise RuntimeError("Export subprocess crashed during wait")
continue
@ -294,17 +287,16 @@ 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 too so users see
# high-level progress (e.g. "Importing Unsloth...",
# "Loading checkpoint: ...") alongside subprocess output.
if message:
self._append_log(
{
@ -315,7 +307,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 +354,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:
@ -488,12 +479,11 @@ class ExportOrchestrator:
def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str, Optional[str]]:
"""Send an export command to the subprocess and wait for result.
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 resolved on-disk dir the worker wrote to (None when the
export only pushed to Hub or failed before writing). Surfaced via
the export route's ``details.output_path`` so the dialog's success
screen shows where the model landed.
"""
with self._lock:
if not self._ensure_subprocess_alive():
@ -527,7 +517,7 @@ class ExportOrchestrator:
"""Cleanup export-related models from memory."""
with self._lock:
if not self._ensure_subprocess_alive():
# No subprocess — just clear local state
# No subprocess — clear local state.
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
@ -542,7 +532,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 +543,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)

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

@ -4,9 +4,9 @@
"""
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.
Each export session runs in a persistent subprocess (mp.get_context("spawn")),
giving a clean 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,
@ -31,41 +31,39 @@ 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.
# 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 -- 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 when actual export work starts and leaves
# it open; the orchestrator spawns a fresh subprocess for the next checkpoint
# load (see orchestrator._spawn_subprocess), resetting 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.
# original stdout/stderr fds so the server console / log file keeps the full
# output for debugging.
_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.
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.
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 stdout/stderr so the
server console keeps receiving the full subprocess output, even
while ``_log_forward_gate`` is closed.
Lines are also echoed back to the original stdout/stderr so the server
console keeps the full subprocess 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
# dup failed (exotic platforms) - give up quietly; export still
# works, just no live log streaming.
return
@ -88,13 +86,13 @@ 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 now 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 bound to the
# (now-redirected) 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 +110,8 @@ 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 +131,9 @@ 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 to the saved
# console fd above; drop the line so the export dialog
# doesn't see import / vendoring noise.
continue
try:
resp_queue.put_nowait(
@ -147,8 +145,8 @@ 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 (full, closed, etc.) - drop the line
# rather than crash the reader thread.
pass
if buf and _log_forward_gate.is_set():
try:
@ -181,7 +179,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 +265,11 @@ 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 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.
_log_forward_gate.set()
output_path: Any = None
@ -372,23 +370,23 @@ 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. This powers the
# live export log stream in the UI.
_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.
# Force unbuffered output from child Python processes (e.g. the GGUF
# converter) so their prints surface in the log stream as they happen,
# not at the end.
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.
# (it isn't -- we redirected fd 1/2 to a pipe), making multi-step
# progress bars look frozen in the export log panel. Force frequent
# flushes so the user sees movement during merge / GGUF conversion. No
# effect on single-step bars (e.g. "Copying 1 files") which only emit
# start/end events anyway.
os.environ.setdefault("TQDM_MININTERVAL", "0.5")
import warnings
@ -419,7 +417,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
@ -433,9 +431,9 @@ 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.
# rationale (torchao -> torch.distributed._functional_collectives crashes
# on Windows ROCm: RCCL backend 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 +509,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__ = [

62
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,9 @@ 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 output buffers so nested blockquotes
# each collect their own content and get prefixed with the correct
# number of ">" markers on close.
self._bq_stack: list[list[str]] = []
# ------------------------------------------------------------------
@ -102,7 +102,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:
@ -117,8 +117,8 @@ class _MarkdownRenderer(HTMLParser):
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 HTML with omitted optional
# end tags (</td>, </tr>) does not lose data.
# ------------------------------------------------------------------
def _finish_cell(self) -> None:
if not self._in_cell:
@ -143,8 +143,8 @@ class _MarkdownRenderer(HTMLParser):
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-level content inside
# an <a> does not produce multiline Markdown link labels.
# ------------------------------------------------------------------
def _finish_link(self) -> None:
text = re.sub(r"\s+", " ", "".join(self._link_text_parts)).strip()
@ -234,13 +234,13 @@ 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 any open cell/row from a previous row that omitted its
# optional </td> or </tr> end tags.
self._finish_cell()
self._finish_row()
elif tag in ("th", "td"):
# Flush any open cell (handles omitted </td>/<th>)
# Flush any open cell (handles omitted </td>/<th>).
self._finish_cell()
self._cell_parts = []
self._in_cell = True
@ -248,8 +248,8 @@ class _MarkdownRenderer(HTMLParser):
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 page text focused on readable content and
# avoids data-URI amplification.
return
def handle_endtag(self, tag: str) -> None:
@ -310,7 +310,7 @@ class _MarkdownRenderer(HTMLParser):
self._finish_row()
elif tag == "table":
# Flush any remaining row (handles omitted </tr>)
# Flush any remaining row (handles omitted </tr>).
self._finish_cell()
self._finish_row()
self._in_table = False
@ -325,15 +325,15 @@ 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 text nodes between table structural elements
# (source indentation) so leading spaces don't break Markdown table
# row alignment.
if self._in_table and not self._in_cell and not text.strip():
return
self._emit(text)
@ -354,9 +354,9 @@ class _MarkdownRenderer(HTMLParser):
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).
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 innermost buffers first so their content propagates outward.
@ -376,7 +376,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)
@ -394,8 +394,8 @@ class _MarkdownRenderer(HTMLParser):
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.
Preserves content inside fenced code blocks verbatim so intentional blank
lines in ``<pre>`` content are not collapsed.
"""
lines = text.split("\n")
out: list[str] = []
@ -411,7 +411,7 @@ def _cleanup(text: str) -> str:
continue
if in_fence:
# Preserve code block content exactly as-is
# Preserve code block content exactly as-is.
out.append(line)
continue
@ -433,11 +433,11 @@ def _cleanup(text: str) -> str:
def html_to_markdown(source_html: str) -> str:
"""Convert an HTML string to Markdown.
Handles headings, links, bold/italic, lists (ordered and unordered),
tables, blockquotes, code blocks, and HTML entities. ``<script>``,
``<style>``, and ``<head>`` sections are stripped entirely.
Handles headings, links, bold/italic, ordered/unordered lists, tables,
blockquotes, code blocks, and HTML entities. ``<script>``, ``<style>``,
and ``<head>`` sections 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)

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

@ -4,7 +4,7 @@
"""
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
@ -46,9 +46,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] = []
@ -104,9 +104,9 @@ 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 part list to preserve text/image interleaving (e.g.
# [text, image, text, image]). tool_result blocks become their
# own OpenAI "tool" role messages.
user_parts: list[dict] = []
has_image = False
tool_results: list[dict] = []
@ -137,8 +137,8 @@ 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 so
# text-only callers keep their simple shape.
text = "\n".join(p["text"] for p in user_parts)
if text:
result.append({"role": "user", "content": text})
@ -181,8 +181,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
@ -209,8 +209,8 @@ def build_anthropic_sse_event(event_type: str, data: dict) -> str:
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
@ -320,8 +320,8 @@ class AnthropicStreamEmitter:
# 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: if a different tool_start arrives while a tool_use
# block is open, close the stale block before starting another.
elif self._open_tool_call_id is not None:
events.append(self._close_block())
self._open_tool_call_id = None
@ -421,10 +421,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:

28
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 (same as
# training; the 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,9 @@ 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 (same pattern as Spark-TTS / BiCodec). The pip
# package has problematic deps; the notebook clones and removes
# 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 +148,8 @@ 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 (matches
# notebook: gguf_model.py under models/, others under outetts/)
remove_paths = [
os.path.join(outetts_pkg, "models", "gguf_model.py"),
os.path.join(outetts_pkg, "interface.py"),
@ -181,9 +181,9 @@ class AudioCodecManager:
"""
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,
strips EOS (128258), redistributes 7-per-frame codes into 3 SNAC layers.
generated_ids: full model output including prompt tokens. 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).
"""
@ -192,7 +192,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]
@ -256,8 +256,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 in
# BiCodec config). Pad with zeros or truncate to 32.
GLOBAL_TOKEN_NUM = 32
if global_matches:
raw = [int(t) for t in global_matches]

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

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

File diff suppressed because it is too large Load diff

340
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,25 @@ 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!``.
gpt-oss emits multi-channel output via special tokens like
``<|channel|>analysis<|message|>...`` / ``<|channel|>final<|message|>...``.
A plain ``TextIteratorStreamer(skip_special_tokens=True)`` drops the tokens
but leaves channel names glued to content, producing garbled output like
``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:
This streamer decodes with ``skip_special_tokens=False`` so the harmony
markup is visible, then uses stateful incremental parsing:
- ``<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
- ``<think>`` once when ``analysis`` channel is first seen
- analysis content streamed incrementally
- ``</think>`` once when ``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.
position as content grows. Implements the same ``put`` / ``end`` / iterator
interface as ``TextIteratorStreamer`` for drop-in use in ``generate_stream``.
"""
import re as _re
@ -87,7 +81,7 @@ 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
@ -102,7 +96,7 @@ class HarmonyTextStreamer:
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 +104,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
@ -125,7 +119,7 @@ class HarmonyTextStreamer:
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)
@ -164,19 +158,19 @@ class HarmonyTextStreamer:
# ------------------------------------------------------------------
def _process_incremental(self, raw: str) -> None:
"""Parse harmony channels and emit deltas per-channel.
"""Parse harmony channels and emit per-channel deltas.
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:
Tracks per-channel content lengths (instead of diffing transformed
whole-text, which breaks when wrapping ``<think>`` tags shift position)
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.
# raw has <|channel|> but no complete channel+message pair yet:
# buffer silently, don't emit partial channel names as text.
has_channel_token = "<|channel|>" in raw
matches = list(self._HARMONY_RE.finditer(raw))
@ -185,8 +179,7 @@ class HarmonyTextStreamer:
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
# No harmony protocol (shouldn't happen for gpt-oss) — emit nothing
return
for m in matches:
@ -228,11 +221,11 @@ 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() calls. Must be a plain
# Lock (NOT RLock): in async FastAPI requests share one event-loop
# thread, so 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 +235,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 +248,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 +256,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,9 +313,9 @@ 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 from local path. base_model is e.g.
# /home/.../Spark-TTS-0.5B/LLM; BiCodec weights live in
# the parent dir (Spark-TTS-0.5B/).
base_path = config.base_model
if os.path.isdir(base_path):
abs_repo_path = os.path.abspath(os.path.dirname(base_path))
@ -348,7 +339,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 +397,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 +426,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 +448,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 +461,19 @@ 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 (e.g.
# GemmaTokenizerFast) instead of a proper Processor for some
# models (e.g. Gemma-3). If so, load the real processor from
# the base model.
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 the base model. Local merged exports:
# read export_metadata.json for the original base model.
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 +513,6 @@ class InferenceBackend:
trust_remote_code = trust_remote_code,
)
# Apply inference optimization
FastLanguageModel.for_inference(model)
self.models[model_name]["model"] = model
@ -530,7 +520,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 +541,25 @@ 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 so the next model gets a fresh one.
# On spawn-based platforms, preserve trainer files so concurrent
# training dataset.map() workers can still import them.
import sys as _sys
from utils.cache_cleanup import clear_unsloth_compiled_cache
@ -593,26 +578,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
# Clear any lingering peft_config. model.unload() can leave a
# peft_config attribute; removing it gives PeftModel.from_pretrained()
# a clean base model without "multiple adapters" warnings.
if hasattr(model, "peft_config"):
del model.peft_config
@ -636,10 +618,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 +630,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 +646,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 +659,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 +670,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 +686,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 +701,8 @@ 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 (replaces the flawed
'enable_adapter')."""
model = self.models[base_model_name].get("model")
try:
logger.info(f"Setting active adapter to: '{adapter_name}'")
@ -733,18 +710,17 @@ 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.
Uses PEFT's disable_adapter_layers() / enable_adapter_layers(), which
toggle a flag (proj.disable_adapters) on each LoRA layer; Unsloth's
fast_linear_forward checks it and skips LoRA when True. Non-destructive,
no unload/reload.
Args:
use_adapter: None = no change, False = disable (base model),
@ -781,7 +757,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,13 +771,12 @@ 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.
in the background generation thread, NOT the event-loop thread. Prevents
the RLock-reentrant race when two async SSE handlers share one
event-loop thread.
Args:
use_adapter: Adapter control (None/False/True/str). See _apply_adapter_state.
@ -833,9 +808,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}``
@ -896,15 +870,14 @@ class InferenceBackend:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Generator[str, None, None]:
"""
Generate response for text or vision models.
"""Generate response for text or vision models.
The generation lock is acquired by the background generation 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.
``tokenizer.apply_chat_template`` so templates that understand them
(Qwen3, Llama 3.1+, gpt-oss harmony, ...) advertise the tool schemas
and reasoning controls to the model.
"""
yield from self._generate_chat_response_inner(
messages = messages,
@ -941,9 +914,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.
@ -960,10 +932,10 @@ class InferenceBackend:
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).
# Vision generation (only with an actual image). Verify the stored
# processor can handle images — FastVisionModel may return a raw
# tokenizer (e.g. GemmaTokenizerFast) instead of a ProcessorMixin
# for some models (e.g. Gemma-3).
from transformers import ProcessorMixin
processor = model_info.get("processor")
@ -991,10 +963,10 @@ 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: training-pipeline approach. 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 +974,13 @@ 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
# Sets the correct template on the tokenizer
tokenizer = get_chat_template(
tokenizer,
chat_template = template_name,
@ -1021,7 +992,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 +1017,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 +1051,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 +1107,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 +1204,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 +1216,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 +1224,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 +1240,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 +1261,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 +1325,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 +1347,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 +1363,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 +1374,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 +1386,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 +1484,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 +1514,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 +1629,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 +1657,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 +1710,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 +1813,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 +1926,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 +1942,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 +1956,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 +1990,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 +2020,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 +2028,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 +2088,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 +2106,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 +2119,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,

24
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,9 @@ 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 is suspicious — a fresh keypair invalidates all in-flight
# ciphertext encrypted against the previous public key. Log loudly so a
# regression that calls init twice is visible.
logger.warning(
"init_key_pair called again — replacing existing RSA keypair "
"(previous fingerprint=%s). Any frontend that cached the old "
@ -111,9 +111,9 @@ 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.
# Surface enough state to tell a key mismatch (wrong public key used on
# encrypt) from a padding/algo mismatch or corrupted bytes. Expected
# RSA-2048 ciphertext length is exactly 256 bytes.
logger.warning(
"decrypt_api_key: RSA decrypt failed (ciphertext_len=%d, expected=256, "
"fingerprint=%s, exc=%s): %s",

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

File diff suppressed because it is too large Load diff

91
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"}),
)
@ -74,10 +73,9 @@ 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`.
(llama-server shorts always start with a letter), strips whitespace, and
normalises attached `-np8` / signed `-np-1` / digit-prefix-junk `-np8x`
to `-np`. Mirrors the CLI's `_expand_attached_np_short`.
"""
token = token.strip()
if not token.startswith("-") or token in {"-", "--"}:
@ -95,9 +93,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 +114,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 +155,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 pass-through numeric
knob Studio's load-time fit logic needs to see.
"""
if not args:
return None
@ -204,10 +201,10 @@ 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 the ``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 reimplement the
conditional and assert against their own reimplementation.
"""
override = parse_ctx_override(args)
return override if override is not None else fallback_n_ctx
@ -216,10 +213,10 @@ 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.
Like parse_ctx_override but for cache type. Recognises -ctk (key) and
-ctv (value); when both appear, returns the last-wins value, treating
key and value flags as one setting because Studio's KV estimate has a
single cache_type_kv knob.
"""
if not args:
return None
@ -256,8 +253,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 +270,9 @@ 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; the
route only strips groups whose first-class field the caller supplied.
"""
shadowing: set[str] = set()
if strip_context:
@ -299,8 +294,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:

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

@ -31,20 +31,20 @@ 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 intact but also keeps surrounding
# quotes on a token. Strip a matched pair so argv reaches the subprocess
# clean ('"C:\\Program Files\\node"' -> C:\\Program Files\\node).
parts = [p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p for p in parts]
return parts
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (and bypass
the python/terminal sandbox), so they are only allowed when the backend
host is the user's own machine. The Tauri desktop app sets
"""stdio MCP servers spawn local processes as the backend user (bypassing
the python/terminal sandbox), so they're allowed only 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."""
self-hosted users can opt in with the same variable. Stays off for Colab
and any network (0.0.0.0) bind."""
return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1"
@ -63,8 +63,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
@ -83,7 +83,7 @@ def _oauth_store():
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.
# 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 +93,11 @@ 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 old
credentials explicitly; otherwise re-registering the same URL silently
reuses the old account's token. The whole body is protected -- store / OAuth
construction failing 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 +123,9 @@ 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's safe default
# env). keep_alive=False tears the subprocess down on exit, so a
# one-shot probe/tool call never leaves an orphan.
return Client(
StdioTransport(
command = parts[0],
@ -192,10 +191,10 @@ 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 HTTP
call is cancelled and a cancellation Error is returned. Polled in parallel
with the tool call via ``asyncio.wait`` so a /cancel POST from the UI
interrupts even mid-network-read.
"""
async def _call() -> Any:
@ -204,14 +203,14 @@ 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
# (reviewer-reproduced race).
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
call_task = asyncio.create_task(_call())

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

@ -65,10 +65,10 @@ class MLXInferenceBackend:
def _configure_memory_limits(self):
"""Apply Metal memory caps before loading a model.
Mirrors MLXTrainer._configure_memory_limits's defaults:
Mirrors MLXTrainer._configure_memory_limits 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.
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 +109,16 @@ 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 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:`) sends
# GGUF traffic to llama-server before reaching the MLX orchestrator.
# Reaching here with is_gguf=True means 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 can't reach the parent's llama-server, so raise loudly to
# give the caller a clear error instead of a cryptic "config.json does
# not exist" from mlx_lm.utils.load_model.
if getattr(config, "is_gguf", False):
raise RuntimeError(
f"MLXInferenceBackend cannot load GGUF model '{model_name}': "
@ -187,9 +185,9 @@ 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 can ship it back
# to the parent and the route layer classifies capabilities the same
# way as the transformers / GGUF paths.
self._populate_chat_template_info(model_name)
logger.info("Model %s loaded successfully", model_name)
@ -198,10 +196,9 @@ 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`` so ``_detect_safetensors_features``
(route layer) sees the template the model actually uses."""
entry = self.models.get(model_name)
if not entry:
return
@ -276,10 +273,9 @@ 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 -- the MLX
# path renders the template via apply_chat_template_for_generation so
# these are honoured the same way as the transformers path.
tools = None,
enable_thinking = None,
reasoning_effort = None,
@ -308,7 +304,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 +385,8 @@ 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
# (1.0 is the no-op value).
logits_processors = None
if repetition_penalty is not None and float(repetition_penalty) not in (
0.0,
@ -425,7 +421,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 — like GPU
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
@ -471,10 +467,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 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.
chat_target = self._processor
if (
getattr(self._processor, "apply_chat_template", None) is None
@ -492,8 +487,8 @@ 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)
# For VLM: always use mlx_vlm's stream_generate, which handles
# pixel_values properly (None for text-only, image for VLM)
images = [image] if image is not None else None
cumulative = ""
@ -503,8 +498,8 @@ 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.
# accepts temp/top_p/top_k/repetition_penalty (building 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).

189
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):
@ -66,9 +64,9 @@ class InferenceOrchestrator:
self._gen_lock = threading.Lock() # Serializes generation — one request at a time
# 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.
# These bypass _gen_lock: they send commands directly and read from
# per-request mailboxes. A dispatcher thread routes resp_queue events
# by request_id.
self._mailboxes: dict[str, queue.Queue] = {}
self._mailbox_lock = threading.Lock() # Protects _mailboxes dict
self._dispatcher_thread: Optional[threading.Thread] = None
@ -92,7 +90,7 @@ class InferenceOrchestrator:
atexit.register(self._cleanup)
logger.info("InferenceOrchestrator initialized (subprocess mode)")
# Kick off background fetch of top models from HF
# Background fetch of top models from HF
threading.Thread(target = self._fetch_top_models, daemon = True, name = "top-models").start()
# ------------------------------------------------------------------
@ -105,10 +103,9 @@ class InferenceOrchestrator:
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 (editorial picks), then HF
# download-ranked models to backfill. Send extras so the frontend
# still has 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 +130,8 @@ 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 - frontend pages through them via infinite
# scroll, so send a deep pool.
gguf_ids = [m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF")][
:40
]
@ -201,7 +198,7 @@ class InferenceOrchestrator:
self._cancel_generation()
time.sleep(0.5) # Brief wait for generation to stop
# 2. Drain stale responses from queue
# 2. Drain stale responses
self._drain_queue()
# 3. Send shutdown command
@ -243,7 +240,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 +274,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 +340,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 +362,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 +416,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 — maybe for a _gen_lock reader or orphaned.
# Can't un-get from mp.Queue, so just log a warning.
if rtype not in ("status",):
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
@ -453,13 +446,13 @@ 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.
Uses a per-request mailbox for tokens, so two compare-mode requests
can be queued in the subprocess at once, avoiding 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.
The subprocess runs commands sequentially from its cmd_queue, so
generation stays serialized at the GPU level — this only avoids
orchestrator-level lock contention.
"""
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess is not running"
@ -532,7 +525,7 @@ 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
@ -574,8 +567,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 +581,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 +611,9 @@ 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
inspect.getsource() failures from a previous model.
"""
from utils.transformers_version import needs_transformers_5
@ -649,9 +642,9 @@ 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 causes
# inspect.getsource() failures on the next load.
if self._ensure_subprocess_alive():
self._cancel_generation()
time.sleep(0.3)
@ -680,7 +673,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 enabled -> 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,7 +682,7 @@ 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 disabled) -> give up
self._shutdown_subprocess(timeout = 5)
raise RuntimeError(
f"Download stalled for '{model_name}' even with "
@ -745,7 +738,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
@ -797,13 +790,13 @@ class InferenceOrchestrator:
"""Generate response, streaming tokens from subprocess.
Optional ``tools`` / ``enable_thinking`` / ``reasoning_effort`` /
``preserve_thinking`` kwargs are forwarded into the worker so
``preserve_thinking`` kwargs are forwarded to the worker so
``tokenizer.apply_chat_template`` can render tool schemas and
reasoning controls when the template understands them.
reasoning controls when the template supports them.
``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 so concurrent streams cannot
read each other's stats.
"""
yield from self._generate_inner(
messages = messages,
@ -847,13 +840,12 @@ 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 this (parent) process,
calling the worker for each generation 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. See
``safetensors_agentic.run_safetensors_tool_loop`` for the protocol.
"""
from core.inference.safetensors_agentic import run_safetensors_tool_loop
from core.inference.tools import execute_tool
@ -861,8 +853,8 @@ class InferenceOrchestrator:
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.
# ``conv`` already carries any system message: the loop appends
# to a list seeded with system+user above.
common_kwargs = dict(
messages = conv,
system_prompt = "",
@ -914,9 +906,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 +938,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"
@ -959,13 +950,13 @@ class InferenceOrchestrator:
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.
# wait for dispatched requests to finish then stop it, so we can read
# from resp_queue directly.
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 — 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.
with self._gen_lock:
yield from self._generate_locked(
messages = messages,
@ -1029,8 +1020,8 @@ 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, so older workers that
# ignore unknown keys still work.
if tools is not None:
cmd["tools"] = tools
if enable_thinking is not None:
@ -1046,8 +1037,8 @@ class InferenceOrchestrator:
yield f"Error: {exc}"
return
# Yield tokens from response queue — we are the only reader
# because _gen_lock is held.
# Yield tokens from response queue — we are the only reader since
# _gen_lock is held.
while True:
resp = self._read_resp(timeout = 30.0)
@ -1071,10 +1062,10 @@ 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
# Wait for the subprocess to ack cancellation
# (gen_done/gen_error) so stale events don't leak into
# the next generation request.
self._drain_until_gen_done(timeout = 5.0)
@ -1117,7 +1108,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 +1233,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,8 +1301,8 @@ 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.
No ML imports — runs locally in the parent process.
"""
if img is None:
return None
@ -1331,26 +1322,26 @@ class InferenceOrchestrator:
return base64.b64encode(buf.getvalue()).decode("ascii")
def get_current_model(self) -> Optional[str]:
"""Get currently active model name."""
"""Currently active model name."""
return self.active_model_name
def is_model_loading(self) -> bool:
"""Check if any model is currently loading."""
"""True if any model is loading."""
return len(self.loading_models) > 0
def get_loading_model(self) -> Optional[str]:
"""Get name of currently loading model."""
"""Name of the currently loading model."""
return next(iter(self.loading_models)) if self.loading_models else None
def check_vision_model_compatibility(self) -> bool:
"""Check if current model supports vision."""
"""True if the current model supports vision."""
if self.active_model_name and self.active_model_name in self.models:
return self.models[self.active_model_name].get("is_vision", False)
return False
def _is_gpt_oss_model(self, model_name: str = None) -> bool:
"""Parent-side gpt-oss detection so the safetensors route can run
the same guard without an IPC round-trip to the subprocess."""
"""Parent-side gpt-oss detection so the safetensors route can run the
same guard without an IPC round-trip to the subprocess."""
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 +1351,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()

85
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,9 @@ 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, longest-first: dated
# snapshots inherit canonical prices, but "claude-opus-4-15" won't
# match "claude-opus-4-1" nor "gpt-5.5-prod" match "gpt-5.5-pro".
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 +105,10 @@ 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) so the
frontend renders a number or tooltip without redoing the math.
Unknown model -> ``priced`` False and USD fields 0.0 (token counts
still report).
"""
prices = _lookup(provider, model)
out: dict[str, float] = {
@ -128,22 +125,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 +147,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 +195,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 +230,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)

121
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,10 @@ 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
# dozens of historical snapshots, fine-tunes, and non-chat models
# (embeddings, TTS, image, moderation) we never want in the chat
# UI. Filter here so the backend is the single source of truth.
"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}$"),
@ -48,9 +47,8 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
],
# 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).
# YYYYMMDD-suffixed variants — same intent as the OpenAI denylist,
# just a different date format (no dashes).
"model_id_denylist": re.compile(r"-\d{8}$"),
"supports_streaming": True,
"supports_vision": True,
@ -68,20 +66,18 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
# 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
# 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:
# Curated lineup -- live ListModels returns dozens of historical /
# experimental / embedding ids. Cap to 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.
# The allowlist below keeps the retired ids from re-appearing via
# live ListModels. Verified against `/v1beta/models` 2026-05-24.
"default_models": [
"gemini-3.1-pro-preview",
"gemini-3.5-flash",
@ -100,8 +96,8 @@ 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 key on `x-goog-api-key`; empty `auth_prefix`
# sends the bare key.
"auth_header": "x-goog-api-key",
"auth_prefix": "",
"openai_compatible": False,
@ -110,20 +106,19 @@ 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.
# Even after the regex match, drop ids Google still returns from
# ListModels but routes via implicit redirect. gemini-3-pro-preview
# was shut down 2026-03-09 and auto-aliased to
# gemini-3.1-pro-preview; surface the canonical id only so users
# don't see two cards for the same model.
"model_id_deny_exact": ("gemini-3-pro-preview",),
# Matches the chat-capable 3.5 / 3.1 / 3 / 2.5 families plus the
# rolling *-latest aliases (which Google rolls forward as new
# generations ship). Image-tier ids (`-image`, `-image-preview`,
# rolling *-latest aliases (rolled 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.
# `responseModalities` path in `_stream_gemini`. Retired 2.0 ids are
# excluded on purpose -- ListModels would re-surface them and they
# 404 on use.
"model_id_allowlist": re.compile(
r"^("
r"gemini-3\.5-(?:flash|pro)(?:-preview)?|"
@ -186,12 +181,12 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"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:
# Listing/overview endpoints:
# 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; everything else (moonshot-v1-*, dated k2 previews) is
# filtered out by model_id_allowlist below.
"default_models": [
"kimi-k2.6",
"kimi-k2.5",
@ -205,8 +200,8 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"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.
# (same for top_p). Strip both fields from the outbound body so the
# server falls back to its required defaults.
"body_omit": ("temperature", "top_p"),
},
"qwen": {
@ -228,9 +223,9 @@ 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 with 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.
"default_models": [
"openai/gpt-oss-120b",
"deepseek-ai/DeepSeek-V3",
@ -249,27 +244,26 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"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.
# catalog (state.org/model[:policy] ids). Remote so users see live
# availability -- the picker has search, and loadModels() merges
# defaults so default_models 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 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.).
"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. /v1/models has no server-side limit or
# popularity sort, so this is just "first N matches"; the
# default_models seed keeps flagship ids near the top regardless of
# the API's order.
"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 layer falls back to
# the payload's base_url when the registry entry has none.
"base_url": "",
"default_models": [],
"supports_streaming": True,
@ -277,15 +271,14 @@ 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 in stream_chat_completion -- vLLM's
# /v1/responses rebuilds messages through the 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.
"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
# /api/providers/registry dropdown -- see list_available_providers.
"hidden": True,
},
"ollama": {
@ -371,10 +364,10 @@ 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 (``"hidden": True``) are filtered out: they exist 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.
"""
result = []
for provider_type, info in PROVIDER_REGISTRY.items():

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

@ -4,17 +4,16 @@
"""
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 (the
``InferenceOrchestrator.generate_chat_response`` pipeline streaming from a
worker subprocess) with the tool-calling, thinking-block, status, and
metadata event protocol used by the GGUF path. The front-end SSE shape
stays identical across backends, so the chat UI is engine-agnostic.
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``.
The GGUF path (``llama_cpp.py``) uses llama-server's structured
``delta.tool_calls`` directly. Native transformers has no such channel, so
this loop parses tool calls from the cumulative text and dispatches them
via ``core.inference.tools``.
"""
import json
@ -41,7 +40,7 @@ from core.inference.tool_call_parser import (
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
@ -105,10 +104,10 @@ def _coerce_arguments(
"""Normalise tool ``arguments`` to a dict.
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).
``heal=True`` a bare string becomes ``{<canonical_key>: ...}`` so a
Hermes-style call without proper JSON still runs. Canonical key per
tool: ``code`` for python, ``command`` for terminal, ``query`` otherwise
(e.g. web_search).
"""
if isinstance(raw_args, dict):
return raw_args
@ -140,28 +139,23 @@ 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"}``
"""
@ -205,7 +199,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
@ -308,7 +302,7 @@ 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):
detect_state = _state_draining
@ -331,9 +325,9 @@ def run_safetensors_tool_loop(
id_offset = next_call_id,
)
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: streaming already emitted content. Skip the
# final=True re-strip so literal "<tool_call>" in prose
# survives when no real tool call parsed.
yield {"type": "status", "text": ""}
return
tool_calls = safety_tc
@ -349,8 +343,8 @@ def run_safetensors_tool_loop(
id_offset = next_call_id,
)
if not tool_calls and auto_heal_tool_calls:
# Parser found nothing -- surface raw content so any
# literal "<tool_call>" prose is preserved.
# Parser found nothing -- surface raw content so literal
# "<tool_call>" prose is preserved.
if content_accum:
yield {"type": "content", "text": content_accum}
if provisional_render_html_started:
@ -436,9 +430,8 @@ def run_safetensors_tool_loop(
render_html_succeeded = True
tool_call_history.append((tc_key, is_error))
# Strip frontend image sentinel from the model's view.
# Cut at the first occurrence so leading and consecutive
# sentinels are both removed.
# Strip frontend image sentinel from the model's view. Cut at the
# first occurrence so leading and consecutive sentinels both go.
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()

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

@ -13,9 +13,9 @@ 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.
# Function-name char set tracks OpenAI's ^[a-zA-Z0-9_-]{1,64}$ so hyphenated
# MCP tool names (e.g. mcp__srv__list-issues) parse the same as built-in
# web_search/python/terminal names.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
@ -72,8 +72,8 @@ _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.
# Parameter names can carry hyphens too (e.g. MCP schemas with
# `issue-number`, `repo-name`); `\w+` alone dropped those keys.
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
_PARAM_CLOSE_TAG = "</parameter>"
@ -124,8 +124,8 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
"""
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
# 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
@ -165,9 +165,9 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
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;
# don't use </function> as body boundary since code values can contain
# that literal.
if not tool_calls:
func_starts = [
fm
@ -190,8 +190,8 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
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 embedded
# </parameter> in code strings is preserved.
pm = param_starts[0]
val = body[pm.end() :]
val = _TC_PARAM_CLOSE_RE.sub("", val)

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

@ -4,7 +4,7 @@
"""
Tool definitions and executors for LLM tool calling.
Supports web search (DuckDuckGo), Python code execution, and terminal commands.
Web search (DuckDuckGo), Python code execution, and terminal commands.
"""
import ast
@ -42,9 +42,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).
# Pre-import _sandbox_preexec modules at module level so the preexec_fn closure
# doesn't trigger imports in the forked child (can deadlock multi-threaded servers).
_libc = None
if sys.platform == "linux":
try:
@ -64,8 +63,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 +121,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 +151,19 @@ _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).
# into their own tokens, so we detect command position even for
# `echo done; rm -rf x` (no whitespace) or quote-split command names
# (`r''m` collapses to `rm` at command position after `;`).
try:
if sys.platform == "win32":
tokens = shlex.split(command, posix = False)
@ -178,9 +175,8 @@ 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.
# shlex may glue meta-chars onto a token (`rm;`); strip leading and
# trailing ones so the basename still matches `rm`.
tok = tok.strip(";&|()`{}")
base = os.path.basename(tok).lower()
stem, ext = os.path.splitext(base)
@ -189,33 +185,32 @@ 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. While a wrapper prefix awaits
# its command (`stdbuf -oL cmd`, `xargs -- cmd`), keep expect_command.
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.
# next non-flag, non-numeric token is the real command. `sudo` is in
# _BLOCKED_COMMANDS, so it's flagged AND we keep walking.
if base in _COMMAND_PREFIXES:
prefix_pending = True
continue
@ -229,10 +224,10 @@ 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: blocked words at command boundaries shlex won't see, e.g. inside
# unquoted $(rm -rf), <(rm), backtick chains, or after a separator with no
# whitespace ("foo;rm"). Anchored to command-position delimiters; no match
# in argument position.
lowered = command.lower()
if _BLOCKED_COMMANDS:
words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS))
@ -243,11 +238,10 @@ 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 'sudo whoami', bash -lc '...',
# bash --login -c '...', cmd /c '...'). On a -c or /c flag, look
# backwards for a shell name (skipping flags like --login, -l, -x)
# 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 +253,9 @@ 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 backwards past flags to find the shell binary. Unix flags
# start with - (skip). Windows flags start with / but so do absolute
# paths, so only skip short single-char /X flags (not /bin/bash).
for j in range(i - 1, -1, -1):
prev = tokens[j]
if prev.startswith("-"):
@ -282,17 +275,15 @@ 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 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 all other parent
vars are absent. HOME points at the sandbox workdir so HF / wandb / aws
SDKs can't read cached credentials from the operator's real ~/.
"""
# 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 with the running interpreter's directory so 'python', 'pip', etc.
# 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 []
@ -354,9 +345,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.
@ -381,10 +372,9 @@ def _sandbox_preexec():
pass
try:
# Default high enough for multi-shard safetensors mmaps + Python's
# own handle count; tunable via env for installs that hit the cap.
# 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).
# handle count; tunable via env. Clamp to the inherited hard limit
# so setrlimit doesn't ValueError where the parent's hard cap is
# below the request (else NOFILE stays at the parent's default).
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 +391,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] = {}
@ -681,10 +670,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).
``timeout``: int = per-call limit in seconds, ``None`` = no limit,
unset (default) = ``_EXEC_TIMEOUT`` (300 s).
``session_id``: optional thread/session ID for per-conversation sandbox isolation.
"""
logger.info(f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}")
@ -726,10 +715,10 @@ def execute_tool(
_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.
# Raw download cap. Larger than _MAX_PAGE_CHARS because SSR pages embed large
# <head> sections (CSS, JS, SVGs) stripped during HTML-to-Markdown conversion.
# 512 KB reaches article content on GitBook / Next.js / Docusaurus pages whose
# <head> alone can be 200 KB.
_MAX_FETCH_BYTES = 512 * 1024
_USER_AGENTS = (
@ -750,14 +739,13 @@ 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 certificate 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 of which fail.
This subclass splits the concerns: TCP connects to the pinned IP (``host``)
while TLS uses ``sni_hostname`` for the ClientHello and cert check.
"""
def __init__(self, host: str, *, sni_hostname: str, **kwargs):
@ -765,8 +753,8 @@ 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 (+ tunnel if a proxy is set
# via set_tunnel, though we don't use one).
http.client.HTTPConnection.connect(self)
# TLS handshake with the real hostname for SNI + cert verification.
self.sock = self._context.wrap_socket(
@ -778,8 +766,8 @@ class _PinnedHTTPSConnection(http.client.HTTPSConnection):
class _SNIHTTPSHandler(urllib.request.HTTPSHandler):
"""HTTPS handler that sends 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
SSRF IP-pinning rewrites URLs to raw IPs, breaking SNI and cert
verification. This handler returns a ``_PinnedHTTPSConnection`` that
connects to the pinned IP but verifies TLS against the original hostname.
"""
@ -798,9 +786,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 +803,14 @@ 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`
# `not ip.is_global` rejects every category the denylist below rejects
# PLUS shared address space (100.64.0.0/10 carrier-grade NAT) and
# benchmarking/documentation/exchange ranges Python marks
# `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.
# The explicit predicates after it give human-readable categories in
# the error message; the non-global check is the source of truth and
# prevents future ranges from leaking.
if (
not ip.is_global
or ip.is_private
@ -872,8 +860,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 the URL
# to the IP and 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
@ -926,7 +914,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)
@ -981,7 +969,7 @@ def _web_search(
def _check_signal_escape_patterns(code: str):
"""
Check if code contains patterns that could escape signal-based timeouts.
Check for 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).
@ -1071,8 +1059,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 +1081,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")
# Maps bare function names to fully-qualified form for
# from-import tracking (e.g. "system" -> "os.system")
self.shell_exec_aliases: dict[str, str] = {}
self.loop_depth = 0
@ -1197,7 +1185,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):
@ -1229,7 +1217,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",
@ -1251,8 +1239,8 @@ 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).
# enabled. Treat any non-literal-False shell= value as
# potentially True (conservative).
_STRING_SHELL_FUNCS = frozenset(
{
"os.system",
@ -1268,7 +1256,7 @@ def _check_signal_escape_patterns(code: str):
shell_safe = shell_node is None or (
isinstance(shell_node, ast.Constant) and shell_node.value is False
)
# Dynamic shell-exec args (chr/format/concat bypasses).
# Dynamic shell-exec args (chr/format/concat bypasses)
if (
shell_func in _STRING_SHELL_FUNCS
or shell_func in _SHELL_EXEC_FUNCS
@ -1311,10 +1299,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.
# except Exception doesn't catch SystemExit/KeyboardInterrupt,
# so it can't suppress timeout enforcement; flagging it would
# false-positive on normal error handling.
if node.type.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
@ -1342,9 +1329,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 instead.
network_calls: list[dict] = []
sensitive_file_reads: list[dict] = []
_NETWORK_FQ_PREFIXES = (
@ -1591,9 +1578,9 @@ def _check_signal_escape_patterns(code: str):
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.
# but should only fire when huggingface_hub / hf_api is imported in the
# snippet; else paramiko.upload_file, boto3 create_commit, etc. would
# false-positive. We pre-scan for the imports.
_HF_IMPORT_MODULES = (
"huggingface_hub",
"hf_api",
@ -1611,8 +1598,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
@ -1635,8 +1622,8 @@ def _check_signal_escape_patterns(code: str):
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.
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
@ -1649,7 +1636,7 @@ def _check_signal_escape_patterns(code: str):
# 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.
# lifted from the parent.
_HF_SENSITIVE_KWARGS = frozenset(
{
"token",
@ -1674,14 +1661,14 @@ 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.
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.
Walking the subtree (not just the root) catches wrapper calls like
`str(os.environ)`, `json.dumps(os.environ)`,
`'-'.join(os.environ.values())`.
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`.
could lift parent env via `printenv` / `env` / `set`.
"""
if node is None:
return False
@ -1729,7 +1716,7 @@ def _check_signal_escape_patterns(code: str):
if node is None:
return False
if isinstance(node, ast.Constant) and isinstance(node.value, (bytes, bytearray)):
return True # inline bytes, no file access
return True # inline bytes; no file access
if isinstance(node, ast.Constant) and isinstance(node.value, str):
return _is_safe_relative_path(node.value)
if isinstance(node, ast.Call):
@ -1751,10 +1738,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 +1800,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 +1934,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 +2018,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 +2074,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):

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

@ -4,9 +4,9 @@
"""
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.
Each inference session runs in a persistent subprocess
(mp.get_context("spawn")), giving a clean Python interpreter with no stale
module state — fully solving the transformers version-switching problem.
The subprocess stays alive while a model is loaded, accepting commands
(generate, load, unload) via mp.Queue. It exits on shutdown or unload.
@ -35,7 +35,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 +96,16 @@ 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.
When *model_names* is provided, only those models' ``blobs/`` dirs are
checked instead of scanning every cached model -- much faster with many
models. Accepts multiple names so LoRA loads can watch both the adapter
repo and the base model repo at once.
*has_incomplete* is True when any ``*.incomplete`` files exist in the
watched blobs directories, indicating that ``huggingface_hub`` is
actively downloading.
watched blobs dirs, indicating ``huggingface_hub`` is actively downloading.
Returns None if the state cannot be determined (import error,
permission error, etc.) so callers can skip stall logic.
Returns None if the state cannot be determined (import error, permission
error, etc.) so callers can skip stall logic.
"""
try:
from huggingface_hub.constants import HF_HUB_CACHE
@ -127,12 +125,12 @@ def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, b
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.
# backslashes, distinguishing them from absolute, relative,
# and Windows paths.
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,14 +163,14 @@ 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.
Monitors the HF Hub cache for download activity. A stall is reported only
when ``*.incomplete`` files are present (``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.
Once the download finishes (no more ``.incomplete`` files), the stall timer
resets, so post-download init (quantization, GPU weight loading) is never
misclassified as a stalled download.
Returns a stop event -- set it to terminate the heartbeat thread.
"""
@ -188,7 +186,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,9 +204,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
# Only fire stall when .incomplete files confirm a download is in
# progress. Once downloads finish (no .incomplete), reset the timer
# so model init time is not counted as a stall.
if not has_incomplete:
last_change = now
@ -224,7 +221,7 @@ def _start_heartbeat(
"ts": time.time(),
},
)
# Only fire once -- the orchestrator will kill us
# Only fire once -- the orchestrator will kill us.
return
_send_response(
@ -249,7 +246,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,7 +277,7 @@ 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.
# 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).
@ -298,12 +295,12 @@ 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)
# Send heartbeats every 30s so the orchestrator knows we're alive
# (download/weight loading can be slow on slow connections).
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 both the model repo and base model repo (for LoRA loads where
# the base model download is the real bottleneck).
watch_repos = [mc.identifier]
base = getattr(mc, "base_model", None)
if base and str(base) != mc.identifier:
@ -328,7 +325,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 +338,8 @@ 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
# without re-entering the subprocess.
try:
_bm = getattr(backend, "models", {}) or {}
_entry = (
@ -397,22 +394,21 @@ 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 shared with the parent. The parent can set it
at any time (e.g. user stops generation, or loads a new model mid-generate)
and generation stops within 1-2 tokens.
"""
request_id = cmd.get("request_id", "")
try:
# Decode image if provided
# 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
# Build generation kwargs.
gen_kwargs = {
"messages": cmd["messages"],
"system_prompt": cmd.get("system_prompt", ""),
@ -426,9 +422,8 @@ 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.
# Optional template/tool plumbing: only forward keys that are present so
# the backend signature can evolve without breaking older payloads.
for opt_key in (
"tools",
"enable_thinking",
@ -438,7 +433,7 @@ 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
# Choose generation path.
use_adapter = cmd.get("use_adapter")
if use_adapter is not None:
generator = backend.generate_with_adapter_control(
@ -451,7 +446,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 +503,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 +537,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)
# Decode audio array from list (numpy arrays can't go through mp.Queue).
audio_array = np.array(cmd["audio_data"], dtype = np.float32)
audio_type = cmd.get("audio_type")
@ -638,12 +633,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 +663,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 +697,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 +755,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 +770,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 +843,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 +868,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
# Load a new model in this subprocess; unload the current one first.
if backend.active_model_name:
backend.unload_model(backend.active_model_name)
_handle_load(backend, cmd, resp_queue)
@ -891,7 +885,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 +901,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
)
elif cmd_type == "status":
# Return current status
# Return current status.
_send_response(
resp_queue,
{
@ -927,7 +921,7 @@ 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
# Unload all models.
for model_name in list(backend.models.keys()):
try:
backend.unload_model(model_name)

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

@ -3,24 +3,24 @@
"""Tool-call XML parsing and stripping helpers.
Extracted verbatim from studio/backend/core/inference/llama_cpp.py so that
Extracted verbatim from studio/backend/core/inference/llama_cpp.py so
external inference servers (llama-server wrappers, llama-swap, custom
shims) can reuse the same logic without importing the inference
orchestrator, structlog, httpx, or the rest of the studio backend.
shims) 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
Regexes and function bodies are byte-for-byte identical to the original
inline implementation in llama_cpp.py; any change must preserve that
equivalence. tests/python/test_tool_healing_extraction_is_exact.py
verifies it with AST comparison.
"""
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
# function/parameter name char-class tracks OpenAI's allowed set so MCP
# tool names (mcp__srv__list-issues) and parameter names (`issue-number`)
# with dashes 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 +46,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
@ -96,11 +96,10 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
# All closing tags optional -- models frequently omit </parameter>,
# </function>, and/or </tool_call>.
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 all <function=name> positions and extract bodies.
# Use only </tool_call> or the next <function= as hard boundaries;
# </function> is avoided since code values can contain it. Trim a
# trailing </function> afterwards if present.
func_starts = list(_TC_FUNC_START_RE.finditer(content))
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
@ -114,18 +113,18 @@ 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)
# Trim the real closing </function> tag if present
body = _TC_FUNC_CLOSE_RE.sub("", body)
# 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 (the common code/command/query case), 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>.
# Single parameter: value is everything after the tag to
# end of body, trimming any trailing </parameter>.
pm = param_starts[0]
val = body[pm.end() :]
val = _TC_PARAM_CLOSE_RE.sub("", val)
@ -134,7 +133,7 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
for pidx, pm in enumerate(param_starts):
param_name = pm.group(1)
val_start = pm.end()
# Value ends at next <parameter= or end of body
# Value ends at the next <parameter= or end of body
next_param = (
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)

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 submodule: backends and trainer classes."""
from .training import TrainingBackend, TrainingProgress, get_training_backend

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

File diff suppressed because it is too large Load diff

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

@ -4,12 +4,12 @@
"""
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).
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.
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.
"""
@ -47,9 +47,9 @@ _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 and any non-numeric-suffix tmp dir are
user-owned and 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 +95,8 @@ 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 — here so the parent process never
imports the heavy ML modules."""
epoch: float = 0
step: int = 0
@ -118,7 +118,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 +136,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 +169,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,8 +244,8 @@ 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 and CPT keep the
# explicit request so 4-bit adapter/raw-text runs stay possible.
if config["training_type"] == "Full Finetuning":
config["load_in_4bit"] = False
@ -296,8 +296,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 — safe: old pump thread confirmed dead, proc.start() succeeded
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
@ -325,7 +324,7 @@ class TrainingBackend:
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
@ -368,7 +367,7 @@ class TrainingBackend:
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)
# (8s covers SQLite's 5s lock timeout plus execution overhead)
if self._pump_thread is not None and self._pump_thread.is_alive():
self._pump_thread.join(timeout = 8.0)
@ -399,7 +398,7 @@ class TrainingBackend:
if p.is_completed or p.error:
return False
# Check status message for activity indicators
# Check status message for activity
status_lower = (p.status_message or "").lower()
if any(
k in status_lower
@ -530,9 +529,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 +540,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; update progress after coercion
_raw_loss = event.get("loss")
_raw_lr = event.get("learning_rate")
try:
@ -580,7 +578,7 @@ class TrainingBackend:
if status:
self._progress.status_message = status
# Update metric histories — reuse sanitized values from above
# Update metric histories — reuse sanitized values above
step = event.get("step", 0)
loss = _safe_loss
lr = _safe_lr
@ -616,7 +614,7 @@ class TrainingBackend:
else:
eval_loss = None
# Buffer metric for DB flush (loss/lr already sanitized above)
# Buffer metric for DB flush (loss/lr already sanitized)
self._metric_buffer.append(
{
"step": step,
@ -630,7 +628,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 +782,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 +829,7 @@ class TrainingBackend:
return events
# ------------------------------------------------------------------
# Plot generation (unchanged from original)
# Plot generation
# ------------------------------------------------------------------
def _create_loss_plot(
@ -954,9 +952,9 @@ 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.
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. No-op placeholder.
"""
logger.info(
"_transfer_to_inference_backend: subprocess training — "

398
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,13 @@ _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.
# Worker subprocesses inherit the parent env but not its os.add_dll_directory
# registrations. Replicate main.py's Windows ROCm DLL setup at module load so
# the first `import torch` finds amdhip64.dll even when HIP_PATH\bin is not on
# PATH. Handles retained at module scope so they aren't GC'd.
_ROCM_DLL_HANDLES: list = []
if sys.platform == "win32":
@ -147,21 +144,20 @@ 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).
BOTH the gcc runtime dir AND ``/usr/include/c++/<N>`` C++ headers, or
``None`` if no match (or non-Linux / non-x86_64).
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::
Ubuntu 24.04 ships gcc-14 runtime objects but not ``/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, 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).
via ``HIPCC_COMPILE_FLAGS_APPEND``. Mirrors the loop ``bbf004c`` added to
``studio/setup.sh`` for the llama.cpp HIP build (PR #5301).
"""
if not sys.platform.startswith("linux"):
return None
@ -300,11 +296,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,
@ -313,15 +308,14 @@ 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
# mamba-ssm / flash-attn source fallbacks) 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).
# Inject --gcc-install-dir for a gcc whose C++ headers exist. Respect any
# pre-existing --gcc-install-dir in HIPCC_COMPILE_FLAGS_APPEND; else
# append. Mirrors fix bbf004c in studio/setup.sh (llama.cpp HIP, 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 +363,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 +481,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 +493,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 +610,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 +684,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 +720,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 +764,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 +795,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 +817,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,9 +850,9 @@ 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.
# Wrap transformers' is_{flash_linear_attention,causal_conv1d}_available so the
# first call (at modeling import) drives the install. Models that query the gate
# get the install; those that never query it (Llama, Gemma, dense Qwen) pay nothing.
# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the legacy substring path.
@ -894,8 +889,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 +932,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 +961,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 +1052,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 +1074,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 +1087,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 +1213,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 +1271,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 +1309,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 +1413,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 +1506,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 +1520,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 +1556,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 +1711,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 +1748,15 @@ 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
# 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.
# Scoped to this subprocess (orchestrator spawns a fresh one).
if "HF_HUB_OFFLINE" not in os.environ:
import socket as _socket
import threading as _threading
@ -1807,7 +1802,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# ── 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.
# activation, causal-conv1d install, and torch imports.
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
@ -1827,7 +1822,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,10 +1855,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).
# NemotronH has transformers config-parsing bugs needing 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).
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_lowered = model_name.lower()
@ -1880,19 +1874,17 @@ 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.
# 1) causal-conv1d ALWAYS runs eagerly via the substring path. Some SSM
# modeling files (nemotron_h, falcon_h1, granitemoehybrid) call
# `lazy_load_kernel("causal-conv1d")` directly and never call
# `is_causal_conv1d_available()`, so the runtime hook wouldn't fire.
# 2) FLA + tilelang: primary gate is the runtime hook on
# `is_flash_linear_attention_available`. Architectures that query it
# auto-trigger the install; others never pay. `_install_fast_path_hooks`
# also wraps `is_causal_conv1d_available` as defence in depth.
# 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.
# 4) `UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1` falls back to the substring
# path for FLA / tilelang.
try:
_ensure_causal_conv1d_fast_path(event_queue, model_name)
if os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1":
@ -1923,12 +1915,11 @@ 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.
# Parent launched us via spawn, but the compiled SFTTrainer checks
# get_start_method() and disables num_proc if not "fork". Linux only: fork
# is the default and safe here (no CUDA context yet). macOS defaults to spawn
# since Python 3.8 (fork unsafe with Metal/MPS, CoreFoundation) -- do NOT
# override there. Windows has no fork.
if sys.platform == "linux":
import multiprocessing as _mp
try:
@ -1951,15 +1942,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.
# Windows ROCm: no RCCL backend). No-op off Windows ROCm. Must run before any
# import of transformers / unsloth_zoo.
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
# ── 1e. Ensure torch.distributed helper attrs are present ──
# Single-GPU training never initialises the process group, so these helpers
# are never called — but transformers/trl import them unconditionally.
# Single-GPU training never inits the process group, so these helpers are
# never called — but transformers/trl import them unconditionally.
_td_stubs = {
"is_initialized": lambda: False,
"is_available": lambda: False,
@ -1989,16 +1980,14 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# 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.
# Root cause: JitDecomp autograd decomposition (NOT torch.compile) dispatches
# _grouped_mm → _fused_adagrad_ → _grouped_mm HIP → null crash.
# TORCHDYNAMO_DISABLE=1 stops the compiler frontend but not JitDecomp, so we
# also override the CUDA dispatch key for _grouped_mm with a 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.
# Fixed in torch==2.11.0+rocm7.13.0: the 3-D batch and grouped (with offs)
# variants now have working HIP kernels on gfx1200. We gate the override on
# HIP < 7.13 so the fixed wheel uses the real GPU kernel.
#
# Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0.
#
@ -2006,17 +1995,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# 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.
# torch is already in sys.modules from section 1e. Module-level
# _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 torch.__version__. AMD SDK /
# Radeon Windows wheels don't 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 and main.py).
_build_version_for_rocm = (
getattr(_torch_for_rocm, "__version__", "").lower()
if _torch_for_rocm is not None
@ -2030,20 +2017,19 @@ 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 below is
# the real fix, but dynamo off avoids any other compile paths).
if "TORCHDYNAMO_DISABLE" not in os.environ:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.info("Windows ROCm: torch.compile (dynamo) disabled")
# BNB auto-detects the HIP version from torch.version.hip and uses
# it to choose which DLL to load (e.g. "7.13" → rocm713.dll).
# AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its
# version suffix does not always match the torch HIP version (e.g.
# torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still
# ships rocm72.dll). We detect the actual DLL name from the installed
# package and override BNB's auto-detection. "72" is a safe fallback
# if detection fails. Callers may override by pre-setting the var.
# BNB auto-detects the HIP version from torch.version.hip to pick a
# DLL (e.g. "7.13" → rocm713.dll). AMD's Windows BNB prerelease wheel
# ships only one rocm DLL whose version suffix may not 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). Detect the actual DLL name from
# the installed package and override auto-detection. "72" is a safe
# fallback if detection fails. Callers may override by pre-setting the var.
if "BNB_ROCM_VERSION" not in os.environ:
_bnb_rocm_ver = None
try:
@ -2064,10 +2050,9 @@ 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.
# Pick the highest numeric suffix so "713" wins over
# "72" when both are present. Glob order isn't guaranteed,
# so always sort rather than take the first match.
if _all_vers:
_bnb_rocm_ver = max(_all_vers, key = lambda v: int(v))
except Exception:
@ -2081,17 +2066,15 @@ 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. torch.version.hip
# can be "7.13.99004", "7.2.0", etc. AMD SDK / Radeon wheels may leave
# it unset and encode the ROCm version in torch.__version__ instead
# (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back to
# that string when version.hip is missing.
def _hip_ver_at_least(major: int, minor: int) -> bool:
_hip_str = getattr(getattr(_torch_for_rocm, "version", None), "hip", None)
if not _hip_str:
# Try the standard "+rocmX.Y.Z" embedded version first
# (e.g. "2.11.0+rocm7.13.0").
# 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 (
@ -2100,12 +2083,11 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
) >= (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.
# explicit rocmX.Y. The rocmsdk format postdates the gfx120X
# null-kernel fix (ROCm 7.13), so any such wheel has working
# HIP kernels. Treat as >= 7.13 rather than returning False
# and installing the Python workaround on a wheel that
# doesn't need it.
if "rocmsdk" in _build_version_for_rocm:
logger.debug(
"Windows ROCm: AMD SDK wheel detected (%r); "
@ -2139,10 +2121,9 @@ 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.
# _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12 (causing
# 0xC0000005); fixed in ROCm 7.13 (torch 2.11+). Install the Python
# fallback only on affected versions so 7.13+ gets the real GPU kernel.
if not _hip_ver_at_least(7, 13):
try:
import warnings as _warnings
@ -2159,24 +2140,22 @@ 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: match the real op — (M, K) x (K, N) ->
# mm, or 3-D batched -> bmm. An unconditional torch.mm
# previously raised "self must be a matrix" on 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 in `self`; mat2 may be 3-D or 2-D.
offs_list = offs.tolist()
pieces = []
prev = 0
@ -2189,7 +2168,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 +2216,23 @@ 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.
# On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can hang
# the HIP driver and freeze the whole system instead of raising. So
# set_per_process_memory_fraction caps the HIP allocator and PyTorch raises
# OutOfMemoryError before the hardware limit, giving the UI a clean error.
# ROCm only -- NVIDIA CUDA has a graceful OOM path.
# 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.
# and system RAM in one pool: 0.90 of 128 GB starves the OS, so use 0.80.
# Primary classifier: gcnArchName (stable within a product family,
# naming-independent). AMD SDK / Radeon wheels may omit it or use a variant
# spelling, so try several attrs then fall back to device-name markers.
# Non-fatal: silently skipped if torch is not importable.
if _hw.IS_ROCM:
try:
import torch as _torch_mem
if _torch_mem.cuda.is_available():
# Classify unified vs discrete via _rocm_classify_unified_memory.
# See that function's docstring for classification priority.
# 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 +2286,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 +2356,9 @@ 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
# (incl. LLM-assisted detection) 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
@ -2451,8 +2427,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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 +2451,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 +2509,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 +2586,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 +2620,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 +2737,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 +2834,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 +2887,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 +3024,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 +3048,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,

24
studio/backend/loggers/config.py
View file

@ -1,15 +1,14 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Logging configuration for structured logging with structlog.
"""Structured logging configuration via structlog.
This module provides centralized logging configuration with environment-specific
formats and processors. Supports both development and production environments
with consistent structured logging.
Centralized config with environment-specific formats and processors for
development and production.
Key Features:
- Environment-specific formatting (JSON for production, console for development)
- Timestamp standardization (ISO format)
Key features:
- Environment-specific formatting (JSON for prod, console for dev)
- ISO timestamps
- Context variable integration
- Log level filtering
- Logger caching for performance
@ -28,8 +27,8 @@ from loggers.handlers import filter_sensitive_data
class LogConfig:
"""Structured logging configuration for the application.
Provides static method to configure structlog with environment-specific
formatting and processors for consistent structured logging.
Static method to configure structlog with environment-specific
formatting and processors.
"""
@staticmethod
@ -41,9 +40,8 @@ class LogConfig:
service_name: Name of the service for logging identification
env: Environment (development/production), affects logging format
"""
# Determine log level from environment
# Log level from environment; fall back to INFO if invalid.
log_level_name = os.getenv("LOG_LEVEL", "INFO").upper()
# Fallback to INFO if an invalid level is provided
log_level = getattr(logging, log_level_name, logging.INFO)
if sys.platform == "win32":
@ -56,13 +54,13 @@ class LogConfig:
structlog.configure(
processors = [
# Reorder processors to control field order
# Ordered to control output field order.
structlog.processors.TimeStamper(fmt = "iso"), # timestamp first
structlog.processors.add_log_level, # level second
structlog.contextvars.merge_contextvars,
structlog.processors.format_exc_info,
filter_sensitive_data,
# Custom processor to flatten the extra field
# Flatten the extra field into the main dict.
lambda logger, method_name, event_dict: {
"timestamp": event_dict.get("timestamp"),
"level": event_dict.get("level"),

21
studio/backend/loggers/handlers.py
View file

@ -3,16 +3,16 @@
"""Logging handlers and middleware for structured logging.
This module provides FastAPI middleware and structlog processors for:
FastAPI middleware + structlog processors for:
- Request/response logging with timing
- Sensitive data filtering in logs
- Structured logging configuration
- Error handling with detailed context
- Sensitive data filtering
- Structured logging config
- Error handling with context
Key Components:
- LoggingMiddleware: FastAPI middleware for request/response logging
- filter_sensitive_data: Structlog processor for data sanitization
- get_logger: Factory function for structured loggers
Key components:
- LoggingMiddleware: request/response logging
- filter_sensitive_data: structlog processor for sanitization
- get_logger: factory for structured loggers
"""
import re
@ -38,7 +38,6 @@ class LoggingMiddleware(BaseHTTPMiddleware):
try:
response = await call_next(request)
# Log response
process_time = (time.time() - start_time) * 1000
EXCLUDED_PATHS = {
@ -108,10 +107,8 @@ def filter_sensitive_data(logger, method_name, event_dict):
def get_logger(name: str) -> structlog.BoundLogger:
"""Get a logger instance for a specific module.
"""Get a bound structured logger for a module.
Args:
name: Usually __name__ of the module
Returns:
A bound structured logger
"""
return structlog.get_logger(name)

185
studio/backend/main.py
View file

@ -9,20 +9,20 @@ import os
import sys
from pathlib import Path as _Path
# Suppress annoying C-level dependency warnings globally
# Suppress C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
if sys.platform == "win32":
# Retained at module scope -- os.add_dll_directory returns a handle that
# Retained at module scope; os.add_dll_directory returns a handle that
# removes the search-path entry when garbage collected.
_ROCM_DLL_HANDLES: list = []
def _add_rocm_dll_dirs() -> None:
candidates = []
# 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer
# 1. HIP_PATH / ROCM_PATH set by the AMD HIP SDK installer
for _var in ("HIP_PATH", "ROCM_PATH"):
_val = os.environ.get(_var)
if _val:
@ -34,7 +34,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"; non-numeric chunks fall back to string
parts = []
for chunk in name.split("."):
try:
@ -62,16 +62,15 @@ if sys.platform == "win32":
del _add_rocm_dll_dirs
# ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
# bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm<ver>.dll
# where <ver> comes from torch.version.hip (e.g. "7.13..." → "713").
# The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without
# this the server process crashes with "Configured ROCm binary not found".
# Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION
# before any import that pulls in bitsandbytes (mirrors worker.py logic).
# Gate on the rocm bnb DLL (the exact file this configures) or HIP_PATH/
# ROCM_PATH, not on torch.version.hip: that needed importing torch on every
# Windows host (NVIDIA/CPU included), adding seconds to startup. Radeon
# wheels without HIP_PATH still ship the rocm bnb DLL, so they are covered.
# bitsandbytes on Windows ROCm loads libbitsandbytes_rocm<ver>.dll where
# <ver> comes from torch.version.hip (e.g. "7.13..." → "713"). The installed
# BNB wheel ships rocm72.dll (not rocm713.dll), so without this the server
# crashes with "Configured ROCm binary not found". Detect the available DLL,
# fall back to "72", and set BNB_ROCM_VERSION before any bitsandbytes import
# (mirrors worker.py). Gate on the rocm bnb DLL (the file this configures) or
# HIP_PATH/ROCM_PATH, not torch.version.hip: that needed importing torch on
# every Windows host (NVIDIA/CPU too), adding startup seconds. Radeon wheels
# without HIP_PATH still ship the rocm bnb DLL, so they're covered.
if "BNB_ROCM_VERSION" not in os.environ:
import glob as _glob
import logging as _logging
@ -82,7 +81,7 @@ if sys.platform == "win32":
try:
import importlib.util as _ilu
_bnb_spec = _ilu.find_spec("bitsandbytes")
# submodule_search_locations (not spec.origin) handles editable installs.
# submodule_search_locations (not spec.origin) handles editable installs
if _bnb_spec and _bnb_spec.submodule_search_locations:
import re as _re_bnb
@ -102,7 +101,7 @@ if sys.platform == "win32":
"Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'",
_e,
)
# rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72").
# rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72")
if _found_rocm_bnb or _hip_env:
_bnb_rocm_ver_final = _bnb_rocm_ver or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
@ -111,8 +110,8 @@ if sys.platform == "win32":
_bnb_rocm_ver_final,
)
# Ensure backend dir is on sys.path so _platform_compat is importable when
# main.py is launched directly (e.g. `uvicorn main:app`).
# Put backend dir on sys.path so _platform_compat is importable when main.py
# is launched directly (e.g. `uvicorn main:app`).
_backend_dir = str(_Path(__file__).parent)
if _backend_dir not in sys.path:
sys.path.insert(0, _backend_dir)
@ -126,14 +125,14 @@ except ValueError as exc:
_raw = os.environ.get("UNSLOTH_CPU_THREADS")
raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}") from None
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# Anaconda/conda-forge Python: seed platform._sys_version_cache before any
# library import triggers attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
import _platform_compat # noqa: F401
# Direct `uvicorn main:app` launches bypass run.py, so re-export here too
# (mirrors run.py). Required BEFORE the unsloth-zoo import below, since
# its LLAMA_CPP_DEFAULT_DIR binding is import-time.
# (mirrors run.py). Required BEFORE the unsloth-zoo import below, whose
# LLAMA_CPP_DEFAULT_DIR binding is import-time.
from utils.paths.storage_roots import studio_root as _studio_root
try:
@ -166,15 +165,13 @@ _STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$")
def _read_studio_install_id() -> str:
"""Per-install opaque id written by install.sh / install.ps1 at
$STUDIO_HOME/share/studio_install_id. Returns "" when the file is
absent (pre-PR install, fresh tree never run through the installer)
or contains anything other than a 64-char lowercase-hex token --
in which case /api/health emits "" and the launcher's _check_health
falls back to the existing "no baked id, accept any healthy
Unsloth backend" path. This intentionally replaces a previous
sha256(resolved_install_path) so the field carries no install-path
information for callers reaching /api/health (relevant when Studio
is run with -H 0.0.0.0)."""
$STUDIO_HOME/share/studio_install_id. Returns "" when the file is absent
(pre-PR install, or fresh tree never run through the installer) or holds
anything other than a 64-char lowercase-hex token; then /api/health emits
"" and the launcher's _check_health falls back to "no baked id, accept any
healthy Unsloth backend". Replaces a previous sha256(resolved_install_path)
so the field carries no install-path info for callers reaching /api/health
(relevant when Studio runs with -H 0.0.0.0)."""
try:
token = (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
except (OSError, ValueError):
@ -186,29 +183,28 @@ _STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()
def _studio_root_id() -> str:
"""Same-install discriminator for /api/health: a per-install opaque
token written once by the installer and read once at module import.
Empty when no installer-written token is present; the launcher
contract treats "" as "no baked id, accept any healthy backend"."""
"""Same-install discriminator for /api/health: a per-install opaque token
written once by the installer and read once at module import. Empty when no
installer token is present; the launcher contract treats "" as "no baked id,
accept any healthy backend"."""
return _STUDIO_ROOT_ID_CACHE
# Fix broken Windows registry MIME types. Some Windows installs map .js to
# "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes
# module reads from the registry, and FastAPI/Starlette's StaticFiles uses
# mimetypes.guess_type() to set Content-Type headers. Browsers enforce strict
# MIME checking for ES module scripts (<script type="module">) and will refuse
# to execute .js files served as text/plain — resulting in a blank page.
# Calling add_type() *before* StaticFiles is instantiated ensures the correct
# types are used regardless of the OS registry.
# Fix broken Windows registry MIME types. Some Windows installs map .js to
# "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes reads
# from the registry, and Starlette's StaticFiles uses mimetypes.guess_type() for
# Content-Type. Browsers strictly MIME-check ES module scripts
# (<script type="module">) and refuse .js served as text/plain, giving a blank
# page. Calling add_type() *before* StaticFiles is instantiated forces the
# correct types regardless of the OS registry.
if sys.platform == "win32":
mimetypes.add_type("application/javascript", ".js")
mimetypes.add_type("text/css", ".css")
# Suppress annoying dependency warnings in production
# Suppress dependency warnings in production
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
# Alternatively, you can be more specific:
# Or be more specific:
# warnings.filterwarnings("ignore", category=DeprecationWarning)
# warnings.filterwarnings("ignore", module="triton.*")
@ -300,7 +296,7 @@ def _desktop_owner() -> dict[str, str] | None:
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
# Clean up any stale compiled cache from previous runs
# Clean up stale compiled cache from previous runs
clear_unsloth_compiled_cache()
# Remove stale .venv_overlay from previous versions — no longer used.
@ -354,8 +350,8 @@ async def lifespan(app: FastAPI):
import structlog
structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc)
# Pre-cache the helper GGUF model for LLM-assisted dataset detection.
# Runs in a background thread so it doesn't block server startup.
# Pre-cache the helper GGUF model for LLM-assisted dataset detection,
# in a background thread so it doesn't block server startup.
import threading
def _precache():
@ -412,7 +408,7 @@ app.add_middleware(LoggingMiddleware)
# Citation favicons load from www.google.com/s2/favicons; *.gstatic.com is
# kept for legacy web-search faviconV2 paths. Everything else is same-origin.
# kept for legacy web-search faviconV2 paths. All else is same-origin.
from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402
from starlette.requests import Request as _StarletteRequest # noqa: E402
@ -421,7 +417,7 @@ _CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
_ARTIFACT_PREVIEW_FRAME_PATH = "/api/inference/artifact-preview-frame"
# /content is Colab's working directory — more reliable than env vars which
# /content is Colab's working directory — more reliable than env vars, which
# aren't always set depending on Colab runtime version.
import importlib.util as _importlib_util
@ -437,19 +433,18 @@ def _build_csp(script_nonce: "str | None" = None) -> str:
if script_nonce:
script_src += f" 'nonce-{script_nonce}'"
# In Colab the parent frame can be colab.research.google.com, a multi-level
# *.prod.colab.dev subdomain (e.g. foo.region.prod.colab.dev — note: CSP
# wildcards only match one level, so *.prod.colab.dev misses these), or a
# sandboxed null-origin output iframe. Use '*' so any ancestor is allowed;
# Colab is already a sandboxed single-user environment.
# *.prod.colab.dev subdomain (e.g. foo.region.prod.colab.dev — CSP wildcards
# only match one level, so *.prod.colab.dev misses these), or a sandboxed
# null-origin output iframe. Use '*' to allow any ancestor; Colab is already
# a sandboxed single-user environment.
frame_ancestors = "*" if _IS_COLAB else "'none'"
# In Colab the frontend is served over the Colab reverse-proxy at an HTTPS
# *.prod.colab.dev URL. Colab's kernel communication layer and the output
# iframe scaffolding inject scripts from *.prod.colab.dev and
# *.googleusercontent.com, and make fetch/WebSocket connections to those
# same origins. Widen script-src and connect-src in Colab mode so those
# requests are not blocked. 'unsafe-inline' for scripts is still omitted;
# our own inline script uses a nonce.
# *.prod.colab.dev URL. Colab's kernel layer and output-iframe scaffolding
# inject scripts from *.prod.colab.dev and *.googleusercontent.com, and
# fetch/WebSocket to those origins. Widen script-src and connect-src in
# Colab mode so those aren't blocked. 'unsafe-inline' for scripts is still
# omitted; our inline script uses a nonce.
if _IS_COLAB:
script_src += " https://*.prod.colab.dev https://*.googleusercontent.com"
connect_src = (
@ -482,7 +477,7 @@ class SecurityHeadersMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: _StarletteRequest, call_next):
response = await call_next(request)
# Strip the internal nonce hand-off header so it never reaches the client.
# Strip the internal nonce hand-off header so it never reaches the client
nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER)
if nonce is not None:
del response.headers[_CSP_SCRIPT_NONCE_HEADER]
@ -505,7 +500,7 @@ app.add_middleware(SecurityHeadersMiddleware)
# Cap request bodies on protected POSTs. Upload routes get explicit multipart
# headroom, while non-upload routes keep the default body cap.
# headroom; non-upload routes keep the default body cap.
import json as _json_for_413 # noqa: E402
from utils.upload_limits import ( # noqa: E402
UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES,
@ -651,7 +646,7 @@ class MaxBodyMiddleware:
if mtype == "http.disconnect":
return
if mtype != "http.request":
# Mid-stream unexpected frame: forwarding would corrupt downstream.
# Mid-stream unexpected frame: forwarding would corrupt downstream
return
body = msg.get("body", b"") or b""
if body:
@ -730,13 +725,13 @@ app.include_router(training_router, prefix = "/api/train", tags = ["training"])
app.include_router(models_router, prefix = "/api/models", tags = ["models"])
app.include_router(chat_history_router, prefix = "/api/chat", tags = ["chat"])
app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
# Studio-only inference endpoints (cancel, etc.) are intentionally NOT
# exposed on the /v1 OpenAI-compat prefix below.
# Studio-only inference endpoints (cancel, etc.) are NOT exposed on the /v1
# OpenAI-compat prefix below.
app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"])
# OpenAI-compatible endpoints: mount the same inference router at /v1
# so external tools (Open WebUI, SillyTavern, etc.) can use the
# standard /v1/chat/completions path.
# OpenAI-compatible endpoints: mount the same inference router at /v1 so
# external tools (Open WebUI, SillyTavern, etc.) can use the standard
# /v1/chat/completions path.
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
app.include_router(settings_router, prefix = "/api/settings", tags = ["settings"])
@ -757,9 +752,9 @@ async def health_check(request: Request):
Unauthenticated callers (Tauri watchdog, frontend bootstrap polls) need
``service`` / ``studio_root_id`` / ``chat_only`` / ``desktop_*`` / ``native_path_leases_supported``
to (a) re-adopt a sibling backend across restarts and (b) gate UI surfaces
before any token is available. None of those leak install path or version.
``version`` / ``studio_version`` / ``device_type`` still require a bearer
because they fingerprint the host.
before a token is available. None of those leak install path or version.
``version`` / ``studio_version`` / ``device_type`` require a bearer since
they fingerprint the host.
"""
base = {
"status": "healthy",
@ -783,7 +778,7 @@ async def health_check(request: Request):
from fastapi.security import HTTPAuthorizationCredentials
creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = auth.split(" ", 1)[1])
# Must await: a bare coroutine is truthy and would skip the auth check.
# Must await: a bare coroutine is truthy and would skip the auth check
subject = await _gcs(creds)
except HTTPException:
return base
@ -819,7 +814,7 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c
"""Gracefully shut down the Unsloth Studio server.
Called by the frontend quit dialog so users can stop the server from the UI
without needing to use the CLI or kill the process manually.
without the CLI or killing the process manually.
"""
import asyncio
@ -842,11 +837,11 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c
async def get_system_info(current_subject: str = Depends(get_current_subject)):
"""Get system information.
Gated behind auth: the response includes platform, Python version,
GPU name, memory total, and ML package set -- enough to fingerprint
a host. Studio's chat-only-mode design assumes only the local user
reaches /api/system; in -H 0.0.0.0 / Colab / Tauri-relayed setups
that assumption breaks unless we require a bearer.
Gated behind auth: the response includes platform, Python version, GPU name,
memory total, and ML package set -- enough to fingerprint a host. Studio's
chat-only design assumes only the local user reaches /api/system; in
-H 0.0.0.0 / Colab / Tauri-relayed setups that breaks unless we require a
bearer.
"""
import platform
import psutil
@ -865,9 +860,8 @@ async def get_system_info(current_subject: str = Depends(get_current_subject)):
return {
"platform": platform.platform(),
"python_version": platform.python_version(),
# Use the centralized _backend_label helper so the /api/system
# endpoint reports "rocm" on AMD hosts instead of "cuda", matching
# the /api/hardware and /api/gpu-visibility endpoints.
# Centralized _backend_label so /api/system reports "rocm" on AMD hosts
# instead of "cuda", matching /api/hardware and /api/gpu-visibility.
"device_backend": _backend_label(get_device()),
"cpu_count": psutil.cpu_count(),
"memory": {
@ -888,8 +882,8 @@ async def get_gpu_visibility(current_subject: str = Depends(get_current_subject)
async def get_hardware_info(current_subject: str = Depends(get_current_subject)):
"""Return GPU name, total VRAM, and key ML package versions.
Gated behind auth alongside /api/system -- same fingerprinting
concern. /api/system/gpu-visibility is also auth-gated already.
Gated behind auth alongside /api/system -- same fingerprinting concern.
/api/system/gpu-visibility is also auth-gated.
"""
from utils.hardware import get_gpu_summary, get_package_versions
return {
@ -904,11 +898,10 @@ async def get_hardware_info(current_subject: str = Depends(get_current_subject))
def _strip_crossorigin(html_bytes: bytes) -> bytes:
"""Remove ``crossorigin`` attributes from script/link tags.
Vite adds ``crossorigin`` by default which forces CORS mode on font
subresource loads. When Studio is served over plain HTTP, Firefox
HTTPS-Only Mode does not exempt CORS font requests -- causing all
@font-face downloads to fail silently. Stripping the attribute
makes them regular same-origin fetches that work on any protocol.
Vite adds ``crossorigin`` by default, forcing CORS mode on font subresource
loads. Over plain HTTP, Firefox HTTPS-Only Mode doesn't exempt CORS font
requests, so all @font-face downloads fail silently. Stripping the attribute
makes them same-origin fetches that work on any protocol.
"""
html = html_bytes.decode("utf-8")
html = _re.sub(r'\s+crossorigin(?:="[^"]*")?', "", html)
@ -917,8 +910,8 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes:
def _inject_bootstrap(html_bytes: bytes, app: FastAPI):
"""Inject bootstrap credentials when password change is pending.
Returns ``(html_bytes, script_nonce_or_None)``; callers forward the
nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so CSP allows the inline script.
Returns ``(html_bytes, script_nonce_or_None)``; callers forward the nonce
via ``_CSP_SCRIPT_NONCE_HEADER`` so CSP allows the inline script.
"""
import json as _json
import secrets as _secrets
@ -949,9 +942,9 @@ _DEFAULT_PORTS = {"http": 80, "https": 443, "ws": 80, "wss": 443}
def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]]:
"""Canonicalise an Origin to ``(scheme, host, port)`` for equality.
Browsers strip default ports (RFC 6454 sec 6.1) and scheme/host are
case-insensitive (RFC 3986), so bare string compare misclassifies
same-origin requests as cross-origin. Returns ``None`` on unparseable
input so callers fall to the safer cross-origin default.
case-insensitive (RFC 3986), so a bare string compare misclassifies
same-origin requests as cross-origin. Returns ``None`` on unparseable input
so callers fall to the safer cross-origin default.
"""
scheme = (scheme or "").strip().lower()
if not scheme or not netloc:
@ -960,7 +953,7 @@ def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]
if "@" in netloc:
netloc = netloc.rsplit("@", 1)[1]
# IPv6 hosts use brackets (RFC 3986 sec 3.2.2): ``[::1]:8902``. Bare
# ``partition(":")`` mis-parses these and breaks ``unsloth studio -H ::1``.
# ``partition(":")`` mis-parses these, breaking ``unsloth studio -H ::1``.
if netloc.startswith("["):
close = netloc.find("]")
if close == -1:
@ -993,7 +986,7 @@ def _is_same_origin_request(request: Request) -> bool:
Top-level same-document GETs omit Origin, so missing counts as same-origin.
Callers must also emit ``Vary: Origin``. Both sides are canonicalised via
:func:`_canonical_origin` so default-port stripping and scheme/host case
do not misclassify same-origin requests as cross-origin.
don't misclassify same-origin requests as cross-origin.
"""
origin = request.headers.get("origin")
if origin is None:
@ -1064,7 +1057,7 @@ def setup_frontend(app: FastAPI, build_path: Path):
file_path = (build_path / full_path).resolve()
# Block path traversal — ensure resolved path stays inside build_path
# Block path traversal — resolved path must stay inside build_path
if not file_path.is_relative_to(build_path.resolve()):
return Response(status_code = 403)

4
studio/backend/models/__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
"""
Pydantic models for API request/response schemas
"""
"""Pydantic models for API request/response schemas."""
from .training import (
TrainingStartRequest,

4
studio/backend/models/auth.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
"""
Pydantic schemas for Authentication API
"""
"""Pydantic schemas for the Authentication API."""
from typing import Optional

4
studio/backend/models/data_recipe.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
"""
Pydantic schemas for Data Recipe (DataDesigner) API.
"""
"""Pydantic schemas for Data Recipe (DataDesigner) API."""
from __future__ import annotations

4
studio/backend/models/datasets.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
"""
Dataset-related Pydantic models for API requests and responses.
"""
"""Dataset Pydantic models for API requests and responses."""
from typing import Any, Dict, List, Optional

4
studio/backend/models/export.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
"""
Pydantic schemas for Export API.
"""
"""Pydantic schemas for Export API."""
from pathlib import Path

181
studio/backend/models/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
"""
Pydantic schemas for Inference API
"""
"""Pydantic schemas for the Inference API."""
from __future__ import annotations
@ -108,11 +106,8 @@ class UnloadRequest(BaseModel):
class ValidateModelRequest(BaseModel):
"""
Lightweight validation request to check whether a model identifier
*can be resolved* into a ModelConfig.
This does NOT actually load weights into GPU memory.
"""Lightweight check whether a model identifier *can be resolved* into a
ModelConfig. Does NOT load weights into GPU memory.
"""
model_path: str = Field(..., description = "Model identifier or local path")
@ -126,8 +121,7 @@ class ValidateModelRequest(BaseModel):
class ValidateModelResponse(BaseModel):
"""
Result of model validation.
"""Result of model validation.
valid == True means ModelConfig.from_identifier() succeeded and basic
introspection (GGUF / LoRA / vision flags) is available.
@ -245,10 +239,9 @@ class UnloadResponse(BaseModel):
class LoadProgressResponse(BaseModel):
"""Progress of the active GGUF load, sampled on demand.
Used by the UI to show a real progress bar during the
post-download warmup window (mmap + CUDA upload), rather than a
generic "Starting model..." spinner that freezes for minutes on
large MoE models.
Drives a real progress bar during the post-download warmup window
(mmap + CUDA upload), instead of a generic "Starting model..." spinner
that freezes for minutes on large MoE models.
"""
phase: Optional[str] = Field(
@ -409,10 +402,10 @@ class InputDocumentContentPart(BaseModel):
Studio-normalised shape. The frontend sends either
``{type:"input_document", file_data:"data:application/pdf;base64,..."}``
or ``{type:"input_document", file_url:"https://..."}``, plus optional
``filename`` and ``media_type``. ``external_provider`` translates this
onto Anthropic's ``document`` block or OpenAI Responses' ``input_file``
block for vision-capable providers; non-vision providers drop the
part entirely (handled in ``_build_external_messages``).
``filename`` and ``media_type``. ``external_provider`` maps this onto
Anthropic's ``document`` block or OpenAI Responses' ``input_file`` block
for vision-capable providers; non-vision providers drop the part
(handled in ``_build_external_messages``).
"""
type: Literal["input_document"]
@ -437,10 +430,10 @@ class InputDocumentContentPart(BaseModel):
class OpenAIReasoningContentPart(BaseModel):
"""OpenAI Responses reasoning item paired with a tool output.
Reasoning models can require the previous ``reasoning`` output item
to be replayed immediately before an ``image_generation_call`` id
when manually managing Responses context. This part is OpenAI-only;
routes strip it for every other provider before proxying.
Reasoning models can require the previous ``reasoning`` output item to be
replayed immediately before an ``image_generation_call`` id when manually
managing Responses context. OpenAI-only; routes strip it for every other
provider before proxying.
"""
type: Literal["reasoning"]
@ -452,12 +445,12 @@ class OpenAIReasoningContentPart(BaseModel):
class ImageGenerationCallContentPart(BaseModel):
"""OpenAI Responses image_generation call reference.
OpenAI accepts prior ``image_generation_call`` items in the next
Responses ``input`` array so follow-up prompts can edit or refine a
generated image without resending the base64 payload. The frontend
forwards this as a synthetic assistant content part when building
the next OpenAI Responses request; ``external_provider`` translates
it back to the provider-specific top-level input item.
OpenAI accepts prior ``image_generation_call`` items in the next Responses
``input`` array so follow-up prompts can edit or refine a generated image
without resending the base64 payload. The frontend forwards this as a
synthetic assistant content part when building the next request;
``external_provider`` maps it back to the provider-specific top-level
input item.
"""
type: Literal["image_generation_call"]
@ -472,14 +465,13 @@ class CompactionContentPart(BaseModel):
"""Anthropic server-side compaction state, attached to an assistant
message for round-tripping on the next turn.
When Anthropic runs compaction during a request, the response
carries a ``{"type": "compaction", "content": "<summary>"}`` block
on the assistant message. The chat-adapter persists it onto the
stored message; the next turn's outbound request must forward it
back so Anthropic recognises the existing compaction state and
doesn't re-summarise the conversation from scratch. See
``external_provider._stream_anthropic`` for the wire-side handling
and https://platform.claude.com/docs/en/build-with-claude/compaction
When Anthropic compacts during a request, the response carries a
``{"type": "compaction", "content": "<summary>"}`` block on the assistant
message. The chat-adapter persists it; the next turn's outbound request
must forward it back so Anthropic recognises the existing compaction state
and doesn't re-summarise from scratch. See
``external_provider._stream_anthropic`` for the wire-side handling and
https://platform.claude.com/docs/en/build-with-claude/compaction
for the upstream contract.
"""
@ -516,9 +508,9 @@ ContentPart = Annotated[
class ChatMessage(BaseModel):
"""Single message in a chat conversation.
``content`` is a string or a list of multimodal content parts. Assistant
messages with only ``tool_calls`` populated may set ``content=None``.
Missing ``tool_call_id`` on ``role="tool"`` is resolved at the
``content`` is a string or list of multimodal content parts. Assistant
messages with only ``tool_calls`` may set ``content=None``. Missing
``tool_call_id`` on ``role="tool"`` is resolved at the
``ChatCompletionRequest`` layer by walking back to the preceding assistant.
"""
@ -571,16 +563,15 @@ class ChatMessage(BaseModel):
class ChatCompletionRequest(BaseModel):
"""
OpenAI-compatible chat completion request.
"""OpenAI-compatible chat completion request.
Extensions (non-OpenAI fields) are marked with 'x-unsloth'.
Non-OpenAI extension fields are marked with 'x-unsloth'.
"""
# Accept unknown fields defensively so future OpenAI fields (seed,
# response_format, logprobs, frequency_penalty, etc.) don't get
# silently dropped by Pydantic before route code runs. Mirrors
# AnthropicMessagesRequest and ResponsesRequest.
# Accept unknown fields so future OpenAI fields (seed, response_format,
# logprobs, frequency_penalty, etc.) aren't silently dropped by Pydantic
# before route code runs. Mirrors AnthropicMessagesRequest and
# ResponsesRequest.
model_config = {"extra": "allow"}
model: str = Field(
@ -738,18 +729,16 @@ class ChatCompletionRequest(BaseModel):
@field_validator("enable_prompt_caching", mode = "before")
@classmethod
def _coerce_enable_prompt_caching(cls, value: Any) -> Any:
"""Preserve the pre-PR coercion: the field used to be Optional[bool],
so callers historically sent JSON strings `"true"` / `"false"` and
Pydantic v1 coerced them. Widening to Optional[Union[bool, str]] for
Gemini cache resource names lets `"false"` slip through as a truthy
string. Coerce the canonical bool literals back so explicit opt-outs
stay opt-out."""
"""Preserve the pre-PR coercion. The field was once Optional[bool], so
callers historically sent JSON strings `"true"` / `"false"` that
Pydantic v1 coerced. Widening to Optional[Union[bool, str]] for Gemini
cache resource names lets `"false"` slip through as truthy; coerce the
canonical bool literals back so explicit opt-outs stay opt-out."""
if isinstance(value, str):
lowered = value.strip().lower()
# Match Pydantic v1's BooleanField coercion table (yes/y/on/t/1
# and no/n/off/f/0) so opt-outs that used to parse still parse.
# Anything else is preserved as a string for Gemini's
# cachedContent resource path.
# and no/n/off/f/0) so old opt-outs still parse. Anything else
# stays a string for Gemini's cachedContent resource path.
if lowered in ("true", "t", "1", "yes", "y", "on"):
return True
if lowered in ("false", "f", "0", "no", "n", "off"):
@ -836,10 +825,10 @@ class ChatCompletionRequest(BaseModel):
OpenAI / Anthropic passthrough require the result id to match the
assistant's tool_calls[].id. Prefer function.name match, else first
unconsumed tool_call; synth random id only if no candidate exists.
Crossing a user turn breaks the lookup.
unconsumed tool_call; synth a random id only if none exists. A user
turn breaks the lookup.
"""
# Pre-mark explicit ids first so a sibling missing-id result does not
# Pre-mark explicit ids first so a sibling missing-id result doesn't
# steal one already claimed by name.
consumed: set[tuple[int, int]] = set()
@ -902,12 +891,10 @@ class ChatCompletionRequest(BaseModel):
class OpenAIContainerRequest(BaseModel):
"""
Shared body for the three OpenAI container endpoints (list / create
/ delete). Carries the encrypted API key + base URL so the route
handler can decrypt it and proxy to the user's OpenAI account.
Same pattern as the inference proxy endpoints — keeps the key off
persistent storage on the backend.
"""Shared body for the three OpenAI container endpoints (list / create /
delete). Carries the encrypted API key + base URL so the route handler
can decrypt it and proxy to the user's OpenAI account. Same pattern as the
inference proxy endpoints — keeps the key off backend persistent storage.
"""
encrypted_api_key: str = Field(
@ -1055,10 +1042,10 @@ class ResponsesOutputTextPart(BaseModel):
"""Assistant ``output_text`` content part replayed on subsequent turns.
When a client (OpenAI Codex CLI, OpenAI Python SDK agents) loops on a
stateless Responses endpoint, prior assistant messages are round-tripped
as ``{"role":"assistant","content":[{"type":"output_text","text":...,
"annotations":[],"logprobs":[]}]}``. We preserve the text and ignore
the annotations/logprobs metadata when flattening into Chat Completions.
stateless Responses endpoint, prior assistant messages round-trip as
``{"role":"assistant","content":[{"type":"output_text","text":...,
"annotations":[],"logprobs":[]}]}``. We keep the text and ignore the
annotations/logprobs metadata when flattening into Chat Completions.
"""
type: Literal["output_text"]
@ -1073,8 +1060,8 @@ class ResponsesUnknownContentPart(BaseModel):
"""Catch-all for content-part types we don't model explicitly.
Keeps validation green when a client sends newer part types (e.g.
``input_audio``, ``input_file``) we haven't mapped; these are silently
skipped during normalisation rather than rejected with a 422.
``input_audio``, ``input_file``) we haven't mapped; these are skipped
during normalisation rather than rejected with a 422.
"""
type: str
@ -1099,16 +1086,15 @@ class ResponsesInputMessage(BaseModel):
# Codex (gpt-5.3-codex+) attaches a `phase` field ("commentary" |
# "final_answer") to assistant messages and requires clients to preserve
# it on subsequent turns. We accept and round-trip it; llama-server does
# not care about it.
# it across turns. We accept and round-trip it; llama-server ignores it.
model_config = {"extra": "allow"}
class ResponsesFunctionCallInputItem(BaseModel):
"""A prior assistant function_call being replayed in a multi-turn Responses input.
"""A prior assistant function_call replayed in a multi-turn Responses input.
The Responses API represents tool calls as top-level input items (not
nested inside assistant messages), correlated across turns by ``call_id``.
nested in assistant messages), correlated across turns by ``call_id``.
"""
type: Literal["function_call"]
@ -1125,7 +1111,7 @@ class ResponsesFunctionCallInputItem(BaseModel):
class ResponsesFunctionCallOutputInputItem(BaseModel):
"""A tool result supplied by the client for a prior function_call.
Replaces Chat Completions' ``role="tool"`` message. Correlated to the
Replaces Chat Completions' ``role="tool"`` message. Correlated to its
originating call by ``call_id``.
"""
@ -1142,10 +1128,9 @@ class ResponsesUnknownInputItem(BaseModel):
"""Catch-all for Responses input item types we don't model explicitly.
Covers ``reasoning`` items (replayed from prior o-series / gpt-5 turns)
and any future item types the client may send. These items are dropped
during normalisation — llama-server-backed GGUFs cannot consume them —
but keeping them in the request-model union stops unrelated turns from
failing validation with a 422.
and any future item types. Dropped during normalisation (llama-server
GGUFs can't consume them), but keeping them in the request-model union
stops unrelated turns from failing validation with a 422.
"""
type: str
@ -1156,12 +1141,11 @@ class ResponsesUnknownInputItem(BaseModel):
def _responses_input_item_discriminator(v: Any) -> str:
"""Route a Responses input item to the correct tagged variant.
Pydantic's default smart-union matching fails when one variant in the
union is tagged with a strict ``Literal`` (``function_call`` /
``function_call_output``) and the incoming dict uses a different
``type`` — the other variants' validation errors are hidden and the
outer ``Union[str, list[...]]`` reports a misleading "Input should be a
valid string" error. An explicit discriminator makes the routing
Pydantic's smart-union matching fails when one union variant is tagged
with a strict ``Literal`` (``function_call`` / ``function_call_output``)
and the incoming dict has a different ``type``: other variants' errors are
hidden and the outer ``Union[str, list[...]]`` reports a misleading "Input
should be a valid string". An explicit discriminator makes routing
deterministic and lets us fall through to the catch-all.
"""
if isinstance(v, dict):
@ -1191,12 +1175,12 @@ ResponsesInputItem = Annotated[
class ResponsesFunctionTool(BaseModel):
"""Flat function-tool definition used by the Responses API request.
"""Flat function-tool definition for the Responses API request.
Unlike Chat Completions (which nests ``{"name": ..., "parameters": ...}``
inside a ``"function"`` key), the Responses API uses a flat shape with
``type``, ``name``, ``description``, ``parameters``, and ``strict`` at the
top level of each tool entry.
under a ``"function"`` key), the Responses API uses a flat shape with
``type``, ``name``, ``description``, ``parameters``, ``strict`` at the top
level of each tool entry.
"""
type: Literal["function"]
@ -1220,11 +1204,11 @@ class ResponsesRequest(BaseModel):
max_output_tokens: Optional[int] = Field(None, ge = 1)
stream: bool = Field(False, description = "Whether to stream the response via SSE")
# OpenAI function-calling fields — forwarded to llama-server via the
# Chat Completions pass-through (see routes/inference.py). Typed as a
# plain list so built-in tool shapes (``web_search``, ``file_search``,
# ``mcp``, ...) round-trip without validation errors — the translator
# picks out only ``type=="function"`` entries for forwarding.
# OpenAI function-calling fields, forwarded to llama-server via the Chat
# Completions pass-through (see routes/inference.py). Plain list so
# built-in tool shapes (``web_search``, ``file_search``, ``mcp``, ...)
# round-trip without validation errors; the translator forwards only
# ``type=="function"`` entries.
tools: Optional[list[dict]] = Field(
None,
description = (
@ -1280,10 +1264,9 @@ class ResponsesOutputMessage(BaseModel):
class ResponsesOutputFunctionCall(BaseModel):
"""A function-call output item in the Responses API response.
Unlike Chat Completions (which nests tool calls inside the assistant
message), the Responses API emits each tool call as its own top-level
``output`` item so clients can correlate results via ``call_id`` on the
next turn.
Unlike Chat Completions (tool calls nested in the assistant message), the
Responses API emits each tool call as its own top-level ``output`` item so
clients can correlate results via ``call_id`` on the next turn.
"""
type: Literal["function_call"] = "function_call"
@ -1439,7 +1422,7 @@ class AnthropicMessagesRequest(BaseModel):
top_k: Optional[int] = None
stop_sequences: Optional[list[str]] = None
metadata: Optional[dict] = None
# [x-unsloth] extensions — mirror the OpenAI endpoint convenience fields
# [x-unsloth] extensions mirroring the OpenAI endpoint convenience fields
min_p: Optional[float] = Field(
None, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold"
)

6
studio/backend/models/models.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
"""
Pydantic schemas for Model Management API
"""
"""Pydantic schemas for Model Management API"""
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any, Literal
@ -56,7 +54,7 @@ class CheckpointListResponse(BaseModel):
class ModelDetails(BaseModel):
"""Detailed model configuration and metadata - can be used for both list and detail views"""
"""Model configuration and metadata; used for both list and detail views"""
id: str = Field(..., description = "Model identifier")
model_name: Optional[str] = Field(

4
studio/backend/models/providers.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
"""
Pydantic schemas for the external LLM providers API.
"""
"""Pydantic schemas for the external LLM providers API."""
from typing import Literal, Optional

6
studio/backend/models/responses.py
View file

@ -1,10 +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
"""
Pydantic response schemas for endpoints that previously returned raw dicts.
These are small response models for training and model management routes.
"""
"""Pydantic response models for training and model management routes
(previously returned as raw dicts)."""
from pydantic import BaseModel, Field
from typing import Optional, List

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

@ -26,7 +26,7 @@ _MAX_LR_VALUE = 1.0
_MAX_LORA_R = 16_384
_MAX_LORA_ALPHA = 32_768
_MIN_VISION_IMAGE_SIZE = 256
# 2048 was the most I could get most llms to work at without getting unstable
# 2048 is the highest most llms stay stable at
_MAX_VISION_IMAGE_SIZE = 2048
@ -343,7 +343,7 @@ class TrainingStartRequest(BaseModel):
@model_validator(mode = "after")
def _check_steps_or_epochs(self) -> "TrainingStartRequest":
# num_epochs and max_steps each accept 0 as a "use the other one"
# sentinel. If both resolve to 0 there's nothing to train against.
# sentinel. Both 0 means nothing to train against.
if (self.max_steps is None or self.max_steps == 0) and self.num_epochs == 0:
raise ValueError("Either num_epochs or max_steps must be > 0; both cannot be 0.")
return self

5
studio/backend/models/users.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
"""Pydantic models for authentication tokens.
This module defines the Token response model used by auth routes.
"""
"""Pydantic models for authentication tokens."""
from pydantic import BaseModel, Field

2
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py
View file

@ -3,5 +3,5 @@
# Intentionally empty. Data-designer loads submodules lazily via qualified names
# (impl_qualified_name / config_qualified_name in plugin.py), so importing this
# package must NOT touch modules that depend on data_designer.engine.* during
# package must NOT touch modules depending on data_designer.engine.* during
# Studio's bootstrap (circular import).

17
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py
View file

@ -17,13 +17,12 @@ from .scraper import ScrapeConfig, materialize_to_jsonl
# In-process cache mapping a stable config signature to the JSONL materialization
# path. A single recipe job invokes the seed reader multiple times (validation,
# preview, per-column sampling), and the default flow re-scrapes the repo on
# every call: for a 2-repo preview that is ~15s of redundant GitHub GraphQL
# traffic before any generation fires. Memoize the materialization so the second
# and third passes reuse the file the first pass wrote. Cache key excludes the
# raw token and uses a short SHA-256 digest so token values never hit memory
# twice and token rotation invalidates cleanly.
# path. A recipe job calls the seed reader multiple times (validation, preview,
# per-column sampling) and the default flow re-scrapes on every call: a 2-repo
# preview is ~15s of redundant GitHub GraphQL traffic before any generation.
# Memoize so later passes reuse the first pass's file. The key excludes the raw
# token and uses a short SHA-256 digest, so token values never hit memory twice
# and rotation invalidates cleanly.
_SCRAPE_CACHE: dict[tuple, str] = {}
_SCRAPE_CACHE_LOCK = threading.Lock()
@ -47,8 +46,8 @@ def _lookup_cached_scrape(key: tuple) -> Optional[str]:
path = _SCRAPE_CACHE.get(key)
if path and Path(path).exists():
return path
# Stale entry (tmp cleanup, user restarted, ...); drop it so the caller
# materializes a fresh file rather than returning a dangling path.
# Stale entry (tmp cleanup, restart, ...); drop it so the caller
# materializes a fresh file instead of returning a dangling path.
if path:
with _SCRAPE_CACHE_LOCK:
_SCRAPE_CACHE.pop(key, None)

11
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py
View file

@ -3,11 +3,10 @@
"""Multi-repo GitHub scraper for the Studio seed plugin.
Drives the GraphQL-based scraper in `scraper_impl/` per repo. Each repo is
scraped with a trial_limits cap so we stop at `limit` items per resource.
After scraping, we read the per-resource JSONL shards and flatten them into
a single unified JSONL with stable columns (`item_type`, `repo`, `number`,
`title`, `body`, ...).
Drives the GraphQL scraper in `scraper_impl/` per repo, capped via trial_limits
to stop at `limit` items per resource. Then reads the per-resource JSONL shards
and flattens them into one unified JSONL with stable columns (`item_type`,
`repo`, `number`, `title`, `body`, ...).
"""
from __future__ import annotations
@ -166,7 +165,7 @@ def scrape(cfg: ScrapeConfig, base_dir: Path):
client = GitHubClient(token = token.value, token_source = token.source)
base_dir.mkdir(parents = True, exist_ok = True)
# Per-resource trial limits. limit <= 0 means "all": use a very large cap.
# Per-resource trial limits. limit <= 0 means "all": use a large cap.
effective_limit = cfg.limit if cfg.limit and cfg.limit > 0 else 1_000_000
trial_limits: dict[str, int] = {}
if "issues" in cfg.item_types:

22
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py
View file

@ -110,12 +110,11 @@ class GitHubClient:
)
def _is_auth_failure(self, r: "requests.Response") -> bool:
"""Distinguish auth failures from rate limiting on 401/403 responses.
"""Tell auth failures apart from rate limiting on 401/403.
- 401: always an auth failure (invalid / expired / wrong-scope token).
- 403: an auth failure UNLESS the response carries a clear rate-limit signal
(Retry-After header, X-RateLimit-Remaining: 0, or GitHub's secondary /
abuse rate-limit response text).
- 403: auth failure UNLESS it carries a rate-limit signal (Retry-After,
X-RateLimit-Remaining: 0, or secondary / abuse rate-limit text).
"""
if r.status_code == 401:
return True
@ -169,7 +168,7 @@ class GitHubClient:
timeout = 120,
)
self.calls_graphql += 1
# Update rate info from response headers
# Update rate info from headers
rem = r.headers.get("X-RateLimit-Remaining")
rst = r.headers.get("X-RateLimit-Reset")
if rem is not None:
@ -190,7 +189,7 @@ class GitHubClient:
if self._is_auth_failure(r):
self._raise_auth_error(r, "GraphQL")
if r.status_code == 403 or r.status_code == 429:
# Check for secondary/abuse
# Secondary/abuse rate limit
retry_after = _retry_after_seconds(r.headers.get("Retry-After"))
if retry_after is not None:
log.warning("Secondary rate limit. Sleep %ds.", retry_after)
@ -204,15 +203,14 @@ class GitHubClient:
r.raise_for_status()
data = r.json()
if "errors" in data and data["errors"]:
# Surface errors but allow partial data
# Surface errors but allow partial data; retry on RATE_LIMITED
errs = data["errors"]
# Retry on RATE_LIMITED
for e in errs:
if e.get("type") == "RATE_LIMITED":
self._sleep_until((self.graphql_reset or int(time.time()) + 60))
break
else:
# No rate-limit error, log and return partial
# No rate-limit error: log and return partial
log.warning("GraphQL errors: %s", json.dumps(errs)[:400])
return data
continue
@ -268,7 +266,7 @@ class GitHubClient:
log.warning("Secondary rate limit on REST. Sleep %ds.", retry_after)
time.sleep(retry_after + 2)
continue
# Check if primary rate
# Primary rate limit
if self.rest_remaining == 0 and self.rest_reset:
self._sleep_until(self.rest_reset)
continue
@ -299,11 +297,11 @@ class GitHubClient:
return
items = r.json()
if isinstance(items, dict):
# Some endpoints return dict with list field
# Some endpoints wrap the list in an "items" field
items = items.get("items", [])
for it in items:
yield it
# Follow link header
# Follow Link header
link = r.headers.get("Link", "")
nxt = None
for part in link.split(","):

6
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py
View file

@ -3,11 +3,11 @@
"""GraphQL queries for GitHub data scraping.
GitHub's GraphQL rejects queries that define unused fragments, so each query
only includes the fragments it actually references.
GitHub's GraphQL rejects queries with unused fragments, so each query includes
only the fragments it references.
"""
# ---- Fragments (kept as raw strings, composed per query) ----
# ---- Fragments (raw strings, composed per query) ----
F_ACTOR = """
fragment ActorFields on Actor {
__typename

14
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py
View file

@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Main scraper orchestration. Collects issues, PRs, discussions, commits, releases, etc.
"""Scraper orchestration: issues, PRs, discussions, commits, releases, etc.
Resumable via state file. Writes JSONL shards under data/{repo}/{resource}.jsonl.
"""
@ -49,9 +49,9 @@ class RepoScraper:
self.base_dir = base_dir
self.client = client
self.trial_limits = trial_limits or {}
# When light=True, use trimmed GraphQL queries (no reviewThreads,
# reviews, commits, timelineItems, files) so PR pages can be much
# larger without blowing GitHub's node-count ceiling.
# light=True uses trimmed GraphQL queries (no reviewThreads,
# reviews, commits, timelineItems, files) so PR pages can be larger
# without hitting GitHub's node-count ceiling.
self.light = light
self.repo_dir = base_dir / f"{owner}__{name}"
self.repo_dir.mkdir(parents = True, exist_ok = True)
@ -225,8 +225,8 @@ class RepoScraper:
page = 0
# Heavy nested PR query is capped at 3 per page (GitHub node-count
# ceiling); light query skips reviewThreads/reviews/commits/etc and
# can safely go to 25 per page. Clamp by trial_limit for small
# previews so limit=1 does not fetch a whole 25-item page.
# goes to 25 per page. Clamp by trial_limit so limit=1 does not
# fetch a whole 25-item page.
page_cap = 25 if self.light else 3
trial_cap = self.trial_limits.get(key)
per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap
@ -370,7 +370,7 @@ class RepoScraper:
f["_owner"] = self.owner
f["_repo"] = self.name
f["_prNumber"] = number
# files don't have id, synthesize one
# files have no id; synthesize one
f["_syntheticId"] = f"{self.owner}/{self.name}#{number}:{f.get('path')}"
self.writers[out_key].write(f)
info = ff.get("pageInfo") or {}

2
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py
View file

@ -65,7 +65,7 @@ class JsonlWriter:
self._lock = threading.Lock()
self._fh = self.path.open("a", buffering = 1)
self._count_seen_keys: set[str] = set()
# Preload seen keys if file exists (for dedup across resumes)
# Preload seen keys for dedup across resumes
if self.path.exists() and self.path.stat().st_size > 0:
try:
with self.path.open() as f:

6
studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py
View file

@ -70,11 +70,11 @@ def build_multi_file_preview_rows(
def _round_robin_preview(rows: list[dict[str, str]], preview_size: int) -> list[dict[str, str]]:
"""Pick preview rows round-robin across source files so every file is represented."""
"""Pick preview rows round-robin across source files so each is represented."""
if not rows or preview_size <= 0:
return []
# Group rows by source_file, preserving order of first appearance
# Group rows by source_file, preserving first-appearance order.
from collections import OrderedDict
grouped: OrderedDict[str, list[dict[str, str]]] = OrderedDict()
@ -147,7 +147,7 @@ def materialize_multi_file_unstructured_seed(
chunk_size: int,
chunk_overlap: int,
) -> tuple[Path, list[dict[str, str]]]:
"""Chunk multiple files and combine into one parquet dataset with source_file column."""
"""Chunk multiple files into one parquet dataset with a source_file column."""
chunk_size, chunk_overlap = resolve_chunking(chunk_size, chunk_overlap)
cache_key = _compute_multi_file_cache_key(file_entries, chunk_size, chunk_overlap)
cached = _CACHE_DIR / f"{cache_key}.parquet"

4
studio/backend/requirements/single-env/patch_metadata.py
View file

@ -2,14 +2,14 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Relax strict metadata pins so pip check matches known working single-env stack.
"""Relax strict metadata pins so pip check passes on the single-env stack.
Why:
- data-designer pins huggingface-hub>=1.0.1 and pyarrow<20.
- unsloth/transformers pins huggingface-hub<1.
- studio datasets pins pyarrow>=21.
Runtime works in this app with hub 0.36.x + pyarrow 23.x, but metadata conflicts.
Runtime works with hub 0.36.x + pyarrow 23.x; only the metadata conflicts.
"""
from __future__ import annotations

44
studio/backend/routes/auth.py
View file

@ -43,17 +43,17 @@ router = APIRouter()
def _reset_password_command() -> str:
"""Shell command shown in the 'incorrect password' hint.
Prefer the ABSOLUTE path to this install's ``unsloth`` launcher (a sibling
of the running interpreter) so the hint works even when the launcher's
directory is not on PATH -- e.g. a terminal opened before install, a stale
Windows PATH, or ``~/.local/bin`` not on PATH (the default on macOS) -- and
regardless of the current working directory.
Prefer the ABSOLUTE path to this install's ``unsloth`` launcher (sibling of
the running interpreter) so the hint works even when the launcher's dir
isn't on PATH -- e.g. a terminal opened before install, a stale Windows
PATH, or ``~/.local/bin`` not on PATH (the macOS default) -- regardless of
the cwd.
On POSIX the path is shell-quoted so spaces are handled. On Windows we only
use the bare absolute path when it has no spaces, because a quoted path needs
different syntax in cmd (``"..."``) vs PowerShell (``& "..."``); when it has
a space we fall back to the PATH-based form to stay unambiguous across
shells. If the launcher can't be located we fall back to the PATH form too.
On POSIX the path is shell-quoted to handle spaces. On Windows we use the
bare absolute path only when it has no spaces, since a quoted path needs
different syntax in cmd (``"..."``) vs PowerShell (``& "..."``); with a
space we fall back to the PATH-based form to stay unambiguous across shells.
We also fall back to the PATH form if the launcher can't be located.
"""
try:
bin_dir = os.path.dirname(os.path.abspath(sys.executable))
@ -72,7 +72,7 @@ def _reset_password_command() -> str:
# Per-(ip, username) bucket + per-IP aggregate. Account bucket stops one user's
# typos from blocking others; the aggregate stops username-rotation spray.
# Single-process only -- multi-worker deployments need a shared store.
# Single-process only; multi-worker deployments need a shared store.
_LOGIN_BUCKETS: dict[tuple[str, str], deque] = {}
_LOGIN_IP_BUCKETS: dict[str, deque] = {}
_LOGIN_BUCKETS_LOCK = threading.Lock()
@ -80,18 +80,18 @@ _LOGIN_WINDOW_SECONDS = 60.0
_LOGIN_MAX_FAILS = 5
_LOGIN_IP_MAX_FAILS = 30
_LOGIN_LOCKOUT_SECONDS = 60
# Bucket-dict cap. On overflow we prune stale entries; if still full the
# failure folds into the per-IP aggregate only.
# Bucket-dict cap. On overflow, prune stale entries; if still full the failure
# folds into the per-IP aggregate only.
_LOGIN_MAX_BUCKETS = 4096
# Unrepresentable as a real username (leading NUL); folds unknown-user attempts
# into one slot so attacker cardinality cannot blow the bucket dict.
# into one slot so attacker cardinality can't blow the bucket dict.
_UNKNOWN_LOGIN_USER = "\x00unknown-user"
def _trust_forwarded_for() -> bool:
"""Honour X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is set.
Off by default so a direct caller cannot spoof the header.
Off by default so a direct caller can't spoof the header.
"""
return os.environ.get("UNSLOTH_STUDIO_TRUST_FORWARDED", "").lower() in (
"1",
@ -112,7 +112,7 @@ def _normalize_forwarded_addr(value: str) -> str:
return ""
host = value[1:end]
elif value.count(":") == 1:
# IPv4:port. Bare IPv6 has multiple colons and takes the else branch.
# IPv4:port. Bare IPv6 has multiple colons → else branch.
head, _, tail = value.rpartition(":")
host = head if tail.isdigit() and head else value
else:
@ -144,7 +144,7 @@ def _client_ip(request: Request | None) -> str:
return normalized
fwd = request.headers.get("forwarded", "")
if fwd:
# First element only -- multi-element headers cannot fork buckets.
# First element only; multi-element headers can't fork buckets.
normalized = _forwarded_for_from_element(fwd.split(",", 1)[0])
if normalized:
return normalized
@ -190,7 +190,7 @@ def _record_login_failure(key: tuple[str, str]) -> int:
_prune_bucket(account_bucket, now)
account_bucket.append(now)
return len(account_bucket)
# Bucket dict is at its cap; per-IP cap still applies via ip_bucket.
# Bucket dict at cap; per-IP cap still applies via ip_bucket.
return len(ip_bucket)
@ -242,16 +242,16 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token:
if blocked_for > 0:
raise HTTPException(
status_code = status.HTTP_429_TOO_MANY_REQUESTS,
# IP is intentionally not interpolated into the body; behind a
# proxy or NAT it is either misleading or an info leak.
# IP not interpolated into the body; behind a proxy/NAT it's
# misleading or an info leak.
detail = (f"Too many failed login attempts. " f"Try again in {blocked_for} seconds."),
headers = {"Retry-After": str(blocked_for)},
)
record = storage.get_user_and_secret(payload.username)
if record is None:
# Record under a single sentinel key per IP so attacker-controlled
# username cardinality does not allocate buckets without bound.
# Record under one sentinel key per IP so attacker-controlled username
# cardinality can't allocate unbounded buckets.
_record_login_failure(unknown_key)
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,

24
studio/backend/routes/chat_history.py
View file

@ -189,21 +189,20 @@ class ChatMessagesBatchResponse(BaseModel):
class ChatImportLedgerResponse(BaseModel):
# Plain list of legacy thread ids. Keeping the payload key-less keeps
# the client diff to a single Set construction.
# Plain list of legacy thread ids; key-less payload keeps the client diff
# to a single Set construction.
threadIds: list[str]
class ChatImportLedgerRecordRequest(BaseModel):
# 10k cap keeps the request body bounded; real users have << 1k threads.
# 10k cap bounds the request body; real users have << 1k threads.
threadIds: list[str] = Field(default_factory = list, max_length = 10_000)
class ChatImportLedgerRecordResponse(BaseModel):
# accepted: deduped non-empty input count. inserted: rows actually new
# (ON CONFLICT DO NOTHING skips already-recorded ids). The client uses
# `accepted >= 0` as the "endpoint exists" signal and ignores the split
# otherwise.
# `accepted >= 0` as the "endpoint exists" signal and ignores the split.
accepted: int
inserted: int
@ -353,7 +352,7 @@ async def batch_thread_messages(
payload: ChatMessagesBatchRequest, current_subject: str = Depends(get_current_subject)
):
"""One round-trip per sidebar/search rebuild instead of N. Unknown thread
ids are returned as empty lists so callers don't need a pre-flight."""
ids return empty lists so callers don't need a pre-flight."""
by_thread: dict[str, list[ChatMessage]] = {tid: [] for tid in payload.threadIds}
for m in list_chat_messages_for_threads(payload.threadIds):
tid = m["threadId"]
@ -444,11 +443,10 @@ async def count_threads(current_subject: str = Depends(get_current_subject)):
@router.get("/import-ledger", response_model = ChatImportLedgerResponse)
async def get_import_ledger(current_subject: str = Depends(get_current_subject)):
"""Legacy-Dexie import ledger. Returns the set of legacy thread ids
already copied into chat_threads / chat_messages. The frontend
uses this on every fresh tab open to decide whether to re-run the
Dexie -> studio.db import. Source of truth lives inside studio.db
so a studio.db wipe makes the import recoverable."""
"""Legacy-Dexie import ledger. Returns the legacy thread ids already copied
into chat_threads / chat_messages. The frontend uses this on every fresh tab
open to decide whether to re-run the Dexie -> studio.db import. Source of
truth lives in studio.db, so a studio.db wipe makes the import recoverable."""
return ChatImportLedgerResponse(threadIds = list_chat_legacy_imports())
@ -480,8 +478,8 @@ async def put_settings(
parsed = ChatSettingsPayload.model_validate(payload)
except ValidationError as exc:
raise HTTPException(status_code = 400, detail = exc.errors()) from exc
# Atomic read + deep-merge + write inside one BEGIN IMMEDIATE so two
# concurrent slider drags can't drop each other's updates.
# Atomic read + deep-merge + write in one BEGIN IMMEDIATE so concurrent
# slider drags can't drop each other's updates.
try:
return ChatSettingsResponse(
settings = upsert_chat_settings_merge(parsed.model_dump(exclude_unset = True))

164
studio/backend/routes/data_recipe/jobs.py
View file

@ -36,14 +36,14 @@ def _resolve_local_v1_endpoint(request: Request) -> str:
"""Return the loopback /v1 URL for the actual backend listen port.
Resolution order:
1. ``app.state.server_port`` - explicitly published by run.py after
the uvicorn server has bound. This is the most reliable source
because it survives reverse proxies, TLS terminators and tunnels.
1. ``app.state.server_port`` - published by run.py after uvicorn
binds. Most reliable; survives reverse proxies, TLS terminators,
and tunnels.
2. ``request.scope["server"]`` - the real (host, port) tuple uvicorn
sets when the request is dispatched. Used when Studio is started
outside ``run_server`` (e.g. ``uvicorn studio.backend.main:app``).
3. ``request.base_url`` parsed - last resort for test fixtures that
do not route through a live uvicorn server.
sets per request. Used when Studio starts outside ``run_server``
(e.g. ``uvicorn studio.backend.main:app``).
3. ``request.base_url`` parsed - last resort for test fixtures with
no live uvicorn server.
"""
port: Any = getattr(request.app.state, "server_port", None)
if not isinstance(port, int) or port <= 0:
@ -76,14 +76,14 @@ def _request_has_desktop_access_token(request: Request) -> bool:
def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]:
"""Return the set of model_aliases that are actually referenced by an
LLM column. Used to narrow the "Chat model loaded" gate so that orphan
model_config nodes on the canvas do not block unrelated recipe runs.
"""Return model_aliases actually referenced by an LLM column. Narrows
the "Chat model loaded" gate so orphan model_config nodes on the canvas
do not block unrelated recipe runs.
The ``llm-`` prefix matches the existing convention in
The ``llm-`` prefix matches the convention in
``core/data_recipe/service.py::_recipe_has_llm_columns`` and covers all
LLM column types emitted by the frontend (llm-text, llm-code,
llm-structured, llm-judge).
frontend LLM column types (llm-text, llm-code, llm-structured,
llm-judge).
"""
aliases: set[str] = set()
for column in recipe.get("columns", []):
@ -180,18 +180,17 @@ def _ensure_selected_local_model_loaded(
def _inject_local_structured_response_format(
recipe: dict[str, Any], local_provider_names: set[str]
) -> None:
"""For each llm-structured column that targets a local-provider model_config,
clone the model_config and inject an OpenAI ``response_format`` with the
column's ``output_format`` JSON schema. The column is rewritten to point at
the clone so llm-text / llm-judge columns that share the same alias keep
free-form sampling.
"""For each llm-structured column targeting a local-provider model_config,
clone the config and inject an OpenAI ``response_format`` with the
column's ``output_format`` JSON schema. The column is repointed at the
clone so llm-text / llm-judge columns sharing the alias keep free-form
sampling.
Without this, data_designer only injects a prompt-level "return JSON in a
```json fence" instruction. Small GGUF models frequently break format,
wasting the full ``max_tokens`` budget per row and then failing to parse.
Forwarding ``response_format`` lets llama-server apply grammar-constrained
sampling from the JSON schema, which guarantees a parseable response and
terminates early.
```json fence" instruction. Small GGUFs often break format, wasting the
full ``max_tokens`` budget per row and then failing to parse. Forwarding
``response_format`` lets llama-server apply grammar-constrained sampling
from the schema, guaranteeing a parseable response and terminating early.
"""
columns = recipe.get("columns")
model_configs = recipe.get("model_configs")
@ -210,8 +209,8 @@ def _inject_local_structured_response_format(
return
# Clone per (alias, column) so each llm-structured column gets its own
# schema without leaking response_format onto other columns that share the
# same base alias.
# schema without leaking response_format onto other columns sharing the
# base alias.
seen_clone_aliases: set[str] = {
mc.get("alias") for mc in model_configs if isinstance(mc.get("alias"), str)
}
@ -243,15 +242,14 @@ def _inject_local_structured_response_format(
if not isinstance(params, dict):
params = {}
clone["inference_parameters"] = params
# data_designer's BaseInferenceParams is a pydantic model with
# extra="forbid", so response_format cannot sit at the top level of
# inference_parameters. It does expose an `extra_body: dict` pass-
# through that the OpenAI client spreads into the request body at the
# top level, which is where llama-server reads response_format from.
# llama.cpp server shape (tools/server/README.md): the schema sits
# directly under response_format, not nested in a json_schema object
# the way OpenAI's Chat Completions API expects. llama-server converts
# the schema to a GBNF grammar and applies it during sampling.
# data_designer's BaseInferenceParams is pydantic extra="forbid", so
# response_format cannot sit at the top level of inference_parameters.
# Its `extra_body: dict` passthrough is spread into the request body
# top level by the OpenAI client, where llama-server reads
# response_format. llama.cpp server shape (tools/server/README.md):
# the schema sits directly under response_format, not nested in a
# json_schema object like OpenAI's Chat Completions API expects.
# llama-server converts the schema to a GBNF grammar for sampling.
extra_body = params.get("extra_body")
if not isinstance(extra_body, dict):
extra_body = {}
@ -269,21 +267,21 @@ def _inject_local_structured_response_format(
def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optional[int]:
"""
Mutate recipe dict in-place: for any provider with is_local=True,
fill in the endpoint pointing at this server and inject a short-lived
internal sk-unsloth-* API key for workflow auth.
Mutate recipe dict in-place: for any provider with is_local=True, set
the endpoint to this server and inject a short-lived internal
sk-unsloth-* API key for workflow auth.
Returns the row id of the minted internal key (so the caller can
revoke it on job completion) or ``None`` when no local provider is
actually reachable from an LLM column.
Returns the row id of the minted internal key (for the caller to revoke
on job completion) or ``None`` when no local provider is reachable from
an LLM column.
"""
providers = recipe.get("model_providers")
if not providers:
return None
# Collect local providers and pop is_local from ALL dicts unconditionally.
# Strict `is True` guard so malformed payloads (is_local: 1,
# is_local: "true") do not accidentally trigger the loopback rewrite.
# Collect local providers and pop is_local from ALL dicts. Strict
# `is True` guard so malformed payloads (is_local: 1, is_local: "true")
# do not trigger the loopback rewrite.
local_indices: list[int] = []
for i, provider in enumerate(providers):
if not isinstance(provider, dict):
@ -297,10 +295,10 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona
endpoint = _resolve_local_v1_endpoint(request)
# Only gate on model-loaded if a local provider is actually reachable
# from an LLM column through a model_config. Orphan model_config nodes
# that reference a local provider but that no LLM column uses should
# not block runs; the recipe would never call /v1 for them.
# Only gate on model-loaded if a local provider is reachable from an LLM
# column via a model_config. Orphan model_config nodes referencing a
# local provider that no LLM column uses should not block runs; the
# recipe would never call /v1 for them.
local_names = {providers[i].get("name") for i in local_indices if providers[i].get("name")}
used_aliases = _used_llm_model_aliases(recipe)
referenced_providers = {
@ -313,19 +311,19 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona
internal_key_id: Optional[int] = None
if local_names & referenced_providers:
# Verify the selected local model is loaded before minting a workflow
# key. This still remains a point-in-time singleton-backend check
# (TOCTOU): a future generation token should bind frontend load and
# job creation, and the inference endpoint returns a clear 400 if the
# model is later unloaded or swapped before the subprocess calls /v1.
# key. Still a point-in-time singleton-backend check (TOCTOU): a
# future generation token should bind frontend load to job creation;
# the inference endpoint returns a clear 400 if the model is unloaded
# or swapped before the subprocess calls /v1.
_ensure_selected_local_model_loaded(recipe, local_names)
from auth import storage # deferred: avoids circular import
# Mint an internal sk-unsloth-* key scoped to this workflow run.
# Uses the unified API-key issuance path (one mint/revoke/verify
# surface instead of a second JWT code path). The key is marked
# internal so it is hidden from the user's API-key list, and the
# caller revokes it when the job terminates.
# Mint an internal sk-unsloth-* key scoped to this workflow run via
# the unified API-key issuance path (one mint/revoke/verify surface
# instead of a second JWT code path). Marked internal so it is hidden
# from the user's API-key list; the caller revokes it when the job
# terminates.
expires_at = (datetime.now(timezone.utc) + timedelta(hours = 24)).isoformat()
token, row = storage.create_api_key(
username = "unsloth",
@ -335,11 +333,11 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona
)
internal_key_id = int(row["id"])
# Defensively strip any stale "external"-only fields the frontend may
# have left on the dict (extra_headers/extra_body/api_key_env). The UI
# hides these inputs in local mode but the payload builder still serializes
# them, so a previously external provider that flipped to local can carry
# invalid JSON or rogue auth headers into the local /v1 call.
# Strip stale "external"-only fields the frontend may have left
# (extra_headers/extra_body/api_key_env). The UI hides these in local
# mode but the payload builder still serializes them, so a provider
# flipped from external to local could carry invalid JSON or rogue auth
# headers into the local /v1 call.
for i in local_indices:
providers[i]["endpoint"] = endpoint
providers[i]["api_key"] = token
@ -348,29 +346,27 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona
providers[i].pop("extra_headers", None)
providers[i].pop("extra_body", None)
# Force skip_health_check on any model_config that references a local
# provider. The frontend now sends the explicit selected local model id,
# but llama-server's /v1/models response can still differ from that id
# for local paths, cache aliases, and GGUF variant loads. The recipe run
# has already gated on a loaded local inference backend above, so the
# data_designer model-list health check would be redundant and can reject
# valid local selections.
# Force skip_health_check on any model_config referencing a local
# provider. The frontend sends the explicit selected local model id, but
# llama-server's /v1/models response can differ from it for local paths,
# cache aliases, and GGUF variant loads. The run already gated on a
# loaded local backend above, so data_designer's model-list health check
# would be redundant and can reject valid local selections.
for mc in recipe.get("model_configs", []):
if not isinstance(mc, dict):
continue
if mc.get("provider") in local_names:
mc["skip_health_check"] = True
# Disable thinking for data-recipe inference on local providers.
# Reasoning models emit a <think>...</think> preamble before the
# answer, which roughly doubles generated token count per row and
# pushes the visible answer past data_designer's json-fence
# regex. Forward chat_template_kwargs={enable_thinking: False}
# through the OpenAI SDK's extra_body passthrough so llama-server
# renders the template without the reasoning preamble. Free-form
# llm-text columns benefit from the latency cut, and structured
# columns also stop leaking think tags into the grammar-
# constrained JSON (llama-server's GBNF path still enforces the
# schema either way).
# Reasoning models emit a <think>...</think> preamble that roughly
# doubles generated tokens per row and pushes the answer past
# data_designer's json-fence regex. Forward
# chat_template_kwargs={enable_thinking: False} via the OpenAI
# SDK's extra_body passthrough so llama-server renders the
# template without the preamble. llm-text columns get the latency
# cut; structured columns stop leaking think tags into the
# grammar-constrained JSON (GBNF still enforces the schema either
# way).
params = mc.get("inference_parameters")
if not isinstance(params, dict):
params = {}
@ -387,7 +383,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona
# Forward each llm-structured column's output_format as an OpenAI
# response_format so llama-server uses grammar-constrained sampling and
# small GGUFs stop wasting the full max_tokens budget on broken JSON.
# small GGUFs stop wasting the max_tokens budget on broken JSON.
_inject_local_structured_response_format(recipe, local_names)
return internal_key_id
@ -446,11 +442,11 @@ def create_job(payload: RecipePayload, request: Request):
log = logger,
) from exc
# Single try block covers get_job_manager() AND mgr.start() so a workflow
# key minted above never outlives the request even when an unexpected
# Single try covers get_job_manager() AND mgr.start() so a minted
# workflow key never outlives the request even on an unexpected
# exception type (TypeError from a stale kwarg, OSError from a queue
# write, etc.) bubbles up. Without the bare except, such exceptions let
# the sk-unsloth-* key live until its 24h TTL.
# write, etc.). Without the bare except, those would let the
# sk-unsloth-* key live until its 24h TTL.
try:
mgr = get_job_manager()
job_id = mgr.start(
@ -488,7 +484,7 @@ def create_job(payload: RecipePayload, request: Request):
def _revoke_internal_api_key_safe(key_id: int) -> None:
"""Best-effort revoke of a workflow-minted key; swallow any error so
that revocation failures never mask the caller's own error path."""
revocation failures never mask the caller's own error path."""
try:
from auth import storage # deferred: avoids circular import
storage.revoke_internal_api_key(key_id)

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

@ -394,7 +394,7 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
def _extract_text_from_file(file_path: Path, ext: str) -> str:
"""Extract text from uploaded file based on extension, converting to markdown where possible."""
"""Extract text from an uploaded file by extension, to markdown where possible."""
if ext in {".txt", ".md"}:
raw = file_path.read_text(encoding = "utf-8", errors = "ignore")
elif ext == ".pdf":
@ -581,7 +581,7 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
seed_source_type = _normalize_optional_text(payload.seed_source_type) or "local"
filename = _sanitize_filename(payload.filename)
ext = Path(filename).suffix.lower()
# Legacy single-file unstructured path only supports .txt/.md
# Legacy single-file unstructured path supports only .txt/.md;
# PDF/DOCX extraction uses the multi-file upload endpoint instead
_LEGACY_UNSTRUCTURED_EXTS = {".txt", ".md"}
if seed_source_type == "unstructured":

18
studio/backend/routes/data_recipe/validate.py
View file

@ -121,9 +121,8 @@ def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
def _patch_local_providers(recipe: dict[str, Any]) -> None:
"""Strip is_local and fill a dummy endpoint so validation doesn't choke.
Uses a strict `is True` check to match _inject_local_providers in
jobs.py - malformed payloads with truthy but non-boolean is_local
values should not be treated as local.
Strict `is True` check matches _inject_local_providers in jobs.py:
truthy but non-boolean is_local values are not treated as local.
"""
for provider in recipe.get("model_providers", []):
if not isinstance(provider, dict):
@ -152,13 +151,12 @@ def validate(payload: RecipePayload) -> ValidateResponse:
build_config_builder(recipe)
except ModuleNotFoundError as exc:
# data_designer is an optional runtime dep. Static validation
# already passed; live access + full config validation are
# deferred to run start (per _GITHUB_VALIDATE_NOTE), so a missing
# optional import at validate time should not block the recipe.
# Restrict the bypass to the data_designer module specifically so
# other ImportErrors (e.g. broken internal imports or missing
# transitive deps after a package upgrade) still surface as
# validation failures instead of being silently swallowed.
# passed; live access + full config validation are deferred to
# run start (per _GITHUB_VALIDATE_NOTE), so a missing optional
# import here should not block the recipe. Restrict the bypass to
# the data_designer module so other ImportErrors (broken internal
# imports, missing transitive deps after an upgrade) still surface
# as validation failures instead of being swallowed.
if not (exc.name or "").startswith("data_designer"):
raise
logger.debug(

78
studio/backend/routes/datasets.py
View file

@ -45,9 +45,9 @@ def _get_dataset_size_cached(repo_id: str) -> int:
def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
"""Pick the most useful on-disk path for a HF cache repo dir.
Mirrors the helper in routes/models.py: prefer the most-recent
snapshot dir, fall back to the cache repo root, return resolved
realpath. Duplicated here to keep routes/datasets.py self-contained.
Mirrors routes/models.py: prefer the most-recent snapshot dir, else the
cache repo root, returned as a resolved realpath. Duplicated here to keep
routes/datasets.py self-contained.
"""
try:
snapshots_dir = repo_dir / "snapshots"
@ -61,12 +61,12 @@ def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
return None
# Add backend directory to path
# Add backend directory to path.
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
# Import dataset utilities
# Import dataset utilities.
from utils.datasets import check_dataset_format
from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label
from auth.authentication import get_current_subject
@ -93,7 +93,7 @@ from utils.paths import (
def _serialize_preview_value(value):
"""make it json safe for client preview ⊂(◉‿◉)つ"""
"""Make a value JSON-safe for the client preview."""
if value is None or isinstance(value, (str, int, float, bool)):
return value
@ -130,10 +130,10 @@ def _serialize_preview_rows(rows):
# --- Endpoints ---
# Recognized data-file extensions for the single-file fallback approach.
# Tabular formats are preferred over archives for Tier 1 preview because
# archives (e.g. images.zip) may be loaded as ImageFolder datasets with
# synthetic columns (image/label) that don't match the real dataset schema.
# Recognized data-file extensions for the single-file fallback.
# Tier 1 preview prefers tabular over archives: archives (e.g. images.zip) can
# load as ImageFolder datasets with synthetic image/label 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
@ -331,9 +331,9 @@ async def upload_dataset(
stored_name = f"{uuid4().hex}_{stem}{ext}"
stored_path = DATASET_UPLOAD_DIR / stored_name
# Stream file to disk in chunks to avoid holding entire file in memory.
# Keep a route-level cap so users get a clear training-dataset-specific
# error and oversized partial files are not left in the Studio uploads directory.
# Stream to disk in chunks to avoid holding the whole file in memory. The
# route-level cap gives a clear training-dataset error and avoids leaving
# oversized partial files in the Studio uploads directory.
upload_limit_bytes = get_upload_limit_bytes()
total_bytes = 0
upload_complete = False
@ -378,12 +378,10 @@ async def get_dataset_download_progress(
"""Return download progress for a HuggingFace dataset repo.
Mirrors ``GET /api/models/download-progress`` but scans the
``datasets--owner--name`` cache directory under HF_HUB_CACHE.
Modern ``datasets``/``huggingface_hub`` caches both raw model and
raw dataset blobs in HF_HUB_CACHE; the ``datasets`` library writes
its processed Arrow shards elsewhere, but the in-progress *download*
bytes are observable here. Returns ``cache_path`` so the UI can
show users where the dataset blobs landed on disk.
``datasets--owner--name`` cache directory under HF_HUB_CACHE. Modern
``datasets``/``huggingface_hub`` cache raw dataset blobs there (processed
Arrow shards live elsewhere, but in-progress *download* bytes are visible
here). Returns ``cache_path`` so the UI can show where blobs landed.
"""
_empty = {
"downloaded_bytes": 0,
@ -433,9 +431,9 @@ async def get_dataset_download_progress(
"cache_path": cache_path,
}
# Same 95% completion threshold as the model endpoint -- HF blob
# dedup makes completed_bytes drift slightly under expected_bytes,
# and inter-file gaps would otherwise look like "done".
# Same 95% completion threshold as the model endpoint -- HF blob dedup
# makes completed_bytes drift slightly under expected_bytes, and
# inter-file gaps would otherwise look "done".
if completed_bytes >= expected_bytes * 0.95:
progress = 1.0
else:
@ -456,15 +454,13 @@ def check_format(request: CheckFormatRequest, current_subject: str = Depends(get
"""
Check if a dataset requires manual column mapping.
Strategy for HuggingFace datasets:
1. list_repo_files → pick the first data file → load_dataset(data_files=[…])
HuggingFace strategy:
1. list_repo_files → first data file → load_dataset(data_files=[…]).
Avoids resolving thousands of files; typically ~2-4 s.
2. Full streaming load_dataset as a last-resort fallback.
Local files are loaded directly.
Using a plain `def` (not async) so FastAPI runs this in a thread-pool,
preventing any blocking IO from freezing the event loop.
Local files load directly. Plain `def` (not async) so FastAPI runs it in a
thread-pool, keeping blocking IO off the event loop.
"""
try:
from itertools import islice
@ -502,14 +498,14 @@ def check_format(request: CheckFormatRequest, current_subject: str = Depends(get
)
data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)]
# Prefer tabular formats over archives (e.g. images.zip → ImageFolder
# with synthetic image/label columns that don't match the real schema).
# Prefer tabular over archives (e.g. images.zip → ImageFolder
# with synthetic image/label columns not in the real schema).
tabular_files = [
f for f in data_files if any(f.endswith(ext) for ext in _TABULAR_EXTS)
]
candidates = tabular_files or data_files
# When a subset is specified, narrow to files whose name matches
# With a subset, narrow to files whose name matches
# (e.g. subset="testmini" → prefer "testmini.parquet").
if request.subset and candidates:
subset_matches = [f for f in candidates if request.subset in Path(f).stem]
@ -560,26 +556,26 @@ def check_format(request: CheckFormatRequest, current_subject: str = Depends(get
preview_slice = Dataset.from_list(rows)
total_rows = None
# Run lightweight format check on the preview slice
# Lightweight format check on the preview slice.
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)}"
)
# Generate preview samples
# Generate preview samples.
preview_samples = None
if not result["requires_manual_mapping"]:
if result.get("suggested_mapping"):
# Heuristic-detected: show raw data so columns match the API response.
# Processing (column stripping) happens at training time, not preview.
# Column stripping happens at training time, not preview.
preview_samples = _serialize_preview_rows(preview_slice)
else:
try:
format_result = format_dataset(
preview_slice,
format_type = "auto",
num_proc = None, # Only 10 preview rows -- no need for multiprocessing
num_proc = None, # Only 10 preview rows -- no multiprocessing needed
)
processed = format_result["dataset"]
preview_samples = _serialize_preview_rows(processed)
@ -589,7 +585,7 @@ def check_format(request: CheckFormatRequest, current_subject: str = Depends(get
else:
preview_samples = _serialize_preview_rows(preview_slice)
# Collect warnings: from check_dataset_format + URL-based image detection
# 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 []):
@ -636,17 +632,17 @@ def ai_assist_mapping(
"""
Run LLM-assisted dataset conversion advisor (user-triggered).
Multi-pass analysis using a 7B helper model:
Pass 1: Classify dataset type from HF card + samples
Pass 2: Generate conversion strategy (system prompt, templates)
Pass 3: Validate conversion quality
Multi-pass analysis with a 7B helper model:
Pass 1: Classify dataset type from HF card + samples.
Pass 2: Generate conversion strategy (system prompt, templates).
Pass 3: Validate conversion quality.
Falls back to simple column classification if the advisor fails.
"""
try:
from utils.datasets.llm_assist import llm_conversion_advisor
# Truncate sample values for the LLM prompt
# Truncate sample values for the LLM prompt.
truncated = [
{col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5]
]

106
studio/backend/routes/export.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
"""
Export API routes: checkpoint discovery and model export operations.
"""
"""Export API routes: checkpoint discovery and model export operations."""
import asyncio
import json
@ -28,7 +26,7 @@ from auth.authentication import get_current_subject
from utils.utils import safe_error_detail
# Import backend functions
# Backend functions
try:
from core.export import get_export_backend
except ImportError:
@ -37,7 +35,7 @@ except ImportError:
sys.path.insert(0, str(parent_backend))
from core.export import get_export_backend
# Import Pydantic models
# Pydantic models
from models import (
LoadCheckpointRequest,
ExportStatusResponse,
@ -56,16 +54,12 @@ logger = get_logger(__name__)
async def load_checkpoint(
request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject)
):
"""
Load a checkpoint into the export backend.
Wraps ExportBackend.load_checkpoint.
"""
"""Load a checkpoint into the export backend (ExportBackend.load_checkpoint)."""
try:
# Version switching is handled automatically by the subprocess-based
# export backend — no need for ensure_transformers_version() here.
# The subprocess-based export backend handles version switching
# automatically — no ensure_transformers_version() needed here.
# Free GPU memory: shut down any running inference/training subprocesses
# Free GPU memory: shut down running inference/training subprocesses
# before loading the export checkpoint (they'd compete for VRAM).
try:
from core.inference import get_inference_backend
@ -87,8 +81,8 @@ async def load_checkpoint(
if trn.is_training_active():
logger.info("Stopping active training to free GPU memory for export")
trn.stop_training()
# Wait for training subprocess to actually exit before proceeding,
# otherwise it may still hold GPU memory when export tries to load.
# Wait for the training subprocess to exit; else it may still
# hold GPU memory when export tries to load.
for _ in range(60): # up to 30s
if not trn.is_training_active():
break
@ -101,8 +95,8 @@ async def load_checkpoint(
backend = get_export_backend()
# load_checkpoint spawns and waits on a subprocess and can take
# minutes. Run it in a worker thread so the event loop stays
# free to serve the live log SSE stream concurrently.
# minutes. Run it in a worker thread so the event loop stays free
# to serve the live log SSE stream concurrently.
success, message = await asyncio.to_thread(
backend.load_checkpoint,
checkpoint_path = request.checkpoint_path,
@ -127,11 +121,7 @@ async def load_checkpoint(
@router.post("/cleanup", response_model = ExportOperationResponse)
async def cleanup_export_memory(current_subject: str = Depends(get_current_subject)):
"""
Cleanup export-related models from memory (GPU/CPU).
Wraps ExportBackend.cleanup_memory.
"""
"""Cleanup export-related models from memory (ExportBackend.cleanup_memory)."""
try:
backend = get_export_backend()
success = await asyncio.to_thread(backend.cleanup_memory)
@ -158,9 +148,7 @@ async def cleanup_export_memory(current_subject: str = Depends(get_current_subje
@router.get("/status", response_model = ExportStatusResponse)
async def get_export_status(current_subject: str = Depends(get_current_subject)):
"""
Get current export backend status (loaded checkpoint, model type, PEFT flag).
"""
"""Get export backend status (loaded checkpoint, model type, PEFT flag)."""
try:
backend = get_export_backend()
return ExportStatusResponse(
@ -177,7 +165,7 @@ async def get_export_status(current_subject: str = Depends(get_current_subject))
def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]:
"""Return the export path relative to exports_root so the install path is not leaked."""
"""Return the export path relative to exports_root, hiding the install path."""
if not output_path:
return None
try:
@ -195,8 +183,7 @@ def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]:
async def export_merged_model(
request: ExportMergedModelRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export a merged PEFT model (e.g., 16-bit or 4-bit) and optionally push to Hub.
"""Export a merged PEFT model (16-bit or 4-bit), optionally pushing to Hub.
Wraps ExportBackend.export_merged_model.
"""
@ -234,8 +221,7 @@ async def export_merged_model(
async def export_base_model(
request: ExportBaseModelRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export a non-PEFT base model and optionally push to Hub.
"""Export a non-PEFT base model, optionally pushing to Hub.
Wraps ExportBackend.export_base_model.
"""
@ -273,8 +259,7 @@ async def export_base_model(
async def export_gguf(
request: ExportGGUFRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export the current model to GGUF format and optionally push to Hub.
"""Export the current model to GGUF format, optionally pushing to Hub.
Wraps ExportBackend.export_gguf.
"""
@ -311,8 +296,7 @@ async def export_gguf(
async def export_lora_adapter(
request: ExportLoRAAdapterRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export only the LoRA adapter (if the loaded model is PEFT).
"""Export only the LoRA adapter (if the loaded model is PEFT).
Wraps ExportBackend.export_lora_adapter.
"""
@ -349,18 +333,18 @@ async def export_lora_adapter(
# Live export log stream (Server-Sent Events)
# ─────────────────────────────────────────────────────────────────────
#
# The export worker subprocess redirects its stdout/stderr into a pipe
# that a reader thread forwards to the orchestrator as log entries (see
# The export worker subprocess redirects its stdout/stderr into a pipe a
# reader thread forwards to the orchestrator as log entries (see
# core/export/worker.py::_setup_log_capture and
# core/export/orchestrator.py::_append_log). This endpoint streams
# those entries to the browser so the export dialog can show a live
# terminal-style output panel while load_checkpoint / export_merged /
# export_gguf / export_lora / export_base run.
# core/export/orchestrator.py::_append_log). This endpoint streams those
# entries to the browser so the export dialog shows a live terminal-style
# panel while load_checkpoint / export_merged / export_gguf / export_lora /
# export_base run.
#
# Shape follows the training progress SSE endpoint
# (routes/training.py::stream_training_progress): each event carries
# `id`, `event`, and `data` fields, the stream starts with a `retry:`
# directive, and `Last-Event-ID` is honored on reconnect.
# (routes/training.py::stream_training_progress): each event carries `id`,
# `event`, `data` fields, the stream starts with a `retry:` directive, and
# `Last-Event-ID` is honored on reconnect.
def _format_sse(
@ -389,28 +373,26 @@ async def stream_export_logs(
current_subject: str = Depends(get_current_subject),
):
"""
Stream live stdout/stderr output from the export worker subprocess
as Server-Sent Events.
Stream live stdout/stderr from the export worker subprocess as
Server-Sent Events.
Events:
- `log` : a single log line (data: {"stream","line","ts"})
- `heartbeat`: periodic keepalive when no new lines are available
- `complete` : emitted once the export worker is idle and no new
lines arrived for ~1 second. Clients should close.
- `complete` : once the worker is idle and no new lines arrived for
~1 second. Clients should close.
- `error` : unrecoverable server-side error
The `id:` field on each event is the log entry's monotonic seq
number so the browser can resume via `Last-Event-ID` on reconnect.
Each event's `id:` field is the log entry's monotonic seq number so the
browser can resume via `Last-Event-ID` on reconnect.
"""
backend = get_export_backend()
# Determine starting cursor. Explicit `since` wins, then
# Last-Event-ID header on reconnect, otherwise start from the
# run-start snapshot captured by clear_logs() so the client sees
# every line emitted since the current run began -- even if the
# SSE connection opened after the POST that kicked off the export.
# Using get_current_log_seq() here would lose the early bootstrap
# lines that arrive in the gap between POST and SSE connect.
# Starting cursor: explicit `since` wins, then Last-Event-ID header on
# reconnect, else the run-start snapshot from clear_logs() so the client
# sees every line since the run began -- even if the SSE connection
# opened after the POST that kicked off the export. get_current_log_seq()
# would lose the early bootstrap lines arriving between POST and connect.
last_event_id = request.headers.get("last-event-id")
if since is None and last_event_id is not None:
try:
@ -425,8 +407,7 @@ async def stream_export_logs(
async def event_generator() -> AsyncGenerator[str, None]:
nonlocal cursor
# Tell the browser to reconnect after 3 seconds if the
# connection drops mid-export.
# Reconnect after 3 seconds if the connection drops mid-export.
yield "retry: 3000\n\n"
last_yield = time.monotonic()
@ -460,9 +441,8 @@ async def stream_export_logs(
yield _format_sse("{}", event = "heartbeat")
last_yield = now
if not backend.is_export_active():
# Give the reader thread a moment to drain any
# trailing lines the worker process printed
# just before signalling done.
# Give the reader thread a moment to drain trailing
# lines the worker printed just before signalling done.
if idle_since is None:
idle_since = now
elif now - idle_since > 1.0:
@ -477,8 +457,8 @@ async def stream_export_logs(
await asyncio.sleep(0.1)
except asyncio.CancelledError:
# Client disconnected mid-yield. Don't re-raise, just end
# the generator cleanly so StreamingResponse finalizes.
# Client disconnected mid-yield. Don't re-raise; end the
# generator cleanly so StreamingResponse finalizes.
return
except Exception as exc:
logger.error("Export log stream failed: %s", exc, exc_info = True)

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

File diff suppressed because it is too large Load diff

28
studio/backend/routes/mcp_servers.py
View file

@ -34,9 +34,9 @@ router = APIRouter()
def _looks_like_command(value: str) -> bool:
"""Whitespace is a one-way signal: a URL can't hold an unencoded space, so a
value with whitespace is definitely a command. No whitespace proves nothing
(a lone token may be a single-arg command or a scheme-less URL)."""
"""Whitespace is a one-way signal: a URL can't hold an unencoded space, so
a value with whitespace is definitely a command. No whitespace proves
nothing (a lone token may be a single-arg command or a scheme-less URL)."""
return any(ch.isspace() for ch in value)
@ -46,8 +46,8 @@ def _validate_url(url: str) -> str:
raise HTTPException(status_code = 400, detail = "url must not be empty")
# When stdio is enabled on this host, a non-HTTP value is a local command.
# Reuse this field so stdio servers ride the existing CRUD/storage with no
# schema change. A lone token (example.com, /usr/bin/srv) is ambiguous, so we
# keep the existing behaviour and treat any non-HTTP value as a command here.
# schema change. A lone token (example.com, /usr/bin/srv) is ambiguous, so
# keep existing behaviour and treat any non-HTTP value as a command here.
if stdio_mcp_enabled() and is_stdio(trimmed):
try:
parts = parse_stdio_command(trimmed)
@ -63,7 +63,7 @@ def _validate_url(url: str) -> str:
raise HTTPException(status_code = 400, detail = "command must not be empty")
if "://" in parts[0]:
# A URL-scheme first token is a mistyped URL, not a command. Reject
# it cleanly instead of exec-ing it (mirrors the frontend check).
# cleanly instead of exec-ing it (mirrors the frontend check).
raise HTTPException(
status_code = 400,
detail = "Enter an http(s):// URL, or a local command whose "
@ -87,7 +87,7 @@ def _validate_url(url: str) -> str:
def _normalize_headers(headers: dict[str, str] | None) -> dict[str, str] | None:
"""Trim header names, drop empties, coerce values to str. None if nothing left."""
"""Trim header names, drop empties, coerce values to str; None if empty."""
if not headers:
return None
out: dict[str, str] = {}
@ -126,7 +126,7 @@ async def create_mcp_server(
url = _validate_url(payload.url)
headers = _normalize_headers(payload.headers)
# OAuth is HTTP-only; force it off for stdio commands so a stale flag can't
# push the probe onto the 305s OAuth timeout. Backend is the enforcer.
# push the probe onto the 305s OAuth timeout. Backend enforces this.
use_oauth = payload.use_oauth and not is_stdio(url)
server_id = uuid.uuid4().hex[:16]
@ -182,7 +182,7 @@ async def update_mcp_server(
if not changes:
raise HTTPException(status_code = 400, detail = "No fields to update")
# headers == HTTP headers (remote) or env vars (stdio). On a transport-type
# switch with no new headers, drop the old ones so env secrets are not
# switch with no new headers, drop the old ones so env secrets aren't
# re-sent as HTTP headers (or vice versa).
if (
"url" in changes
@ -190,9 +190,9 @@ async def update_mcp_server(
and "headers_json" not in changes
):
changes["headers_json"] = None
# Clear persisted OAuth tokens when the URL changes or OAuth is
# disabled; fastmcp keys tokens by URL and would otherwise let a
# re-pointed server silently inherit the old account's credentials.
# Clear persisted OAuth tokens when the URL changes or OAuth is disabled;
# fastmcp keys tokens by URL and would otherwise let a re-pointed server
# silently inherit the old account's credentials.
if bool(old.get("use_oauth")) and (
("url" in changes and changes["url"] != old["url"]) or changes.get("use_oauth") is False
):
@ -248,8 +248,8 @@ async def test_mcp_server(
payload: McpServerTestRequest, current_subject: str = Depends(get_current_subject)
):
# URL/header validation must surface as 400 like create/update so the
# frontend's create-form pre-flight gets the same error semantics as
# the actual save call. Only catch transport/timeout errors below.
# frontend's create-form pre-flight gets the same error semantics as the
# save call. Only catch transport/timeout errors below.
url = _validate_url(payload.url)
headers = _normalize_headers(payload.headers)
try:

487
studio/backend/routes/models.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
"""
Model Management API routes
"""
"""Model management API routes."""
import hashlib
import json
@ -28,14 +26,13 @@ def _is_valid_repo_id(repo_id: str) -> bool:
def _safe_is_dir(path) -> bool:
"""``Path.is_dir()`` that returns ``False`` instead of raising.
"""``Path.is_dir()`` returning ``False`` instead of raising.
On Python >= 3.12 ``is_dir()``'s ``os.stat`` only suppresses
"not found"-class errors and now propagates ``PermissionError``
(EACCES); on Python <= 3.11 it returned ``False``. The folder-scan
endpoints probe well-known system locations (e.g. a root-owned,
mode-700 ``/usr/share/ollama/.ollama/models``) and must treat an
un-stat-able path as "not a directory", never 500.
Python >= 3.12's ``is_dir()`` propagates ``PermissionError``
(EACCES); <= 3.11 returned ``False``. Folder-scan endpoints probe
system locations (e.g. root-owned mode-700
``/usr/share/ollama/.ollama/models``) and must treat an un-stat-able
path as "not a directory", never 500.
"""
try:
return Path(path).is_dir()
@ -43,7 +40,7 @@ def _safe_is_dir(path) -> bool:
return False
# Add backend directory to path
# Add backend dir to path
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
@ -79,7 +76,7 @@ try:
resolve_export_dir,
)
except ImportError:
# Fallback: try to import from parent directory
# Fallback: import from parent directory
parent_backend = backend_path.parent / "backend"
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
@ -167,15 +164,14 @@ def _resolve_hf_cache_dir() -> Path:
def _is_model_directory(d: Path) -> bool:
"""Return ``True`` when *d* looks like a model directory.
A model directory must have **both** a config file (``config.json`` or
``adapter_config.json``) **and** actual model weight files. Both
conditions are required: a bare directory with only loose ``.gguf``
files (no config) might be a mixed collection, and a ``config.json``
alone (no weights) is not a model directory.
Requires **both** a config (``config.json`` or
``adapter_config.json``) **and** weight files: loose ``.gguf`` files
without config may be a mixed collection, and config without weights
is not a model dir.
Excludes ``mmproj`` GGUF files (vision projectors) and non-weight
``.bin`` files (``tokenizer.bin``, ``vocab.bin``, etc.) from the
weight check to avoid false positives.
``.bin`` files (``tokenizer.bin``, ``vocab.bin``, etc.) to avoid
false positives.
"""
def _is_weight_file(f: Path) -> bool:
@ -239,8 +235,8 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
or any(child.glob("*.gguf"))
)
except OSError:
# Skip individual children that are unreadable (permissions, broken
# symlinks, etc.) rather than failing the entire scan.
# Skip unreadable children (permissions, broken symlinks)
# rather than failing the whole scan.
continue
if not has_model_files:
continue
@ -257,7 +253,7 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
updated_at = updated_at,
),
)
# Also scan for standalone .gguf files directly in the models directory
# Also scan standalone .gguf files in the models directory
if limit is None or len(found) < limit:
for gguf_file in models_dir.glob("*.gguf"):
if limit is not None and len(found) >= limit:
@ -315,16 +311,16 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
"""Scan an LM Studio models directory for model files.
LM Studio uses a ``publisher/model-name`` folder structure containing
GGUF files, or standalone GGUF files at the top level.
LM Studio uses a ``publisher/model-name`` folder structure with GGUF
files, or standalone GGUF files at the top level.
"""
if not lm_dir.exists() or not lm_dir.is_dir():
return []
# If the directory itself is a model directory (has config AND weight
# files), it is not an LM Studio publisher structure -- return it as a
# single model entry. We cannot skip it silently because this function
# is the only scanner called for default LM Studio roots.
# If lm_dir is itself a model directory (config AND weights), it's
# not an LM Studio publisher structure -- return it as a single
# entry. Can't skip it silently: this is the only scanner called for
# default LM Studio roots.
if _is_model_directory(lm_dir):
try:
updated_at = lm_dir.stat().st_mtime
@ -360,9 +356,9 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
)
continue
# If the child directory itself looks like a model directory
# (has config AND weight files), surface it directly instead
# of descending into it as a publisher.
# If the child looks like a model directory (config AND
# weights), surface it directly instead of descending into
# it as a publisher.
if _is_model_directory(child):
try:
updated_at = child.stat().st_mtime
@ -379,7 +375,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
)
continue
# child is a publisher directory -- scan its sub-directories
# child is a publisher directory -- scan its subdirectories
for model_dir in child.iterdir():
try:
if model_dir.is_dir():
@ -430,11 +426,11 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
"""Return a writable directory for Ollama ``.gguf`` symlinks.
Prefers ``<ollama_dir>/.studio_links/`` so the links sit next to the
blobs they point at. Falls back to a per-ollama-dir namespace under
Studio's own cache when the models directory is read-only (common
for system installs under ``/usr/share/ollama`` or ``/var/lib/ollama``)
so we still surface Ollama models in those environments.
Prefers ``<ollama_dir>/.studio_links/`` so links sit next to their
blobs. Falls back to a per-ollama-dir namespace under Studio's cache
when the models dir is read-only (common for system installs under
``/usr/share/ollama`` or ``/var/lib/ollama``) so Ollama models still
surface there.
"""
from utils.paths.storage_roots import cache_root
@ -449,9 +445,8 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
e,
)
# Fallback: 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.
# Fallback: namespace by a hash of ollama_dir so two Ollama roots
# don't collide. Cache path, not a security boundary.
try:
digest = hashlib.sha256(str(ollama_dir.resolve()).encode()).hexdigest()[:12]
except OSError:
@ -472,34 +467,29 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[LocalModelInfo]:
"""Scan an Ollama models directory for downloaded models.
Ollama stores models in a content-addressable layout::
Ollama uses a content-addressable layout::
<ollama_dir>/manifests/<host>/<namespace>/<model>/<tag>
<ollama_dir>/blobs/sha256-...
The default host is ``registry.ollama.ai`` with namespace
``library`` (official models), but users can pull from custom
namespaces (``mradermacher/llama3``) or entirely different hosts
(``hf.co/org/repo:tag``). We iterate all manifest files via
``rglob`` so every layout depth is discovered.
Default host is ``registry.ollama.ai`` namespace ``library``
(official), but users can pull from custom namespaces
(``mradermacher/llama3``) or other hosts (``hf.co/org/repo:tag``).
We ``rglob`` all manifest files so every layout depth is found.
Each manifest is JSON with a ``layers`` array. The layer with
``mediaType == "application/vnd.ollama.image.model"`` contains the
GGUF weights. Vision models also have a projector layer
(``application/vnd.ollama.image.projector``). We read the config
layer to extract family/size info.
Each manifest is JSON with a ``layers`` array. The
``application/vnd.ollama.image.model`` layer holds the GGUF weights;
vision models add an ``application/vnd.ollama.image.projector``
layer. We read the config layer for family/size info.
Since Ollama blobs lack a ``.gguf`` extension (which the GGUF
loading pipeline requires), we create ``.gguf``-named links
pointing at the blobs so the existing ``detect_gguf_model`` and
``llama-server -m`` paths work unchanged. Each model gets its
own subdirectory under the links dir (keyed by a short hash of
the manifest path) so that ``detect_mmproj_file`` only sees the
projector for *that* model. Links are created as symlinks when
possible, falling back to hardlinks (Windows without Developer
Mode) as a last resort. The link dir lives under
``<ollama_dir>/.studio_links/`` when writable, otherwise under
Studio's own cache directory.
Ollama blobs lack the ``.gguf`` extension the loading pipeline
requires, so we create ``.gguf``-named links to them so
``detect_gguf_model`` and ``llama-server -m`` work unchanged. Each
model gets its own subdir under the links dir (keyed by a short hash
of the manifest path) so ``detect_mmproj_file`` only sees that
model's projector. Links are symlinks when possible, falling back to
hardlinks (Windows without Developer Mode). The link dir is
``<ollama_dir>/.studio_links/`` when writable, else Studio's cache.
"""
manifests_root = ollama_dir / "manifests"
if not manifests_root.is_dir():
@ -518,10 +508,9 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
def _make_link(link_dir: Path, link_name: str, target: Path) -> Optional[str]:
"""Create a .gguf-named link to an Ollama blob.
Tries symlink first, then hardlink (works on Windows without
Developer Mode when target is on the same filesystem). Skips
the model if neither works -- a full file copy of a multi-GB
GGUF inside a synchronous API request would block the backend.
Tries symlink, then hardlink (Windows without Developer Mode,
same filesystem). Skips the model if neither works -- a full
multi-GB copy in a sync API request would block the backend.
Idempotent: skips recreation when a valid link already exists.
"""
@ -529,9 +518,9 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
link_path = link_dir / link_name
resolved = target.resolve()
# Skip if the link already points at the exact same blob.
# Only use samefile -- size-based checks can reuse stale links
# after `ollama pull` updates a tag to a same-sized blob.
# Skip if the link already points at the same blob. Use samefile
# only -- size checks can reuse stale links after `ollama pull`
# updates a tag to a same-sized blob.
try:
if link_path.exists() and os.path.samefile(str(link_path), str(resolved)):
return str(link_path)
@ -700,9 +689,9 @@ async def list_local_models(
hf_default = hf_default_cache_dir()
lm_dirs = lmstudio_model_dirs()
# Validate models_dir against an allowlist of trusted directories.
# Only the trusted Path objects are used for filesystem access -- the
# user-supplied string is only used for matching, never for path construction.
# Validate models_dir against an allowlist of trusted dirs. Only the
# trusted Path objects are used for FS access; the user string is
# only used for matching, never for path construction.
allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir]
if legacy_hf.is_dir():
allowed_roots.append(legacy_hf)
@ -719,7 +708,7 @@ async def list_local_models(
for root in allowed_roots:
root_str = os.path.realpath(str(root))
if requested == root_str or requested.startswith(root_str + os.sep):
models_root = root # Use the trusted root, not the user-supplied path
models_root = root # trusted root, not the user-supplied path
break
if models_root is None:
raise HTTPException(
@ -734,7 +723,7 @@ async def list_local_models(
if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve():
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ when env vars are overridden)
# Scan HF system default cache (may differ under env overrides)
if (
hf_default.is_dir()
and hf_default.resolve() != hf_cache_dir.resolve()
@ -746,7 +735,7 @@ async def list_local_models(
for lm_dir in lm_dirs:
local_models += _scan_lmstudio_dir(lm_dir)
# Scan user-added custom folders (cap per-folder to avoid unbounded scans)
# Scan user-added custom folders (per-folder cap)
from storage.studio_db import list_scan_folders
_MAX_MODELS_PER_FOLDER = 200
@ -758,9 +747,9 @@ async def list_local_models(
for folder in custom_folders:
folder_path = Path(folder["path"])
try:
# Ollama scanner creates .studio_links/ with .gguf symlinks.
# Filter those from the generic scanners to avoid duplicates
# and leaking internal paths into the UI.
# Ollama scanner creates .studio_links/ with .gguf
# symlinks. Filter those from generic scanners to avoid
# duplicates and leaking internal paths into the UI.
_generic = [
m
for m in (
@ -781,10 +770,10 @@ async def list_local_models(
continue
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate models, but always keep custom folder entries so they
# appear in the "Custom Folders" UI section even when the same model
# also exists in the HF cache or default models directory. Use a
# (id, source) key for custom entries to avoid collisions.
# Deduplicate, but always keep custom folder entries so they show
# in the "Custom Folders" UI section even when the same model is
# also in the HF cache or default models dir. Custom entries use
# an (id, source) key to avoid collisions.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
@ -855,11 +844,10 @@ async def remove_scan_folder_endpoint(
async def get_recommended_folders(current_subject: str = Depends(get_current_subject)):
"""Return well-known model directories that exist on this machine.
Lightweight alternative to ``browse-folders`` for showing quick-pick
chips without the overhead of enumerating a directory tree. Returns
paths that actually exist on disk (HF cache, LM Studio, Ollama,
``~/models``, etc.) so the frontend can offer them as one-click
"Recommended" shortcuts in the Custom Folders section.
Lightweight alternative to ``browse-folders`` for quick-pick chips
without enumerating a directory tree. Returns paths that exist on
disk (HF cache, LM Studio, Ollama, ``~/models``, etc.) for the
frontend's one-click "Recommended" shortcuts.
"""
from utils.paths.storage_roots import lmstudio_model_dirs
@ -900,26 +888,25 @@ async def get_recommended_folders(current_subject: str = Depends(get_current_sub
return {"folders": folders}
# Heuristic ceiling on how many children to stat when checking whether a
# directory "looks like" it contains models. Keeps the browser snappy
# even when a directory has thousands of unrelated entries.
# Max children to stat when checking if a directory "looks like" it
# holds models. Keeps the browser snappy on dirs with thousands of
# unrelated entries.
_BROWSE_MODEL_HINT_PROBE = 64
# Hard cap on how many subdirectory entries we send back. Pointing the
# browser at something like ``/usr/lib`` or ``/proc`` must not stat-storm
# the process or send tens of thousands of rows to the client.
# Hard cap on subdirectory entries returned. Browsing ``/usr/lib`` or
# ``/proc`` must not stat-storm the process or send tens of thousands of
# rows to the client.
_BROWSE_ENTRY_CAP = 2000
def _count_model_files(directory: Path, cap: int = 200) -> int:
"""Count GGUF/safetensors files immediately inside *directory*.
Used to surface a count-hint on the response so the UI can tell
users that a leaf directory (no subdirs, only weights) is a valid
"Use this folder" target.
Bounded by *visited entries*, not by *match count*: in directories
with many non-model files (or many subdirectories) the scan still
stops after ``cap`` entries so a UI hint never costs more than a
bounded directory walk.
Surfaces a count-hint so the UI can mark a leaf directory (no
subdirs, only weights) as a valid "Use this folder" target.
Bounded by *visited entries*, not match count: the scan stops after
``cap`` entries even in dirs full of non-model files/subdirs, so a UI
hint never costs more than a bounded directory walk.
"""
n = 0
visited = 0
@ -945,10 +932,10 @@ def _count_model_files(directory: Path, cap: int = 200) -> int:
def _has_direct_model_signal(directory: Path) -> bool:
"""Return True if *directory* has an immediate child that signals
it holds a model: a GGUF/safetensors/config.json file, or a
`models--*` subdir (HF hub cache). Bounded by
``_BROWSE_MODEL_HINT_PROBE`` to stay fast."""
"""Return True if *directory* has an immediate child signalling a
model: a GGUF/safetensors/config.json file, or a `models--*` subdir
(HF hub cache). Bounded by ``_BROWSE_MODEL_HINT_PROBE`` to stay
fast."""
try:
it = directory.iterdir()
except OSError:
@ -975,23 +962,22 @@ def _has_direct_model_signal(directory: Path) -> bool:
def _looks_like_model_dir(directory: Path) -> bool:
"""Bounded heuristic used by the folder browser to flag directories
worth exploring. False negatives are fine; the real scanner is
"""Bounded heuristic to flag directories worth exploring in the
folder browser. False negatives are fine; the real scanner is
authoritative.
Three signals, cheapest first:
1. Directory name itself: ``models--*`` is the HuggingFace hub cache
layout (``blobs``/``refs``/``snapshots`` children wouldn't match
the file-level probes below).
2. An immediate child is a weight file or config (handled by
1. Directory name ``models--*`` (HuggingFace hub cache layout; its
``blobs``/``refs``/``snapshots`` children wouldn't match the
file-level probes below).
2. An immediate child is a weight file or config (via
:func:`_has_direct_model_signal`).
3. A grandchild has a direct signal -- this catches the
``publisher/model/weights.gguf`` layout used by LM Studio and
Ollama. We probe at most the first
``_BROWSE_MODEL_HINT_PROBE`` child directories, each of which is
checked with a bounded :func:`_has_direct_model_signal` call,
so the total cost stays O(PROBE^2) worst-case.
3. A grandchild has a direct signal -- catches the
``publisher/model/weights.gguf`` layout of LM Studio and Ollama.
Probes at most the first ``_BROWSE_MODEL_HINT_PROBE`` child dirs,
each via a bounded :func:`_has_direct_model_signal`, so total cost
is O(PROBE^2) worst-case.
"""
if directory.name.startswith("models--"):
return True
@ -1011,7 +997,7 @@ def _looks_like_model_dir(directory: Path) -> bool:
continue
except OSError:
continue
# Fast name check first
# Fast name check first.
if child.name.startswith("models--"):
return True
if _has_direct_model_signal(child):
@ -1022,16 +1008,15 @@ def _looks_like_model_dir(directory: Path) -> bool:
def _build_browse_allowlist() -> list[Path]:
"""Return the list of root directories the folder browser is allowed
to walk. The same list is used to seed the sidebar suggestion chips,
so chip targets are always reachable.
"""Return the root directories the folder browser may walk. The same
list seeds the sidebar suggestion chips, so chip targets are always
reachable.
Roots include the current user's HOME, the resolved HF cache dirs,
Studio's own outputs/exports/studio root, registered scan folders,
and well-known third-party local-LLM dirs (LM Studio, Ollama,
`~/models`). Each is added only if it currently resolves to a real
directory, so we never produce a "dead" sandbox boundary the user
can't navigate into.
Roots: the user's HOME, resolved HF cache dirs, Studio's
outputs/exports/studio root, registered scan folders, and well-known
third-party local-LLM dirs (LM Studio, Ollama, `~/models`). Each is
added only if it resolves to a real directory, so we never produce a
"dead" sandbox boundary.
"""
from utils.paths import (
hf_default_cache_dir,
@ -1100,9 +1085,8 @@ def _build_browse_allowlist() -> list[Path]:
def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
"""Return True if *target* equals or is a descendant of any allowed
root. The comparison uses ``os.path.realpath`` so symlinks cannot be
used to escape the sandbox.
"""Return True if *target* equals or descends from any allowed root.
Uses ``os.path.realpath`` so symlinks can't escape the sandbox.
"""
try:
target_real = os.path.realpath(str(target))
@ -1262,30 +1246,28 @@ async def browse_folders(
"""
List immediate subdirectories of *path* for the Custom Folders picker.
The frontend uses this to render a modal folder browser without needing
a native OS dialog (Studio is served over HTTP, so the browser can't
reveal absolute paths on the host). The endpoint is read-only and does
not create, move, or delete anything. It simply enumerates visible
subdirectories so the user can click their way to a folder and hand
the resulting string back to POST `/api/models/scan-folders`.
Lets the frontend render a modal folder browser without a native OS
dialog (Studio is served over HTTP, so the browser can't reveal host
paths). Read-only: it only enumerates visible subdirectories so the
user can click to a folder and hand the string back to POST
`/api/models/scan-folders`.
Sandbox: requests are bounded to the allowlist returned by
Sandbox: bounded to the allowlist from
:func:`_build_browse_allowlist` (HOME, HF cache, Studio dirs,
registered scan folders, well-known model dirs). Paths outside the
allowlist return 403 so users cannot probe ``/etc``, ``/proc``,
``/root`` (when not HOME), or other sensitive system locations
even if the server process can read them. Symlinks are resolved
via ``os.path.realpath`` before the check, so symlink traversal
cannot escape the sandbox either.
registered scan folders, well-known model dirs). Paths outside it
return 403 so users can't probe ``/etc``, ``/proc``, ``/root`` (when
not HOME), or other sensitive locations even if the process can read
them. Symlinks are resolved via ``os.path.realpath`` first, so
traversal can't escape the sandbox.
Sorting: directories that look like they hold models come first, then
plain directories, then hidden entries (if `show_hidden=true`).
Sorting: model-bearing dirs first, then plain dirs, then hidden
entries (if `show_hidden=true`).
"""
from utils.paths import hf_default_cache_dir, well_known_model_dirs
from storage.studio_db import list_scan_folders
# Build the allowlist once -- both the sandbox check below and the
# suggestion chips use the same set, so chips are always navigable.
# Build the allowlist once -- the sandbox check and suggestion chips
# share it, so chips are always navigable.
allowed_roots = _build_browse_allowlist()
try:
@ -1301,7 +1283,7 @@ async def browse_folders(
raise
# Enumerate immediate subdirectories with a bounded cap so a stray
# query against ``/usr/lib`` or ``/proc`` can't stat-storm the process.
# query against ``/usr/lib`` or ``/proc`` can't stat-storm us.
entries: list[BrowseEntry] = []
truncated = False
visited = 0
@ -1321,13 +1303,12 @@ async def browse_folders(
try:
for child in it:
# Bound by *visited entries*, not by *appended entries*: in
# directories full of files (or hidden subdirs when
# ``show_hidden=False``) the cap on ``len(entries)`` would
# never trigger and we'd still stat every child. Counting
# visits keeps the worst-case work to ``_BROWSE_ENTRY_CAP``
# iterdir/is_dir calls regardless of how many of them
# survive the filters below.
# Bound by *visited*, not *appended*: in dirs full of files
# (or hidden subdirs when ``show_hidden=False``) a cap on
# ``len(entries)`` would never trigger and we'd stat every
# child. Counting visits caps worst-case work at
# ``_BROWSE_ENTRY_CAP`` iterdir/is_dir calls regardless of
# how many survive the filters below.
visited += 1
if visited > _BROWSE_ENTRY_CAP:
truncated = True
@ -1366,10 +1347,10 @@ async def browse_folders(
entries.sort(key = _sort_key)
# Parent is None at the filesystem root (`p.parent == p`) AND when
# the parent would step outside the sandbox -- otherwise the up-row
# would 403 on click. Users can still hop to other allowed roots
# via the suggestion chips below.
# Parent is None at the filesystem root (`p.parent == p`) and when
# the parent would leave the sandbox -- otherwise the up-row would
# 403 on click. Users can still hop to other allowed roots via the
# suggestion chips below.
parent: Optional[str]
if target.parent == target or not _is_path_inside_allowlist(target.parent, allowed_roots):
parent = None
@ -1393,27 +1374,25 @@ async def browse_folders(
seen_sug.add(resolved)
suggestions.append(resolved)
# Home always comes first -- it's the safe fallback when everything
# else is cold.
# Home first -- the safe fallback when everything else is cold.
_add_sug(Path.home())
# The HF cache root the process is actually using.
try:
_add_sug(hf_default_cache_dir())
except Exception:
pass
# Already-registered scan folders (what the user has curated).
# Already-registered scan folders (user-curated).
try:
for folder in list_scan_folders():
_add_sug(Path(folder.get("path", "")))
except Exception as exc:
logger.debug("browse-folders: could not load scan folders: %s", exc)
# Directories commonly used by other local-LLM tools: LM Studio
# (`~/.lmstudio/models` + legacy `~/.cache/lm-studio/models` +
# user-configured downloadsFolder from LM Studio's settings.json),
# Ollama (`~/.ollama/models` + common system paths + OLLAMA_MODELS
# env var), and generic user-choice spots (`~/models`, `~/Models`).
# Each helper only returns paths that currently exist so we never
# show dead chips.
# Dirs used by other local-LLM tools: LM Studio
# (`~/.lmstudio/models` + legacy `~/.cache/lm-studio/models` + LM
# Studio settings.json downloadsFolder), Ollama (`~/.ollama/models`
# + common system paths + OLLAMA_MODELS env var), and generic spots
# (`~/models`, `~/Models`). Each helper returns only existing paths
# so we never show dead chips.
try:
for p in well_known_model_dirs():
_add_sug(p)
@ -1441,18 +1420,12 @@ def _looks_like_mlx_repo(model_id: str) -> bool:
@router.get("/list")
async def list_models(current_subject: str = Depends(get_current_subject)):
"""
List available models (default models and loaded models).
This endpoint returns the default models and any currently loaded models.
"""
"""List available models: default plus currently loaded."""
try:
inference_backend = get_inference_backend()
# Get default models
default_models = inference_backend.default_models
# Get loaded models
loaded_models = []
for model_name, model_data in inference_backend.models.items():
_is_vision = model_data.get("is_vision", False)
@ -1470,7 +1443,7 @@ async def list_models(current_subject: str = Depends(get_current_subject)):
)
loaded_models.append(model_info)
# Include active GGUF model (loaded via llama-server)
# Include active GGUF model (loaded via llama-server).
from routes.inference import get_llama_cpp_backend
llama_backend = get_llama_cpp_backend()
@ -1486,12 +1459,11 @@ async def list_models(current_subject: str = Depends(get_current_subject)):
)
)
# Combine default and loaded models
# Combine default and loaded models.
all_models = []
seen_ids = set()
# Prefer loaded entries for duplicate ids so runtime flags
# (is_mlx, is_vision, is_audio, ...) are not lost.
# Prefer loaded entries for duplicate ids so runtime flags survive.
loaded_by_id = {model_info.id: model_info for model_info in loaded_models}
# Add default models
@ -1506,7 +1478,7 @@ async def list_models(current_subject: str = Depends(get_current_subject)):
all_models.append(model_info)
seen_ids.add(model_id)
# Add loaded models
# Add loaded models.
for model_info in loaded_models:
if model_info.id not in seen_ids:
all_models.append(model_info)
@ -1525,7 +1497,7 @@ async def list_models(current_subject: str = Depends(get_current_subject)):
def _get_max_position_embeddings(config) -> Optional[int]:
"""Extract max_position_embeddings from a model config, checking text_config fallback."""
"""Extract max_position_embeddings from a config, with text_config fallback."""
if hasattr(config, "max_position_embeddings"):
return config.max_position_embeddings
if hasattr(config, "text_config") and hasattr(config.text_config, "max_position_embeddings"):
@ -1534,7 +1506,7 @@ def _get_max_position_embeddings(config) -> Optional[int]:
def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Optional[int]:
"""Get total size of model weight files from HF Hub."""
"""Total size of model weight files from HF Hub."""
try:
from huggingface_hub import HfApi
@ -1562,11 +1534,7 @@ async def get_model_config(
hf_token: Optional[str] = Query(None),
current_subject: str = Depends(get_current_subject),
):
"""
Get configuration for a specific model.
This endpoint wraps the backend load_model_defaults function.
"""
"""Get configuration for a specific model (wraps load_model_defaults)."""
try:
if not is_local_path(model_name):
resolved = resolve_cached_repo_id_case(model_name)
@ -1581,15 +1549,14 @@ async def get_model_config(
logger.info(f"Getting model config for: {model_name}")
from utils.models.model_config import detect_audio_type
# Load model defaults from backend
config_dict = load_model_defaults(model_name)
# Detect model capabilities (pass HF token for gated models)
# Detect capabilities (pass HF token for gated models).
is_vision = is_vision_model(model_name, hf_token = hf_token)
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
audio_type = detect_audio_type(model_name, hf_token = hf_token)
# Check if it's a LoRA adapter
# Check if it's a LoRA adapter.
is_lora = False
base_model = None
max_position_embeddings = None
@ -1601,7 +1568,7 @@ async def get_model_config(
except Exception:
pass
# Fallback: try AutoConfig directly if not found yet
# Fallback: try AutoConfig directly if not found yet.
if max_position_embeddings is None:
try:
from transformers import AutoConfig as _AutoConfig
@ -1653,18 +1620,17 @@ async def scan_loras(
),
current_subject: str = Depends(get_current_subject),
):
"""
Scan for trained LoRA adapters and exported models.
"""Scan for trained LoRA adapters and exported models.
Returns both training outputs (from outputs_dir) and exported models
(from exports_dir) in a single list, distinguished by source field.
Returns training outputs (outputs_dir) and exported models
(exports_dir) in one list, distinguished by the source field.
"""
try:
resolved_outputs_dir = str(resolve_output_dir(outputs_dir))
resolved_exports_dir = str(resolve_export_dir(exports_dir))
lora_list = []
# Scan training outputs
# Scan training outputs.
trained_models = scan_trained_models(outputs_dir = resolved_outputs_dir)
for display_name, model_path, model_type in trained_models:
base_model = get_base_model_from_checkpoint(model_path)
@ -1712,7 +1678,7 @@ def _is_path_under(path: Path, root: Path) -> bool:
def _is_path_under_lexically(path: Path, root: Path) -> bool:
"""Check containment without resolving the final path's symlink target."""
"""Check containment without resolving the final path's symlink."""
try:
absolute_path = Path(os.path.abspath(str(path)))
absolute_root = Path(os.path.abspath(str(root)))
@ -1744,10 +1710,10 @@ def _loading_model_matches_deleted_path(loading_model: object, deleted_path: Pat
def _prune_empty_parents(start: Path, stop_at: Path) -> None:
"""Remove empty ancestor directories of ``start`` up to (but not including) ``stop_at``.
"""Remove empty ancestors of ``start`` up to (not including) ``stop_at``.
Used after deleting a model checkpoint so the enclosing run directory does
not linger as an empty entry in scan results.
Used after deleting a checkpoint so the enclosing run dir doesn't
linger as an empty entry in scan results.
"""
try:
stop_resolved = stop_at.resolve()
@ -1799,8 +1765,8 @@ async def delete_finetuned_model(
):
"""Delete a Studio-trained or exported model from disk.
Only paths under Studio's outputs/exports roots are accepted. Exported
GGUF entries can delete one quantization variant at a time.
Only paths under Studio's outputs/exports roots are accepted.
Exported GGUF entries can delete one quant variant at a time.
"""
if source not in {"training", "exported"}:
raise HTTPException(
@ -2126,18 +2092,16 @@ async def get_gguf_variants(
hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"),
current_subject: str = Depends(get_current_subject),
):
"""
List available GGUF quantization variants for a HuggingFace repo
or a local directory (e.g. LM Studio model folder).
"""List GGUF quantization variants for a HuggingFace repo or 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.
Returns all variants (Q4_K_M, Q8_0, BF16, etc.) with file sizes,
whether the model supports vision, and the recommended default.
"""
try:
from utils.models.model_config import is_local_path, list_local_gguf_variants
# Local directory path (e.g. LM Studio models) — scan filesystem
# 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)
@ -2160,26 +2124,25 @@ async def get_gguf_variants(
default_variant = default_variant,
)
# Remote HuggingFace repo — query HF API
# Remote HuggingFace repo — query HF API.
variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token)
# Determine default variant
# Determine default variant.
filenames = [v.filename for v in variants]
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
# Check which variants are fully downloaded in the HF cache.
# For split GGUFs, ALL shards must be present -- sum cached bytes
# per variant and compare against the expected total.
# HF cache dir uses the exact case from the repo_id at download time,
# which may differ from the canonical HF repo_id, so do a
# case-insensitive match.
# Which variants are fully downloaded in the HF cache. For split
# GGUFs, ALL shards must be present -- sum cached bytes per
# variant vs. the expected total. The HF cache dir uses the
# repo_id casing from download time, which may differ from the
# canonical repo_id, so match case-insensitively.
cached_bytes_by_quant: dict[str, int] = {}
try:
import re as _re
from huggingface_hub import constants as hf_constants
# Sanitize repo_id: must be "owner/name" with safe chars only
# Sanitize repo_id: "owner/name" with safe chars only.
if not _is_valid_repo_id(repo_id):
raise ValueError(f"Invalid repo_id format: {repo_id}")
@ -2203,7 +2166,7 @@ async def get_gguf_variants(
cached = cached_bytes_by_quant.get(variant.quant, 0)
if cached == 0 or variant.size_bytes == 0:
return False
# Allow small rounding tolerance (symlinks vs real sizes)
# Small rounding tolerance (symlinks vs real sizes).
return cached >= variant.size_bytes * 0.99
return GgufVariantsResponse(
@ -2236,10 +2199,10 @@ async def get_gguf_download_progress(
expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
current_subject: str = Depends(get_current_subject),
):
"""Return download progress by checking cached GGUF files for a specific variant.
"""Download progress from cached GGUF files for a specific variant.
Tracks completed shard downloads in snapshots and in-progress downloads
in the blobs directory (incomplete files).
Tracks completed shards in snapshots and in-progress (.incomplete)
downloads in the blobs directory.
"""
try:
if not _is_valid_repo_id(repo_id):
@ -2258,12 +2221,12 @@ async def get_gguf_download_progress(
in_progress_bytes = 0
for entry in cache_dir.iterdir():
if entry.name.lower() == target:
# Count completed .gguf files matching this variant in snapshots
# Completed .gguf files for this variant in snapshots.
for f in _iter_gguf_paths(entry):
fname = f.name.lower().replace("-", "").replace("_", "")
if not variant_lower or variant_lower in fname:
downloaded_bytes += f.stat().st_size
# Check blobs for in-progress downloads (.incomplete files)
# In-progress downloads (.incomplete) in blobs.
blobs_dir = entry / "blobs"
if blobs_dir.is_dir():
for f in blobs_dir.iterdir():
@ -2273,7 +2236,7 @@ async def get_gguf_download_progress(
total_progress_bytes = downloaded_bytes + in_progress_bytes
progress = min(total_progress_bytes / expected_bytes, 0.99) if expected_bytes > 0 else 0
# Only report 1.0 when all bytes are in completed files (not in-progress)
# Report 1.0 only when all bytes are in completed files.
if expected_bytes > 0 and downloaded_bytes >= expected_bytes:
progress = 1.0
return {
@ -2288,9 +2251,9 @@ async def get_gguf_download_progress(
def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
"""Pick the most useful on-disk path for a HF cache repo.
Prefers the most-recent snapshot dir (what `from_pretrained` actually
points at). Falls back to the cache repo root. Returns the resolved
realpath so symlinks under snapshots/ are followed back to blobs/.
Prefers the most-recent snapshot dir (what `from_pretrained` points
at), falling back to the cache repo root. Returns the resolved
realpath so symlinks under snapshots/ follow back to blobs/.
"""
try:
snapshots_dir = repo_dir / "snapshots"
@ -2312,11 +2275,10 @@ async def get_download_progress(
"""Return download progress for any HuggingFace model repo.
Checks the local HF cache for completed blobs and in-progress
(.incomplete) downloads. Uses the HF API to determine the expected
total size on the first call, then caches it for subsequent polls.
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.
(.incomplete) downloads. Gets the expected total size from the HF API
on the first call, then caches it for later polls. Also returns
``cache_path``: the realpath of the snapshot dir (or cache repo root
if no snapshot yet) so the UI can show where weights live on disk.
"""
_empty = {
"downloaded_bytes": 0,
@ -2356,10 +2318,10 @@ async def get_download_progress(
if downloaded_bytes == 0:
return {**_empty, "cache_path": cache_path}
# Get expected size from HF API (cached per repo_id)
# Expected size from HF API (cached per repo_id).
expected_bytes = _get_repo_size_cached(repo_id)
if expected_bytes <= 0:
# Cannot determine total; report bytes only, no percentage
# Total unknown; report bytes only, no percentage.
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": 0,
@ -2367,11 +2329,10 @@ async def get_download_progress(
"cache_path": cache_path,
}
# Use 95% threshold for completion (blob deduplication can make
# completed_bytes differ slightly from expected_bytes).
# Do NOT use "no .incomplete files" as a completion signal --
# HF downloads files sequentially, so between files there are
# no .incomplete files even though the download is far from done.
# 95% completion threshold (blob dedup can make completed_bytes
# differ slightly from expected_bytes). Do NOT treat "no
# .incomplete files" as done -- HF downloads sequentially, so
# between files none exist even though it's far from finished.
if completed_bytes >= expected_bytes * 0.95:
progress = 1.0
else:
@ -2406,14 +2367,14 @@ def _get_repo_size_cached(repo_id: str) -> int:
def _all_hf_cache_scans():
"""Return scan_cache_dir results for the active, legacy, and default HF caches."""
"""scan_cache_dir results for the active, legacy, and default HF caches."""
from huggingface_hub import scan_cache_dir
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir
scans = [scan_cache_dir()]
seen: set[str] = set()
try:
# Resolve the active cache dir so we can dedup
# Resolve the active cache dir for deduplication.
from huggingface_hub.constants import HF_HUB_CACHE
seen.add(str(Path(HF_HUB_CACHE).resolve()))
except Exception:
@ -2435,14 +2396,14 @@ def _is_gguf_filename(name: str) -> bool:
def _is_mmproj_filename(name: str) -> bool:
"""Match GGUF vision-adapter (mmproj) files. Kept consistent with
"""Match GGUF vision-adapter (mmproj) files. Consistent with
``utils.models.model_config._is_mmproj``."""
return "mmproj" in name.lower()
def _is_main_gguf_filename(name: str) -> bool:
"""A GGUF file that is a primary weight artifact, not an mmproj
vision adapter."""
"""A GGUF file that is a primary weight, not an mmproj vision
adapter."""
return _is_gguf_filename(name) and not _is_mmproj_filename(name)
@ -2453,17 +2414,16 @@ def _iter_gguf_paths(root: Path):
def _repo_gguf_size_bytes(repo_info) -> int:
"""Return the total on-disk size of primary GGUF weight files across
all revisions, excluding mmproj vision-adapter files.
"""Total on-disk size of primary GGUF weight files across all
revisions, excluding mmproj vision-adapter files.
Hugging Face hardlinks blobs shared between revisions, so this
deduplicates by blob path (or, as a fallback, by revision commit
hash + filename) to avoid double-counting the same bytes. Files
with an unknown size (``size_on_disk is None``, e.g. a partial or
interrupted download) are treated as zero bytes. mmproj files are
excluded so that repos whose only ``.gguf`` artifact is a vision
adapter are not classified as GGUF repos: the variant selector
filters mmproj out and would otherwise show zero pickable variants.
deduplicates by blob path (or revision commit hash + filename as a
fallback) to avoid double-counting. Unknown sizes (``size_on_disk is
None``, e.g. a partial download) count as zero. mmproj files are
excluded so repos whose only ``.gguf`` artifact is a vision adapter
aren't classed as GGUF repos: the variant selector filters mmproj
out and would otherwise show zero pickable variants.
"""
unique_blobs: dict[str, int] = {}
for revision in repo_info.revisions:
@ -2480,9 +2440,9 @@ def _repo_gguf_size_bytes(repo_info) -> int:
def _repo_has_gguf_files(repo_info) -> bool:
"""Return True when any revision in a cached repo contains a
primary GGUF weight file. Repos whose only ``.gguf`` artifact is
an mmproj vision adapter are not treated as GGUF here."""
"""True when any revision in a cached repo has a primary GGUF weight
file. Repos whose only ``.gguf`` artifact is an mmproj vision adapter
are not treated as GGUF here."""
return _repo_gguf_size_bytes(repo_info) > 0
@ -2576,14 +2536,14 @@ async def delete_cached_model(
):
"""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.
With *variant*, only GGUF files matching that quant label are removed
(e.g. ``UD-Q4_K_XL``); otherwise the whole 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")
# Check if model is currently loaded
# Refuse if the model is currently loaded.
try:
from routes.inference import get_llama_cpp_backend
llama_backend = get_llama_cpp_backend()
@ -2641,7 +2601,7 @@ async def delete_cached_model(
quant = _extract_quant_label(f.file_name)
if quant.lower() != variant.lower():
continue
# Delete the blob (actual data) and the snapshot symlink
# Delete the blob (data) and the snapshot symlink.
try:
blob = Path(f.blob_path)
snap = Path(f.file_path)
@ -2700,8 +2660,7 @@ async def list_checkpoints(
),
current_subject: str = Depends(get_current_subject),
):
"""
List available checkpoints in the outputs directory.
"""List checkpoints in the outputs directory.
Scans the outputs folder for training runs and their checkpoints.
"""

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

@ -4,12 +4,12 @@
"""
API routes for external LLM provider management.
Provides endpoints for:
- Discovering available provider types (registry)
Endpoints:
- Discover available provider types (registry)
- CRUD for saved provider configurations (no API keys stored)
- Fetching the RSA public key for API key encryption
- Testing provider connectivity
- Listing models from a provider
- Fetch the RSA public key for API key encryption
- Test provider connectivity
- List models from a provider
"""
import uuid
@ -54,11 +54,10 @@ router = APIRouter()
async def get_public_key(current_subject: str = Depends(get_current_subject)):
"""Return the RSA public key PEM for client-side API key encryption.
The ``fingerprint`` field is a short SHA256 of the PEM and is meant
purely for diagnostics — a mismatch between what the frontend
captured at encrypt time and what the server reports here is a
clear signal that the keypair rotated mid-flight (e.g. the server
re-ran ``init_key_pair`` for any reason).
The ``fingerprint`` field is a short SHA256 of the PEM, for diagnostics: a
mismatch between what the frontend captured at encrypt time and what the
server reports here signals the keypair rotated mid-flight (e.g. the server
re-ran ``init_key_pair``).
"""
return {
"public_key": get_public_key_pem(),
@ -80,10 +79,10 @@ async def list_registry(current_subject: str = Depends(get_current_subject)):
@router.get("/pricing")
async def get_pricing_snapshot(current_subject: str = Depends(get_current_subject)):
"""Static per-MTok pricing table the frontend uses to convert
upstream usage chunks into a per-turn USD cost. See
``core/inference/pricing.py`` for sourcing notes; values reflect
the published prices as of the file's last update."""
"""Static per-MTok pricing table the frontend uses to convert upstream
usage chunks into a per-turn USD cost. See ``core/inference/pricing.py``
for sourcing notes; values reflect the published prices as of the file's
last update."""
return pricing_snapshot()
@ -196,7 +195,7 @@ async def test_provider(
Test connectivity to an external provider.
Makes a lightweight GET /models call to verify the API key works.
The encrypted_api_key is decrypted server-side and never stored.
encrypted_api_key is decrypted server-side and never stored.
"""
info = get_provider_info(payload.provider_type)
if info is None:
@ -267,7 +266,7 @@ async def list_provider_models(
"""
List models available from an external provider.
The encrypted_api_key is decrypted server-side and never stored.
encrypted_api_key is decrypted server-side and never stored.
"""
info = get_provider_info(payload.provider_type)
if info is None:
@ -308,16 +307,15 @@ async def list_provider_models(
try:
models = await client.list_models()
# Registry-level model-id filters are scoped to the canonical
# native Gemini base. A custom Gemini OAI-compatible proxy
# (LiteLLM, deployment gateway) returns IDs like
# `google/gemini-2.5-flash`, `gemini/gemini-2.5-flash`, or
# team-prefixed deployment aliases; the native allowlist regex
# would strip those out and leave the picker empty even though
# the chat path now routes them via the OAI-compatible
# dispatcher (the same gate ExternalProviderClient applies for
# request building). Match the host check here so the model
# list and chat dispatch agree on what counts as "native".
# Registry-level model-id filters are scoped to the canonical native
# Gemini base. A custom Gemini OAI-compatible proxy (LiteLLM, deployment
# gateway) returns IDs like `google/gemini-2.5-flash`,
# `gemini/gemini-2.5-flash`, or team-prefixed deployment aliases; the
# native allowlist regex would strip those and leave the picker empty,
# even though the chat path routes them via the OAI-compatible
# dispatcher (the same gate ExternalProviderClient applies for request
# building). Match the host check here so the model list and chat
# dispatch agree on what counts as "native".
apply_registry_model_filters = True
if payload.provider_type == "gemini":
try:
@ -344,11 +342,10 @@ async def list_provider_models(
denylist = info.get("model_id_denylist")
if denylist is not None:
models = [m for m in models if not denylist.search(m.get("id", ""))]
# Apply an optional cap after filtering so registry entries with a
# large remote catalog (e.g. HF Inference Providers) can stay
# picker-sized. No popularity sort happens server-side, so this is
# "first N matches" — pair with default_models for any must-have
# flagship ids.
# Optional cap after filtering so registry entries with a large remote
# catalog (e.g. HF Inference Providers) stay picker-sized. No
# server-side popularity sort, so this is "first N matches" — pair with
# default_models for any must-have flagship ids.
limit = info.get("model_id_limit")
if isinstance(limit, int) and limit > 0:
models = models[:limit]

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

@ -16,13 +16,12 @@ import asyncio
from datetime import datetime
import uuid as _uuid
# Add backend directory to path
# The backend code should be in the same directory structure
# Add backend directory to path (same directory structure).
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
# Import backend functions
# Import backend functions.
try:
from core.training import get_training_backend
from core.training.resume import (
@ -34,7 +33,7 @@ try:
from utils.models.model_config import load_model_defaults
from utils.paths import resolve_dataset_path
except ImportError:
# Fallback: try to import from parent directory
# Fallback: import from the parent directory.
parent_backend = backend_path.parent / "backend"
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
@ -95,10 +94,9 @@ def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset"
@router.get("/hardware")
async def get_hardware_utilization(current_subject: str = Depends(get_current_subject)):
"""
Get a live snapshot of GPU hardware utilization.
Live snapshot of GPU hardware utilization for the active backend.
Designed to be polled by the frontend during training.
Returns live GPU memory usage information for the active backend.
Polled by the frontend during training.
"""
from utils.hardware import get_gpu_utilization
return get_gpu_utilization()
@ -117,8 +115,8 @@ async def start_training(
"""
Start a training job.
This endpoint initiates training in the background and returns immediately.
Use the /status endpoint to check training progress.
Initiates training in the background and returns immediately. Use /status
to check progress.
"""
try:
logger.info(f"Starting training job with model: {request.model_name}")
@ -129,7 +127,7 @@ async def start_training(
backend = get_training_backend()
# Check if training is already active (before mutating any state)
# Check if training is already active (before mutating state).
if backend.is_training_active():
existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
return TrainingJobResponse(
@ -142,11 +140,11 @@ async def start_training(
error = "Training already active",
)
# Generate job ID — passed into start_training() which sets it on the
# backend only after confirming the old pump thread is dead.
# Job ID — passed to start_training(), which sets it on the backend only
# after confirming the old pump thread is dead.
job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}"
# Validate dataset paths if provided
# Validate dataset paths if provided.
if request.local_datasets:
request.local_datasets = _validate_local_dataset_paths(
request.local_datasets, "Local dataset"
@ -178,7 +176,7 @@ async def start_training(
)
request.resume_from_checkpoint = resume_checkpoint
# Convert request to kwargs for backend
# Convert request to backend kwargs.
training_kwargs = {
"model_name": request.model_name,
"training_type": request.training_type,
@ -242,8 +240,8 @@ async def start_training(
}
# Training page has no trust_remote_code toggle — the value comes from
# YAML model defaults applied when the user selects a model. As a safety
# net, consult the YAML directly so models that need it always get it.
# YAML model defaults on model select. As a safety net, consult the YAML
# directly so models that need it always get it.
if not training_kwargs["trust_remote_code"]:
model_defaults = load_model_defaults(request.model_name)
yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
@ -252,7 +250,7 @@ async def start_training(
training_kwargs["trust_remote_code"] = True
# Free GPU memory: shut down any running inference/export subprocesses
# before training starts (they'd compete for VRAM otherwise)
# before training (they'd compete for VRAM otherwise).
try:
from core.inference import get_inference_backend
inf_backend = get_inference_backend()
@ -279,7 +277,7 @@ async def start_training(
except Exception as e:
logger.warning("Could not shut down export subprocess: %s", e)
# start_training now spawns a subprocess (non-blocking)
# start_training spawns a subprocess (non-blocking).
success = backend.start_training(job_id = job_id, **training_kwargs)
if not success:
@ -334,7 +332,6 @@ async def stop_training(
status = "idle", message = "No training job is currently running"
)
# Call backend stop method
backend.stop_training(save = body.save)
return TrainingStopResponse(
@ -354,16 +351,14 @@ async def stop_training(
@router.post("/reset")
async def reset_training(current_subject: str = Depends(get_current_subject)):
"""
Reset training state so the user can return to configuration.
"""
"""Reset training state so the user can return to configuration."""
try:
backend = get_training_backend()
is_active = backend.is_training_active()
if is_active:
if backend._cancel_requested:
# Cancel (save=False) was requested — force-terminate so we can reset immediately
# Cancel (save=False) requested — force-terminate to reset immediately.
logger.info("Force-terminating subprocess for immediate reset (cancel path)")
backend.force_terminate()
else:
@ -460,7 +455,7 @@ async def get_training_status(current_subject: str = Depends(get_current_subject
if output_dir:
details["output_dir"] = output_dir
# Build metric history for chart recovery after SSE reconnection
# Metric history for chart recovery after SSE reconnection.
metric_history = None
if backend.step_history:
metric_history = {
@ -540,16 +535,15 @@ async def stream_training_progress(
request: Request, current_subject: str = Depends(get_current_subject)
):
"""
Stream training progress updates using Server-Sent Events (SSE).
Stream training progress via Server-Sent Events (SSE).
This endpoint provides real-time updates on training progress.
Supports reconnection via the SSE spec:
- Sends `id:` with each event so the browser tracks position.
- Sends `retry:` to control reconnection interval.
- Sends named `event:` types (progress, heartbeat, complete, error).
- Reads `Last-Event-ID` header on reconnect to replay missed steps.
Real-time progress with reconnection support per the SSE spec:
- `id:` per event so the browser tracks position.
- `retry:` to control reconnection interval.
- Named `event:` types (progress, heartbeat, complete, error).
- Reads `Last-Event-ID` on reconnect to replay missed steps.
"""
# Read Last-Event-ID header for reconnection resume
# Read Last-Event-ID header for reconnection resume.
last_event_id = request.headers.get("last-event-id")
resume_from_step: Optional[int] = None
if last_event_id is not None:
@ -580,7 +574,7 @@ async def stream_training_progress(
else:
progress_percent = float(step) / float(total) * 100.0 if total > 0 else 0.0
# Get actual values from progress object if available
# Pull values from the progress object if available.
elapsed_seconds = getattr(progress, "elapsed_seconds", None) if progress else None
eta_seconds = getattr(progress, "eta_seconds", None) if progress else None
grad_norm = grad_norm_override
@ -622,7 +616,7 @@ async def stream_training_progress(
return "\n".join(lines)
# ── Retry directive ──────────────────────────────────────
# Tell the browser to reconnect after 3 seconds if the connection drops
# Reconnect after 3 seconds if the connection drops.
yield "retry: 3000\n\n"
# ── Replay missed steps on reconnect ─────────────────────
@ -707,7 +701,7 @@ async def stream_training_progress(
# ── Live polling loop ────────────────────────────────────
last_step = resume_from_step if resume_from_step is not None else -1
no_update_count = 0
max_no_updates = 1800 # Timeout after 30 minutes (large models need time for compilation)
max_no_updates = 1800 # Timeout after 30 min (large models need compile time)
while backend.is_training_active():
try:
@ -721,7 +715,7 @@ async def stream_training_progress(
)
current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None
# Only send if step changed
# Only send if the step changed.
if current_step != last_step:
progress_payload = build_progress(
current_step,
@ -740,7 +734,7 @@ async def stream_training_progress(
no_update_count = 0
else:
no_update_count += 1
# Send heartbeat every 10 seconds
# Heartbeat every 10 seconds.
if no_update_count % 10 == 0:
heartbeat_payload = build_progress(
current_step,
@ -756,11 +750,11 @@ async def stream_training_progress(
event_id = current_step,
)
else:
# No steps yet, but training is active (model loading, etc.)
# No steps yet, but training is active (model loading, etc.).
no_update_count += 1
if no_update_count % 5 == 0:
# Pull total_steps and status from trainer so
# the frontend can show "Tokenizing…" etc.
# Pull total_steps + status from trainer so the frontend
# can show "Tokenizing…" etc.
tp_prep = getattr(
getattr(backend, "trainer", None),
"training_progress",

266
studio/backend/run.py
View file

@ -1,9 +1,9 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Run script for Unsloth UI Backend.
Works independently and can be moved to any directory.
"""Run script for Unsloth UI Backend.
Self-contained; can be moved to any directory.
"""
import os
@ -11,10 +11,10 @@ import sys
from pathlib import Path
from typing import Optional
# Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked)
# Suppress C-level dependency warnings globally (e.g. SwigPyPacked).
os.environ["PYTHONWARNINGS"] = "ignore"
# Add the backend directory to Python path early so local modules are importable
# Add the backend dir to sys.path early so local modules import.
backend_dir = Path(__file__).parent
if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
@ -27,8 +27,8 @@ except ValueError as exc:
configured = os.environ.get("UNSLOTH_CPU_THREADS")
raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}") from None
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# Anaconda/conda-forge Python: seed platform._sys_version_cache before
# imports that trigger attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
import _platform_compat # noqa: F401
@ -39,18 +39,17 @@ logger = get_logger(__name__)
def _resolve_external_ip() -> str:
"""
Resolve the machine's external IP address.
"""Resolve the machine's external IP address.
Tries (in order):
1. GCE metadata server (instant, works on Google Cloud VMs)
2. ifconfig.me (works anywhere with internet)
Tries, in order:
1. GCE metadata server (instant on Google Cloud VMs)
2. ifconfig.me (anywhere with internet)
3. LAN IP via UDP socket trick (fallback)
"""
import urllib.request
import socket
# 1. Try GCE metadata server (responds in <10ms on GCE, times out fast elsewhere)
# 1. GCE metadata server (<10ms on GCE, times out fast elsewhere).
try:
req = urllib.request.Request(
"http://metadata.google.internal/computeMetadata/v1/instance/network-interfaces/0/access-configs/0/external-ip",
@ -63,7 +62,7 @@ def _resolve_external_ip() -> str:
except Exception:
pass
# 2. Try public IP service
# 2. Public IP service.
try:
with urllib.request.urlopen("https://ifconfig.me", timeout = 3) as resp:
ip = resp.read().decode().strip()
@ -84,9 +83,10 @@ def _resolve_external_ip() -> str:
def _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> None:
"""Rewrite Uvicorn's startup log line: swap wildcard bind for the
externally-reachable address, replace the CTRL+C suffix with our Mac-aware
stop hint, and rename the prefix to "Unsloth Studio running on"."""
"""Rewrite Uvicorn's startup log line: swap the wildcard bind for the
externally-reachable address, replace the CTRL+C suffix with our
Mac-aware stop hint, and rename the prefix to "Unsloth Studio running
on"."""
import logging
import re
@ -137,7 +137,7 @@ def _local_port_open(
port: int,
timeout: float = 1.0,
) -> bool:
"""Return True iff a TCP connection to (host, port) succeeds within timeout."""
"""True iff a TCP connection to (host, port) succeeds within timeout."""
import socket
try:
with socket.create_connection((host, port), timeout = timeout):
@ -147,8 +147,8 @@ def _local_port_open(
def _working_local_url(port: int) -> "str | None":
"""Return a working loopback URL on this machine, or None if neither
127.0.0.1 nor ::1 responds. Used as a fallback when external reachability fails."""
"""A working loopback URL on this machine, or None if neither
127.0.0.1 nor ::1 responds. Fallback when external reachability fails."""
if _local_port_open("127.0.0.1", port):
return f"http://127.0.0.1:{port}"
if _local_port_open("::1", port):
@ -157,13 +157,12 @@ def _working_local_url(port: int) -> "str | None":
def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None":
"""Return the IPv4 loopback URL when localhost will not reach 127.0.0.1.
"""Return the IPv4 loopback URL when localhost won't reach 127.0.0.1.
Local Studio intentionally binds to 127.0.0.1. On hosts where localhost
resolves to IPv6 only (::1), a browser pointed at http://localhost:<port>
fails -- or worse, reaches a different process listening on ::1 -- even
though http://127.0.0.1:<port> works. Return the IPv4 URL so the caller can
tell the user which address to open.
Local Studio binds to 127.0.0.1. Where localhost resolves to IPv6
only (::1), http://localhost:<port> fails -- or worse, hits a
different process on ::1 -- even though http://127.0.0.1:<port> works.
Return the IPv4 URL so the caller can tell the user what to open.
"""
import socket
@ -172,7 +171,7 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None":
ipv4_url = f"http://127.0.0.1:{port}"
# Only warn once Studio is confirmed answering on the IPv4 loopback.
# Only warn once Studio is confirmed answering on IPv4 loopback.
if _working_local_url(port) != ipv4_url:
return None
@ -194,11 +193,11 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None":
if host == "::1":
has_ipv6_loopback = True
# A successful connection to ::1 is NOT evidence that Studio is reachable
# there: Studio binds 127.0.0.1 only, so anything answering on ::1 is a
# different process -- which is exactly when the user must be steered to
# 127.0.0.1. Dual-stack localhost is fine (browsers fall back to 127.0.0.1
# when ::1 refuses), so only the IPv6-only case strands the user.
# A connection to ::1 is NOT evidence Studio is reachable there:
# Studio binds 127.0.0.1 only, so anything on ::1 is a different
# process -- exactly when to steer the user to 127.0.0.1. Dual-stack
# localhost is fine (browsers fall back to 127.0.0.1 when ::1
# refuses), so only the IPv6-only case strands the user.
if has_ipv6_loopback and not has_ipv4_loopback:
return ipv4_url
return None
@ -231,11 +230,11 @@ def _print_localhost_ipv6_mismatch_warning(local_url: str, port: int) -> None:
def _verify_global_reachability(display_host: str, port: int) -> None:
"""Probe check-host.net to confirm display_host:port is reachable from the
public internet. Synchronous so the caller can render output between the
banner URL section and the trailing stop hint. Bounded at ~15s; failures
are swallowed (the verifier failing is not Studio failing). Only meaningful
when bound to a wildcard host."""
"""Probe check-host.net to confirm display_host:port is reachable
from the public internet. Synchronous so the caller can render output
between the banner URL section and the trailing stop hint. Bounded at
~15s; failures are swallowed (verifier failing != Studio failing).
Only meaningful when bound to a wildcard host."""
import ipaddress
import json
import time
@ -256,7 +255,7 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
url = f"http://{display_host}:{port}"
# Private / loopback / link-local addresses are not globally routable.
# Private/loopback/link-local addresses aren't globally routable.
try:
addr = ipaddress.ip_address(display_host)
if addr.is_loopback or addr.is_private or addr.is_link_local:
@ -304,7 +303,7 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
continue
if results and all(v is not None for v in results.values()):
break
# Two decisive nodes is enough; stop polling early.
# Two decisive nodes is enough; stop early.
decisive = [
v
for v in results.values()
@ -373,7 +372,7 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
f"{dim} then open http://localhost:{port}/ in your browser.{reset}",
flush = True,
)
# Only offer the local URL if loopback actually answers.
# Only offer the local URL if loopback answers.
local_url = _working_local_url(port)
if local_url:
print(
@ -388,7 +387,7 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
flush = True,
)
except urllib.error.URLError:
# Outbound HTTPS blocked; skip silently.
# Outbound HTTPS blocked; skip.
pass
except Exception:
pass
@ -397,16 +396,16 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
def _emit_startup_output(host: str, port: int, display_host: str) -> None:
"""Print the access banner plus any post-startup warnings.
Extracted from ``_run`` so the banner/warning wiring is unit-testable. The
``localhost``-to-::1 mismatch warning and the wildcard reachability check
are mutually exclusive (the mismatch helper returns None for any non
127.0.0.1 bind, and wildcard binds are never 127.0.0.1), so the trailing
stop hint is emitted exactly once.
Extracted from ``_run`` so the banner/warning wiring is testable. The
``localhost``-to-::1 mismatch warning and the wildcard reachability
check are mutually exclusive (the mismatch helper returns None for any
non-127.0.0.1 bind, and wildcard binds are never 127.0.0.1), so the
trailing stop hint is emitted exactly once.
"""
wildcard_bind = host in ("0.0.0.0", "::")
localhost_mismatch_url = _localhost_ipv6_mismatch_url(host, port)
# For wildcard binds, run the reachability check between the URL section
# and the stop hint so the stop hint stays last on screen.
# For wildcard binds, run the reachability check between the URL
# section and the stop hint so the stop hint stays last.
print_studio_access_banner(
port = port,
bind_host = host,
@ -422,12 +421,11 @@ def _emit_startup_output(host: str, port: int, display_host: str) -> None:
def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
"""Return (pid, process_name) of the process listening on *port*, or None.
"""Return (pid, process_name) listening on *port*, or None.
Uses psutil when available. Falls back gracefully to None so callers
can still report the port conflict without process details.
Works on Windows, macOS, and Linux wherever psutil is installed.
Uses psutil when available, falling back to None so callers can still
report the conflict without process details. Works on Windows, macOS,
and Linux wherever psutil is installed.
"""
try:
import psutil
@ -444,7 +442,7 @@ def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
except (psutil.NoSuchProcess, psutil.AccessDenied):
return (conn.pid, "<unknown>")
except (psutil.AccessDenied, OSError) as e:
# psutil.net_connections() needs elevated privileges on some platforms
# net_connections() needs elevated privileges on some platforms.
logger.debug("Failed to scan network connections for port %s: %s", port, e)
return None
@ -452,19 +450,17 @@ def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
def _is_port_free(host: str, port: int) -> bool:
"""Check if a port is available for binding.
When *host* is ``0.0.0.0`` (wildcard), we also check whether anything
is already listening on ``127.0.0.1`` (and ``::1`` when IPv6 is
available). An SSH tunnel or similar process may hold the loopback
address while our wildcard bind still succeeds, making Unsloth Studio
unreachable via ``localhost``.
When *host* is ``0.0.0.0`` (wildcard), also check whether anything is
already listening on ``127.0.0.1`` (and ``::1`` when IPv6 exists). An
SSH tunnel may hold the loopback address while our wildcard bind
succeeds, making Studio unreachable via ``localhost``.
Works on Windows, macOS, and Linux.
"""
import socket
# 1. Can we bind to the requested address?
# Use getaddrinfo so both IPv4 ("0.0.0.0") and IPv6 ("::") hosts
# resolve to the correct address family automatically.
# 1. Can we bind to the requested address? getaddrinfo resolves both
# IPv4 ("0.0.0.0") and IPv6 ("::") to the right address family.
try:
addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
family, socktype, proto, _, sockaddr = addr_info[0]
@ -474,9 +470,9 @@ def _is_port_free(host: str, port: int) -> bool:
except OSError:
return False
# 2. When binding to all interfaces, verify that localhost is not
# already claimed by another process (e.g. an SSH -L tunnel).
# We attempt a TCP connect -- if it succeeds something is listening.
# 2. When binding to all interfaces, verify localhost isn't already
# claimed by another process (e.g. an SSH -L tunnel). A successful
# TCP connect means something is listening.
if host in ("0.0.0.0", "::"):
for loopback, family in [
("127.0.0.1", socket.AF_INET),
@ -486,10 +482,10 @@ def _is_port_free(host: str, port: int) -> bool:
with socket.socket(family, socket.SOCK_STREAM) as s:
s.settimeout(1)
if s.connect_ex((loopback, port)) == 0:
# Connection succeeded -- port is taken on loopback
# Port is taken on loopback.
return False
except OSError:
# IPv6 disabled or other OS-level restriction -- skip
# IPv6 disabled or other OS-level restriction -- skip.
continue
return True
@ -500,7 +496,7 @@ def _find_free_port(
start: int,
max_attempts: int = 20,
) -> int:
"""Find a free port starting from `start`, trying up to max_attempts ports."""
"""Find a free port from `start`, trying up to max_attempts ports."""
for offset in range(max_attempts):
candidate = start + offset
if _is_port_free(host, candidate):
@ -514,7 +510,7 @@ _PID_FILE = _studio_root() / "studio.pid"
# Direct backend launches bypass the CLI's env re-export; do it here for
# real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR
# picks up the custom build. Skip for legacy-default to avoid flipping
# picks up the custom build. Skip legacy-default to avoid flipping
# default-mode installs into env-override.
try:
_LEGACY_STUDIO_ROOT = (Path.home() / ".unsloth" / "studio").resolve()
@ -552,20 +548,20 @@ def _remove_pid_file():
def _graceful_shutdown(server = None):
"""Explicitly shut down all subprocess backends and the uvicorn server.
"""Shut down all subprocess backends and the uvicorn server.
Called from signal handlers to ensure child processes are cleaned up
before the parent exits. This is critical on Windows where atexit
handlers are unreliable after Ctrl+C.
Called from signal handlers to clean up child processes before the
parent exits. Critical on Windows where atexit handlers are
unreliable after Ctrl+C.
"""
_remove_pid_file()
logger.info("Graceful shutdown initiated — cleaning up subprocesses...")
# 1. Shut down uvicorn server (releases the listening socket)
# 1. Shut down uvicorn (releases the listening socket).
if server is not None:
server.should_exit = True
# 2. Clean up inference subprocess (if instantiated)
# 2. Clean up inference subprocess (if instantiated).
try:
from core.inference.orchestrator import _inference_backend
if _inference_backend is not None:
@ -573,7 +569,7 @@ def _graceful_shutdown(server = None):
except Exception as e:
logger.warning("Error shutting down inference subprocess: %s", e)
# 3. Clean up export subprocess (if instantiated)
# 3. Clean up export subprocess (if instantiated).
try:
from core.export.orchestrator import _export_backend
if _export_backend is not None:
@ -581,7 +577,7 @@ def _graceful_shutdown(server = None):
except Exception as e:
logger.warning("Error shutting down export subprocess: %s", e)
# 4. Clean up training subprocess (if active)
# 4. Clean up training subprocess (if active).
try:
from core.training.training import _training_backend
if _training_backend is not None:
@ -589,7 +585,7 @@ def _graceful_shutdown(server = None):
except Exception as e:
logger.warning("Error shutting down training subprocess: %s", e)
# 5. Kill llama-server subprocess (if loaded)
# 5. Kill llama-server subprocess (if loaded).
try:
from routes.inference import _llama_cpp_backend
if _llama_cpp_backend is not None:
@ -601,10 +597,10 @@ def _graceful_shutdown(server = None):
# The uvicorn server instance -- set by run_server(), used by callers
# that need to tell the server to exit (e.g. signal handlers).
# that tell the server to exit (e.g. signal handlers).
_server = None
# Shutdown event -- used to wake the main loop on signal
# Shutdown event -- wakes the main loop on signal.
_shutdown_event = None
@ -615,8 +611,8 @@ def _iter_frontend_fallback_candidates() -> "list[Path]":
"""Yield `studio/frontend/dist` paths to try when the default is missing.
Covers PATH-shadowed binaries whose __file__ resolves into a
site-packages tree that never received a vite build (e.g. plain
`pip install unsloth` from PyPI).
site-packages tree with no vite build (e.g. plain `pip install
unsloth` from PyPI).
"""
import ast
import re
@ -642,9 +638,9 @@ def _iter_frontend_fallback_candidates() -> "list[Path]":
src = finder.read_text(encoding = "utf-8")
except OSError:
continue
# Tolerate single- or multi-line dict literals; [^}]* still
# rejects nested dicts, which the setuptools template never
# emits for editable installs.
# Tolerate single- or multi-line dict literals; [^}]*
# still rejects nested dicts, which the setuptools
# template never emits for editable installs.
m = re.search(r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S)
if not m:
continue
@ -652,8 +648,8 @@ def _iter_frontend_fallback_candidates() -> "list[Path]":
mapping = ast.literal_eval(m.group(1))
except (SyntaxError, ValueError):
continue
# Defensive: literal_eval can return a set / list / None if the
# matched literal is not a dict (regex captures `{...}`).
# Defensive: literal_eval can return a set/list/None if
# the matched `{...}` literal isn't a dict.
if not isinstance(mapping, dict):
continue
studio_pkg = mapping.get("studio")
@ -663,10 +659,10 @@ def _iter_frontend_fallback_candidates() -> "list[Path]":
def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Path]]:
"""Pick a frontend dir that actually contains `index.html`.
"""Pick a frontend dir that contains `index.html`.
Returns (chosen, attempted). `chosen` is None if nothing servable was
found; `attempted` is the full ordered list for diagnostics.
found; `attempted` is the ordered list for diagnostics.
"""
attempted: list[Path] = []
seen: set[Path] = set()
@ -706,25 +702,25 @@ def run_server(
port: Port to bind to (auto-increments if in use)
frontend_path: Path to frontend build directory (optional)
silent: Suppress startup messages
api_only: Run API server only, no frontend serving (for Tauri desktop app)
llama_parallel_slots: Number of parallel slots for llama-server
api_only: API server only, no frontend (for Tauri desktop app)
llama_parallel_slots: parallel slots for llama-server
Note:
Signal handlers are NOT registered here so that embedders
(e.g. Colab notebooks) keep their own interrupt semantics.
Standalone callers should register handlers after calling this.
Signal handlers are NOT registered here so embedders (e.g. Colab
notebooks) keep their own interrupt semantics. Standalone callers
should register handlers after calling this.
"""
global _server, _shutdown_event
# On Windows the default console encoding (cp1252) cannot encode emoji.
# Reconfigure stdout to UTF-8 so startup messages do not crash the server.
# Windows console encoding (cp1252) can't encode emoji. Reconfigure
# stdout to UTF-8 so startup messages don't crash the server.
if sys.platform == "win32" and hasattr(sys.stdout, "reconfigure"):
try:
sys.stdout.reconfigure(encoding = "utf-8", errors = "replace")
except Exception:
pass
# Set env var BEFORE importing main so CORS middleware picks it up
# Set env var BEFORE importing main so CORS middleware picks it up.
if api_only:
os.environ["UNSLOTH_API_ONLY"] = "1"
@ -739,10 +735,10 @@ def run_server(
from main import app, setup_frontend, _IS_COLAB
from utils.paths import ensure_studio_directories
# Create all standard directories on startup
# Create all standard directories on startup.
ensure_studio_directories()
# Auto-find free port if requested port is in use
# Auto-find a free port if the requested one is in use.
if not _is_port_free(host, port):
original_port = port
blocker = _get_pid_on_port(port)
@ -760,14 +756,14 @@ def run_server(
print("=" * 50)
print("")
# Setup frontend if path provided (skip in api-only mode).
# Falls back through alternate locations if the default lacks a built
# dist; errors out loudly rather than silently serving 404 on `/`.
# Setup frontend if path provided (skip in api-only mode). Falls back
# through alternate locations if the default lacks a built dist;
# errors loudly rather than serving 404 on `/`.
if frontend_path and not api_only:
chosen, attempted = _resolve_frontend_path(Path(frontend_path))
if chosen is not None and setup_frontend(app, chosen):
if not silent:
# Resolve so logs always show an absolute path for support.
# Resolve so logs show an absolute path for support.
try:
display = chosen.resolve()
except OSError:
@ -779,8 +775,8 @@ def run_server(
or os.environ.get("STUDIO_HOME")
or str(Path.home() / ".unsloth" / "studio")
)
# Windows ships the user-facing shim at $STUDIO_HOME/bin/unsloth.exe
# (a hardlink to the venv exe); Linux/macOS use the venv binary
# Windows ships the shim at $STUDIO_HOME/bin/unsloth.exe (a
# hardlink to the venv exe); Linux/macOS use the venv binary
# at $STUDIO_HOME/unsloth_studio/bin/unsloth.
home = Path(home_str).expanduser()
if sys.platform == "win32":
@ -803,7 +799,7 @@ def run_server(
" - reinstall: curl -fsSL https://unsloth.ai/install.sh | sh"
)
# Resolve once; shared by the log rewrite and the banner.
# Resolve once; shared by the log rewrite and banner.
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
_install_uvicorn_startup_log_rewrite(host, display_host)
@ -825,11 +821,11 @@ def run_server(
access_log = False,
server_header = False,
)
# Only in Colab: trust X-Forwarded-* from Colab's reverse proxy so the app
# sees the real https origin. forwarded_allow_ips="*" is fine inside Colab's
# single-user sandbox, but would be an unwanted security relaxation for a
# normal local/standalone Studio, so leave uvicorn's safe defaults
# (forwarded headers trusted from loopback only) in place there.
# Colab only: trust X-Forwarded-* from Colab's reverse proxy so the
# app sees the real https origin. forwarded_allow_ips="*" is fine in
# Colab's single-user sandbox but an unwanted relaxation for a normal
# local/standalone Studio, so leave uvicorn's safe defaults
# (forwarded headers trusted from loopback only) elsewhere.
if _IS_COLAB:
config_kwargs["proxy_headers"] = True
config_kwargs["forwarded_allow_ips"] = "*"
@ -837,17 +833,16 @@ def run_server(
_server = _ReadyServer(config)
_shutdown_event = Event()
# Expose the actual bound port so request-handling code can build
# loopback URLs that point at the real backend, not whatever port a
# reverse proxy or tunnel exposed in the request URL. Only publish
# an explicit value when we know the concrete port; for ephemeral
# binds (port==0) leave it unset and let request handlers fall back
# to the ASGI request scope or request.base_url.
# Expose the actual bound port so request handlers build loopback
# URLs pointing at the real backend, not whatever port a proxy/tunnel
# exposed in the request URL. Only publish a concrete port; for
# ephemeral binds (port==0) leave it unset so handlers fall back to
# the ASGI request scope or request.base_url.
app.state.server_port = port if port and port > 0 else None
app.state.llama_parallel_slots = llama_parallel_slots
# Expose a shutdown callable via app.state before the server can accept
# requests so /api/shutdown is available as soon as readiness is published.
# Expose a shutdown callable via app.state before the server accepts
# requests so /api/shutdown is ready as soon as readiness publishes.
def _trigger_shutdown():
_graceful_shutdown(_server)
if _shutdown_event is not None:
@ -855,11 +850,10 @@ def run_server(
app.state.trigger_shutdown = _trigger_shutdown
# Run server in a daemon thread.
# Use an explicit new_event_loop() + run_until_complete() instead of
# asyncio.run() to avoid nest_asyncio's global patches to asyncio.run
# interfering when called from a thread while Colab/IPython already has
# a running loop on the main thread.
# Run server in a daemon thread. Use explicit new_event_loop() +
# run_until_complete() rather than asyncio.run() so nest_asyncio's
# global patches to asyncio.run don't interfere when called from a
# thread while Colab/IPython already runs a loop on the main thread.
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
@ -876,9 +870,9 @@ def run_server(
thread = Thread(target = _run, daemon = True)
thread.start()
# Wait until uvicorn has completed lifespan startup and bound sockets, or
# until the server exits/fails before startup. This intentionally has no
# correctness deadline: a slow but live startup should remain in progress.
# Wait until uvicorn finishes lifespan startup and binds sockets, or
# until the server exits/fails before startup. No correctness
# deadline: a slow but live startup should remain in progress.
try:
while not ready_event.is_set():
if startup_failed.is_set() or not thread.is_alive():
@ -898,8 +892,8 @@ def run_server(
atexit.register(_remove_pid_file)
# Output port for Tauri to parse when in api-only mode. Emit only after
# uvicorn sockets are bound and FastAPI lifespan/startup has completed.
# Output port for Tauri to parse in api-only mode. Emit only after
# uvicorn sockets are bound and FastAPI startup completed.
if api_only:
print(f"TAURI_PORT={port}", flush = True)
@ -909,13 +903,13 @@ def run_server(
return app
# For direct execution (also invoked by CLI via os.execvp / subprocess)
# For direct execution (also invoked by CLI via os.execvp / subprocess).
if __name__ == "__main__":
import argparse
import signal
import traceback
# Ensure stderr can handle Unicode on Windows (tracebacks with non-ASCII paths)
# Ensure stderr handles Unicode on Windows (non-ASCII path tracebacks).
if sys.platform == "win32" and hasattr(sys.stderr, "reconfigure"):
try:
sys.stderr.reconfigure(encoding = "utf-8", errors = "replace")
@ -942,7 +936,7 @@ if __name__ == "__main__":
help = "API server only, no frontend (for Tauri)",
)
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1
# applies only to direct backend launches; `unsloth studio run`
# applies to direct backend launches only; `unsloth studio run`
# always passes its own value (4) explicitly.
_PARALLEL_MIN = 1
_PARALLEL_MAX = 64
@ -985,7 +979,7 @@ if __name__ == "__main__":
sys.stderr.flush()
sys.exit(1)
# Signal handler -- ensures subprocess cleanup on Ctrl+C
# Signal handler -- ensures subprocess cleanup on Ctrl+C.
def _signal_handler(signum, frame):
_graceful_shutdown(_server)
_shutdown_event.set()
@ -993,13 +987,13 @@ if __name__ == "__main__":
signal.signal(signal.SIGINT, _signal_handler)
signal.signal(signal.SIGTERM, _signal_handler)
# On Windows, some terminals send SIGBREAK for Ctrl+C / Ctrl+Break
# On Windows, some terminals send SIGBREAK for Ctrl+C / Ctrl+Break.
if hasattr(signal, "SIGBREAK"):
signal.signal(signal.SIGBREAK, _signal_handler)
# Keep running until shutdown signal.
# NOTE: Event.wait() without a timeout blocks at the C level on Linux,
# which prevents Python from delivering SIGINT (Ctrl+C). Using a
# short timeout in a loop lets the interpreter process pending signals.
# preventing Python from delivering SIGINT (Ctrl+C). A short timeout
# in a loop lets the interpreter process pending signals.
while not _shutdown_event.is_set():
_shutdown_event.wait(timeout = 1)

10
studio/backend/startup_banner.py
View file

@ -3,7 +3,7 @@
"""Terminal banner for Studio startup.
Stdlib only — safe to import without the rest of the backend (no structlog/uvicorn).
Stdlib only — safe to import without the rest of the backend.
"""
from __future__ import annotations
@ -34,7 +34,7 @@ def print_port_in_use_notice(original_port: int, new_port: int) -> None:
def print_studio_stop_hint() -> None:
"""Print the trailing stop hint + closing divider. Separate from the main
"""Print the trailing stop hint + closing divider. Separate from the
banner so callers can interleave content (e.g. a reachability check)."""
use_color = stdout_supports_color()
dim = "\033[38;5;245m"
@ -67,7 +67,7 @@ def print_studio_access_banner(
display_host: str,
include_stop_hint: bool = True,
) -> None:
"""Pretty-print URLs after the server is listening. Set
"""Pretty-print URLs once the server is listening. Set
``include_stop_hint=False`` to omit the trailing stop block; pair with
:func:`print_studio_stop_hint` after inserting your own content."""
use_color = stdout_supports_color()
@ -96,8 +96,8 @@ def print_studio_access_banner(
listen_all = bind_host in ("0.0.0.0", "::")
loopback_bind = bind_host in ("127.0.0.1", "localhost", "::1")
# Use loopback URL only when the server is reachable on loopback;
# otherwise show the actual bound address.
# Use the loopback URL only when reachable on loopback; otherwise show
# the actual bound address.
primary_url = loopback_url if listen_all or loopback_bind else external_url
tip_url = alt_local if listen_all or loopback_bind else external_url
api_base = primary_url

2
studio/backend/storage/mcp_servers_db.py
View file

@ -28,7 +28,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
)
"""
)
# use_oauth was added after the first release; backfill for pre-existing DBs.
# use_oauth added after the first release; backfill for pre-existing DBs.
cols = {r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall()}
if "use_oauth" not in cols:
conn.execute("ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0")

6
studio/backend/storage/providers_db.py
View file

@ -4,8 +4,8 @@
"""
SQLite storage for external LLM provider configurations.
Follows the same pattern as studio_db.py — module-level functions,
raw sqlite3, WAL mode, per-function connections.
Same pattern as studio_db.py: module-level functions, raw sqlite3, WAL
mode, per-function connections.
NOTE: API keys are NOT stored here. They live only in the browser
(localStorage) and are sent encrypted per-request.
@ -26,7 +26,7 @@ _schema_ready = False
def _ensure_schema(conn: sqlite3.Connection) -> None:
"""Create the llm_providers table if it doesn't exist. Called once per process."""
"""Create the llm_providers table if absent. Called once per process."""
conn.execute("PRAGMA journal_mode=WAL")
conn.execute(
"""

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

@ -4,8 +4,8 @@
"""
SQLite storage for training run history and metrics.
Follows the same pattern as auth/storage.py — module-level functions,
raw sqlite3, per-function connections. Enhancements over auth:
Same pattern as auth/storage.py — module-level functions, raw sqlite3,
per-function connections. Enhancements over auth:
- WAL mode for concurrent read/write access
- PRAGMA foreign_keys = ON for CASCADE deletes
"""
@ -34,8 +34,8 @@ def _denied_path_prefixes() -> list[str]:
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.
# realpath() resolves /etc -> /private/etc, /tmp -> /private/tmp on
# macOS, so include the /private variants to avoid bypasses.
return [
"/System",
"/Library",
@ -174,9 +174,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
"""
)
conn.execute("CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)")
# Use COLLATE NOCASE on Windows so C:\Models and c:\models dedup via the
# UNIQUE constraint. On Linux/macOS (case-sensitive FS) keep the default
# BINARY collation so /Models and /models remain distinct.
# COLLATE NOCASE on Windows so C:\Models and c:\models dedup via the
# UNIQUE constraint. On Linux/macOS (case-sensitive FS) keep the default
# BINARY collation so /Models and /models stay distinct.
collation = "COLLATE NOCASE" if platform.system() == "Windows" else ""
conn.execute(
f"""
@ -287,15 +287,14 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
)
"""
)
# Server-side import ledger so a studio.db wipe correctly re-triggers
# the legacy Dexie import. The previous boolean localStorage sentinel
# (`unsloth_chat_legacy_imported_to_studio_db`) is non-recoverable:
# if studio.db is recreated while the browser keeps the flag, legacy
# Dexie threads are silently hidden from the sidebar. The ledger
# lives inside studio.db so it disappears together with the data it
# is supposed to track, which is the recovery the boolean lacked.
# Keyed by legacy thread id; per-thread is sufficient because Dexie
# is read-only after this PR (a thread's message set does not grow).
# Server-side import ledger so a studio.db wipe re-triggers the legacy
# Dexie import. The old boolean localStorage sentinel
# (`unsloth_chat_legacy_imported_to_studio_db`) was non-recoverable: if
# studio.db is recreated while the browser keeps the flag, legacy Dexie
# threads are silently hidden from the sidebar. The ledger lives inside
# studio.db so it disappears with the data it tracks -- the recovery the
# boolean lacked. Keyed by legacy thread id; per-thread suffices because
# Dexie is read-only after this PR (a thread's message set doesn't grow).
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_legacy_imports (
@ -313,7 +312,7 @@ def get_connection() -> sqlite3.Connection:
ensure_dir(db_path.parent)
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
# foreign_keys is session-scoped, must be set per connection
# foreign_keys is session-scoped; set per connection
conn.execute("PRAGMA foreign_keys=ON")
if not _schema_ready:
with _schema_lock:
@ -701,9 +700,9 @@ def add_scan_folder(path: str) -> dict:
if not os.access(normalized, os.R_OK | os.X_OK):
raise ValueError("Path is not readable")
# On Windows, use normcase for denylist comparison but store the
# original-cased path so downstream consumers see the native
# drive-letter casing the user expects (e.g. C:\Models, not c:\models).
# On Windows, normcase for the denylist comparison but store the
# original-cased path so downstream consumers see the native drive-letter
# casing the user expects (e.g. C:\Models, not c:\models).
is_win = platform.system() == "Windows"
check = os.path.normcase(normalized) if is_win else normalized
for prefix in _denied_path_prefixes():
@ -713,8 +712,8 @@ def add_scan_folder(path: str) -> dict:
conn = get_connection()
try:
now = datetime.now(timezone.utc).isoformat()
# On Windows, use case-insensitive lookup so C:\Models and c:\models
# dedup correctly while preserving the originally-stored casing.
# On Windows, case-insensitive lookup so C:\Models and c:\models dedup
# while preserving the originally-stored casing.
if is_win:
existing = conn.execute(
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE",
@ -735,8 +734,8 @@ def add_scan_folder(path: str) -> dict:
conn.commit()
except sqlite3.IntegrityError:
pass # duplicate -- fall through to SELECT
# Use the same collation as the pre-check so we find the row even
# when a concurrent writer stored it with different casing (Windows).
# Same collation as the pre-check so we find the row even when a
# concurrent writer stored it with different casing (Windows).
fallback_sql = (
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE"
if is_win
@ -1408,8 +1407,8 @@ def _deep_merge_settings(current: dict[str, Any], updates: dict[str, Any]) -> di
def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]:
"""Atomic read-merge-write under BEGIN IMMEDIATE so two concurrent writers
cannot drop one another's updates."""
"""Atomic read-merge-write under BEGIN IMMEDIATE so concurrent writers
cannot drop each other's updates."""
if not updates:
return list_chat_settings()
conn = get_connection()
@ -1457,8 +1456,8 @@ def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]:
def list_chat_legacy_imports() -> list[str]:
"""Return the legacy_thread_id of every thread already imported.
Cheap: scans a single small PK-only table. The frontend stuffs the
result into a Set before walking Dexie, so the diff is O(|Dexie|).
Cheap: scans one small PK-only table. The frontend puts the result in a
Set before walking Dexie, so the diff is O(|Dexie|).
"""
conn = get_connection()
try:
@ -1472,13 +1471,13 @@ def upsert_chat_legacy_imports(legacy_thread_ids: list[str]) -> tuple[int, int]:
"""Mark each given legacy thread id as imported. Idempotent.
Returns (accepted, inserted):
- accepted: number of non-empty deduped input ids
- inserted: number of rows that were actually new (not already in ledger)
- accepted: count of non-empty deduped input ids
- inserted: count of rows that were actually new (not already in ledger)
ON CONFLICT DO NOTHING keeps the existing imported_at when an id is
recorded twice. INSERT...RETURNING reports only the rows that were
actually inserted, so callers can distinguish first-time imports
from idempotent re-runs without an extra SELECT.
recorded twice. INSERT...RETURNING reports only newly-inserted rows, so
callers distinguish first-time imports from idempotent re-runs without an
extra SELECT.
"""
ids = list(dict.fromkeys(tid for tid in legacy_thread_ids if tid))
if not ids:

42
studio/backend/tests/conftest.py
View file

@ -7,19 +7,19 @@ Shared pytest configuration for the backend test suite.
Responsibilities:
1. Put the backend root on sys.path so `from models.inference import ...`
(and similar flat imports) resolve in test modules — mirrors how the
app itself is launched.
app is launched.
2. Provide a hybrid ``studio_server`` session fixture for end-to-end tests
(see ``test_studio_api.py``). The fixture supports two invocation modes:
(see ``test_studio_api.py``), supporting two invocation modes:
a. **External server.** If ``UNSLOTH_E2E_BASE_URL`` is set, tests point
at an already-running Studio instance. ``UNSLOTH_E2E_API_KEY`` must
also be set. This is the fast-iteration mode: start the server once
with ``unsloth studio run ...``, then run pytest against it many
times with no per-run GGUF load cost.
at an already-running Studio instance (``UNSLOTH_E2E_API_KEY`` must
also be set). Fast-iteration mode: start the server once with
``unsloth studio run ...``, then run pytest against it many times with
no per-run GGUF load cost.
b. **Fixture-managed server.** Otherwise, the fixture launches a fresh
server via ``_start_server`` and tears it down at session end. This
is the one-shot mode for CI or a clean-slate verification run.
b. **Fixture-managed server.** Otherwise the fixture launches a fresh
server via ``_start_server`` and tears it down at session end. One-shot
mode for CI or a clean-slate verification run.
The model / variant for mode (b) come from ``--unsloth-model`` /
``--unsloth-gguf-variant`` pytest options, then ``UNSLOTH_E2E_MODEL`` /
@ -80,16 +80,16 @@ def studio_server(request):
Resolution order:
1. If ``UNSLOTH_E2E_BASE_URL`` is set → point at that server,
require ``UNSLOTH_E2E_API_KEY`` alongside (skip if missing).
2. Otherwise → start a fresh ``unsloth studio run`` subprocess via
the existing ``_start_server`` helper in ``test_studio_api.py``
and tear it down on session teardown.
1. If ``UNSLOTH_E2E_BASE_URL`` is set → point at that server, requiring
``UNSLOTH_E2E_API_KEY`` alongside (skip if missing).
2. Otherwise → start a fresh ``unsloth studio run`` subprocess via the
``_start_server`` helper in ``test_studio_api.py`` and tear it down on
session teardown.
Session-scoped so the expensive GGUF load happens at most once per
pytest invocation. Lazily instantiated — tests that don't request
the fixture (e.g. the unit tests in ``test_anthropic_messages.py``
or ``test_help_output``) do not trigger server startup.
Session-scoped so the expensive GGUF load happens at most once per pytest
invocation. Lazy — tests that don't request the fixture (e.g. the unit
tests in ``test_anthropic_messages.py`` or ``test_help_output``) don't
trigger server startup.
"""
external_url = os.environ.get("UNSLOTH_E2E_BASE_URL")
if external_url:
@ -103,9 +103,9 @@ def studio_server(request):
yield external_url, api_key
return
# Lazy import: pytest has already loaded test_studio_api into
# sys.modules by the time any test requests this fixture, so this
# is a cache hit, not a re-execution.
# Lazy import: pytest has already loaded test_studio_api into sys.modules
# by the time any test requests this fixture, so this is a cache hit, not
# a re-execution.
import test_studio_api as _e2e
model = (

32
studio/backend/tests/test_anthropic_cache_ttl.py
View file

@ -1,18 +1,17 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Unit tests for the prompt_cache_ttl threading on the Anthropic path.
"""Unit tests for prompt_cache_ttl threading on the Anthropic path.
Anthropic accepts an optional ``ttl`` on each ``cache_control`` marker:
the default is the 5-minute ephemeral pool; ``ttl:"1h"`` writes into
the 1-hour pool instead. The 1h pool is the right pick when
conversations span multiple short bursts more than 5 minutes apart --
1h writes are billed at 2x base input vs 1.25x for 5m, but reads stay
at 0.1x for both, so one extra read pays off the premium.
default is the 5-minute ephemeral pool; ``ttl:"1h"`` writes into the
1-hour pool. The 1h pool fits conversations spanning short bursts more
than 5 minutes apart -- 1h writes bill at 2x base input vs 1.25x for 5m,
but reads stay at 0.1x for both, so one extra read pays the premium.
These tests pin the outbound body shape: when prompt_cache_ttl="1h"
both cache_control markers carry ``ttl:"1h"``; default omits the field
entirely so the 5m pool is used; garbage values are silently dropped.
These tests pin the outbound body shape: prompt_cache_ttl="1h" puts
``ttl:"1h"`` on both markers; default omits the field (5m pool); garbage
values are silently dropped.
"""
import asyncio
@ -128,13 +127,12 @@ def test_1h_ttl_writes_into_1h_pool(monkeypatch):
def test_1h_ttl_does_not_send_extended_cache_ttl_beta_header(monkeypatch):
# The `extended-cache-ttl-2025-04-11` beta header that originally
# gated 1h cache TTL has been promoted to GA: verified live against
# api.anthropic.com on 2026-05-22 -- a request with
# `cache_control:{type:"ephemeral", ttl:"1h"}` and NO beta header
# returns 200 and populates `ephemeral_1h_input_tokens`. Pin the
# contract so we don't reintroduce the gate by accident; a future
# regression that re-adds the header would surface here.
# The `extended-cache-ttl-2025-04-11` beta header that originally gated
# 1h cache TTL is now GA: verified live against api.anthropic.com on
# 2026-05-22 -- a request with `cache_control:{type:"ephemeral",
# ttl:"1h"}` and NO beta header returns 200 and populates
# `ephemeral_1h_input_tokens`. Pin the contract so a regression that
# re-adds the header surfaces here.
captured = _capture(monkeypatch, ttl = "1h")
beta = captured["headers"].get("anthropic-beta", "")
assert "extended-cache-ttl-2025-04-11" not in beta, beta
@ -156,7 +154,7 @@ def test_unknown_ttl_silently_dropped(monkeypatch, bogus):
assert len(ccs) == 2, ccs
for cc in ccs:
# Bogus TTLs must NOT round-trip; marker stays at the default
# (no `ttl` key, which means the 5m pool upstream).
# (no `ttl` key = 5m pool upstream).
assert cc == {"type": "ephemeral"}, cc

9
studio/backend/tests/test_anthropic_citations.py
View file

@ -3,11 +3,10 @@
"""Tests for Anthropic ``citations_delta`` handling in the streaming proxy.
Verifies the proxy injects inline ``[N]`` markers after cited text,
dedupes by type-specific anchor (char_location, page_location,
content_block_location, search_result_location), forwards a synthetic
``document_citations`` tool_event at message_stop, and stays inert when
no citations_delta events fire. See
Verifies the proxy injects inline ``[N]`` markers after cited text, dedupes by
type-specific anchor (char_location, page_location, content_block_location,
search_result_location), forwards a synthetic ``document_citations`` tool_event
at message_stop, and stays inert when no citations_delta events fire. See
https://platform.claude.com/docs/en/build-with-claude/citations
"""

77
studio/backend/tests/test_anthropic_citations_edge.py
View file

@ -50,10 +50,9 @@ def _capture(
messages: list[dict] | None = None,
captured_body: dict | None = None,
) -> list[str]:
"""Drive ``stream_chat_completion`` against a mocked Anthropic
response and return the SSE lines. Pass ``captured_body`` to also
capture the outgoing request body for assertions on the translated
Anthropic shape.
"""Drive ``stream_chat_completion`` against a mocked Anthropic response
and return the SSE lines. Pass ``captured_body`` to also capture the
outgoing request body for assertions on the translated Anthropic shape.
"""
def handler(request: httpx.Request) -> httpx.Response:
@ -147,8 +146,8 @@ def _joined(lines: list[str]) -> str:
def _citation_payload(body: str) -> dict:
"""Pull the ``document_citations`` synthetic tool_event from the
SSE body and return its payload. Raises if absent."""
"""Return the ``document_citations`` synthetic tool_event payload from
the SSE body. Raises if absent."""
assert "document_citations" in body, body
for line in body.splitlines():
if not line.startswith("data: "):
@ -167,8 +166,8 @@ def _citation_payload(body: str) -> dict:
def test_citation_with_no_preceding_text_still_emits_marker(monkeypatch):
"""citations_delta before any text_delta must not crash; marker
lands at the start of the block."""
"""citations_delta before any text_delta must not crash; marker lands
at the start of the block."""
cit = {
"type": "char_location",
"document_index": 0,
@ -195,8 +194,8 @@ def test_citation_with_no_preceding_text_still_emits_marker(monkeypatch):
def test_citations_delta_with_non_dict_citation_is_ignored(monkeypatch):
"""Non-dict ``delta.citation`` must not crash, emit a marker, or
poison the document_citations list."""
"""Non-dict ``delta.citation`` must not crash, emit a marker, or poison
the document_citations list."""
lines = _capture(
monkeypatch,
[
@ -220,8 +219,8 @@ def test_citations_delta_with_non_dict_citation_is_ignored(monkeypatch):
def test_citations_delta_with_missing_citation_field_is_ignored(monkeypatch):
"""Missing ``citation`` field is treated like a non-dict citation:
skip without crashing."""
"""Missing ``citation`` field is treated like a non-dict citation: skip
without crashing."""
lines = _capture(
monkeypatch,
[
@ -245,8 +244,8 @@ def test_citations_delta_with_missing_citation_field_is_ignored(monkeypatch):
def test_char_location_with_reversed_indices_does_not_crash(monkeypatch):
"""Malformed char_location with reversed indices must not crash;
the dedup key accepts any int pair and still surfaces a footnote."""
"""Malformed char_location with reversed indices must not crash; the
dedup key accepts any int pair and still surfaces a footnote."""
cit = {
"type": "char_location",
"document_index": 0,
@ -275,8 +274,8 @@ def test_char_location_with_reversed_indices_does_not_crash(monkeypatch):
def test_page_location_missing_document_index_does_not_crash(monkeypatch):
"""page_location missing ``document_index`` still produces a
footnote; dedup key falls back to ``None`` for the missing field."""
"""page_location missing ``document_index`` still produces a footnote;
dedup key falls back to ``None`` for the missing field."""
cit = {
"type": "page_location",
"document_title": "Untitled PDF",
@ -303,8 +302,8 @@ def test_page_location_missing_document_index_does_not_crash(monkeypatch):
def test_content_block_location_with_non_int_block_index_does_not_crash(monkeypatch):
"""content_block_location with string block indices must not crash;
dedup key tolerates non-int values."""
"""content_block_location with string block indices must not crash; dedup
key tolerates non-int values."""
cit = {
"type": "content_block_location",
"document_index": 0,
@ -332,8 +331,8 @@ def test_content_block_location_with_non_int_block_index_does_not_crash(monkeypa
def test_unknown_citation_type_falls_back_to_stringified_key(monkeypatch):
"""Unknown citation ``type`` (forward-compat) still dedupes:
identical ones collapse, differing ones get distinct numbers."""
"""Unknown citation ``type`` (forward-compat) still dedupes: identical
ones collapse, differing ones get distinct numbers."""
cit_a = {
"type": "future_shape_location",
"anchor": "abc",
@ -369,8 +368,8 @@ def test_unknown_citation_type_falls_back_to_stringified_key(monkeypatch):
def test_mixed_citation_types_same_document_get_distinct_keys(monkeypatch):
"""char_location and page_location on the same document_index are
distinct shapes; dedup key uses citation type as its first slot."""
"""char_location and page_location on the same document_index are distinct
shapes; dedup key uses citation type as its first slot."""
cit_char = {
"type": "char_location",
"document_index": 0,
@ -406,9 +405,9 @@ def test_mixed_citation_types_same_document_get_distinct_keys(monkeypatch):
def test_cited_text_is_preserved_in_synthetic_event(monkeypatch):
"""``cited_text`` must survive into the synthetic event so the
Sources panel can render it as a tooltip. Anthropic does not bill
cited_text against output tokens, so preserving it is free."""
"""``cited_text`` must survive into the synthetic event so the Sources
panel can render it as a tooltip. Anthropic does not bill cited_text
against output tokens, so preserving it is free."""
cit = {
"type": "char_location",
"document_index": 0,
@ -435,8 +434,8 @@ def test_cited_text_is_preserved_in_synthetic_event(monkeypatch):
def test_internal_key_field_never_leaks_to_client(monkeypatch):
"""The internal ``_key`` dedup sentinel must be stripped before
the synthetic event is forwarded; it is not an Anthropic field."""
"""The internal ``_key`` dedup sentinel must be stripped before the
synthetic event is forwarded; it is not an Anthropic field."""
cit = {
"type": "char_location",
"document_index": 0,
@ -465,8 +464,8 @@ def test_internal_key_field_never_leaks_to_client(monkeypatch):
def test_citation_across_multiple_content_blocks_numbers_continue(monkeypatch):
"""Footnote numbering is per-message, not per-content-block:
citations across separate blocks emit [1] then [2]."""
"""Footnote numbering is per-message, not per-content-block: citations
across separate blocks emit [1] then [2]."""
cit_a = {
"type": "char_location",
"document_index": 0,
@ -509,9 +508,9 @@ def test_citation_across_multiple_content_blocks_numbers_continue(monkeypatch):
def test_inline_marker_lands_after_text_run(monkeypatch):
"""Inline ``[N]`` must land AFTER the cited text run: Anthropic
streams text then citation, so the proxy emits ``"...green.[1]"``
not ``"[1]green"``."""
"""Inline ``[N]`` must land AFTER the cited text run: Anthropic streams
text then citation, so the proxy emits ``"...green.[1]"`` not
``"[1]green"``."""
cit = {
"type": "char_location",
"document_index": 0,
@ -541,8 +540,8 @@ def test_inline_marker_lands_after_text_run(monkeypatch):
def test_no_synthetic_event_when_only_text_deltas(monkeypatch):
"""No citations_delta means no synthetic ``document_citations``
event; Sources panel relies on absence to suppress the section."""
"""No citations_delta means no synthetic ``document_citations`` event;
Sources panel relies on absence to suppress the section."""
lines = _capture(
monkeypatch,
[
@ -561,9 +560,9 @@ def test_no_synthetic_event_when_only_text_deltas(monkeypatch):
def test_input_document_translation_enables_citations(monkeypatch):
"""``input_document`` must translate to an Anthropic ``document``
block carrying ``citations: {enabled: true}`` (both base64 and url
source branches) so upstream emits citations_delta."""
"""``input_document`` must translate to an Anthropic ``document`` block
carrying ``citations: {enabled: true}`` (both base64 and url source
branches) so upstream emits citations_delta."""
captured_b64: dict = {}
_capture(
monkeypatch,
@ -631,8 +630,8 @@ def test_input_document_translation_enables_citations(monkeypatch):
def test_cited_text_truncated_in_synthetic_event(monkeypatch):
"""``cited_text`` is capped server-side so multi-KB spans do not
balloon the SSE payload."""
"""``cited_text`` is capped server-side so multi-KB spans don't balloon
the SSE payload."""
from core.inference.external_provider import _CITED_TEXT_MAX_LEN
long_quote = "x" * (_CITED_TEXT_MAX_LEN + 4000)

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

@ -6,20 +6,18 @@ Unit tests for Anthropic's server-side `code_execution_20250825` tool
translation in `_stream_anthropic`.
Covers:
- Request body: when ``enabled_tools=["code_execution"]``, the outbound
``tools`` array carries ``{"type": "code_execution_20250825", "name":
"code_execution"}`` and the ``anthropic-beta`` header includes
``code-execution-2025-08-25``.
- Combined request: ``enabled_tools=["web_search", "code_execution"]``
sends both tool entries; the beta header still merges the code-exec
flag onto whatever the registry contributed.
- SSE translation: a `bash_code_execution` server_tool_use +
- Request body: ``enabled_tools=["code_execution"]`` puts ``{"type":
"code_execution_20250825", "name": "code_execution"}`` in the outbound
``tools`` and ``code-execution-2025-08-25`` in the ``anthropic-beta``
header.
- Combined: ``["web_search", "code_execution"]`` sends both tool entries;
the beta header merges the code-exec flag onto the registry's.
- SSE: a `bash_code_execution` server_tool_use +
`bash_code_execution_tool_result` pair emits one tool_start and one
tool_end ``_toolEvent`` chunk with the expected arguments and result.
- SSE translation: a `text_editor_code_execution` create + result emits
a tool_start with ``kind="text_editor"`` + parsed args, and tool_end
with ``"Created"`` (or ``"Updated"``) based on the ``is_file_update``
flag.
tool_end ``_toolEvent`` chunk with the expected args and result.
- SSE: a `text_editor_code_execution` create + result emits a tool_start
with ``kind="text_editor"`` + parsed args, and tool_end with
``"Created"`` (or ``"Updated"``) per the ``is_file_update`` flag.
- Error path: a ``bash_code_execution_tool_result_error`` with
``error_code="container_expired"`` renders as ``"Error:
container_expired"`` in the tool_end ``result``.
@ -116,13 +114,13 @@ def test_code_execution_tool_appended_to_request_body(monkeypatch):
body = captured["body"]
tools = body.get("tools") or []
# Opus 4.7 gets the newer date-pinned variant (`_20260120`) that
# supports REPL state persistence + programmatic tool calling.
# Opus 4.7 gets the newer date-pinned variant (`_20260120`) with REPL
# state persistence + programmatic tool calling.
assert {"type": "code_execution_20260120", "name": "code_execution"} in tools
# No web_search entry when only code_execution is enabled.
assert all("web_search" not in (t.get("type") or "") for t in tools)
# Beta header still carries the documented flag; both `_20250825`
# and `_20260120` are unlocked by the same header per upstream docs.
# Beta header still carries the flag; both `_20250825` and `_20260120`
# are unlocked by the same header per upstream docs.
beta_header = captured["headers"].get("anthropic-beta", "")
assert "code-execution-2025-08-25" in beta_header
@ -193,11 +191,11 @@ def test_no_code_execution_tool_when_pill_off(monkeypatch):
_drive(run())
tools = captured["body"].get("tools") or []
# Pill off -- neither the legacy nor the new code_execution variant
# may appear on the wire.
# Pill off -- neither the legacy nor new code_execution variant may
# appear on the wire.
assert all("code_execution" not in (t.get("type") or "") for t in tools)
# Beta header must NOT mention code-execution when the tool isn't on
# -- that flag is opt-in only.
# Beta header must NOT mention code-execution when the tool is off --
# that flag is opt-in only.
assert "code-execution-2025-08-25" not in captured["headers"].get("anthropic-beta", "")
@ -270,8 +268,8 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch):
assert start["type"] == "tool_start"
assert start["tool_name"] == "code_execution"
assert start["tool_call_id"] == "srvtoolu_1"
# `_server_tool: True` marks this as a provider-side synthetic
# tool card for the frontend's history serializer.
# `_server_tool: True` marks a provider-side synthetic tool card for
# the frontend's history serializer.
assert start["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True}
assert end["type"] == "tool_end"

86
studio/backend/tests/test_anthropic_compaction.py
View file

@ -4,15 +4,15 @@
"""Unit tests for Anthropic server-side context compaction wiring.
Compaction is a beta feature (header ``compact-2026-01-12``) gated to
Opus 4.6, Opus 4.7, Sonnet 4.6, and Mythos preview. When enabled,
Studio attaches ``context_management.edits[{type:"compact_20260112",
trigger:{type:"input_tokens", value:N}}]`` to the outbound body. The
minimum upstream-accepted threshold is 50k tokens; lower values are
clamped to 50k so the request doesn't 400.
Opus 4.6, Opus 4.7, Sonnet 4.6, and Mythos preview. When enabled, Studio
attaches ``context_management.edits[{type:"compact_20260112",
trigger:{type:"input_tokens", value:N}}]`` to the outbound body. Minimum
upstream threshold is 50k tokens; lower values are clamped to 50k so the
request doesn't 400.
These tests pin: the body shape per model, the beta header merge with
the existing code-execution beta, threshold clamping, and silent no-op
on unsupported models.
These tests pin: body shape per model, the beta header merge with the
code-execution beta, threshold clamping, and silent no-op on unsupported
models.
"""
import asyncio
@ -168,10 +168,9 @@ def test_omitted_threshold_no_body_field(monkeypatch):
def test_chat_completion_request_accepts_sub_50k_compaction_threshold():
# Codex P1 caught that ge=50_000 on the field caused FastAPI to
# 422 the request before the in-helper clamp could fire. The
# schema must accept any positive int and let _stream_anthropic
# clamp upward.
# Codex P1 caught that ge=50_000 on the field made FastAPI 422 the
# request before the in-helper clamp could fire. The schema must
# accept any positive int and let _stream_anthropic clamp upward.
from models.inference import ChatCompletionRequest
req = ChatCompletionRequest.model_validate(
@ -209,13 +208,12 @@ def test_chat_completion_request_accepts_sub_50k_compaction_threshold():
def test_message_delta_iterations_array_aggregates_compaction_tokens(monkeypatch, capsys):
# When Anthropic compacts mid-stream, the SSE message_delta usage
# payload carries `iterations: [{type:"compaction", ...}, ...]`.
# The top-level input_tokens / output_tokens only account for the
# `message` iteration, so the cost surface needs the compaction
# totals exposed separately. The stream helper folds them into
# last_usage as `compaction_input_tokens` / `compaction_output_tokens`
# and surfaces them in the closing summary log so an operator can
# eyeball "did compaction cost us 180k tokens this turn?".
# carries `iterations: [{type:"compaction", ...}, ...]`. The top-level
# input_tokens / output_tokens only cover the `message` iteration, so
# compaction totals need to be exposed separately. The stream helper
# folds them into last_usage as `compaction_input_tokens` /
# `compaction_output_tokens` and surfaces them in the closing summary
# log so an operator can see "did compaction cost us 180k tokens?".
def http_handler(request: httpx.Request) -> httpx.Response:
body = (
@ -258,8 +256,8 @@ def test_message_delta_iterations_array_aggregates_compaction_tokens(monkeypatch
_drive(run())
# structlog renders the closing summary through the stdlib bridge,
# which lands on stdout. Capture and check the rendered line.
# structlog renders the closing summary via the stdlib bridge onto
# stdout. Capture and check the rendered line.
out = capsys.readouterr().out
summary = next(
(line for line in out.splitlines() if "Anthropic stream complete" in line),
@ -271,8 +269,8 @@ def test_message_delta_iterations_array_aggregates_compaction_tokens(monkeypatch
def test_message_delta_no_iterations_leaves_compaction_keys_unset(monkeypatch, capsys):
# Re-applying a previous compaction block does NOT emit a fresh
# iterations array. The helper must not invent compaction keys
# in that case (would otherwise double-bill).
# iterations array. The helper must not invent compaction keys then
# (would otherwise double-bill).
def http_handler(request: httpx.Request) -> httpx.Response:
body = (
b"event: message_delta\n"
@ -331,19 +329,17 @@ def _async_collect(agen):
def test_compaction_block_emitted_as_tool_event(monkeypatch):
# Codex P1: once context_management is enabled and Anthropic runs
# compaction during a turn, the response carries a
# `{type:"compaction", content:"<summary>"}` block. The translator
# must surface it so the chat-adapter can persist it onto the
# assistant message; otherwise the next turn loses the state and
# Anthropic re-compacts from scratch.
# Codex P1: once context_management is enabled and Anthropic compacts
# during a turn, the response carries a `{type:"compaction",
# content:"<summary>"}` block. The translator must surface it so the
# chat-adapter persists it onto the assistant message; otherwise the
# next turn loses the state and Anthropic re-compacts from scratch.
def http_handler(request: httpx.Request) -> httpx.Response:
# Anthropic ships compaction blocks as a content_block_start
# with `type:"compaction"`, then either includes the summary
# on that start event AND/OR streams it via text_delta events
# on the same block index. Test the streamed-delta path since
# it's the harder case.
# Anthropic ships compaction blocks as a content_block_start with
# `type:"compaction"`, then includes the summary on that start
# event AND/OR streams it via text_delta events on the same block
# index. Test the streamed-delta path since it's the harder case.
body = (
b"event: message_start\n"
b'data: {"type":"message_start","message":{"usage":{}}}\n\n'
@ -411,8 +407,8 @@ def test_compaction_block_emitted_as_tool_event(monkeypatch):
except json.JSONDecodeError:
continue
# tool_event payloads ride inside chat.completion.chunk.choices[0].delta.content
# as a JSON-encoded string. The simpler path: look for the
# marker substring anywhere in the chunk.
# as a JSON-encoded string. Simpler: look for the marker substring
# anywhere in the chunk.
if "compaction_block" in raw:
events.append(raw)
assert events, f"no compaction_block tool event found in {lines}"
@ -445,10 +441,10 @@ def test_compaction_block_emitted_as_tool_event(monkeypatch):
def test_compaction_block_round_trips_through_outbound_messages(monkeypatch):
# Once the prior turn persisted a compaction block onto the
# assistant message, the next turn's outbound body must forward
# the {type:"compaction", content:"..."} block to Anthropic
# verbatim so the API recognises the existing state.
# Once the prior turn persisted a compaction block onto the assistant
# message, the next turn's outbound body must forward the
# {type:"compaction", content:"..."} block to Anthropic verbatim so the
# API recognises the existing state.
captured: dict = {}
def http_handler(request: httpx.Request) -> httpx.Response:
@ -527,10 +523,10 @@ def test_compaction_content_part_accepted_by_chat_message_schema():
def test_build_external_messages_passes_compaction_for_anthropic_only():
# Compaction is an Anthropic-only synthetic content part. The
# builder MUST gate it on provider_type=="anthropic"; every other
# provider would 400 on the unknown content type via generic
# /chat/completions passthrough (Codex P1 follow-up).
# Compaction is an Anthropic-only synthetic content part. The builder
# MUST gate it on provider_type=="anthropic"; every other provider
# would 400 on the unknown content type via generic /chat/completions
# passthrough (Codex P1 follow-up).
from models.inference import ChatMessage
from routes.inference import _build_external_messages
@ -554,7 +550,7 @@ def test_build_external_messages_passes_compaction_for_anthropic_only():
def test_build_external_messages_strips_compaction_for_non_anthropic_providers():
# Provider switch (or reused history) hands compaction blocks to a
# non-Anthropic provider. Those land on generic /chat/completions
# non-Anthropic provider, landing on generic /chat/completions
# passthrough where the unknown content type fails the upstream
# validator. Builder must strip the part for every non-anthropic
# provider, including OpenAI/DeepSeek/Mistral/Gemini/Kimi/OpenRouter.

16
studio/backend/tests/test_anthropic_fast_mode_and_refusal.py
View file

@ -3,9 +3,9 @@
"""Tests for Anthropic fast-mode wiring and streaming refusal handling.
fast_mode=True on Opus 4.6/4.7 attaches the ``fast-mode-2026-02-01``
beta header and sets ``speed: "fast"``; unsupported models drop both.
Streaming ``stop_reason: "refusal"`` surfaces a user notice before the
fast_mode=True on Opus 4.6/4.7 attaches the ``fast-mode-2026-02-01`` beta
header and sets ``speed: "fast"``; unsupported models drop both. Streaming
``stop_reason: "refusal"`` surfaces a user notice before the
``content_filter`` finish chunk.
https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals
"""
@ -63,7 +63,7 @@ def _capture(
sse: bytes = b"",
**kwargs,
) -> tuple[dict, list[str]]:
"""Install a MockTransport, drive one streamed call, return body+lines."""
"""Install a MockTransport, drive one streamed call; return body+lines."""
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
@ -157,12 +157,12 @@ def test_refusal_emits_user_facing_notice_and_content_filter_finish(monkeypatch)
def test_refusal_emits_tool_event_for_chat_adapter_drop(monkeypatch):
"""Refused turns emit an out-of-band `_toolEvent` that the chat-adapter
latches into assistant `metadata.custom.anthropicRefusal`, driving
the next-request prune. Tool event (not text) prevents spoofing.
latches into assistant `metadata.custom.anthropicRefusal`, driving the
next-request prune. Tool event (not text) prevents spoofing.
"""
_, lines = _capture(monkeypatch, sse = _refusal_sse())
body = "\n".join(lines)
assert '"_toolEvent": {"type": "anthropic_refusal"}' in body, body
# Visible refusal text must not embed a sentinel that could spoof
# a context reset if echoed by another assistant message.
# Visible refusal text must not embed a sentinel that could spoof a
# context reset if echoed by another assistant message.
assert "studio:anthropic-refusal" not in body, body

17
studio/backend/tests/test_anthropic_fast_mode_edge.py
View file

@ -1,12 +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
"""Edge-case coverage for the Anthropic fast-mode + refusal wiring.
"""Edge-case coverage for Anthropic fast-mode + refusal wiring.
Complements ``test_anthropic_fast_mode_and_refusal.py`` (happy path)
with dated snapshots, strict opt-in (future Opus families do not
auto-enable), multi-beta header merging, refusal stream ordering, and
the non-destruction guarantee for unset/None fast_mode.
Complements ``test_anthropic_fast_mode_and_refusal.py`` (happy path) with
dated snapshots, strict opt-in (future Opus families do not auto-enable),
multi-beta header merging, refusal stream ordering, and the
non-destruction guarantee for unset/None fast_mode.
"""
import asyncio
@ -128,7 +128,7 @@ def test_fast_mode_attaches_on_dated_opus_4_6_snapshot(monkeypatch):
# ──────────────────────────── strict opt-in semantics ────────────────────────────
def test_fast_mode_does_not_auto_enable_on_future_opus_4_8(monkeypatch):
"""Future ``claude-opus-4-8`` must not auto-enable; opt-in per family."""
"""Future ``claude-opus-4-8`` must not auto-enable; per-family opt-in."""
cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-8")
assert "speed" not in cap["body"], cap["body"]
assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "")
@ -248,7 +248,7 @@ def test_fast_mode_unset_is_byte_identical_to_omitted(monkeypatch):
_drive(run())
assert cap_none["body"] == captured["body"], (cap_none["body"], captured["body"])
# Headers can vary by httpx-injected fields (host, connection); compare
# Headers vary by httpx-injected fields (host, connection); compare
# the load-bearing ones.
for key in ("anthropic-version", "x-api-key", "content-type"):
assert cap_none["headers"].get(key) == captured["headers"].get(key), key
@ -323,8 +323,7 @@ def test_refusal_chunk_is_proper_openai_delta_shape(monkeypatch):
assert notice_chunk is not None, lines
choice = notice_chunk["choices"][0]
assert "delta" in choice and "content" in choice["delta"], notice_chunk
# Must NOT carry a finish_reason itself -- that comes on the next
# chunk.
# Must NOT carry a finish_reason itself -- that comes on the next chunk.
assert choice.get("finish_reason") in (None,), notice_chunk
# Refusal text is plain-spoken; no embedded sentinel.
assert "studio:anthropic-refusal" not in choice["delta"]["content"]

45
studio/backend/tests/test_anthropic_messages.py
View file

@ -2,7 +2,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""
Tests for the Anthropic Messages API schemas and translation layer.
Tests for Anthropic Messages API schemas and translation layer.
No running server or GPU required.
"""
@ -1187,14 +1187,14 @@ class TestAnthropicRequestedStudioTools:
def test_bare_name_without_type_is_not_treated_as_server_tool(self):
# Anthropic dispatches server tools by `type`; bare-name matching
# would let a malformed client tool (e.g. user forgot input_schema)
# silently flip the request into server-execution mode.
# would let a malformed client tool (missing input_schema) silently
# flip the request into server-execution mode.
tools = [{"name": "python"}]
assert _anthropic_requested_studio_tools(tools) == set()
def test_client_tool_named_python_is_not_misclassified(self):
# input_schema is the client-tool discriminator; presence of it
# must prevent the name from being treated as a Studio alias.
# input_schema is the client-tool discriminator; its presence must
# prevent the name from being treated as a Studio alias.
tools = [
{
"name": "python",
@ -1239,8 +1239,8 @@ class _ToolPathCalled(Exception):
def _mock_backend(monkeypatch, **overrides):
"""Install a minimal stub backend on routes.inference.
Generation methods raise sentinel exceptions so the caller can assert
which path the route entered.
Generation methods raise sentinels so the caller can assert which path
the route entered.
"""
import routes.inference as inf_mod
@ -1299,10 +1299,10 @@ class TestAnthropicMessagesToolRouting:
assert "Mixing Anthropic server tools" in exc.value.detail
def test_mixed_rejected_when_client_tool_name_collides_with_server_alias(self, monkeypatch):
# Regression: a client tool sharing a name with a mapped server
# tool (e.g. user defines their own "web_search") must still
# trigger the mixed-mode 400 — the post-name filter would
# otherwise drop the client tool and silently route to server-only.
# Regression: a client tool sharing a name with a mapped server tool
# (e.g. a custom "web_search") must still trigger the mixed-mode 400;
# otherwise the post-name filter drops the client tool and silently
# routes to server-only.
_mock_backend(monkeypatch)
payload = _basic_payload(
tools = [
@ -1329,10 +1329,10 @@ class TestAnthropicMessagesToolRouting:
def test_client_tool_missing_name_rejected_with_400(self, monkeypatch):
# Regression: AnthropicTool.name was relaxed to Optional for server
# tools, so a client-tool payload that has input_schema but omits
# `name` (e.g. typo) now parses successfully but would be silently
# dropped by anthropic_tools_to_openai, leaving the request with
# tool calling disabled. Reject at the boundary instead.
# tools, so a client-tool payload with input_schema but no `name`
# (typo) now parses but would be silently dropped by
# anthropic_tools_to_openai, leaving tool calling disabled. Reject at
# the boundary instead.
_mock_backend(monkeypatch)
payload = _basic_payload(
tools = [{"input_schema": {"type": "object"}}],
@ -1346,7 +1346,7 @@ class TestAnthropicMessagesToolRouting:
def test_client_tool_empty_name_rejected_with_400(self, monkeypatch):
# Same silent-disable class as missing-name: `name: ""` passes the
# isinstance check but is dropped by anthropic_tools_to_openai's
# `if not name` guard. Reject at the boundary so the typo surfaces.
# `if not name` guard. Reject at the boundary so the typo shows.
_mock_backend(monkeypatch)
payload = _basic_payload(
tools = [{"name": "", "input_schema": {"type": "object"}}],
@ -1358,10 +1358,10 @@ class TestAnthropicMessagesToolRouting:
assert "name" in exc.value.detail
def test_alias_named_client_tool_without_schema_rejected_with_400(self, monkeypatch):
# Regression: a typo'd client tool whose name happens to collide
# with a Studio alias (e.g. user meant a custom "python" tool but
# forgot input_schema) must surface a 400, not silently switch
# the request into Studio's built-in python execution.
# Regression: a typo'd client tool whose name collides with a Studio
# alias (e.g. a custom "python" tool missing input_schema) must
# surface a 400, not silently switch into Studio's built-in python
# execution.
_mock_backend(monkeypatch)
payload = _basic_payload(tools = [{"name": "python"}])
@ -1380,9 +1380,8 @@ class TestAnthropicMessagesToolRouting:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
def test_disable_tools_policy_overrides_server_tool_alias(self, monkeypatch):
# CLI `unsloth run --disable-tools` sets policy=False. A request
# carrying a Studio server-tool alias must NOT enter the agentic
# loop in that configuration.
# CLI `unsloth run --disable-tools` sets policy=False. A request with
# a Studio server-tool alias must NOT enter the agentic loop then.
_mock_backend(monkeypatch)
set_tool_policy(False)
payload = _basic_payload(

45
studio/backend/tests/test_anthropic_thinking_translation.py
View file

@ -2,23 +2,19 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Unit tests for the Anthropic extended-thinking translation in
external_provider.
Unit tests for the Anthropic extended-thinking translation in external_provider.
Covers:
- Adaptive-mode request body nests effort under
``output_config: {effort: "<level>"}`` per the Messages API
reference (a top-level ``effort`` field 400s with
"effort: Extra inputs are not permitted").
- Streaming SSE: ``content_block_delta`` with
``delta.type == "thinking_delta"`` is translated into inline
``<think>...</think>`` chat-completion chunks so the frontend's
reasoning-panel pipeline lifts it correctly.
- The ``<think>`` tag closes when the first ``text_delta`` arrives,
on ``content_block_stop``, on ``message_delta``, or on
``message_stop``.
- Thinking is paired with ``temperature=1`` and no ``top_p`` /
``top_k`` on the wire (Anthropic extended-thinking contract).
``output_config: {effort: "<level>"}`` per the Messages API reference (a
top-level ``effort`` field 400s with "effort: Extra inputs are not permitted").
- Streaming SSE: ``content_block_delta`` with ``delta.type == "thinking_delta"``
is translated into inline ``<think>...</think>`` chat-completion chunks so the
frontend's reasoning-panel pipeline lifts it correctly.
- The ``<think>`` tag closes when the first ``text_delta`` arrives, on
``content_block_stop``, on ``message_delta``, or on ``message_stop``.
- Thinking is paired with ``temperature=1`` and no ``top_p`` / ``top_k`` on the
wire (Anthropic extended-thinking contract).
"""
import asyncio
@ -113,9 +109,8 @@ def test_adaptive_thinking_body_uses_output_config_effort_shape(monkeypatch):
# display=summarized is set explicitly so Opus 4.7 (which defaults to
# "omitted") still emits thinking_delta events for the reasoning panel.
assert body["thinking"] == {"type": "adaptive", "display": "summarized"}
# Documented shape: effort is nested under output_config.
# A top-level `effort` field produces a 400:
# "effort: Extra inputs are not permitted".
# Documented shape: effort is nested under output_config. A top-level
# `effort` field produces a 400: "effort: Extra inputs are not permitted".
assert body["output_config"] == {"effort": "medium"}
assert "effort" not in body
# Extended-thinking contract: temperature=1, no top_p / top_k.
@ -258,10 +253,10 @@ def test_manual_thinking_body_uses_budget_tokens_on_4_5(monkeypatch):
body = captured["body"]
assert body["thinking"] == {"type": "enabled", "budget_tokens": 4096}
# max_tokens must be strictly greater than budget_tokens; we shipped 1024
# and budget is 4096, so the wrapper should bump max_tokens.
# and budget is 4096, so the wrapper must bump max_tokens.
assert body["max_tokens"] > body["thinking"]["budget_tokens"]
# Manual-thinking path does not use output_config / effort — those are
# the adaptive-mode controls (Claude 4.6 / 4.7).
# Manual-thinking path does not use output_config / effort — those are the
# adaptive-mode controls (Claude 4.6 / 4.7).
assert "effort" not in body
assert "output_config" not in body
@ -338,8 +333,8 @@ def test_thinking_delta_wrapped_in_think_tags(monkeypatch):
if isinstance(p, dict) and p["choices"][0]["delta"]
)
# Reasoning text should be wrapped in <think>...</think>, followed by the
# answer text, and the stream should terminate with [DONE].
# Reasoning text is wrapped in <think>...</think>, then the answer text, and
# the stream terminates with [DONE].
assert "<think>First I plan.</think>" in combined
assert combined.endswith("Answer.")
# signature_delta is intentionally dropped — no leaked signature text.
@ -350,9 +345,9 @@ def test_thinking_delta_wrapped_in_think_tags(monkeypatch):
def test_thinking_only_turn_closes_tag_without_text_delta(monkeypatch):
"""display=omitted on Claude 4.7 emits a signature_delta and no text.
The <think> open is still triggered by the (synthetic) thinking_delta;
we want content_block_stop to close it cleanly so the tag never leaks
into the next chunk."""
The <think> open is still triggered by the (synthetic) thinking_delta; we
want content_block_stop to close it cleanly so the tag never leaks into the
next chunk."""
def handler(request: httpx.Request) -> httpx.Response:
events = [

27
studio/backend/tests/test_anthropic_tool_versions.py
View file

@ -7,13 +7,12 @@ helpers in ``core.inference.external_provider``.
Anthropic ships date-pinned tool versions per model family. The newer
``_20260209`` web_search / web_fetch and ``_20260120`` code_execution
variants only run on a subset of models; sending them to an older
model returns a 400 from upstream, and sending the older
``_20250305`` / ``_20250910`` / ``_20250825`` variants to a newer
model misses dynamic filtering and the free-when-paired pricing. The
helpers below decide which version goes out per model; this test pins
the dispatch matrix so future model launches keep working without
silently regressing the newer-version path.
variants only run on a subset of models; sending them to an older model
returns a 400 from upstream, and sending the older ``_20250305`` /
``_20250910`` / ``_20250825`` variants to a newer model misses dynamic
filtering and the free-when-paired pricing. The helpers below pick the
version per model; this test pins the dispatch matrix so future model
launches keep working without regressing the newer-version path.
Covers:
- ``_anthropic_web_search_version`` / ``_anthropic_web_fetch_version``
@ -23,10 +22,10 @@ Covers:
- ``_anthropic_code_execution_version`` picks ``_20260120`` for the
Opus 4.5+ / Sonnet 4.5+ / Opus 4.7 / Sonnet 4.6 family and falls back
to ``_20250825`` everywhere else (Haiku 4.5, 4.1, 4.0).
- ``_stream_anthropic`` body integration: when ``enabled_tools=
["web_search", "code_execution"]`` is set on Opus 4.7, the outbound
body carries the newer pinned versions; the same payload on Haiku
4.5 falls back to the legacy versions.
- ``_stream_anthropic`` body integration: with ``enabled_tools=
["web_search", "code_execution"]`` on Opus 4.7, the outbound body
carries the newer pinned versions; the same payload on Haiku 4.5 falls
back to the legacy versions.
- The ``anthropic-beta: code-execution-2025-08-25`` header is sent
unchanged for both code-execution variants (no header rev needed).
"""
@ -165,9 +164,9 @@ def test_outbound_body_uses_new_versions_on_opus_4_7(monkeypatch):
assert "code_execution_20260120" in tool_types
assert "web_search_20250305" not in tool_types
assert "code_execution_20250825" not in tool_types
# Beta header for code execution stays on the existing flag for
# both _20250825 and _20260120; the API uses one header to gate
# the feature, not the date.
# Beta header for code execution stays on the existing flag for both
# _20250825 and _20260120; the API gates the feature by one header, not
# the date.
assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "")

60
studio/backend/tests/test_anthropic_web_fetch.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
"""
Unit tests for Anthropic's `web_fetch_20250910` / `web_fetch_20260209`
"""Unit tests for Anthropic's `web_fetch_20250910` / `web_fetch_20260209`
translation in ``_stream_anthropic``. Covers request body emission
(version picked by ``_anthropic_web_fetch_version``: ``_20260209`` for
Opus 4.6/4.7 + Sonnet 4.6, ``_20250910`` otherwise), combined tool
requests, off-by-default behavior, and SSE translation of success and
(version from ``_anthropic_web_fetch_version``: ``_20260209`` for Opus
4.6/4.7 + Sonnet 4.6, else ``_20250910``), combined tool requests,
off-by-default behavior, and SSE translation of success and
``url_not_accessible`` error paths into ``tool_start`` / ``tool_end``.
"""
@ -103,7 +102,7 @@ def test_web_fetch_tool_appended_to_request_body(monkeypatch):
body = captured["body"]
tools = body.get("tools") or []
# claude-opus-4-7 routes web_fetch to _20260209 (dynamic filtering).
# claude-opus-4-7 routes web_fetch to _20260209.
assert {"type": "web_fetch_20260209", "name": "web_fetch", "max_uses": 5} in tools
# web_fetch is GA; no beta header is required.
assert "web-fetch" not in captured["headers"].get("anthropic-beta", "")
@ -140,13 +139,13 @@ def test_web_fetch_combined_with_web_search_and_code_execution(monkeypatch):
tools = captured["body"].get("tools") or []
tool_types = [t.get("type") for t in tools]
# claude-opus-4-7 routes web_search and web_fetch to _20260209
# and code_execution to _20260120 (per PR 5679 dispatch).
# claude-opus-4-7 routes web_search/web_fetch to _20260209 and
# code_execution to _20260120 (per PR 5679 dispatch).
assert "web_search_20260209" in tool_types, tool_types
assert "web_fetch_20260209" in tool_types, tool_types
assert "code_execution_20260120" in tool_types, tool_types
# Code-execution still adds its beta flag; web_fetch must not
# have accidentally stripped it.
# Code-execution still adds its beta flag; web_fetch must not have
# stripped it.
assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "")
@ -263,8 +262,8 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch):
assert start["type"] == "tool_start"
assert start["tool_name"] == "web_fetch"
assert start["tool_call_id"] == "srvtoolu_wf1"
# `_server_tool: True` marks this as a provider-side synthetic
# tool card for the frontend's history serializer.
# `_server_tool: True` marks this a provider-side synthetic tool card
# for the frontend's history serializer.
assert start["arguments"] == {"url": "https://example.com/article", "_server_tool": True}
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_wf1"
@ -345,9 +344,10 @@ def test_web_fetch_error_renders_error_code(monkeypatch):
def _finish_reasons(lines: list[str]) -> list:
"""Return non-null finish_reason fields from each chat.completion.chunk.
Mid-stream content deltas carry ``finish_reason: None`` and are skipped
(the refusal path emits a notice delta before the content_filter chunk)."""
"""Non-null finish_reason fields from each chat.completion.chunk.
Mid-stream content deltas carry ``finish_reason: None`` and are
skipped (refusal emits a notice delta before the content_filter
chunk)."""
out: list = []
for line in lines:
if not line.startswith("data:"):
@ -369,12 +369,12 @@ def _finish_reasons(lines: list[str]) -> list:
def test_pause_turn_does_not_emit_finish_reason_chunk(monkeypatch):
# `pause_turn` is what Anthropic emits when a long server-tool turn
# (typically web_search / web_fetch) pauses and will resume on the
# next request. Treating it as finish_reason="stop" makes the
# OpenAI-formatted client truncate the rendered assistant message.
# The adapter must skip the chunk so the stream ends cleanly with
# [DONE] and no terminal finish_reason.
# Anthropic emits `pause_turn` when a long server-tool turn
# (typically web_search / web_fetch) pauses and resumes on the next
# request. Mapping it to finish_reason="stop" makes the
# OpenAI-formatted client truncate the assistant message. The adapter
# must skip the chunk so the stream ends cleanly with [DONE] and no
# terminal finish_reason.
sse_events = [
{"type": "message_start", "message": {"usage": {}}},
{
@ -408,15 +408,14 @@ def test_pause_turn_does_not_emit_finish_reason_chunk(monkeypatch):
)
lines = _drive(run())
# No finish_reason chunk for pause_turn -- the only completion
# signal is the [DONE] line.
# No finish_reason chunk for pause_turn -- only [DONE] signals
# completion.
assert _finish_reasons(lines) == [], lines
assert any(line.strip() == "data: [DONE]" for line in lines), lines
def test_end_turn_still_emits_stop_finish_reason(monkeypatch):
# Sanity: the pause_turn -> None mapping must not regress normal
# end_turn handling.
# Sanity: pause_turn -> None mapping must not regress end_turn.
sse_events = [
{"type": "message_start", "message": {"usage": {}}},
{
@ -489,12 +488,11 @@ def test_refusal_maps_to_content_filter(monkeypatch):
def test_web_fetch_titleless_document_falls_back_to_url(monkeypatch):
# Anthropic may omit `document.title` on pages where the HTML
# provides nothing usable. Without a fallback the formatter would
# emit `URL: ...\nSnippet: ...` only, and the frontend's
# parseSourcesFromResult skips entries that lack a `Title:` line,
# so the source pill silently disappears. Verify the formatter
# mirrors the web_search behaviour and falls back to the URL.
# Anthropic may omit `document.title` when the HTML has nothing
# usable. Without a fallback the formatter emits `URL: ...\nSnippet:
# ...` only, and the frontend's parseSourcesFromResult skips entries
# lacking a `Title:` line, so the source pill disappears. Verify the
# formatter mirrors web_search and falls back to the URL.
sse_events = [
{"type": "message_start", "message": {"usage": {}}},
{

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

@ -1,8 +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
"""Tests for tokenizer-based audio_type detection patterns, covering both
Gemma 3n (<audio_soft_token>) and Gemma 4 (<|audio|>) audio-input tokens."""
"""Tests for tokenizer-based audio_type detection, covering Gemma 3n
(<audio_soft_token>) and Gemma 4 (<|audio|>) audio-input tokens."""
from __future__ import annotations
@ -10,8 +10,7 @@ from utils.models.model_config import _AUDIO_TOKEN_PATTERNS, is_audio_input_type
def _classify(tokens: list[str]) -> str | None:
"""Mirror _detect_audio_from_tokenizer._check_token_patterns: first match
in dict order wins."""
"""Mirror _check_token_patterns: first match in dict order wins."""
for audio_type, check in _AUDIO_TOKEN_PATTERNS.items():
if check(tokens):
return audio_type

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

@ -9,8 +9,7 @@ from pathlib import Path
import pytest
from fastapi import HTTPException
# Keep this test runnable in lightweight environments where optional logging
# deps are not installed.
# Keep runnable in lightweight environments lacking optional logging deps.
if "structlog" not in sys.modules:
class _DummyLogger:

4
studio/backend/tests/test_cache_case_resolution.py
View file

@ -111,10 +111,10 @@ def test_resolve_cached_repo_id_case_late_cache_population(tmp_path, monkeypatch
first = resolve_cached_repo_id_case("org/model")
assert first == "org/model"
# Simulate cache being populated after first miss (e.g. another code path/download).
# Cache populated after first miss (e.g. another code path/download).
_mk_cache_repo(tmp_path, "Org/Model")
second = resolve_cached_repo_id_case("org/model")
# Desired behavior: second lookup should pick up the now-existing variant.
# Second lookup should pick up the now-existing variant.
assert second == "Org/Model"

32
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -7,8 +7,8 @@ import types
from pathlib import Path
from types import SimpleNamespace
# Keep this test runnable in lightweight environments where optional logging
# deps are not installed.
# Keep this test runnable in lightweight environments without optional
# logging deps.
if "structlog" not in sys.modules:
class _DummyLogger:
@ -212,7 +212,7 @@ def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypa
def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypatch, tmp_path):
"""Mirror of the _skips_ test: the mixed repo should still surface in
cached-gguf so the picker can show it as a GGUF download."""
cached-gguf so the picker shows it as a GGUF download."""
mixed = _repo(
"Org/MixedRepo",
[
@ -241,8 +241,8 @@ def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypa
def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path):
"""A partial/interrupted GGUF download has ``size_on_disk = None``. The
route must treat the unknown bytes as zero instead of raising TypeError
out of ``sum()`` and wiping the entire response."""
route must treat unknown bytes as zero instead of raising TypeError from
``sum()`` and wiping the whole response."""
partial = _repo(
"Org/PartialDownload",
[_file("Q4_K_M.gguf", None), _file("Q6_K.gguf", 5_000)],
@ -268,7 +268,7 @@ def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path):
def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypatch, tmp_path):
"""One repo raising during classification must not poison the response
for every other repo in the scan."""
for the other repos in the scan."""
class _ExplodingRepo:
repo_id = "Org/Broken"
@ -303,9 +303,9 @@ def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypat
def test_list_cached_gguf_skips_repo_with_only_mmproj_gguf(monkeypatch, tmp_path):
"""A repo whose only ``.gguf`` artifact is an mmproj vision adapter
must not be classified as a GGUF repo: the variant selector filters
mmproj out and the picker would otherwise show zero variants."""
"""A repo whose only ``.gguf`` artifact is an mmproj vision adapter must
not be classified as a GGUF repo: the variant selector filters mmproj
out and the picker would otherwise show zero variants."""
mmproj_only = _repo(
"Org/MmprojOnly",
[
@ -327,9 +327,9 @@ def test_list_cached_gguf_skips_repo_with_only_mmproj_gguf(monkeypatch, tmp_path
def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp_path):
"""Mirror of the cached-gguf skip: a safetensors repo with an
auxiliary mmproj vision adapter must still surface in cached-models
so the user can load it as a normal model."""
"""Mirror of the cached-gguf skip: a safetensors repo with an auxiliary
mmproj vision adapter must still surface in cached-models so the user
can load it as a normal model."""
mmproj_aux = _repo(
"Org/MmprojAux",
[
@ -351,10 +351,10 @@ def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp
def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path):
"""A vision-capable GGUF repo (main weight + mmproj adapter) is still
a GGUF repo. The reported size is the main weight size; mmproj is
excluded from the GGUF-size accounting because it is filtered out at
classification time."""
"""A vision-capable GGUF repo (main weight + mmproj adapter) is still a
GGUF repo. Reported size is the main weight size; mmproj is excluded
from GGUF-size accounting since it is filtered out at classification
time."""
vision_repo = _repo(
"Org/VisionGguf",
[

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

@ -30,9 +30,9 @@ def _reset_studio_db(
def workspace_projects_home(tmp_path):
"""Projects root outside the platform delete denylist.
tmp_path resolves under /private/tmp on macOS, which the workspace
delete guard refuses by design. Linux/Windows tmp is not denied and is
used as-is; only the denied case falls back to a home subdir.
tmp_path resolves under /private/tmp on macOS, which the workspace delete
guard refuses by design. Linux/Windows tmp is not denied and is used as-is;
only the denied case falls back to a home subdir.
"""
candidate = tmp_path / "Projects"
resolved = str(candidate.resolve())
@ -367,8 +367,8 @@ def test_legacy_imports_dedups_input(tmp_path, monkeypatch):
accepted, inserted = studio_db.upsert_chat_legacy_imports(
["x", "x", "y", "x"],
)
# accepted is the deduped non-empty input size; inserted is the rows
# actually new in the ledger after ON CONFLICT DO NOTHING.
# accepted is the deduped non-empty input size; inserted is the rows newly
# added to the ledger after ON CONFLICT DO NOTHING.
assert accepted == 2
assert inserted == 2
assert set(studio_db.list_chat_legacy_imports()) == {"x", "y"}

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

@ -16,11 +16,11 @@ if str(_BACKEND_ROOT) not in sys.path:
@pytest.fixture
def outputs_setup(tmp_path, monkeypatch):
"""Point outputs_root() at a temp dir so cleanup is allowed to run on it.
"""Point outputs_root() at a temp dir so cleanup may run on it.
The training module binds ``outputs_root`` at import time
(``from utils.paths import outputs_root``), so we have to patch
the symbol on the importer module, not on storage_roots.
(``from utils.paths import outputs_root``), so patch the symbol on the
importer module, not on storage_roots.
"""
from core.training import training as training_mod
@ -36,8 +36,8 @@ def _mk_dir(parent: Path, name: str) -> Path:
def test_completed_checkpoints_are_preserved(outputs_setup):
"""The big regression: prior to this fix, every completed
checkpoint-N/ was rmtree'd on Cancel, destroying resume points."""
"""The big regression: before this fix, every completed checkpoint-N/
was rmtree'd on Cancel, destroying resume points."""
from core.training.training import _cleanup_cancelled_checkpoints
out = outputs_setup / "run-1"

4
studio/backend/tests/test_cpu_threads.py
View file

@ -116,8 +116,8 @@ def _ast_line_of_platform_compat_import(source: str) -> int:
raise AssertionError("_platform_compat import not found")
# AST-based ordering: configure_cpu_threads() must precede _platform_compat
# in both run.py and main.py. Robust to formatting / line shifts.
# AST ordering: configure_cpu_threads() must precede _platform_compat in both
# run.py and main.py. Robust to formatting / line shifts.
@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY])
def test_cpu_thread_configuration_runs_before_backend_imports(entry_point):
source = entry_point.read_text()

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