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

Merge branch 'main' into docker-blackwell-build

Browse source
This commit is contained in:
Daniel Han 2026-07-18 08:11:41 +00:00
commit c3ba8c3801
68 changed files with 5692 additions and 644 deletions
Download patch file Download diff file Expand all files Collapse all files

377
.github/workflows/release-desktop.yml vendored
View file

@ -19,6 +19,19 @@ on:
permissions:
contents: read
env:
DESKTOP_RELEASE_NOTES: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
concurrency:
group: release-desktop-${{ github.repository }}
cancel-in-progress: false
@ -295,14 +308,6 @@ jobs:
PY
build:
# TODO: split into a "build (no secrets)" + "publish (secrets)" job pair
# with actions/upload-artifact handoff so the matrix build cannot
# publish a Release on its own. The current matrix runs across
# Linux/macOS/Windows in a single job, so the split needs artefact
# collection across the OS matrix and is out of scope for this
# hardening pass.
permissions:
contents: write # tauri-apps/tauri-action creates / uploads a GitHub Release
strategy:
fail-fast: false
max-parallel: 1
@ -311,15 +316,21 @@ jobs:
- platform: macos-latest
args: '--target aarch64-apple-darwin'
label: macOS (Apple Silicon)
artifact: macos-aarch64
release_arch: aarch64
# - platform: macos-latest
# args: '--target x86_64-apple-darwin'
# label: macOS (Intel)
- platform: ubuntu-22.04
args: ''
label: Linux (x64)
artifact: linux-x64
release_arch: x64
- platform: windows-latest
args: ''
label: Windows (x64)
artifact: windows-x64
release_arch: x64
name: Build ${{ matrix.label }}
needs: prepare-version
@ -465,41 +476,18 @@ jobs:
if (chmodIdx !== -1 && sha256Idx > chmodIdx) {
throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x');
}
const releaseBodies = [];
for (let i = 0; i < lines.length; i += 1) {
const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/);
if (!match) continue;
const baseIndent = match[1].length;
const bodyLines = [];
i += 1;
for (; i < lines.length; i += 1) {
const line = lines[i];
if (line.trim() === '') {
bodyLines.push('');
continue;
}
const indent = line.match(/^\s*/)[0].length;
if (indent <= baseIndent) {
i -= 1;
break;
}
bodyLines.push(line.slice(baseIndent + 2));
}
releaseBodies.push(bodyLines.join('\n'));
const releaseBody = process.env.DESKTOP_RELEASE_NOTES;
if (!releaseBody) {
throw new Error('DESKTOP_RELEASE_NOTES must not be empty');
}
if (releaseBodies.length === 0) {
throw new Error('Expected at least one desktop release body');
if (/\brpm\b|\.rpm/i.test(releaseBody)) {
throw new Error('Desktop release body must not advertise RPM packages');
}
for (const body of releaseBodies) {
if (/\brpm\b|\.rpm/i.test(body)) {
throw new Error('Desktop release body must not advertise RPM packages');
}
if (/AppImage.*universal|universal.*AppImage/i.test(body)) {
throw new Error('Desktop release body must not advertise AppImage as universal');
}
if (!/AppImage.*experimental/i.test(body)) {
throw new Error('Desktop release body must mark AppImage as experimental');
}
if (/AppImage.*universal|universal.*AppImage/i.test(releaseBody)) {
throw new Error('Desktop release body must not advertise AppImage as universal');
}
if (!/AppImage.*experimental/i.test(releaseBody)) {
throw new Error('Desktop release body must mark AppImage as experimental');
}
JS
@ -644,48 +632,33 @@ jobs:
dest="$tools_dir/linuxdeploy-x86_64.AppImage"
curl -fsSL "$LINUXDEPLOY_URL" -o "$dest"
# Verify the digest BEFORE the binary is ever marked executable. The
# next step builds the AppImage with the Tauri signing key and a
# contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy
# that ran here could exfiltrate signing material or tamper with
# published release artifacts. Fail closed on any mismatch.
# next step builds the AppImage with the Tauri signing key, so a
# substituted linuxdeploy that ran here could exfiltrate signing
# material or tamper with release artifacts. Fail closed on any
# mismatch.
echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c -
chmod +x "$dest"
# ── Linux: build + sign + upload ──
# ── Linux: build + sign ──
- name: Build Linux app
id: build_linux
if: matrix.platform == 'ubuntu-22.04'
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
args: -v ${{ matrix.args }}
# ── macOS: build + sign + notarize + upload ──
# ── macOS: build + sign + notarize ──
- name: Build macOS app
id: build_macos
if: matrix.platform == 'macos-latest'
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
APPLE_SIGNING_IDENTITY: ${{ secrets.APPLE_SIGNING_IDENTITY }}
@ -695,29 +668,14 @@ jobs:
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
args: -v ${{ matrix.args }}
# ── Windows: build + sign + upload ──
# ── Windows: build + sign ──
- name: Build Windows app
id: build_windows
if: matrix.platform == 'windows-latest'
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }}
@ -728,35 +686,83 @@ jobs:
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
args: -v ${{ matrix.args }}
# Release process note: only non-draft workflow runs advance the public
# desktop-latest updater channel. Draft builds are for private review; if a
# draft is manually published later, this channel intentionally remains
# unchanged until a narrow manual channel-publish flow is added or a public
# desktop release is created by running this workflow with draft=false.
publish-updater-channel:
name: Publish desktop updater channel
- name: Stage release assets
shell: bash
env:
ARTIFACT_PATHS: ${{ steps.build_linux.outputs.artifactPaths || steps.build_macos.outputs.artifactPaths || steps.build_windows.outputs.artifactPaths }}
RELEASE_ARCH: ${{ matrix.release_arch }}
run: |
set -euo pipefail
if command -v python3 >/dev/null 2>&1; then
PYTHON=python3
else
PYTHON=python
fi
"$PYTHON" <<'PY'
import json
import os
import pathlib
import re
import shutil
import sys
import unicodedata
raw_paths = os.environ.get('ARTIFACT_PATHS', '')
try:
artifact_paths = json.loads(raw_paths)
except json.JSONDecodeError as error:
sys.exit(f'Invalid tauri-action artifactPaths output: {error}')
if not isinstance(artifact_paths, list) or not artifact_paths:
sys.exit('tauri-action did not return any release artifacts')
destination = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets')
destination.mkdir(parents=True, exist_ok=True)
staged = []
for raw_path in artifact_paths:
source = pathlib.Path(raw_path)
if not source.is_file():
continue
name = source.name
for extension in ('.app.tar.gz.sig', '.app.tar.gz'):
if name.endswith(extension):
name = f'{name[:-len(extension)]}_{os.environ["RELEASE_ARCH"]}{extension}'
break
name = unicodedata.normalize('NFD', name)
name = ''.join(character for character in name if not unicodedata.combining(character))
name = re.sub(r'[ ()\[\]{}]', '.', name)
while '..' in name:
name = name.replace('..', '.')
target = destination / name
if target.exists():
sys.exit(f'Duplicate staged release asset name: {name}')
shutil.copy2(source, target)
staged.append(name)
if not staged:
sys.exit('No release files were staged')
print('Staged release assets:')
print('\n'.join(sorted(staged)))
PY
- name: Upload signed release assets
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: desktop-release-${{ matrix.artifact }}
path: ${{ runner.temp }}/desktop-release-assets/*
if-no-files-found: error
compression-level: 0
retention-days: 1
# Only this job gets write access; builds hand off signed files via artifacts.
# Draft runs do not advance the public desktop-latest channel.
publish-release:
name: Publish desktop release
needs: [prepare-version, build]
if: ${{ !inputs.draft }}
runs-on: ubuntu-latest
permissions:
contents: write
contents: write # create the versioned Release and replace updater-channel metadata
env:
GH_REPO: ${{ github.repository }}
APP_VERSION: ${{ needs.prepare-version.outputs.app_version }}
@ -765,7 +771,164 @@ jobs:
DESKTOP_PRERELEASE: ${{ needs.prepare-version.outputs.prerelease }}
steps:
- name: Harden runner (audit)
uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1
with:
egress-policy: audit
- name: Download signed release assets
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
pattern: desktop-release-*
path: ${{ runner.temp }}/desktop-release-assets
merge-multiple: true
- name: Validate release asset set
shell: bash
run: |
set -euo pipefail
python3 <<'PY'
import pathlib
import os
import sys
asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets')
files = [path for path in asset_dir.iterdir() if path.is_file()]
required_suffixes = (
'.dmg',
'.app.tar.gz',
'.app.tar.gz.sig',
'.deb',
'.AppImage',
'.AppImage.sig',
'-setup.exe',
'-setup.exe.sig',
)
for suffix in required_suffixes:
matches = [path for path in files if path.name.endswith(suffix)]
if len(matches) != 1:
sys.exit(f'Expected exactly one {suffix} release asset, found {len(matches)}')
if any(path.name == 'latest.json' for path in files):
sys.exit('Build artifacts must not supply latest.json')
print('\n'.join(sorted(path.name for path in files)))
PY
- name: Create or validate versioned release
shell: bash
env:
GH_TOKEN: ${{ github.token }}
RELEASE_DRAFT: ${{ inputs.draft }}
run: |
set -euo pipefail
notes_file="$RUNNER_TEMP/desktop-release-notes.md"
printf '%s\n' "$DESKTOP_RELEASE_NOTES" > "$notes_file"
release_json="$RUNNER_TEMP/versioned-release.json"
# REST tag lookup omits drafts; `gh release view` also checks pending tags.
if gh release view "$DESKTOP_RELEASE_TAG" \
--json tagName,isDraft,isPrerelease > "$release_json" 2>/dev/null; then
python3 <<'PY'
import json
import os
import pathlib
import sys
release = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'versioned-release.json').read_text())
expected_draft = os.environ['RELEASE_DRAFT'].lower() == 'true'
expected_prerelease = os.environ['DESKTOP_PRERELEASE'].lower() == 'true'
if release.get('tagName') != os.environ['DESKTOP_RELEASE_TAG']:
sys.exit('Existing desktop release tag does not match the requested tag')
if bool(release.get('isDraft')) != expected_draft:
sys.exit('Existing desktop release draft state does not match the workflow input')
if bool(release.get('isPrerelease')) != expected_prerelease:
sys.exit('Existing desktop release prerelease state does not match the requested version')
PY
else
release_flags=(
--title "Unsloth Studio (Desktop) ${STUDIO_VERSION}"
--notes-file "$notes_file"
--target "$GITHUB_SHA"
)
if [ "$RELEASE_DRAFT" = "true" ]; then
release_flags+=(--draft)
fi
if [ "$DESKTOP_PRERELEASE" = "true" ]; then
release_flags+=(--prerelease)
fi
gh release create "$DESKTOP_RELEASE_TAG" "${release_flags[@]}"
fi
- name: Publish versioned release assets
shell: bash
env:
GH_TOKEN: ${{ github.token }}
run: |
set -euo pipefail
gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/desktop-release-assets"/* --clobber
- name: Generate and publish versioned updater metadata
shell: bash
env:
GH_TOKEN: ${{ github.token }}
run: |
set -euo pipefail
python3 <<'PY'
import datetime
import json
import os
import pathlib
import sys
import urllib.parse
asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets')
files = [path for path in asset_dir.iterdir() if path.is_file()]
def exactly_one(suffix: str) -> pathlib.Path:
matches = [path for path in files if path.name.endswith(suffix)]
if len(matches) != 1:
sys.exit(f'Expected exactly one {suffix} updater asset, found {len(matches)}')
return matches[0]
def entry(signature_suffix: str) -> dict[str, str]:
signature_path = exactly_one(signature_suffix)
bundle_name = signature_path.name.removesuffix('.sig')
bundle_path = asset_dir / bundle_name
if not bundle_path.is_file():
sys.exit(f'Missing updater bundle for {signature_path.name}: {bundle_name}')
encoded_tag = urllib.parse.quote(os.environ['DESKTOP_RELEASE_TAG'], safe='')
encoded_name = urllib.parse.quote(bundle_name, safe='')
return {
'signature': signature_path.read_text(),
'url': (
f'https://github.com/{os.environ["GITHUB_REPOSITORY"]}/releases/download/'
f'{encoded_tag}/{encoded_name}'
),
}
darwin = entry('.app.tar.gz.sig')
linux = entry('.AppImage.sig')
windows = entry('.exe.sig')
notes = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-notes.md').read_text()
metadata = {
'version': os.environ['APP_VERSION'],
'notes': notes,
'pub_date': datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
'platforms': {
'darwin-aarch64': darwin,
'darwin-aarch64-app': darwin,
'linux-x86_64': linux,
'linux-x86_64-appimage': linux,
'windows-x86_64': windows,
'windows-x86_64-nsis': windows,
},
}
output = pathlib.Path(os.environ['RUNNER_TEMP'], 'latest.json')
output.write_text(json.dumps(metadata, indent=2) + '\n')
PY
gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/latest.json" --clobber
- name: Download versioned updater metadata
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}
@ -790,6 +953,7 @@ jobs:
test -s "$RUNNER_TEMP/desktop-updater/latest.json"
- name: Validate versioned updater metadata
if: ${{ !inputs.draft }}
shell: bash
run: |
python3 <<'PY'
@ -849,6 +1013,7 @@ jobs:
PY
- name: Ensure desktop updater channel release
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}
@ -881,6 +1046,7 @@ jobs:
PY
- name: Prevent updater channel downgrade
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}
@ -971,6 +1137,7 @@ jobs:
PY
- name: Publish desktop updater channel metadata
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}

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

@ -729,6 +729,7 @@ jobs:
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": prompt}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": enabled,
"session_id": f"{session}-att{attempt_i}",
"temperature": TOOL_PROBE_TEMP,
@ -818,6 +819,7 @@ jobs:
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["web_search"],
"session_id": "ci-tool-calling-web",
"temperature": 0.0,

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

@ -612,6 +612,7 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["python"],
"session_id": "ci-tool-calling-py",
"temperature": TEMP,
@ -647,6 +648,7 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["web_search"],
"session_id": "ci-tool-calling-web",
"temperature": TEMP,

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

@ -791,6 +791,7 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["python"],
"session_id": "ci-tool-calling-py",
"temperature": TEMP,
@ -816,6 +817,7 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Use the terminal tool to run `echo hello-bash-tool` and tell me the exact output."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["terminal"],
"session_id": "ci-tool-calling-bash",
"temperature": TEMP,
@ -840,6 +842,7 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["web_search"],
"session_id": "ci-tool-calling-web",
"temperature": TEMP,

34
studio/MCP.md Normal file
View file

@ -0,0 +1,34 @@
# Unsloth Studio MCP server
Studio can expose a local MCP server so an MCP client can inspect models and
GPU state, validate recipes, start or stop training, inspect recipe output, and
export a loaded model.
The server is disabled by default. Enable it for a local Studio process with:
```bash
UNSLOTH_STUDIO_ENABLE_MCP=1 \
UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \
unsloth studio
```
The endpoint is `http://127.0.0.1:8888/mcp/` when Studio uses its default port
(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Studio
port when it is configured differently.
The high-impact tools are:
- `studio_status` and `list_local_models` for discovery
- `get_training_status`, `start_training`, `stop_training`, and `list_training_runs`
- `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset`
- `load_checkpoint` and `export_gguf`
`start_training` accepts the same fields as the Studio `TrainingStartRequest`.
The request is validated by the existing Pydantic model before a subprocess is
started. Export paths use the existing Studio validation as well.
The endpoint always requires `UNSLOTH_STUDIO_MCP_TOKEN` and checks an exact
Bearer token for both HTTP and WebSocket connections. Keep it on localhost
unless the deployment has an authenticated reverse proxy. The MCP endpoint is
intentionally opt-in because tools can consume GPU memory, write model
artifacts, and stop active work.

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

@ -10,10 +10,242 @@ native-chat-template fallback used by the transformers and MLX backends.
import copy
import json
import logging
from dataclasses import dataclass
from typing import Optional
_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"
_GEMMA_CHANNEL_START = "<|channel>"
_GEMMA_THOUGHT_OPEN = "<|channel>thought"
_GEMMA_THOUGHT_CLOSE = "<channel|>"
_GEMMA_TEMPLATE_OPENERS = (
_GEMMA_THOUGHT_OPEN + "\n",
_GEMMA_THOUGHT_OPEN + "\\n",
_GEMMA_THOUGHT_OPEN + _GEMMA_THOUGHT_CLOSE,
)
def _tokenizer_objects(tokenizer) -> tuple:
"""Return a processor/tokenizer and its distinct nested tokenizer."""
if tokenizer is None:
return ()
nested = getattr(tokenizer, "tokenizer", None)
return (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested)
def _selected_template_strings_from_value(
template,
tools = None,
*,
prefer_tool_use: bool = True,
) -> tuple[str, ...]:
"""Return the named chat template matching HF's default selection rules."""
tools = tools or None
if isinstance(template, str):
return (template,)
if not isinstance(template, dict):
return ()
if prefer_tool_use and tools and isinstance(template.get("tool_use"), str):
return (template["tool_use"],)
if isinstance(template.get("default"), str):
return (template["default"],)
values = tuple(value for value in template.values() if isinstance(value, str))
return values if len(values) == 1 else ()
def _selected_chat_template_strings(tokenizer, tools = None) -> tuple[str, ...]:
"""Return the active chat template selected for this request."""
tools = tools or None
getter = getattr(tokenizer, "get_chat_template", None)
if callable(getter):
for kwargs in ({"chat_template": None, "tools": tools}, {"tools": tools}, {}):
try:
selected = getter(**kwargs)
except Exception:
continue
if isinstance(selected, str):
return (selected,)
# ProcessorMixin.apply_chat_template does not switch to "tool_use" implicitly;
# it uses "default" unless chat_template= names another template.
is_processor = getattr(tokenizer, "tokenizer", None) is not None and callable(
getattr(tokenizer, "apply_chat_template", None)
)
return _selected_template_strings_from_value(
getattr(tokenizer, "chat_template", None),
tools,
prefer_tool_use = not is_processor,
)
def _detect_reasoning_channel_markers_from_templates(
templates: tuple[str, ...],
) -> Optional[tuple[str, str]]:
"""Return Gemma native reasoning markers only when a template emits them."""
if any(opener in template for template in templates for opener in _GEMMA_TEMPLATE_OPENERS):
return _GEMMA_THOUGHT_OPEN, _GEMMA_THOUGHT_CLOSE
return None
def detect_reasoning_channel_markers(tokenizer, tools = None) -> Optional[tuple[str, str]]:
"""Return native Gemma thought-channel markers supported by a tokenizer.
Detection uses the active chat template rather than model names or vocabulary
membership. Some models expose Gemma control tokens without using the native
thought-channel response protocol, and those must keep normal
``skip_special_tokens`` streaming.
"""
for obj in _tokenizer_objects(tokenizer):
templates = _selected_chat_template_strings(obj, tools)
if templates:
return _detect_reasoning_channel_markers_from_templates(templates)
return None
def detect_reasoning_channel_markers_from_template(
template, tools = None
) -> Optional[tuple[str, str]]:
"""Return native Gemma thought-channel markers from a raw template value."""
return _detect_reasoning_channel_markers_from_templates(
_selected_template_strings_from_value(template, tools)
)
def detect_reasoning_channel_markers_from_model_info(
tokenizer,
model_info: Optional[dict] = None,
tools = None,
) -> Optional[tuple[str, str]]:
"""Return reasoning markers from the active or cached native template."""
markers = detect_reasoning_channel_markers(tokenizer, tools = tools)
if markers is not None or not isinstance(model_info, dict):
return markers
native_templates = (
model_info.get("native_chat_template"),
(model_info.get("chat_template_info") or {}).get("template"),
)
for template in native_templates:
markers = detect_reasoning_channel_markers_from_template(template, tools)
if markers is not None:
return markers
return None
@dataclass(frozen = True)
class ChatTemplateRenderResult:
"""Prompt plus response-protocol metadata selected by the renderer."""
prompt: str
reasoning_channel_markers: Optional[tuple[str, str]] = None
def _split_partial_marker(text: str, marker: str) -> tuple[str, str]:
"""Hold the longest suffix that may become ``marker`` in the next chunk."""
for length in range(min(len(text), len(marker) - 1), 0, -1):
if text.endswith(marker[:length]):
return text[:-length], text[-length:]
return text, ""
class ReasoningChannelNormalizer:
"""Incrementally convert one native reasoning channel to ``<think>``.
The parser follows mlx-vlm's streaming boundary behavior but emits Studio's
established canonical text contract. Only the configured opening and
closing markers are consumed; tool-call and other control markers remain
available to downstream parsers.
"""
def __init__(self, opening_marker: str, closing_marker: str):
self._opening_marker = opening_marker
self._closing_marker = closing_marker
self._buffer = ""
self._in_reasoning = False
self._reasoning_done = False
self._skip_opening_newline = False
def feed(self, text: str) -> str:
"""Consume a raw text delta and return the stable canonical delta."""
self._buffer += text or ""
output: list[str] = []
while self._buffer:
if self._reasoning_done:
output.append(self._buffer)
self._buffer = ""
break
if self._in_reasoning and self._skip_opening_newline:
if self._buffer.startswith("\n"):
self._buffer = self._buffer[1:]
self._skip_opening_newline = False
if not self._buffer:
break
marker = self._closing_marker if self._in_reasoning else self._opening_marker
index = self._buffer.find(marker)
if index < 0:
stable, self._buffer = _split_partial_marker(self._buffer, marker)
output.append(stable)
break
output.append(self._buffer[:index])
self._buffer = self._buffer[index + len(marker) :]
if self._in_reasoning:
output.append(_THINK_CLOSE)
self._in_reasoning = False
self._reasoning_done = True
else:
output.append(_THINK_OPEN)
self._in_reasoning = True
self._skip_opening_newline = True
return "".join(output)
def finish(self) -> str:
"""Flush a naturally completed stream and close an open think block."""
output = self.drain()
if self._in_reasoning:
output += _THINK_CLOSE
self._in_reasoning = False
self._reasoning_done = True
return output
def drain(self) -> str:
"""Flush buffered literal text without synthesizing a closing tag."""
output = self._buffer
self._buffer = ""
return output
def normalize_reasoning_snapshots(
stream,
tokenizer = None,
cancel_event = None,
markers: Optional[tuple[str, str]] = None,
tools = None,
):
"""Normalize a prefix-monotonic cumulative text stream when supported."""
markers = markers or detect_reasoning_channel_markers(tokenizer, tools = tools)
if markers is None:
yield from stream
return
normalizer = ReasoningChannelNormalizer(*markers)
raw_output = ""
normalized_output = ""
for snapshot in stream:
if not snapshot.startswith(raw_output):
raise RuntimeError("Reasoning normalization requires cumulative text snapshots")
delta = normalizer.feed(snapshot[len(raw_output) :])
raw_output = snapshot
if delta:
normalized_output += delta
yield normalized_output
cancelled = cancel_event is not None and cancel_event.is_set()
tail = normalizer.drain() if cancelled else normalizer.finish()
if tail:
normalized_output += tail
yield normalized_output
def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str:
@ -166,7 +398,8 @@ def render_native_template(
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> Optional[str]:
return_metadata: bool = False,
):
"""Render ``messages`` + ``tools`` with the model's NATIVE chat template.
Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit
@ -175,7 +408,9 @@ def render_native_template(
tool-calling syntax. It is loaded straight from the repo (bypassing any
override on the live tokenizer) and cached on ``model_info``. Returns the
rendered prompt only if the native template actually emits the tools (render
differs with vs without tools); otherwise ``None``.
differs with vs without tools); otherwise ``None``. With ``return_metadata``,
returns ``ChatTemplateRenderResult`` so callers can stream with the response
protocol selected by this request's template.
``hf_token`` is the token the model was loaded with -- passed to the repo load
so a gated/private model's native template can still be fetched (otherwise the
@ -261,7 +496,16 @@ def render_native_template(
exc,
)
return None
return with_tools if with_tools != no_tools else None
if with_tools == no_tools:
return None
if return_metadata:
return ChatTemplateRenderResult(
with_tools,
_detect_reasoning_channel_markers_from_templates(
_selected_template_strings_from_value(native_tpl, tools)
),
)
return with_tools
def render_with_native_template_fallback(
@ -277,7 +521,8 @@ def render_with_native_template_fallback(
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> str:
return_metadata: bool = False,
):
"""Return ``formatted_prompt``, swapping in a native-template render when an
override template dropped the ``tools`` schema.
@ -285,9 +530,27 @@ def render_with_native_template_fallback(
them (detected by comparison, robust against tool names in the system prompt),
re-render with the model's native template. Shared by the transformers and MLX
backends so both advertise tools consistently. ``hf_token`` is forwarded so a
gated/private model's native template can still be fetched."""
gated/private model's native template can still be fetched. With
``return_metadata``, returns the selected prompt plus reasoning-channel markers
for the exact template used by this request."""
live_markers = detect_reasoning_channel_markers(tokenizer, tools = tools)
def _result(prompt: str, markers = live_markers):
if return_metadata:
return ChatTemplateRenderResult(prompt, markers)
return prompt
if not tools:
return formatted_prompt
# Gemma 4 can emit its native reasoning protocol even when a generation-time
# Unsloth override rendered a marker-free prompt. Preserve the live-verified
# no-tools thinking behavior without letting cached native metadata describe
# unrelated tool prompts that kept the active override.
markers = live_markers
if markers is None:
markers = detect_reasoning_channel_markers_from_model_info(
tokenizer, model_info, tools = None
)
return _result(formatted_prompt, markers)
if apply_fn is None:
apply_fn = apply_chat_template_for_generation
# Probe whether the live template dropped the schema. A tools-requiring template
@ -307,9 +570,9 @@ def render_with_native_template_fallback(
active_model_name,
exc,
)
return formatted_prompt
return _result(formatted_prompt)
if formatted_prompt != probe_no_tools:
return formatted_prompt # template already emits the tools schema
return _result(formatted_prompt) # template already emits the tools schema
native_prompt = render_native_template(
model_info = model_info,
active_model_name = active_model_name,
@ -320,6 +583,7 @@ def render_with_native_template_fallback(
preserve_thinking = preserve_thinking,
apply_fn = apply_fn,
hf_token = hf_token,
return_metadata = return_metadata,
)
if native_prompt:
logger.info(
@ -328,4 +592,4 @@ def render_with_native_template_fallback(
active_model_name,
)
return native_prompt
return formatted_prompt
return _result(formatted_prompt)

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

@ -5,7 +5,7 @@
from unsloth import FastLanguageModel, FastVisionModel
from unsloth.chat_templates import get_chat_template
from transformers import TextStreamer
from transformers import TextIteratorStreamer, TextStreamer
from peft import PeftModel, PeftModelForCausalLM
import json
@ -32,6 +32,11 @@ from core.inference.chat_eos import (
chat_eos_repair,
resolve_chat_turn_end_eos_ids_using,
)
from core.inference.chat_template_helpers import (
ReasoningChannelNormalizer,
detect_reasoning_channel_markers,
detect_think_prefill,
)
from core.inference.presence_penalty import _make_presence_penalty_processor
from io import StringIO
import structlog
@ -187,6 +192,53 @@ class HarmonyTextStreamer:
self._queue.put(new_content)
class ReasoningTextIteratorStreamer(TextIteratorStreamer):
"""TextIteratorStreamer that preserves native channel tokens until parsed."""
def __init__(
self,
tokenizer,
*,
markers: tuple[str, str],
skip_prompt: bool = True,
timeout: float = 0.2,
cancel_event = None,
**decode_kwargs,
):
decode_kwargs["skip_special_tokens"] = False
super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs)
self._normalizer = ReasoningChannelNormalizer(*markers)
self._cancel_event = cancel_event
self._aborted = False
def abort(self):
"""Mark generation as failed so ``end`` drains without closing."""
self._aborted = True
def on_finalized_text(
self,
text: str,
stream_end: bool = False,
):
"""Queue canonical deltas, closing only on natural stream completion."""
delta = self._normalizer.feed(text)
if delta:
self.text_queue.put(delta, timeout = self.timeout)
if stream_end:
cancelled = self._aborted or (
self._cancel_event is not None and self._cancel_event.is_set()
)
tail = self._normalizer.drain() if cancelled else self._normalizer.finish()
if tail:
self.text_queue.put(tail, timeout = self.timeout)
self.text_queue.put(self.stop_signal, timeout = self.timeout)
class _GenerationThreadError(RuntimeError):
"""Generation worker failures that should propagate through stream routes."""
class InferenceBackend:
"""Unified inference backend supporting text, vision, and LoRA models"""
@ -836,6 +888,7 @@ class InferenceBackend:
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
presence_penalty: float = 0.0,
reasoning_prefilled: bool = False,
):
"""Run an agentic tool loop on top of ``generate_chat_response``.
@ -889,6 +942,7 @@ class InferenceBackend:
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
reasoning_prefilled = reasoning_prefilled,
)
def generate_chat_response(
@ -960,8 +1014,7 @@ class InferenceBackend:
thread can toggle adapters under the generation lock.
"""
if not self.active_model_name:
yield "Error: No active model"
return
raise RuntimeError("No active model")
model_info = self.models[self.active_model_name]
is_vision = model_info.get("is_vision", False)
@ -1049,6 +1102,7 @@ class InferenceBackend:
template_messages = [{"role": "system", "content": system_prompt}] + messages
else:
template_messages = messages
reasoning_channel_markers_resolved = False
try:
if not (hasattr(tokenizer, "chat_template") and tokenizer.chat_template):
raise ValueError(
@ -1058,6 +1112,7 @@ class InferenceBackend:
f"Please use a model that includes a chat template, or manually set "
f"one via tokenizer.chat_template before inference."
)
reasoning_channel_markers = None
formatted_prompt = self._apply_chat_template_for_generation(
tokenizer,
template_messages,
@ -1073,7 +1128,7 @@ class InferenceBackend:
render_with_native_template_fallback,
)
formatted_prompt = render_with_native_template_fallback(
render_result = render_with_native_template_fallback(
formatted_prompt = formatted_prompt,
tokenizer = tokenizer,
model_info = model_info,
@ -1085,13 +1140,19 @@ class InferenceBackend:
preserve_thinking = preserve_thinking,
apply_fn = self._apply_chat_template_for_generation,
hf_token = model_info.get("hf_token"),
return_metadata = True,
)
formatted_prompt = render_result.prompt
reasoning_channel_markers = render_result.reasoning_channel_markers
reasoning_channel_markers_resolved = True
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
except Exception as e:
logger.error(f"Error applying chat template: {e}")
# Fall back to manual formatting
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
reasoning_channel_markers = None
reasoning_channel_markers_resolved = True
# Step 3: generate
yield from self.generate_stream(
@ -1105,6 +1166,8 @@ class InferenceBackend:
cancel_event = cancel_event,
_adapter_state = _adapter_state,
presence_penalty = presence_penalty,
reasoning_channel_markers = reasoning_channel_markers,
reasoning_channel_markers_resolved = reasoning_channel_markers_resolved,
)
def _generate_vision_response(
@ -1190,21 +1253,27 @@ class InferenceBackend:
# Stream with TextIteratorStreamer + background thread
try:
from core.inference.chat_template_helpers import detect_think_prefill
# Re-emit an open <think> prefill swallowed by skip_prompt (see
# generate_stream).
think_prefix = detect_think_prefill(
prompt_text, getattr(raw_tokenizer, "all_special_tokens", None)
)
from transformers import TextIteratorStreamer
import threading
streamer = TextIteratorStreamer(
streamer = self._make_text_streamer(
raw_tokenizer,
protocol_source = processor,
# The text-only VLM fallback above did not render with the
# processor template, so its native markers do not describe
# this request's response protocol.
reasoning_channel_markers = detect_reasoning_channel_markers(processor)
if image
else None,
reasoning_channel_markers_resolved = True,
skip_prompt = True,
skip_special_tokens = True,
timeout = 0.2,
cancel_event = cancel_event,
use_harmony = self._is_gpt_oss_model(),
)
generation_kwargs = dict(
@ -1226,6 +1295,10 @@ class InferenceBackend:
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
stopping_criteria = self._cancel_stopping_criteria(cancel_event)
if stopping_criteria is not None:
generation_kwargs["stopping_criteria"] = stopping_criteria
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
err: dict[str, str] = {}
@ -1235,6 +1308,8 @@ class InferenceBackend:
model.generate(**generation_kwargs)
except Exception as e:
err["msg"] = str(e)
if hasattr(streamer, "abort"):
streamer.abort()
logger.error(f"Vision generation error in thread: {e}")
finally:
try:
@ -1251,12 +1326,17 @@ class InferenceBackend:
if think_prefix:
yield think_prefix
from queue import Empty
import time
generation_complete = False
cancel_deadline = None
try:
while True:
if cancel_event is not None and cancel_event.is_set():
break
if cancel_deadline is None:
cancel_deadline = time.monotonic() + 10
elif time.monotonic() >= cancel_deadline:
break
try:
new_token = next(streamer)
except StopIteration:
@ -1265,27 +1345,48 @@ class InferenceBackend:
except Empty:
if not thread.is_alive():
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
if cancel_deadline is not None:
remaining = cancel_deadline - time.monotonic()
if remaining <= 0:
break
thread.join(timeout = remaining)
if thread.is_alive():
break
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
continue
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
output, cleaned = self._append_stream_delta(
output, new_token, active_stop_token_ids
)
yield cleaned
finally:
if cancel_event is not None and not generation_complete:
cancel_event.set()
thread.join(timeout = 10)
join_timeout = 10
if cancel_deadline is not None:
join_timeout = max(0, cancel_deadline - time.monotonic())
thread.join(timeout = join_timeout)
if thread.is_alive():
logger.warning(
"Vision generation thread did not exit after cancel/join timeout"
)
if err.get("msg"):
yield f"Error: {err['msg']}"
raise _GenerationThreadError(err["msg"])
except _GenerationThreadError:
raise
except Exception as e:
logger.error(f"Vision generation error: {e}")
yield f"Error: {str(e)}"
raise
def generate_audio_input_response(
self,
@ -1410,11 +1511,13 @@ class InferenceBackend:
)
if err.get("msg"):
yield f"Error: {err['msg']}"
raise _GenerationThreadError(err["msg"])
except _GenerationThreadError:
raise
except Exception as e:
logger.error(f"Audio input generation error: {e}")
yield f"Error: {str(e)}"
raise
def generate_whisper_response(
self,
@ -1447,6 +1550,86 @@ class InferenceBackend:
from utils.datasets import is_gpt_oss_model_name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
def _make_text_streamer(
self,
tokenizer,
*,
protocol_source = None,
reasoning_channel_markers = None,
reasoning_channel_markers_resolved: bool = False,
skip_prompt: bool = True,
timeout: float = 0.2,
cancel_event = None,
use_harmony: bool = False,
):
"""Create the streamer matching this model's native response protocol."""
if use_harmony:
try:
return HarmonyTextStreamer(
tokenizer,
skip_prompt = skip_prompt,
timeout = timeout,
)
except Exception as e:
logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}")
return TextIteratorStreamer(
tokenizer,
skip_prompt = skip_prompt,
skip_special_tokens = True,
timeout = timeout,
)
markers = (
reasoning_channel_markers
if reasoning_channel_markers_resolved
else reasoning_channel_markers
or detect_reasoning_channel_markers(protocol_source or tokenizer)
)
if markers is not None:
return ReasoningTextIteratorStreamer(
tokenizer,
markers = markers,
skip_prompt = skip_prompt,
timeout = timeout,
cancel_event = cancel_event,
)
return TextIteratorStreamer(
tokenizer,
skip_prompt = skip_prompt,
skip_special_tokens = True,
timeout = timeout,
)
def _append_stream_delta(
self,
output: str,
new_token: str,
stop_token_ids = None,
):
"""Append a streamer delta and apply response-boundary cleanup."""
output += new_token
return output, self._clean_generated_text(output, stop_token_ids = stop_token_ids)
def _drain_streamer_tail(
self,
streamer,
output: str,
stop_token_ids = None,
):
"""Drain queued streamer text after the producer exits."""
while True:
try:
new_token = next(streamer)
except StopIteration:
return output
except Exception:
return output
if new_token:
output, cleaned = self._append_stream_delta(
output, new_token, stop_token_ids = stop_token_ids
)
yield cleaned
def generate_stream(
self,
prompt: str,
@ -1459,6 +1642,8 @@ class InferenceBackend:
cancel_event = None,
_adapter_state = None,
presence_penalty: float = 0.0,
reasoning_channel_markers = None,
reasoning_channel_markers_resolved: bool = False,
) -> Generator[str, None, None]:
"""Generate a streaming text response (text models only).
@ -1467,8 +1652,7 @@ class InferenceBackend:
``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it).
"""
if not self.active_model_name:
yield "Error: No active model"
return
raise RuntimeError("No active model")
model_info = self.models[self.active_model_name]
model = model_info["model"]
@ -1481,9 +1665,7 @@ class InferenceBackend:
try:
inputs = tokenizer(prompt, return_tensors = "pt").to(model.device)
from transformers import TextIteratorStreamer
import threading
from core.inference.chat_template_helpers import detect_think_prefill
# skip_prompt swallows an open <think> prefilled by the template;
# re-emit it so the frontend can render the thinking block.
@ -1494,30 +1676,16 @@ class InferenceBackend:
else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None))
)
# gpt-oss models: HarmonyTextStreamer parses the multi-channel
# harmony protocol into <think> tags
if self._is_gpt_oss_model():
try:
streamer = HarmonyTextStreamer(
tokenizer,
skip_prompt = True,
timeout = 0.2,
)
except Exception as e:
logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}")
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt = True,
skip_special_tokens = True,
timeout = 0.2,
)
else:
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt = True,
skip_special_tokens = True,
timeout = 0.2,
)
streamer = self._make_text_streamer(
tokenizer,
protocol_source = model_info.get("tokenizer"),
reasoning_channel_markers = reasoning_channel_markers,
reasoning_channel_markers_resolved = reasoning_channel_markers_resolved,
skip_prompt = True,
timeout = 0.2,
cancel_event = cancel_event,
use_harmony = self._is_gpt_oss_model(),
)
generation_kwargs = dict(
**inputs,
@ -1535,27 +1703,16 @@ class InferenceBackend:
if tokenizer.pad_token_id is None
else tokenizer.pad_token_id,
)
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
# Presence penalty (GGUF parity); prompt_len excludes prompt tokens.
_pp = _make_presence_penalty_processor(
presence_penalty, int(inputs["input_ids"].shape[1])
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
if cancel_event is not None:
from transformers.generation.stopping_criteria import (
StoppingCriteria,
StoppingCriteriaList,
)
class _CancelCriteria(StoppingCriteria):
def __init__(self, ev):
self.ev = ev
def __call__(self, input_ids, scores, **kwargs):
return self.ev.is_set()
generation_kwargs["stopping_criteria"] = StoppingCriteriaList(
[_CancelCriteria(cancel_event)]
)
stopping_criteria = self._cancel_stopping_criteria(cancel_event)
if stopping_criteria is not None:
generation_kwargs["stopping_criteria"] = stopping_criteria
def generate_fn():
with self._generation_lock:
@ -1565,6 +1722,8 @@ class InferenceBackend:
model.generate(**generation_kwargs)
except Exception as e:
err["msg"] = str(e)
if hasattr(streamer, "abort"):
streamer.abort()
logger.error(f"Generation error: {e}")
finally:
try:
@ -1582,12 +1741,17 @@ class InferenceBackend:
if think_prefix:
yield think_prefix
from queue import Empty
import time
generation_complete = False
cancel_deadline = None
try:
while True:
if cancel_event is not None and cancel_event.is_set():
break
if cancel_deadline is None:
cancel_deadline = time.monotonic() + 10
elif time.monotonic() >= cancel_deadline:
break
try:
new_token = next(streamer)
except StopIteration:
@ -1596,11 +1760,27 @@ class InferenceBackend:
except Empty:
if not thread.is_alive():
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
if cancel_deadline is not None:
remaining = cancel_deadline - time.monotonic()
if remaining <= 0:
break
thread.join(timeout = remaining)
if thread.is_alive():
break
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
continue
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
output, cleaned = self._append_stream_delta(
output, new_token, active_stop_token_ids
)
yield cleaned
finally:
# Set cancel_event only on early exit (user cancel), NOT on
@ -1609,16 +1789,21 @@ class InferenceBackend:
# 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)
join_timeout = 10
if cancel_deadline is not None:
join_timeout = max(0, cancel_deadline - time.monotonic())
thread.join(timeout = join_timeout)
if thread.is_alive():
logger.warning("Generation thread did not exit after cancel/join timeout")
if err.get("msg"):
yield f"Error: {err['msg']}"
raise _GenerationThreadError(err["msg"])
except _GenerationThreadError:
raise
except Exception as e:
logger.error(f"Error during generation: {e}")
yield f"Error: {str(e)}"
raise
# ── Audio (TTS) Generation ────────────────────────────────────
@ -2107,8 +2292,42 @@ class InferenceBackend:
return img.resize(new_size, Image.Resampling.LANCZOS)
return img
def _clean_generated_text(self, text: str) -> str:
"""Strip leaked special tokens using the tokenizer's own token list."""
def _generation_stop_token_ids(self, model, generation_kwargs: dict):
"""Return the stop-token ids active for a ``generate`` call."""
if "eos_token_id" in generation_kwargs:
return generation_kwargs.get("eos_token_id")
generation_config = getattr(model, "generation_config", None)
eos_token_id = getattr(generation_config, "eos_token_id", None)
if eos_token_id is not None:
return eos_token_id
config = getattr(model, "config", None)
return getattr(config, "eos_token_id", None)
def _cancel_stopping_criteria(self, cancel_event):
"""Build a Transformers stopping criteria list for user cancellation."""
if cancel_event is None:
return None
from transformers.generation.stopping_criteria import (
StoppingCriteria,
StoppingCriteriaList,
)
class _CancelCriteria(StoppingCriteria):
def __init__(self, ev):
self.ev = ev
def __call__(self, input_ids, scores, **kwargs):
return self.ev.is_set()
return StoppingCriteriaList([_CancelCriteria(cancel_event)])
def _clean_generated_text(
self,
text: str,
*,
stop_token_ids = None,
) -> str:
"""Strip leaked response-boundary tokens after streaming."""
if self._is_gpt_oss_model():
# HarmonyTextStreamer emits clean <think>...</think>. Strip any
# harmony protocol tokens and other gpt-oss tokens (e.g.
@ -2118,10 +2337,28 @@ class InferenceBackend:
return text.strip()
tokenizer = self.models.get(self.active_model_name, {}).get("tokenizer")
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
if tokenizer:
for token in getattr(tokenizer, "all_special_tokens", []):
if token in text:
text = text.replace(token, "")
if stop_token_ids is None:
stop_token_ids = self.models.get(self.active_model_name, {}).get(
"chat_turn_end_eos_ids"
)
if isinstance(stop_token_ids, int):
stop_token_ids = (stop_token_ids,)
for token_id in stop_token_ids or ():
try:
token = tokenizer.convert_ids_to_tokens(int(token_id))
except Exception:
token = None
if isinstance(token, str) and token and text.endswith(token):
text = text[: -len(token)]
elif (
isinstance(token, str)
and token
and text.endswith("</think>")
and text[: -len("</think>")].endswith(token)
):
text = text[: -len("</think>") - len(token)] + "</think>"
return text.strip()
def _load_chat_template_info(self, model_name: str):

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

@ -9,6 +9,7 @@ OpenAI-compatible /v1/chat/completions endpoint.
import atexit
import contextlib
import functools
import json
import os
import re
@ -82,6 +83,7 @@ from core.inference.tool_call_parser import (
)
from core.inference.tool_loop_controller import (
ToolLoopController,
append_deferred_nudges,
tool_event_provenance,
)
from state.tool_approvals import (
@ -998,6 +1000,283 @@ def _cached_colocated_split_main(
return None
def _cached_variant_resolution(repo_id: str, hf_variant: str) -> tuple[Optional[str], list[str]]:
"""Find a cached main GGUF and its shards for a variant."""
candidate = next(_cached_variant_candidates(repo_id, hf_variant), None)
if candidate is None:
return None, []
_, main, shards, _ = candidate
return main, shards
def _cached_variant_candidates(
repo_id: str,
hf_variant: str,
*,
require_mmproj: bool = False,
) -> Generator[tuple[str, str, list[str], Path], None, None]:
"""Yield complete cached variant copies in snapshot preference order."""
try:
from utils.models.model_config import _iter_hf_cache_snapshots
for snap in _iter_hf_cache_snapshots(repo_id):
cached_files = _gguf_snapshot_files(snap)
matches = _gguf_files_for_variant(cached_files, hf_variant)
if not matches:
continue
main = matches[0]
shards = _gguf_extra_shards(matches, main)
split = _SHARD_FULL_RE.match(main)
if split:
numbers = {
int(match.group(2))
for path in [main, *shards]
if (match := _SHARD_FULL_RE.match(path))
}
if numbers != set(range(1, int(split.group(3)) + 1)):
continue
main_path = snap.joinpath(*main.replace("\\", "/").split("/"))
if not main_path.is_file() or not _snapshot_has_all_shards(
str(main_path), main, shards, {}
):
continue
if require_mmproj and not _pick_mmproj(cached_files):
continue
yield str(main_path), main, shards, snap
except Exception as e:
logger.debug(f"Cache lookup for variant failed: {e}")
def _cached_candidate_matches_revision_size(
repo_id: str, candidate: tuple[str, str, list[str], Path], hf_token: Optional[str]
) -> bool:
"""Check cached byte sizes against the snapshot's own Hub revision.
A snapshot pointer is normally published only after its blob is complete.
When the old revision is still queryable, also compare every weight file's
size so a manually truncated cache entry is not treated as reusable. If
metadata cannot be reached, retain the cache's normal offline semantics.
"""
main_path, main, shards, snap = candidate
paths = [main, *shards]
try:
from huggingface_hub import get_paths_info
infos = list(
get_paths_info(
repo_id,
paths,
revision = snap.name,
token = hf_token,
)
)
except Exception as e:
logger.debug(
"Could not size-check cached GGUF %s at revision %s: %s",
repo_id,
snap.name,
e,
)
return True
if not infos:
# The Hub answers an unknown (e.g. force-pushed away) revision with an
# empty result, not an error; treat it like unreachable metadata.
return True
expected_sizes = {info.path: info.size for info in infos if info.size is not None}
if any(path not in expected_sizes for path in paths):
return False
try:
if os.path.getsize(main_path) < expected_sizes[main]:
return False
except OSError:
return False
return _snapshot_has_all_shards(main_path, main, shards, expected_sizes)
def _cached_complete_candidate(
repo_id: str, gguf_filename: Optional[str], shards: list[str]
) -> Optional[tuple[str, str, list[str], Path]]:
"""Return one complete exact-filename cache candidate with snapshot context."""
if not gguf_filename:
return None
if shards:
main_path = _cached_colocated_split_main(repo_id, gguf_filename, shards, {})
else:
m = _SHARD_FULL_RE.match(gguf_filename)
if m and int(m.group(3)) > 1:
return None
main_path = _cached_hf_snapshot_file(repo_id, gguf_filename)
if main_path is None:
return None
snap = _snapshot_dir_of(main_path)
if snap is None:
return None
return main_path, gguf_filename, shards, snap
def cached_gguf_for_load(
hf_repo: str,
hf_variant: Optional[str],
*,
require_mmproj: bool = False,
verify_sizes: bool = False,
hf_token: Optional[str] = None,
) -> Optional[str]:
"""Return a cached GGUF that can be loaded without downloading."""
if not hf_variant:
return None
hf_repo = _resolve_repo_id_casing(hf_repo)
for candidate in _cached_variant_candidates(
hf_repo,
hf_variant,
require_mmproj = require_mmproj,
):
if verify_sizes and not _cached_candidate_matches_revision_size(
hf_repo, candidate, hf_token
):
continue
return candidate[0]
return None
def _snapshot_dir_of(path: str) -> Optional[Path]:
"""Return the HF cache snapshot containing path, if any."""
try:
p = Path(os.path.abspath(path))
except OSError:
return None
for ancestor in p.parents:
if ancestor.parent.name == "snapshots":
return ancestor
return None
def _companion_snapshot_sibling(
near_path: str, pick: Callable[[list[str]], Optional[str]]
) -> Optional[str]:
"""Find a companion in the same snapshot as near_path."""
snap = _snapshot_dir_of(near_path)
if snap is None:
return None
try:
sibling = pick(_gguf_snapshot_files(snap))
except Exception:
return None
if not sibling:
return None
candidate = snap / sibling
return str(candidate) if candidate.is_file() else None
def _pick_mmproj(candidates: list[str]) -> Optional[str]:
mmproj_files = sorted(
f for f in candidates if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower()
)
if not mmproj_files:
return None
return next((f for f in mmproj_files if f.lower().endswith("-f16.gguf")), mmproj_files[0])
def _hub_download_in_flight(hf_repo: str) -> bool:
try:
from hub.utils.download_registry import get_models_registry
return bool(get_models_registry().active_job_refs(hf_repo))
except Exception:
return False
def _hub_download_blocks_gguf_load(
hf_repo: str,
hf_variant: Optional[str],
*,
require_mmproj: bool = False,
hf_token: Optional[str] = None,
) -> bool:
"""Whether an active Hub job makes this GGUF load unsafe.
Same-variant jobs can reclaim the stale snapshot a load would reuse, so
they always block. Other jobs block only when this load lacks a complete
cached copy and would write to the shared cache itself.
"""
try:
from hub.utils.download_registry import get_models_registry
registry = get_models_registry()
if not registry.active_job_refs(hf_repo):
return False
if registry.has_active_variant(hf_repo, hf_variant):
return True
except Exception:
return False
return (
cached_gguf_for_load(
hf_repo,
hf_variant,
require_mmproj = require_mmproj,
verify_sizes = True,
hf_token = hf_token,
)
is None
)
# Active GGUF loads by normalized repo ID.
_LOADS_IN_FLIGHT: dict[str, int] = {}
_LOADS_IN_FLIGHT_LOCK = threading.Lock()
@contextlib.contextmanager
def gguf_load_in_flight(hf_repo: Optional[str]):
"""Track an HF GGUF load until the context exits."""
key = (hf_repo or "").strip().lower()
if not key:
yield
return
with _LOADS_IN_FLIGHT_LOCK:
_LOADS_IN_FLIGHT[key] = _LOADS_IN_FLIGHT.get(key, 0) + 1
try:
yield
finally:
with _LOADS_IN_FLIGHT_LOCK:
remaining = _LOADS_IN_FLIGHT.get(key, 1) - 1
if remaining <= 0:
_LOADS_IN_FLIGHT.pop(key, None)
else:
_LOADS_IN_FLIGHT[key] = remaining
def hf_gguf_load_in_flight(hf_repo: str) -> bool:
"""Return whether a GGUF load is active for hf_repo."""
key = (hf_repo or "").strip().lower()
if not key:
return False
with _LOADS_IN_FLIGHT_LOCK:
return _LOADS_IN_FLIGHT.get(key, 0) > 0
def _with_gguf_load_marker(load: Callable):
"""Keep an HF repo marked for the full synchronous load call."""
@functools.wraps(load)
def wrapped(self, *args, **kwargs):
hf_repo = kwargs.get("hf_repo")
with gguf_load_in_flight(hf_repo):
if hf_repo and _hub_download_blocks_gguf_load(
hf_repo,
kwargs.get("hf_variant"),
require_mmproj = bool(
kwargs.get("is_vision")
and not extra_args_disable_mmproj(kwargs.get("extra_args"))
),
hf_token = kwargs.get("hf_token"),
):
raise RuntimeError(
f"'{hf_repo}' is currently being downloaded by the download manager"
)
return load(self, *args, **kwargs)
return wrapped
def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]:
m = _SHARD_FULL_RE.match(first_shard)
if not m:
@ -1541,6 +1820,31 @@ def _is_external_link(path: Path) -> bool:
return False
# Inkling's template takes a numeric thinking-effort dial (0..0.99) and its
# float() coercion turns unrecognized named levels into 0, i.e. no thinking.
# Map OpenAI-style names to the values the model was trained on. Module-level
# so duck-typed engine stand-ins in tests do not need the attribute.
_INKLING_REASONING_EFFORT = {
"none": 0.0,
"minimal": 0.1,
"low": 0.2,
"medium": 0.7,
"high": 0.9,
"xhigh": 0.99,
"max": 0.99,
}
def _coerce_reasoning_effort(architecture, kwargs: dict) -> dict:
if architecture == "inkling":
effort = kwargs.get("reasoning_effort")
if isinstance(effort, str):
mapped = _INKLING_REASONING_EFFORT.get(effort.strip().lower())
if mapped is not None:
kwargs["reasoning_effort"] = mapped
return kwargs
class LlamaCppBackend:
"""Manages a llama-server subprocess for GGUF model inference.
@ -1941,36 +2245,15 @@ class LlamaCppBackend:
def reasoning_default(self) -> bool:
return self._reasoning_default
# Inkling's template takes a numeric thinking-effort dial (0..0.99) and its
# float() coercion turns unrecognized named levels into 0, i.e. no thinking.
# Map OpenAI-style names to the values the model was trained on.
_INKLING_REASONING_EFFORT = {
"none": 0.0,
"minimal": 0.2,
"low": 0.2,
"medium": 0.7,
"high": 0.9,
"xhigh": 0.99,
"max": 0.99,
}
def _coerce_reasoning_effort(self, kwargs: dict) -> dict:
if getattr(self, "_architecture", None) == "inkling":
effort = kwargs.get("reasoning_effort")
if isinstance(effort, str):
mapped = self._INKLING_REASONING_EFFORT.get(effort.strip().lower())
if mapped is not None:
kwargs["reasoning_effort"] = mapped
return kwargs
def _reasoning_kwargs(self, enable_thinking: bool) -> dict:
if self._reasoning_style == "enable_thinking_effort":
# GLM-5.2-style: enable_thinking is the on/off gate; when on, leave
# the template's default effort (max) in place.
return {"enable_thinking": enable_thinking}
if self._reasoning_style == "reasoning_effort":
return self._coerce_reasoning_effort(
{"reasoning_effort": "high" if enable_thinking else "low"}
return _coerce_reasoning_effort(
getattr(self, "_architecture", None),
{"reasoning_effort": "high" if enable_thinking else "low"},
)
return {"enable_thinking": enable_thinking}
@ -2019,7 +2302,7 @@ class LlamaCppBackend:
kwargs["enable_thinking"] = enable_thinking
if self._supports_preserve_thinking and preserve_thinking is not None:
kwargs["preserve_thinking"] = preserve_thinking
self._coerce_reasoning_effort(kwargs)
_coerce_reasoning_effort(getattr(self, "_architecture", None), kwargs)
return kwargs or None
@property
@ -3801,13 +4084,13 @@ class LlamaCppBackend:
hf_repo: str,
free_bytes: int,
hf_token: Optional[str] = None,
) -> Optional[tuple[str, int]]:
) -> Optional[tuple[str, int, list[str]]]:
"""Find the smallest GGUF variant (including all shards) that fits.
Groups split shards by variant prefix and sums their sizes (e.g.
UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total).
Returns (first_shard_filename, total_size_bytes) or None.
Returns (first_shard_filename, total_size_bytes, extra_shards) or None.
"""
try:
from huggingface_hub import get_paths_info, list_repo_files
@ -3843,9 +4126,13 @@ class LlamaCppBackend:
# Smallest that fits
variant_sizes.sort(key = lambda x: x[1])
for first_file, total_size, _ in variant_sizes:
for first_file, total_size, shard_files in variant_sizes:
if total_size > 0 and total_size <= free_bytes:
return first_file, total_size
return (
first_file,
total_size,
[path for path in sorted(shard_files) if path != first_file],
)
return None
except Exception:
@ -4441,42 +4728,52 @@ class LlamaCppBackend:
except Exception as e:
logger.warning(f"Could not list repo files: {e}")
# Offline: resolve variant -> filename from the local HF cache.
# The heuristic below assumes filenames echo the repo name, which
# breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file). Match
# against the rel path (not just basename) so subdir layouts like
# ``BF16/foo.gguf`` are findable.
# Fall back to the local cache when the repo listing is unavailable.
if not gguf_filename:
try:
from utils.models.model_config import _iter_hf_cache_snapshots
for snap in _iter_hf_cache_snapshots(hf_repo):
cached_files = _gguf_snapshot_files(snap)
matches = _gguf_files_for_variant(cached_files, hf_variant)
if not matches:
continue
gguf_filename = matches[0]
gguf_extra_shards = _gguf_extra_shards(matches, gguf_filename)
logger.info(
"Resolved variant %s -> %s from local HF cache",
hf_variant,
gguf_filename,
)
break
except Exception as e:
logger.debug(f"Offline cache lookup for variant failed: {e}")
cached_name, cached_shards = _cached_variant_resolution(hf_repo, hf_variant)
if cached_name:
gguf_filename = cached_name
gguf_extra_shards = cached_shards
logger.info(
"Resolved variant %s -> %s from local HF cache",
hf_variant,
gguf_filename,
)
if not gguf_filename:
repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
gguf_filename = f"{repo_name}-{hf_variant}.gguf"
# Prefer the existing model. Updates use force=True to fetch a new revision.
if not force:
if hf_variant:
# Resolve by variant so a newer revision's filename does not hide
# the complete older copy. Size-check against that older snapshot's
# own revision when its metadata remains available.
cached_main = cached_gguf_for_load(
hf_repo,
hf_variant,
verify_sizes = True,
hf_token = hf_token,
)
else:
candidate = _cached_complete_candidate(hf_repo, gguf_filename, gguf_extra_shards)
cached_main = (
candidate[0]
if candidate is not None
and _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token)
else None
)
if cached_main is not None:
logger.info(f"Reusing cached GGUF: {cached_main}")
return cached_main
# Check disk space; fall back to a smaller variant if needed
all_gguf_files = [gguf_filename] + gguf_extra_shards
expected_sizes: dict[str, int] = {}
try:
from huggingface_hub import get_paths_info, try_to_load_from_cache
path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token))
expected_sizes = {p.path: p.size for p in path_infos if p.size}
total_bytes = sum((p.size or 0) for p in path_infos)
# Subtract bytes already in the HF cache so we only preflight
@ -4485,25 +4782,10 @@ class LlamaCppBackend:
# cold whenever free disk is below the full weight footprint,
# even though nothing needs downloading.
already_cached_bytes = 0
# Cross-snapshot / case-variant cache reuse is offline-only (see the download
# path below); online, hf_hub_download fetches the current revision and
# resumes partials, so an old snapshot must not be counted as cached here or
# the preflight would under-count the download and skip the disk fallback.
# Count only files that can resume this download.
offline = _hf_env_offline()
# A split GGUF whose shards are not co-located in a single snapshot is
# refetched as a whole set later, so it must not be counted as cached here.
split_needs_refetch = False
if offline and not force and gguf_extra_shards:
# Scan all snapshots for one that holds the whole set co-located, so a
# newer snapshot with only the first shard does not mask an older
# complete one and needlessly trip the disk fallback.
if (
_cached_colocated_split_main(
hf_repo, gguf_filename, gguf_extra_shards, expected_sizes
)
is None
):
split_needs_refetch = True
# Offline split sets are reusable only when every shard shares a snapshot.
split_needs_refetch = bool(offline and not force and gguf_extra_shards)
if not force and not split_needs_refetch:
for p in path_infos:
if not p.size:
@ -4564,32 +4846,26 @@ class LlamaCppBackend:
hf_token,
)
if smaller:
fallback_file, fallback_size = smaller
fallback_file, fallback_size, fallback_shards = smaller
logger.info(
f"Selected variant too large ({total_gb:.1f} GB), "
f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
)
gguf_filename = fallback_file
_m = _SHARD_RE.match(gguf_filename)
_prefix = _m.group(1) if _m else None
if _prefix:
prefix_lower = _prefix.lower()
gguf_extra_shards = sorted(
f
for f in all_gguf_files
if f.lower().startswith(prefix_lower)
and f != gguf_filename
and not _is_companion_gguf_path(f)
gguf_extra_shards = fallback_shards
# The selected fallback is a new load target. Apply the
# same any-revision reuse policy before starting a fetch.
fallback_candidate = _cached_complete_candidate(
hf_repo, gguf_filename, gguf_extra_shards
)
if fallback_candidate is not None and (
_cached_candidate_matches_revision_size(
hf_repo, fallback_candidate, hf_token
)
else:
gguf_extra_shards = []
# Record the fallback's size so the later cache-reuse probe can
# size-verify it; only for a single-file fallback, since
# _find_smallest_fitting_variant returns the whole-variant size
# and using that as the first shard's expected size would reject
# a valid cached first shard of a split fallback.
if not gguf_extra_shards:
expected_sizes[fallback_file] = fallback_size
):
logger.info(f"Reusing cached fallback GGUF: {fallback_candidate[0]}")
return fallback_candidate[0]
else:
raise RuntimeError(
f"Not enough disk space to download any variant. "
@ -4609,45 +4885,25 @@ class LlamaCppBackend:
raise RuntimeError("Cancelled")
dl_start = time.monotonic()
# Xet primary, HTTP fallback on stall; per-file so finished shards stay cached.
local_path = None
# Reuse a cached copy from another snapshot / case-variant repo dir only when
# offline. Online, fall through to hf_hub_download so its revision/etag check
# fetches the current file (and resumes a partial) instead of serving a stale
# same-name blob from an older revision.
if not force and _hf_env_offline():
if gguf_extra_shards:
# A split GGUF must load every shard from one snapshot; reuse only a
# snapshot that holds the whole set co-located, scanning past a newer
# snapshot that has just the first shard while an older one is complete.
local_path = _cached_colocated_split_main(
hf_repo, gguf_filename, gguf_extra_shards, expected_sizes
)
else:
local_path = _cached_hf_snapshot_file(
hf_repo,
gguf_filename,
expected_size = expected_sizes.get(gguf_filename),
)
if local_path is None:
local_path = hf_hub_download_with_xet_fallback(
local_path = hf_hub_download_with_xet_fallback(
hf_repo,
gguf_filename,
hf_token,
cancel_event = cancel_event,
on_status = lambda m: logger.info(m),
force_download = force,
)
for shard in gguf_extra_shards:
if cancel_event.is_set():
raise RuntimeError("Cancelled")
logger.info(f"Resolving GGUF shard: {shard}")
hf_hub_download_with_xet_fallback(
hf_repo,
gguf_filename,
shard,
hf_token,
cancel_event = cancel_event,
on_status = lambda m: logger.info(m),
force_download = force,
)
for shard in gguf_extra_shards:
if cancel_event.is_set():
raise RuntimeError("Cancelled")
logger.info(f"Resolving GGUF shard: {shard}")
hf_hub_download_with_xet_fallback(
hf_repo,
shard,
hf_token,
cancel_event = cancel_event,
force_download = force,
)
except Exception as e:
if isinstance(e, RuntimeError) and "Cancelled" in str(e):
raise
@ -4670,10 +4926,12 @@ class LlamaCppBackend:
pick: Callable[[list[str]], Optional[str]],
label: str,
cancel_event: Optional[threading.Event] = None,
near_path: Optional[str] = None,
) -> Optional[str]:
"""Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name.
Tries the live repo file list, then the local HF cache snapshots
Prefers a companion co-located with ``near_path``'s cache snapshot,
then tries the live repo file list, then the local HF cache snapshots
(offline, same fallback as _download_gguf), then hf_hub_download.
Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so
an /unload between the main download and here skips the fetch.
@ -4683,6 +4941,17 @@ class LlamaCppBackend:
if cancel_event.is_set():
return None
# Keep companion files in the main GGUF's snapshot.
if near_path:
cached = _companion_snapshot_sibling(near_path, pick)
if cached:
logger.info("Reusing cached %s: %s", label, cached)
return cached
if _hub_download_in_flight(hf_repo):
logger.info("Skipping %s download while a hub download is active", label)
return None
target: Optional[str] = None
from huggingface_hub import list_repo_files
@ -4755,33 +5024,23 @@ class LlamaCppBackend:
hf_repo: str,
hf_token: Optional[str] = None,
cancel_event: Optional[threading.Event] = None,
near_path: Optional[str] = None,
) -> Optional[str]:
"""Download the mmproj (vision projection) file from a GGUF repo.
Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local
path, or None if none exists. ``cancel_event`` overrides
``self._cancel_event`` (defaults to it).
``self._cancel_event`` (defaults to it). ``near_path`` prefers a
copy co-located with the main GGUF's cache snapshot.
"""
def _pick_mmproj(candidates: list[str]) -> Optional[str]:
mmproj_files = sorted(
f
for f in candidates
if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower()
)
if not mmproj_files:
return None
for f in mmproj_files:
if f.lower().endswith("-f16.gguf"):
return f
return mmproj_files[0]
return self._download_companion_gguf(
hf_repo = hf_repo,
hf_token = hf_token,
pick = _pick_mmproj,
label = "mmproj",
cancel_event = cancel_event,
near_path = near_path,
)
def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
@ -4812,6 +5071,7 @@ class LlamaCppBackend:
*,
hf_repo: str,
hf_token: Optional[str] = None,
near_path: Optional[str] = None,
) -> Optional[str]:
"""Download the separate MTP drafter (speculative head) from a GGUF repo.
@ -4823,16 +5083,6 @@ class LlamaCppBackend:
are intentionally skipped. Returns the local path, or None.
"""
# Offline, reuse any drafter already on disk (a fresh copy can't be
# fetched). Online, _download_companion_gguf/hf_hub_download reuse the
# current cached file and refetch a changed one, so skip the probe here
# rather than pair new weights with a stale draft.
if _hf_env_offline():
cached = self._cached_repo_mtp_drafter(hf_repo)
if cached:
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
return cached
def _pick_mtp(candidates: list[str]) -> Optional[str]:
# Root-level only: MTP/ subdir copies now share the mtp- prefix but
# are explicit-selection, not auto-fetch (they'd sort ahead of root).
@ -4845,11 +5095,28 @@ class LlamaCppBackend:
)
return mtp_files[0] if mtp_files else None
if near_path:
cached = _companion_snapshot_sibling(near_path, _pick_mtp)
if cached:
logger.info("Reusing cached MTP drafter: %s", cached)
return cached
# Offline, reuse any drafter already on disk (a fresh copy can't be
# fetched). Online, _download_companion_gguf/hf_hub_download reuse the
# current cached file and refetch a changed one, so skip the probe here
# rather than pair new weights with a stale draft.
if _hf_env_offline():
cached = self._cached_repo_mtp_drafter(hf_repo)
if cached:
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
return cached
return self._download_companion_gguf(
hf_repo = hf_repo,
hf_token = hf_token,
pick = _pick_mtp,
label = "MTP drafter",
near_path = near_path,
)
def _resolve_launch_mmproj_path(
@ -5431,6 +5698,7 @@ class LlamaCppBackend:
)
self._stdout_thread.start()
@_with_gguf_load_marker
def load_model(
self,
*,
@ -5576,6 +5844,7 @@ class LlamaCppBackend:
mmproj_path = self._download_mmproj(
hf_repo = hf_repo,
hf_token = hf_token,
near_path = model_path,
)
# Auto-download the separate MTP drafter (e.g. Gemma) when
# the requested spec mode can use it. Repos with the head
@ -5593,6 +5862,7 @@ class LlamaCppBackend:
mtp_draft_path = self._download_mtp(
hf_repo = hf_repo,
hf_token = hf_token,
near_path = model_path,
)
elif gguf_path:
if not Path(gguf_path).is_file():
@ -8227,6 +8497,11 @@ class LlamaCppBackend:
if not is_ours:
continue
# A live parent means a running Studio (or the user's
# shell) still owns it -- not an orphan.
if LlamaCppBackend._pid_parent_is_alive(proc.info["pid"]):
continue
proc.kill()
killed += 1
logger.info(
@ -8279,6 +8554,9 @@ class LlamaCppBackend:
if not owned:
continue
if LlamaCppBackend._pid_parent_is_alive(pid):
continue
try:
os.kill(pid, signal.SIGKILL)
killed += 1
@ -10107,6 +10385,9 @@ class LlamaCppBackend:
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
# Collect no-op nudges and flush them after the batch, so a no-op
# doesn't abort it and drop the parallel calls that follow.
deferred_noop_msgs: list = []
# The text-path provisional card uses the parser's default id ("call_0");
# a Mistral-style call carries its own id and would open a duplicate. Reuse
@ -10149,14 +10430,14 @@ class LlamaCppBackend:
"provenance": decision.provenance,
}
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
deferred_noop_msgs.append(completion.model_message())
if _forced_tool_call_pending:
_forced_tool_call_pending = False
logger.info(
"Suppressed local GGUF tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
break
continue
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
@ -10275,6 +10556,8 @@ class LlamaCppBackend:
if _forced_tool_call_pending:
_forced_tool_call_pending = False
append_deferred_nudges(conversation, deferred_noop_msgs)
# Close provisional cards not resolved by execution/no-op handling.
for _pid, _pname in provisional_started_tool_calls.items():
if _pid not in resolved_provisional_tool_call_ids:

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

@ -11,6 +11,10 @@ import threading
from typing import Optional, Generator
from core.inference.message_content import content_to_text
from core.inference.runtime_context import runtime_context_length
from core.inference.chat_template_helpers import (
ReasoningChannelNormalizer,
normalize_reasoning_snapshots,
)
from loggers import get_logger
logger = get_logger(__name__)
@ -533,7 +537,7 @@ class MLXInferenceBackend:
break
if self._is_vlm:
yield from self._generate_vlm(
stream = self._generate_vlm(
full_messages,
image,
temperature,
@ -550,7 +554,7 @@ class MLXInferenceBackend:
presence_penalty = presence_penalty,
)
else:
yield from self._generate_text(
stream = self._generate_text(
full_messages,
temperature,
top_p,
@ -565,6 +569,7 @@ class MLXInferenceBackend:
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
yield from stream
def _generate_text(
self,
@ -609,7 +614,7 @@ class MLXInferenceBackend:
# probe and native render share a renderer. (VLM renders via the
# processor for image tokens and is not wired here.)
model_info = self.models.get(self.active_model_name, {})
prompt = render_with_native_template_fallback(
render_result = render_with_native_template_fallback(
formatted_prompt = prompt,
tokenizer = self._tokenizer,
model_info = model_info,
@ -620,7 +625,10 @@ class MLXInferenceBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
hf_token = model_info.get("hf_token"),
return_metadata = True,
)
prompt = render_result.prompt
reasoning_channel_markers = render_result.reasoning_channel_markers
# An open <think> prefilled by the template lives in the prompt, not
# the generated tokens; re-emit it so the frontend renders the block.
@ -654,7 +662,17 @@ class MLXInferenceBackend:
if not logits_processors:
logits_processors = None
preserve_native_channels = reasoning_channel_markers is not None
token_ids = []
normalizer = (
ReasoningChannelNormalizer(*reasoning_channel_markers)
if reasoning_channel_markers is not None
else None
)
# MLX consumers diff cumulative snapshots. Keep a prompt-prefilled
# <think> prefix on every native-protocol snapshot just as the normal
# decoding path does below.
normalized_output = think_prefix
logger.info(
"Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s",
len(prompt),
@ -678,12 +696,19 @@ class MLXInferenceBackend:
**gen_kwargs,
):
final_response = response
token_ids.append(response.token)
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
)
yield think_prefix + cumulative
if preserve_native_channels:
piece = getattr(response, "text", None) or ""
delta = normalizer.feed(piece)
if delta:
normalized_output += delta
yield normalized_output
else:
token_ids.append(response.token)
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
)
yield think_prefix + cumulative
if cancel_event and cancel_event.is_set():
break
@ -700,6 +725,12 @@ class MLXInferenceBackend:
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
if normalizer is not None:
cancelled = cancel_event is not None and cancel_event.is_set()
tail = normalizer.drain() if cancelled else normalizer.finish()
if tail:
normalized_output += tail
yield normalized_output
def _generate_vlm(
self,
@ -858,31 +889,37 @@ class MLXInferenceBackend:
elif _rep_active:
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
with self._generation_lock:
final_response = None
try:
for response in vlm_stream(
self._model,
self._processor,
prompt,
images,
**vlm_kwargs,
):
final_response = response
token_text = response.text if hasattr(response, "text") else str(response)
cumulative += token_text
yield cumulative
if cancel_event and cancel_event.is_set():
break
finally:
# mlx_vlm exposes the same stats fields as mlx_lm.
if final_response is not None:
self.last_generation_stats = _build_generation_stats(
getattr(final_response, "prompt_tokens", 0),
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
def _stream_vlm_snapshots():
nonlocal cumulative
with self._generation_lock:
final_response = None
try:
for response in vlm_stream(
self._model,
self._processor,
prompt,
images,
**vlm_kwargs,
):
final_response = response
token_text = response.text if hasattr(response, "text") else str(response)
cumulative += token_text
yield cumulative
if cancel_event and cancel_event.is_set():
break
finally:
# mlx_vlm exposes the same stats fields as mlx_lm.
if final_response is not None:
self.last_generation_stats = _build_generation_stats(
getattr(final_response, "prompt_tokens", 0),
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
yield from normalize_reasoning_snapshots(
_stream_vlm_snapshots(), chat_target, cancel_event, tools = tools
)
def generate_with_adapter_control(
self,

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

@ -59,7 +59,32 @@ class GenStreamError(str):
"Error:" by checking isinstance(chunk, GenStreamError).
"""
__slots__ = ()
__slots__ = ("public",)
def __new__(
cls,
value,
*,
public: bool = False,
):
obj = str.__new__(cls, value)
obj.public = bool(public)
return obj
class GenStreamErrorRaised(RuntimeError):
"""Internal exception form of ``GenStreamError`` for generator boundaries."""
__slots__ = ("public",)
def __init__(
self,
value,
*,
public: bool = False,
):
super().__init__(value)
self.public = bool(public)
class InferenceOrchestrator:
@ -531,13 +556,19 @@ class InferenceOrchestrator:
initial_resp_queue = self._resp_queue
while True:
if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue:
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
yield GenStreamError(
f"Error: {self._subprocess_crash_message(crash_context)}",
public = True,
)
return
resp = read_one(read_timeout)
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
yield GenStreamError(
f"Error: {self._subprocess_crash_message(crash_context)}",
public = True,
)
return
continue
@ -689,11 +720,11 @@ class InferenceOrchestrator:
GPU work stays serialized; this only avoids orchestrator lock contention.
"""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running")
yield GenStreamError("Error: Inference subprocess is not running", public = True)
return
if not self.active_model_name:
yield GenStreamError("Error: No active model")
yield GenStreamError("Error: No active model", public = True)
return
# Latch the target model so the recheck below can detect a switch that completed
# between _start_dispatcher and mailbox registration (mirrors the locked path's
@ -704,7 +735,7 @@ class InferenceOrchestrator:
# so without this early-out a compare request would enqueue a generate on the
# outgoing model and delay the switch.
if self._unload_pending:
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
@ -776,7 +807,7 @@ class InferenceOrchestrator:
# _stop_dispatcher joins the dispatcher, which itself takes that lock.
if orphaned_dispatcher:
self._stop_dispatcher()
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
try:
@ -1376,6 +1407,7 @@ class InferenceOrchestrator:
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
reasoning_prefilled: bool = False,
**_unused,
):
"""Run the safetensors agentic tool loop in the parent process,
@ -1414,12 +1446,27 @@ class InferenceOrchestrator:
presence_penalty = presence_penalty,
)
if use_adapter is not None:
yield from self.generate_with_adapter_control(
stream = self.generate_with_adapter_control(
use_adapter = use_adapter,
**common_kwargs,
)
else:
yield from self.generate_chat_response(**common_kwargs)
stream = self.generate_chat_response(**common_kwargs)
close_stream = False
try:
for chunk in stream:
if isinstance(chunk, GenStreamError):
close_stream = True
raise GenStreamErrorRaised(str(chunk), public = chunk.public)
yield chunk
finally:
if close_stream:
close = getattr(stream, "close", None)
if callable(close):
try:
close()
except Exception:
logger.debug("failed to close errored generation stream", exc_info = True)
initial = list(messages)
if system_prompt:
@ -1441,6 +1488,7 @@ class InferenceOrchestrator:
confirm_tool_calls = confirm_tool_calls,
bypass_permissions = bypass_permissions,
permission_mode = permission_mode,
reasoning_prefilled = reasoning_prefilled,
)
def generate_with_adapter_control(
@ -1489,11 +1537,11 @@ class InferenceOrchestrator:
readers don't consume each other's tokens off the shared resp_queue.
"""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running")
yield GenStreamError("Error: Inference subprocess is not running", public = True)
return
if not self.active_model_name:
yield GenStreamError("Error: No active model")
yield GenStreamError("Error: No active model", public = True)
return
expected_model = self.active_model_name
@ -1510,7 +1558,7 @@ class InferenceOrchestrator:
# so we never generate on the wrong one.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
request_id = str(uuid.uuid4())
image_b64 = self._pil_to_base64(image) if image is not None else None
@ -1695,10 +1743,10 @@ class InferenceOrchestrator:
) -> Generator[str, None, None]:
"""Shared inner logic for audio input generation (Whisper + ASR)."""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running")
yield GenStreamError("Error: Inference subprocess is not running", public = True)
return
if not self.active_model_name:
yield GenStreamError("Error: No active model")
yield GenStreamError("Error: No active model", public = True)
return
expected_model = self.active_model_name
@ -1707,7 +1755,7 @@ class InferenceOrchestrator:
# cleared or swapped the model while we waited.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
request_id = str(uuid.uuid4())

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

@ -50,6 +50,7 @@ from core.inference.tool_call_parser import (
# pattern lists, so the safetensors streaming strip stays aligned with the parser.
from core.tool_healing import (
_REHEARSAL_TAIL_STRIP_RE,
_THINK_CLOSE_RE,
_strip_bracket_tag_calls,
_think_spans_outside_tool_markup,
apply_tool_strip_patterns,
@ -57,6 +58,7 @@ from core.tool_healing import (
)
from core.inference.tool_loop_controller import (
ToolLoopController,
append_deferred_nudges,
coerce_tool_arguments,
status_for_tool,
tool_event_provenance,
@ -303,6 +305,45 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str:
return status_for_tool(tool_name, arguments)
def _reprompt_intent_text(text: str, *, reasoning_prefilled: bool = False) -> str:
"""Return visible answer text for the plan-without-action classifier.
Safetensors reasoning shares the cumulative text channel with the answer.
Forward-looking phrases inside ``<think>`` / ``[THINK]`` are private
planning, not a user-visible promise to call a tool. Match GGUF's behavior:
classify visible content when present and fall back to reasoning only for a
reasoning-only stall.
"""
prefilled_reasoning = ""
if reasoning_prefilled:
close = _THINK_CLOSE_RE.search(text)
if close is None:
return text.strip()
prefilled_reasoning = text[: close.end()].strip()
text = text[close.end() :].strip()
if not text:
return prefilled_reasoning
spans = _think_spans_outside_tool_markup(text)
if not spans:
return text.strip()
visible: list[str] = []
reasoning: list[str] = []
cursor = 0
for start, end in spans:
visible.append(text[cursor:start])
reasoning.append(text[start:end])
cursor = end
visible.append(text[cursor:])
visible_text = "".join(visible).strip()
reasoning_text = "".join(reasoning).strip()
if visible_text:
return visible_text
return "\n".join(part for part in (prefilled_reasoning, reasoning_text) if part).strip()
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool:
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
probe = strip_llama3_leading_sentinels(text.lstrip())
@ -447,6 +488,7 @@ def run_safetensors_tool_loop(
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
permission_mode: Optional[str] = None,
reasoning_prefilled: bool = False,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
@ -955,7 +997,10 @@ def run_safetensors_tool_loop(
# (GGUF loop parity). The retry is gated on nudge_tool_calls so
# Studio callers (which send True) always nudge, while API callers
# who omit the flag keep today's no-reprompt behavior (opt-in).
stripped_answer = content_accum.strip()
intent_text = _reprompt_intent_text(
content_accum,
reasoning_prefilled = reasoning_prefilled,
)
if (
auto_heal_tool_calls
and nudge_tool_calls
@ -964,7 +1009,7 @@ def run_safetensors_tool_loop(
and not rag_autoinjected
and not tool_denied
and not any(record.executed for record in tool_controller.history)
and is_short_intent_without_action(stripped_answer)
and is_short_intent_without_action(intent_text)
):
reprompt_count += 1
logger.info(
@ -972,9 +1017,9 @@ def run_safetensors_tool_loop(
"calling tools (%d chars)",
reprompt_count,
MAX_ACT_REPROMPTS,
len(stripped_answer),
len(intent_text),
)
conversation.append({"role": "assistant", "content": stripped_answer})
conversation.append({"role": "assistant", "content": intent_text})
tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool"
conversation.append(
{
@ -1099,6 +1144,9 @@ def run_safetensors_tool_loop(
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
# Collect no-op nudges and flush them after the batch, so a no-op doesn't
# abort it and drop the parallel calls that follow.
deferred_noop_msgs: list = []
for tc in tool_calls or []:
func = tc.get("function", {}) or {}
@ -1127,12 +1175,12 @@ def run_safetensors_tool_loop(
"provenance": decision.provenance,
}
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
deferred_noop_msgs.append(completion.model_message())
logger.info(
"Suppressed local safetensors tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
break
continue
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
@ -1243,6 +1291,8 @@ def run_safetensors_tool_loop(
yield completion.tool_end_event()
conversation.append(completion.tool_message())
append_deferred_nudges(conversation, deferred_noop_msgs)
# Clear the status badge before the next turn.
yield {"type": "status", "text": ""}

20
studio/backend/core/inference/tool_loop_controller.py
View file

@ -266,6 +266,17 @@ def strip_result_for_model(result: str) -> str:
return result
def append_deferred_nudges(conversation: list, msgs: Sequence[dict]) -> None:
"""Append a batch's no-op nudges as one deduped ``role=user`` message.
Deferred to after the batch's tool results so a no-op never splits an
assistant's ``tool_calls`` from their ``role=tool`` results.
"""
contents = list(dict.fromkeys(msg["content"] for msg in msgs))
if contents:
conversation.append({"role": "user", "content": "\n\n".join(contents)})
def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
function = tool.get("function")
if not isinstance(function, Mapping):
@ -277,8 +288,9 @@ def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
def _noop_result(reason: NoopReason, tool_name: str) -> str:
if reason == "duplicate":
return (
"The previous tool request was not executed because this exact "
"tool call already completed successfully. Do not repeat the same "
f"One earlier request to call tool '{tool_name}' in this batch was "
"not executed because an identical call had already completed "
"successfully. Do not repeat the same "
"tool call. Continue with a different enabled tool if that would "
"materially help, or provide the final answer if you have enough "
"information."
@ -291,8 +303,8 @@ def _noop_result(reason: NoopReason, tool_name: str) -> str:
"the requested final note or answer."
)
return (
f"The previous tool request was not executed because tool "
f"'{tool_name}' is not enabled for this request. Provide the "
f"One earlier request to call tool '{tool_name}' in this batch was "
"not executed because that tool is not enabled for this request. Provide the "
"final answer now without calling more tools."
)

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

@ -3600,14 +3600,18 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion)
# Raw download cap > _MAX_PAGE_CHARS since SSR pages embed large <head> sections
# stripped during conversion; 512 KB still reaches article content.
_MAX_FETCH_BYTES = 512 * 1024
# PDF cross-reference data lives at EOF, so extraction needs the whole body.
_MAX_PDF_FETCH_BYTES = 10 * 1024 * 1024
_MAX_WEB_PDF_PAGES = 50
# Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs).
# Binary when they exceed 12.5%, after allowing 16 minor encoding glitches.
_BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]")
_MIN_BINARY_CHARS = 16
_BINARY_CHAR_DIVISOR = 8
# Common binary signatures that can otherwise look text-heavy when mislabeled.
_PDF_MAGIC = b"%PDF-"
_BINARY_MAGIC = (
b"%PDF-", # PDF
_PDF_MAGIC,
b"PK\x03\x04", # zip / docx / xlsx / pptx / epub / jar
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", # OLE / legacy Office
b"\x89PNG\r\n\x1a\n", # PNG
@ -3641,14 +3645,22 @@ def _looks_binary(text: str) -> bool:
)
def _has_binary_magic(data: bytes) -> bool:
"""Whether a common binary signature follows optional BOM or whitespace."""
def _magic_head(data: bytes) -> bytes:
head = data[:1024].lstrip()
for bom, _codec in _UNICODE_BOM_CODECS:
if head.startswith(bom):
head = head.removeprefix(bom).lstrip()
break
return head.startswith(_BINARY_MAGIC)
return head
def _has_pdf_magic(data: bytes) -> bool:
return _magic_head(data).startswith(_PDF_MAGIC)
def _has_binary_magic(data: bytes) -> bool:
"""Whether a common binary signature follows optional BOM or whitespace."""
return _magic_head(data).startswith(_BINARY_MAGIC)
def _has_single_byte_text_evidence(data: bytes) -> bool:
@ -3659,6 +3671,45 @@ def _has_single_byte_text_evidence(data: bytes) -> bool:
return ascii_text_bytes / len(data) >= _MIN_SINGLE_BYTE_ASCII_RATIO
def _extract_pdf_text(data: bytes) -> str:
"""Extract page-delimited text with the same parser used by RAG ingestion."""
from ..rag.parsers import parse_pdf_bytes
pages, total_pages = parse_pdf_bytes(data, max_pages = _MAX_WEB_PDF_PAGES)
page_limit_reached = total_pages > _MAX_WEB_PDF_PAGES
parts: list[str] = []
length = 0
text_limited = False
for page in pages:
page_text = page.text.strip()
if not page_text:
continue
section = f"## Page {page.page_number}\n\n{page_text}"
piece = ("\n\n" if parts else "") + section
remaining = _MAX_PAGE_CHARS - length
if len(piece) > remaining:
parts.append(piece[:remaining])
text_limited = True
break
parts.append(piece)
length += len(piece)
text = "".join(parts).rstrip()
if not text:
if page_limit_reached:
return f"(PDF contains no extractable text in the first {_MAX_WEB_PDF_PAGES} pages)"
return ""
limits = []
if text_limited:
limits.append(f"text limited to {_MAX_PAGE_CHARS:,} characters")
if page_limit_reached:
limits.append(f"page processing capped at {_MAX_WEB_PDF_PAGES} pages")
if limits:
marker = f"\n\n... (PDF extraction {'; '.join(limits)})"
text = text[: _MAX_PAGE_CHARS - len(marker)].rstrip() + marker
return text
_USER_AGENTS = (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
@ -4054,29 +4105,71 @@ def _fetch_url_raw(
return reason2, "", ""
current_host = rp.hostname
continue
# get_content_type() defaults to "text/plain" when the header is
# absent (RFC 2045); report "" instead so callers can tell a missing
# header apart from a server that really declared text/plain.
if resp.headers.get("Content-Type") is None:
content_type = ""
else:
content_type = (resp.headers.get_content_type() or "").lower()
# Success: read the capped body enforcing the budget between chunks
# (see _read_capped_body), so a slow-drip server can't stretch a
# single resp.read past the deadline.
declared_pdf = content_type == "application/pdf"
read_limit = _MAX_PDF_FETCH_BYTES + 1 if declared_pdf else max_bytes
body_error, raw_bytes = _read_capped_body(
resp,
max_bytes,
read_limit,
timeout,
deadline,
cancel_event,
)
if body_error is not None:
return body_error, "", ""
# A missing or wrong PDF MIME type is common: once the initial text-sized
# read identifies PDF magic, finish the bounded download to reach the EOF xref.
if not declared_pdf and len(raw_bytes) == max_bytes and _has_pdf_magic(raw_bytes):
tail_error, tail = _read_capped_body(
resp,
_MAX_PDF_FETCH_BYTES - max_bytes + 1,
timeout,
deadline,
cancel_event,
)
if tail_error is not None:
return tail_error, "", ""
raw_bytes += tail
break
else:
return "Failed to fetch URL: too many redirects.", "", ""
# get_content_type() defaults to "text/plain" when the header is
# absent (RFC 2045); report "" instead so callers can tell a missing
# header apart from a server that really declared text/plain.
if resp.headers.get("Content-Type") is None:
content_type = ""
else:
content_type = (resp.headers.get_content_type() or "").lower()
is_pdf = declared_pdf or _has_pdf_magic(raw_bytes)
if is_pdf:
if len(raw_bytes) > _MAX_PDF_FETCH_BYTES:
return (
"(PDF content exceeds the download limit; not readable as text)",
"",
content_type,
)
budget_error = _fetch_budget_exceeded(deadline, cancel_event)
if budget_error is not None:
return budget_error, "", content_type
try:
pdf_text = _extract_pdf_text(raw_bytes)
except Exception as exc:
logger.debug("web PDF text extraction failed (%s)", type(exc).__name__)
return "(PDF content could not be read as text)", "", content_type
budget_error = _fetch_budget_exceeded(deadline, cancel_event)
if budget_error is not None:
return budget_error, "", content_type
if not pdf_text:
pdf_text = "(PDF contains no extractable text)"
# Report the true type even for a mislabeled body so the caller's "html"
# check routes the extracted text to the plain-text path, not html_to_markdown.
return None, pdf_text, "application/pdf"
# Reject known-binary MIME types before decoding. Binary is returned as the
# error string so the caller surfaces the placeholder, not replacement chars.

55
studio/backend/core/rag/parsers.py
View file

@ -103,7 +103,7 @@ def _markdown_incomplete(markdown: str, plain: str) -> bool:
return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters
def _pdf_markdown(doc) -> list[str] | None:
def _pdf_markdown(doc, pages: range | None = None) -> list[str] | None:
"""Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index
i maps to page i+1. Returns None when the lib is missing, extraction fails, or the
page count does not line up, so the caller falls back to plain PyMuPDF text."""
@ -112,28 +112,44 @@ def _pdf_markdown(doc) -> list[str] | None:
except Exception:
return None
try:
chunks = pymupdf4llm.to_markdown(
doc,
page_chunks = True,
show_progress = False,
)
kwargs = {"page_chunks": True, "show_progress": False}
if pages is not None:
kwargs["pages"] = list(pages)
chunks = pymupdf4llm.to_markdown(doc, **kwargs)
except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion
logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True)
return None
if not isinstance(chunks, list) or len(chunks) != doc.page_count:
expected_pages = doc.page_count if pages is None else len(pages)
if not isinstance(chunks, list) or len(chunks) != expected_pages:
return None
return [str(c.get("text") or "") for c in chunks]
def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
def _pdf(
source: str | bytes,
want_images: bool,
max_pages: int | None = None,
) -> tuple[list[Page], list[ParsedImage], int]:
import fitz # PyMuPDF
pages: list[Page] = []
images: list[ParsedImage] = []
doc = fitz.open(path)
doc = (
fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source)
)
try:
md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
for i, page in enumerate(doc):
if doc.needs_pass:
raise ValueError("encrypted PDF requires a password")
total_pages = doc.page_count
page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages))
if not config.PDF_MARKDOWN:
md = None
elif max_pages is None:
md = _pdf_markdown(doc)
else:
md = _pdf_markdown(doc, page_numbers)
for i, page_number in enumerate(page_numbers):
page = doc[page_number]
plain = page.get_text("text") or ""
candidate = md[i] if md else ""
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval),
@ -147,7 +163,7 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
text = candidate
else:
text = plain
pages.append(_page(text, i + 1))
pages.append(_page(text, page_number + 1))
if want_images:
for img in page.get_images(full = True):
xref = img[0]
@ -161,13 +177,22 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
images.append(
ParsedImage(
image_bytes = image_bytes,
page_number = i + 1,
page_number = page_number + 1,
xref = xref,
)
)
finally:
doc.close()
return pages, images
return pages, images, total_pages
def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]:
"""Extract PDF pages from an in-memory download using the ingestion parser.
Returns the (capped) pages plus the document's full page count, so a caller
that set ``max_pages`` can tell a fully-read short PDF from a truncated one."""
pages, _images, total_pages = _pdf(data, want_images = False, max_pages = max_pages)
return pages, total_pages
def _merge_rects(boxes: list) -> list:
@ -416,7 +441,7 @@ def parse(path: str, *, want_images: bool = False):
ext = os.path.splitext(path)[1].lower()
if ext == ".pdf":
pages, images = _pdf(path, want_images)
pages, images, _total = _pdf(path, want_images)
return (pages, images) if want_images else pages
if ext == ".docx":

270
studio/backend/hub/services/download_lifecycle.py
View file

@ -8,6 +8,7 @@ import os
import signal
import subprocess
import sys
import time
import threading
from pathlib import Path
from typing import Callable, Optional
@ -210,7 +211,9 @@ def finalize_worker_exit(
repo_type: Optional[RepoType] = None,
repo_id: Optional[str] = None,
transport: Optional[str] = None,
) -> None:
cancel_marker_transport: Optional[str] = None,
defer_error: bool = False,
) -> str:
"""Block until *proc* exits, then record the job's terminal state in
*registry*. Drains and scrubs stderr first, then classifies the exit code.
A no-op when the process was already dropped (e.g. superseded).
@ -222,7 +225,7 @@ def finalize_worker_exit(
rc = proc.wait()
cancel_requested = registry.cancel_requested(key)
if not registry.drop_process(key, proc):
return
return "idle"
stderr_text = download_registry.scrub_secrets(
(stderr_data or b"").decode("utf-8", "replace").strip(),
hf_token = hf_token,
@ -230,6 +233,8 @@ def finalize_worker_exit(
state = classify_exit(rc, cancel_requested = cancel_requested)
if state == "complete":
registry.set_job(key, "complete")
if transport == download_registry.TRANSPORT_HTTP:
registry.update_job_transport(key, download_registry.TRANSPORT_HTTP)
if stderr_text:
if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text:
logger.warning(
@ -262,18 +267,226 @@ def finalize_worker_exit(
metadata.variant
if metadata is not None and metadata.variant
else download_registry.variant_from_key(key),
transport,
cancel_marker_transport or transport,
logger = logger,
)
else:
registry.set_job(
key,
"error",
stderr_text or f"worker exited with code {rc}",
)
if not defer_error:
registry.set_job(
key,
"error",
stderr_text or f"worker exited with code {rc}",
)
logger.error(
f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}",
)
return state
def _set_retry_failure_state(
registry: download_registry.DownloadRegistry,
key: str,
error: str,
*,
repo_type: RepoType,
repo_id: str,
fallback_variant: Optional[str],
fallback_transport: Optional[str],
logger,
) -> str:
state, metadata = registry.set_error_unless_cancelled(key, error)
if state == "cancelled":
download_registry.persist_cancel_marker(
repo_type,
repo_id,
metadata.variant if metadata is not None and metadata.variant else fallback_variant,
metadata.transport
if metadata is not None and metadata.transport
else fallback_transport,
logger = logger,
)
return state
def _try_http_retry(
registry: download_registry.DownloadRegistry,
key: str,
*,
hf_token: Optional[str],
label: str,
log_prefix: str,
logger,
repo_type: RepoType,
repo_id: str,
watch_name: str,
) -> bool:
"""Reclaim *key* with HTTP transport and spawn a recovery worker.
Returns ``True`` when the HTTP worker was successfully registered.
Caller is responsible for ensuring this is only called when: the job is
in ``"error"`` state, the original transport was XET, and HTTP is available.
Derives variant and blob-hash metadata from the registry entry written by
the original XET claim so callers do not re-construct worker arguments.
Re-queries peer protection hashes at spawn time to reflect any concurrent
sibling changes between the XET failure and this call.
"""
original_metadata = registry.get_job_metadata(key)
if original_metadata is None:
logger.debug("%s XET retry skipped for %s; metadata unavailable", log_prefix, label)
_set_retry_failure_state(
registry,
key,
"XET retry skipped: metadata unavailable",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = download_registry.variant_from_key(key),
fallback_transport = download_registry.TRANSPORT_XET,
logger = logger,
)
return False
if original_metadata.transport != download_registry.TRANSPORT_XET:
logger.debug(
"%s XET retry skipped for %s; original transport was %s",
log_prefix,
label,
original_metadata.transport,
)
_set_retry_failure_state(
registry,
key,
f"XET retry skipped: original transport was {original_metadata.transport}",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = original_metadata.variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
variant = original_metadata.variant
blob_hashes = original_metadata.blob_hashes
progress_blob_hashes = original_metadata.progress_blob_hashes
completed_baseline_bytes = (
download_registry.completed_blob_bytes(
repo_type,
repo_id,
progress_blob_hashes,
)
if progress_blob_hashes
else 0
)
generation = registry.current_generation(key)
registry.release_active_slot(key)
while True:
if registry.cancel_requested(key):
_set_retry_failure_state(
registry,
key,
"HTTP retry cancelled before reclaiming the download slot",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
claimed, conflict_state = registry.claim(
key,
download_registry.TRANSPORT_HTTP,
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
blob_hashes = blob_hashes,
progress_blob_hashes = progress_blob_hashes,
completed_baseline_bytes = completed_baseline_bytes,
generation = generation,
replace_active = True,
cancel_marker_transport = original_metadata.transport,
)
if claimed:
break
if conflict_state == "deleting":
logger.debug(
"%s XET retry claim rejected for %s; repo is being deleted",
log_prefix,
label,
)
_set_retry_failure_state(
registry,
key,
"HTTP retry could not reclaim the download slot",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
logger.debug(
"%s XET retry claim blocked for %s by active sibling state %s; waiting",
log_prefix,
label,
conflict_state,
)
time.sleep(0.05)
args: list[str] = ["--repo-id", repo_id]
if repo_type == "dataset":
args.append("--dataset")
elif variant:
args.extend(["--variant", variant])
# Re-query at spawn time: sibling state may have changed since XET failed.
peer_hashes = registry.peer_blob_hashes(key) if variant else frozenset()
logger.warning(
"%s XET worker failed for %s; retrying over HTTP",
log_prefix,
label,
)
try:
proc = spawn_worker(
args,
hf_token,
use_xet = False,
protected_blob_hashes = peer_hashes or None,
)
except Exception as exc:
scrubbed = download_registry.scrub_secrets(str(exc), hf_token = hf_token)
logger.error(
"%s HTTP retry spawn failed for %s: %s",
log_prefix,
label,
scrubbed,
)
registry.update_job_transport(key, original_metadata.transport)
_set_retry_failure_state(
registry,
key,
scrubbed,
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
return register_worker(
registry,
key,
proc,
hf_token = hf_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = download_registry.TRANSPORT_HTTP,
cancel_marker_transport = original_metadata.transport,
watch_name = watch_name,
)
def kill_and_reap_process(
@ -309,6 +522,7 @@ def register_worker(
repo_type: RepoType,
repo_id: str,
transport: str,
cancel_marker_transport: Optional[str] = None,
watch_name: str,
) -> bool:
if not registry.register_process(key, proc):
@ -319,7 +533,14 @@ def register_worker(
def _watch() -> None:
try:
finalize_worker_exit(
can_retry_http = (
transport == download_registry.TRANSPORT_XET
and download_registry.download_transport_unavailable_reason(
download_registry.TRANSPORT_HTTP
)
is None
)
state = finalize_worker_exit(
registry,
key,
proc,
@ -330,7 +551,25 @@ def register_worker(
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
cancel_marker_transport = cancel_marker_transport,
defer_error = can_retry_http,
)
# XET-to-HTTP recovery: when a non-cancelled XET worker fails and
# HTTP is available, attempt one automatic retry over HTTP. The
# transport check is the recursion guard: an HTTP worker that errors
# never satisfies `transport == TRANSPORT_XET`, so it stays terminal.
if can_retry_http and state == "error":
_try_http_retry(
registry,
key,
hf_token = worker_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
watch_name = watch_name,
)
except Exception:
# finalize_worker_exit is the only thing that clears running/cancelling;
# if it raises, force a terminal state so claim() isn't blocked until restart.
@ -426,8 +665,19 @@ def cancel_worker(
return "cancelling"
return registry.get_job(key).state
# Worker already exited; let its watcher classify the real return code.
# Arming a pending cancel here could mislabel a genuine failure as a cancel.
if proc.poll() is not None:
get_metadata = getattr(registry, "get_job_metadata", None)
metadata = get_metadata(key) if get_metadata is not None else None
can_retry_http = (
metadata is not None
and metadata.transport == download_registry.TRANSPORT_XET
and download_registry.download_transport_unavailable_reason(
download_registry.TRANSPORT_HTTP
)
is None
)
if can_retry_http and registry.mark_pending_cancel(key, generation):
return "cancelling"
return registry.get_job(key).state
if not registry.request_cancel(key, proc, generation):

30
studio/backend/hub/services/models/downloads.py
View file

@ -60,6 +60,30 @@ def _job_status(
return DownloadJobStatus(state = state, error = error, generation = generation)
def _load_in_flight(repo_id: str) -> bool:
try:
from core.inference.llama_cpp import hf_gguf_load_in_flight
return hf_gguf_load_in_flight(repo_id)
except Exception:
return False
def _load_in_flight_error(repo_id: str) -> HTTPException:
return HTTPException(
status_code = 409,
detail = (
f"A model load for '{repo_id}' is in progress and may be "
"downloading it. Wait for the load to finish (or cancel it), "
"then start the download."
),
)
def _reject_if_load_in_flight(repo_id: str) -> None:
if _load_in_flight(repo_id):
raise _load_in_flight_error(repo_id)
def _spawn_download_worker(
repo_id: str,
variant: Optional[str],
@ -89,6 +113,9 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
# Canonicalize so two different-cased paste-ins share one job + cache dir.
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
# Avoid concurrent writers to the same HF cache files.
_reject_if_load_in_flight(repo_id)
variant = (body.gguf_variant or "").strip() or None
if variant is not None and not _is_valid_gguf_variant(variant):
raise HTTPException(
@ -147,9 +174,12 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
blob_hashes = variant_blob_hashes,
progress_blob_hashes = variant_progress_blob_hashes,
completed_baseline_bytes = completed_baseline_bytes,
admission_check = lambda: not _load_in_flight(repo_id),
)
generation = _registry.current_generation(key)
if not claimed:
if claim_state == "admission_blocked":
raise _load_in_flight_error(repo_id)
# claim_state is the blocking job's state. The client can attach only
# when the blocker is this key's own in-flight job (adoptable); a
# cross-variant conflict or in-progress delete is not accepted.

158
studio/backend/hub/tests/test_download_lifecycle.py
View file

@ -1,27 +1,147 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import io
import logging
from hub.services import download_lifecycle
from hub.utils import download_registry, state_dir
def _set_xet_reason(monkeypatch, reason):
class _Proc:
pid = 4242
def __init__(
self,
rc,
stderr = b"",
):
self.rc = rc
self.stderr = io.BytesIO(stderr)
self.waited = False
def poll(self):
return self.rc if self.waited else None
def wait(self, timeout = None):
self.waited = True
return self.rc
def kill(self):
pass
class _ImmediateThread:
def __init__(self, *, target, **_kwargs):
self.target = target
def start(self):
self.target()
def test_resolve_effective_use_xet(monkeypatch):
for requested, unavailable_reason, expected in (
(False, "unused", False),
(True, None, True),
(True, "hf_xet is not installed", False),
):
monkeypatch.setattr(
download_lifecycle.download_registry,
"download_transport_unavailable_reason",
lambda _transport, reason = unavailable_reason: reason,
)
assert download_lifecycle.resolve_effective_use_xet(requested) is expected
def test_xet_failure_retries_over_http_for_model_and_dataset(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread)
register_worker = download_lifecycle.register_worker
for repo_type, repo_id, variant, expected_args in (
("model", "Org/Model", "Q4_K_M", ["--repo-id", "Org/Model", "--variant", "Q4_K_M"]),
("dataset", "Org/Data", None, ["--repo-id", "Org/Data", "--dataset"]),
):
registry = download_registry.DownloadRegistry()
key = download_registry.normalize_job_key(f"{repo_id}::{variant}" if variant else repo_id)
assert registry.claim(
key,
download_registry.TRANSPORT_XET,
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
blob_hashes = frozenset({"blob"}),
)[0]
generation = registry.current_generation(key)
spawned = []
def fake_spawn(
args,
_token,
*,
use_xet,
protected_blob_hashes = None,
):
spawned.append((args, use_xet, protected_blob_hashes))
return _Proc(0)
def fake_retry_register(*_args, **kwargs):
assert kwargs["transport"] == download_registry.TRANSPORT_HTTP
return True
monkeypatch.setattr(download_lifecycle, "spawn_worker", fake_spawn)
monkeypatch.setattr(download_lifecycle, "register_worker", fake_retry_register)
assert register_worker(
registry,
key,
_Proc(1, b"xet failed"),
hf_token = None,
label = repo_id,
log_prefix = "Download",
logger = logging.getLogger("test"),
repo_type = repo_type,
repo_id = repo_id,
transport = download_registry.TRANSPORT_XET,
watch_name = f"{repo_type}-watch",
)
metadata = registry.get_job_metadata(key)
assert spawned == [(expected_args, False, None)]
assert metadata.transport == download_registry.TRANSPORT_HTTP
assert metadata.blob_hashes == frozenset({"blob"})
assert registry.current_generation(key) == generation
def test_http_failure_remains_terminal(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread)
register_worker = download_lifecycle.register_worker
registry = download_registry.DownloadRegistry()
key = download_registry.normalize_repo_key("Org/Data")
assert registry.claim(
key,
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "Org/Data",
)[0]
monkeypatch.setattr(
download_lifecycle.download_registry,
"download_transport_unavailable_reason",
lambda _transport: reason,
download_lifecycle,
"register_worker",
lambda *_args, **_kwargs: (_ for _ in ()).throw(
AssertionError("HTTP failures must not retry")
),
)
def test_resolve_effective_use_xet_keeps_http_when_not_requested(monkeypatch):
_set_xet_reason(monkeypatch, "should not be consulted")
assert download_lifecycle.resolve_effective_use_xet(False) is False
def test_resolve_effective_use_xet_keeps_xet_when_available(monkeypatch):
_set_xet_reason(monkeypatch, None)
assert download_lifecycle.resolve_effective_use_xet(True) is True
def test_resolve_effective_use_xet_downgrades_when_xet_unavailable(monkeypatch):
_set_xet_reason(monkeypatch, "Xet transport is unavailable because hf_xet is not installed.")
assert download_lifecycle.resolve_effective_use_xet(True) is False
assert register_worker(
registry,
key,
_Proc(1, b"http failed"),
hf_token = None,
label = "Org/Data",
log_prefix = "Download",
logger = logging.getLogger("test"),
repo_type = "dataset",
repo_id = "Org/Data",
transport = download_registry.TRANSPORT_HTTP,
watch_name = "dataset-watch",
)
assert registry.get_job(key).state == "error"

200
studio/backend/hub/utils/download_registry.py
View file

@ -45,9 +45,9 @@ import sys
import threading
import time
import weakref
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Iterator, Literal, Optional
from typing import Callable, Iterator, Literal, Optional
from loggers import get_logger
@ -126,6 +126,9 @@ def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMeta
"repo_id": metadata.repo_id if metadata is not None else None,
"variant": metadata.variant if metadata is not None else None,
"transport": metadata.transport if metadata is not None else None,
"cancel_marker_transport": metadata.cancel_marker_transport
if metadata is not None
else None,
}
tmp = path.with_name(f".{path.name}.tmp-{pid}")
try:
@ -305,7 +308,7 @@ def reap_orphan_workers() -> None:
data.get("repo_type"),
repo_id,
data.get("variant"),
data.get("transport"),
data.get("cancel_marker_transport") or data.get("transport"),
)
except Exception as exc:
logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc)
@ -699,6 +702,7 @@ class DownloadMetadata:
repo_id: str
variant: Optional[str]
transport: Optional[str]
cancel_marker_transport: Optional[str] = None
# GGUF variant main/writable hashes, identifying the variant-specific shards
# for concurrency decisions.
blob_hashes: frozenset[str] = field(default_factory = frozenset)
@ -801,6 +805,7 @@ class DownloadRegistry:
self._processes: dict[str, subprocess.Popen] = {}
self._repo_active: dict[str, set[str]] = {}
self._metadata: dict[str, DownloadMetadata] = {}
self._cancel_marker_transports: dict[str, str] = {}
self._pending_cancel: dict[str, Optional[int]] = {}
self._generations: dict[str, int] = {}
# Monotonic across keys so an evicted then re-claimed key never reuses a
@ -839,6 +844,7 @@ class DownloadRegistry:
if state in TERMINAL_STATES:
self._put_terminal_job_locked(key, state, error)
self._pending_cancel.pop(key, None)
self._cancel_marker_transports.pop(key, None)
repo = _repo_of_key(key)
active = self._repo_active.get(repo)
if active is not None:
@ -848,6 +854,57 @@ class DownloadRegistry:
else:
self._jobs[key] = DownloadState(state, error)
def set_error_unless_cancelled(
self, key: str, error: str
) -> tuple[JobState, Optional[DownloadMetadata]]:
key = normalize_job_key(key)
with self._lock:
current = self._jobs.get(key, DownloadState("idle")).state
has_pending_cancel = key in self._pending_cancel
pending_generation = self._pending_cancel.get(key)
metadata = self._metadata.get(key)
should_cancel = current == "cancelling" or (
has_pending_cancel and self._generation_matches_locked(key, pending_generation)
)
terminal_state: JobState = "cancelled" if should_cancel else "error"
marker_transport = self._cancel_marker_transports.pop(key, None)
if marker_transport is None and metadata is not None:
marker_transport = metadata.cancel_marker_transport
self._put_terminal_job_locked(
key,
terminal_state,
None if should_cancel else error,
)
self._pending_cancel.pop(key, None)
repo = _repo_of_key(key)
active = self._repo_active.get(repo)
if active is not None:
active.discard(key)
if not active:
self._repo_active.pop(repo, None)
if should_cancel and metadata is not None and marker_transport is not None:
metadata = replace(metadata, transport = marker_transport)
return terminal_state, metadata
def update_job_transport(self, key: str, transport: str) -> None:
key = normalize_job_key(key)
with self._lock:
metadata = self._metadata.get(key)
if metadata is None or metadata.transport == transport:
return
self._metadata[key] = replace(metadata, transport = transport)
def release_active_slot(self, key: str) -> None:
key = normalize_job_key(key)
repo = _repo_of_key(key)
with self._lock:
active = self._repo_active.get(repo)
if active is None:
return
active.discard(key)
if not active:
self._repo_active.pop(repo, None)
def get_job(self, key: str) -> DownloadState:
key = normalize_job_key(key)
with self._lock:
@ -884,6 +941,14 @@ class DownloadRegistry:
):
self._put_terminal_job_locked(key, "cancelled")
metadata_to_persist = self._metadata.pop(key, None)
marker_transport = self._cancel_marker_transports.pop(key, None)
if marker_transport is None and metadata_to_persist is not None:
marker_transport = metadata_to_persist.cancel_marker_transport
if metadata_to_persist is not None and marker_transport is not None:
metadata_to_persist = replace(
metadata_to_persist,
transport = marker_transport,
)
repo = _repo_of_key(key)
active = self._repo_active.get(repo)
if active is not None:
@ -963,12 +1028,24 @@ class DownloadRegistry:
blob_hashes: Optional[frozenset[str]] = None,
progress_blob_hashes: Optional[frozenset[str]] = None,
completed_baseline_bytes: int = 0,
admission_check: Optional[Callable[[], bool]] = None,
generation: Optional[int] = None,
replace_active: bool = False,
metadata_transport: Optional[str] = None,
cancel_marker_transport: Optional[str] = None,
) -> tuple[bool, str]:
key = normalize_job_key(key)
repo = _repo_of_key(key)
requested_hashes = blob_hashes or frozenset()
requested_progress_hashes = progress_blob_hashes or frozenset()
with self._lock:
# Run the final external admission check while the registry lock is
# held, immediately before inspecting and publishing active state.
# The GGUF load path establishes its marker before calling
# its active-job probe, so either this claim observes that marker
# or the load's later probe observes this claim.
if admission_check is not None and not admission_check():
return False, "admission_blocked"
deleting_scopes = self._deleting.get(repo)
if deleting_scopes is not None and (
None in deleting_scopes or variant_from_key(key) in deleting_scopes
@ -1007,10 +1084,13 @@ class DownloadRegistry:
if conflict_state is not None:
return False, conflict_state
current = self._jobs.get(key, DownloadState("idle")).state
if current in _ACTIVE_STATES:
if current in _ACTIVE_STATES and not replace_active:
return False, current
self._generation_seq += 1
self._generations[key] = self._generation_seq
if generation is None:
self._generation_seq += 1
self._generations[key] = self._generation_seq
else:
self._generations[key] = generation
self._jobs[key] = DownloadState("running")
self._repo_active.setdefault(repo, active).add(key)
if repo_type and repo_id:
@ -1018,7 +1098,8 @@ class DownloadRegistry:
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
transport = transport,
transport = metadata_transport if metadata_transport is not None else transport,
cancel_marker_transport = cancel_marker_transport,
blob_hashes = requested_hashes,
progress_blob_hashes = requested_progress_hashes,
completed_baseline_bytes = max(
@ -1026,8 +1107,13 @@ class DownloadRegistry:
int(completed_baseline_bytes or 0),
),
)
if cancel_marker_transport is not None:
self._cancel_marker_transports[key] = cancel_marker_transport
else:
self._cancel_marker_transports.pop(key, None)
else:
self._metadata.pop(key, None)
self._cancel_marker_transports.pop(key, None)
return True, "running"
def adoptable(self, key: str) -> bool:
@ -1053,7 +1139,8 @@ class DownloadRegistry:
download. A variant delete conflicts only with that same variant or a
whole-repo download writing the shared snapshot; other quantizations
download concurrently and never block it."""
for key in self._repo_active.get(repo_id, set()):
active_keys = self._repo_active.get(repo_id, set())
for key in active_keys:
job = self._jobs.get(key)
if job is None or job.state not in _ACTIVE_STATES:
continue
@ -1062,6 +1149,16 @@ class DownloadRegistry:
other_variant = self._active_job_variant_locked(key)
if other_variant is None or other_variant == variant:
return True
for key, job in self._jobs.items():
if key in active_keys or _repo_of_key(key) != repo_id:
continue
if job.state not in _ACTIVE_STATES:
continue
if variant is None:
return True
other_variant = self._active_job_variant_locked(key)
if other_variant is None or other_variant == variant:
return True
return False
def peer_blob_hashes(self, key: str) -> frozenset[str]:
@ -1108,6 +1205,16 @@ class DownloadRegistry:
candidate_keys = list(self._repo_active.get(repo_key, set()))
else:
candidate_keys = [key for active in self._repo_active.values() for key in active]
# An XET->HTTP retry handoff briefly drops its key from _repo_active
# while its job stays active; include those released-but-active jobs
# so the waiting retry still lists and can be adopted or cancelled.
seen = set(candidate_keys)
for key, job in self._jobs.items():
if key in seen or job.state not in _ACTIVE_STATES:
continue
if repo_key is not None and _repo_of_key(key) != repo_key:
continue
candidate_keys.append(key)
refs: list[ActiveDownloadRef] = []
for key in candidate_keys:
job = self._jobs.get(key)
@ -1123,6 +1230,23 @@ class DownloadRegistry:
)
return refs
def has_active_variant(self, repo_id: str, variant: Optional[str]) -> bool:
"""Whether an active model job targets this exact GGUF variant.
Scans the job table rather than only ``_repo_active`` so an XET-to-HTTP
retry handoff remains visible while it has temporarily released its
active slot.
"""
repo_key = normalize_repo_key(repo_id)
target = (variant or "").strip().lower() or None
with self._lock:
for key, job in self._jobs.items():
if _repo_of_key(key) != repo_key or job.state not in _ACTIVE_STATES:
continue
if self._active_job_variant_locked(key) == target:
return True
return False
def begin_delete(
self,
repo_id: str,
@ -1169,12 +1293,25 @@ class DownloadRegistry:
repo_id = normalize_repo_key(repo_id)
target = (variant or "").strip().lower() or None
with self._lock:
for key in self._repo_active.get(repo_id, set()):
active_keys = self._repo_active.get(repo_id, set())
for key in active_keys:
job = self._jobs.get(key)
if job is None or job.state not in _ACTIVE_STATES:
continue
if self._active_job_variant_locked(key) != target:
return True
# An XET->HTTP retry peer between release_active_slot() and its reclaim
# is briefly absent from _repo_active while its job stays active and
# still owns the shared companion; mirror the released-but-active scan
# used by _delete_blocked_by_active_locked so it still blocks companion
# deletion of a different variant.
for key, job in self._jobs.items():
if key in active_keys or _repo_of_key(key) != repo_id:
continue
if job.state not in _ACTIVE_STATES:
continue
if self._active_job_variant_locked(key) != target:
return True
return False
def request_cancel(
@ -1198,17 +1335,58 @@ class DownloadRegistry:
return True
def terminate_all(self, kind: str = "download") -> None:
settled_no_proc: list[Optional[DownloadMetadata]] = []
with self._lock:
live = [
(key, proc, self._metadata.get(key))
for key, proc in self._processes.items()
if proc.poll() is None
]
live_keys = {key for key, _proc, _metadata in live}
# Flag as an intentional stop so the watcher's exit classification
# reports them cancelled rather than an OOM/crash once SIGKILL lands.
for key, _proc, _metadata in live:
if self._jobs.get(key, DownloadState("idle")).state == "running":
self._jobs[key] = DownloadState("cancelling")
# Settle active jobs without a live worker too. Two cases: an
# XET->HTTP retry parked in the reclaim wait loop has dropped its
# worker and slot guard, so it is absent from `live`; and a
# registered worker that already exited with an error but whose
# watcher has not yet run would otherwise stay `running` and spawn an
# HTTP retry after this shutdown snapshot. Skip a registered worker
# that exited cleanly (rc == 0): it completed and the watcher will
# mark it done, so marking it cancelling would strand a stale marker.
for key, job in list(self._jobs.items()):
if job.state not in _ACTIVE_STATES or key in live_keys:
continue
proc = self._processes.get(key)
if proc is not None:
if proc.poll() == 0:
continue
# A registered worker that exited nonzero on its own over HTTP
# is a genuine terminal download failure, not a shutdown cancel
# and not retry-capable: leave its error status intact rather
# than persisting a cancel marker that would read as
# cancelled/resumable after restart. Only an exited XET worker
# could still spawn a post-shutdown HTTP retry, so only that
# needs settling here.
metadata = self._metadata.get(key)
if metadata is not None and metadata.transport == TRANSPORT_HTTP:
continue
self._pending_cancel[key] = self._generations.get(key)
self._jobs[key] = DownloadState("cancelling")
settled_no_proc.append(self._metadata.get(key))
# Persist a cancel marker for each settled no-live-worker job outside the
# lock (mirroring the reaped path) so shutdown records resumable/cancelled
# state even if it returns before the daemon watcher wakes to do so.
for metadata in settled_no_proc:
if metadata is not None:
persist_cancel_marker(
metadata.repo_type,
metadata.repo_id,
metadata.variant,
metadata.cancel_marker_transport or metadata.transport,
)
reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = []
for key, proc, metadata in live:
try:
@ -1222,7 +1400,7 @@ class DownloadRegistry:
metadata.repo_type,
metadata.repo_id,
metadata.variant,
metadata.transport,
metadata.cancel_marker_transport or metadata.transport,
)
continue
reaped.append((key, proc, metadata))
@ -1242,7 +1420,7 @@ class DownloadRegistry:
metadata.repo_type,
metadata.repo_id,
metadata.variant,
metadata.transport,
metadata.cancel_marker_transport or metadata.transport,
)

17
studio/backend/main.py
View file

@ -612,6 +612,22 @@ app = FastAPI(
lifespan = lifespan,
)
# The MCP surface is opt-in because it can start GPU jobs and write model
# artifacts. Mount it only when explicitly enabled by the Studio process.
if os.environ.get("UNSLOTH_STUDIO_ENABLE_MCP") == "1":
from fastmcp.utilities.lifespan import combine_lifespans
from mcp_server import BearerTokenMiddleware, create_studio_mcp
_studio_mcp_app = create_studio_mcp().http_app(path = "/")
_studio_mcp_lifespan = _studio_mcp_app.lifespan
_mcp_token = os.environ.get("UNSLOTH_STUDIO_MCP_TOKEN")
if not _mcp_token:
raise RuntimeError("UNSLOTH_STUDIO_MCP_TOKEN is required when MCP is enabled")
_studio_mcp_app = BearerTokenMiddleware(_studio_mcp_app, _mcp_token)
app.router.lifespan_context = combine_lifespans(lifespan, _studio_mcp_lifespan)
app.mount("/mcp", _studio_mcp_app)
from loggers.config import LogConfig
from loggers.handlers import LoggingMiddleware
@ -752,6 +768,7 @@ _BODY_PROTECTED_PREFIXES = (
"/api/settings",
"/api/train",
"/api/export",
"/mcp",
)
_DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload"
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = (

259
studio/backend/mcp_server.py Normal file
View file

@ -0,0 +1,259 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Curated MCP tools for driving an Unsloth Studio instance.
The MCP surface deliberately wraps the existing Studio services instead of
duplicating training or export logic. It is opt-in because several tools can
start GPU work or write model artifacts.
"""
from __future__ import annotations
import hmac
from typing import Any
from fastmcp import FastMCP
class BearerTokenMiddleware:
"""Require an exact bearer token when Studio MCP is exposed remotely."""
def __init__(self, app: Any, token: str) -> None:
if not token or not token.strip():
raise ValueError("Studio MCP bearer token must be a non-empty value")
if not token.isascii():
# A non-ASCII token cannot be sent in an HTTP header; reject it here.
raise ValueError("Studio MCP bearer token must contain ASCII characters only")
self.app = app
# Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII
# input, which would surface as a 500 instead of a clean 401.
self.expected = token.encode("utf-8")
async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
scope_type = scope.get("type")
if scope_type not in ("http", "websocket"):
await self.app(scope, receive, send)
return
headers = dict(scope.get("headers", []))
raw_auth = headers.get(b"authorization", b"")
scheme, _, supplied = raw_auth.partition(b" ")
if scheme.lower() != b"bearer" or not hmac.compare_digest(supplied, self.expected):
await _send_unauthorized(send, scope_type)
return
await self.app(scope, receive, send)
async def _send_unauthorized(send: Any, scope_type: str) -> None:
if scope_type == "websocket":
await send({"type": "websocket.close", "code": 4401})
return
await send(
{
"type": "http.response.start",
"status": 401,
"headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")],
}
)
await send(
{
"type": "http.response.body",
"body": b'{"detail":"MCP bearer token required"}',
}
)
def _dump(value: Any) -> Any:
"""Convert Pydantic responses to plain JSON values for MCP clients."""
if hasattr(value, "model_dump"):
return value.model_dump(mode = "json")
return value
def _clamp(value: int, low: int, high: int) -> int:
"""Clamp an MCP-supplied integer into an inclusive range.
MCP tools call the Studio route functions directly, which skips FastAPI's
Query(ge=, le=) validation, so we re-apply the same bounds here.
"""
return max(low, min(value, high))
def create_studio_mcp() -> FastMCP:
"""Create the Studio MCP server and register the high-value tools."""
mcp = FastMCP(
"Unsloth Studio",
instructions = (
"Use read tools to inspect the local Studio state before starting GPU work. "
"Training and export tools can consume substantial VRAM and write files. "
"Never expose tokens or local paths from tool results unless the user asks."
),
)
@mcp.tool
async def studio_status() -> dict[str, Any]:
"""Return the current training, export, inference, and GPU state."""
from routes.export import get_export_status
from routes.inference import get_status as get_inference_status
from routes.training import get_training_status
from utils.hardware import get_gpu_utilization
training, export, inference = await _gather_status(
get_training_status(current_subject = "mcp"),
get_export_status(current_subject = "mcp"),
get_inference_status(current_subject = "mcp"),
)
return {
"training": _dump(training),
"export": _dump(export),
"inference": _dump(inference),
"hardware": get_gpu_utilization(),
}
@mcp.tool
async def list_local_models(models_dir: str = "./models") -> dict[str, Any]:
"""List local and cached models available to Studio."""
from routes.models import list_local_models as list_models
return _dump(await list_models(models_dir = models_dir, current_subject = "mcp"))
@mcp.tool
async def get_training_status() -> dict[str, Any]:
"""Read the active training job, phase, progress, and recent metrics."""
from routes.training import get_training_status as get_status
return _dump(await get_status(current_subject = "mcp"))
@mcp.tool
async def start_training(config: dict[str, Any]) -> dict[str, Any]:
"""Start a validated Studio training job from a TrainingStartRequest-shaped object.
The config is validated by the same Pydantic model used by the Studio UI.
Call get_training_status first and do not start work while another job runs.
"""
from models import TrainingStartRequest
from routes.training import start_training as start
request = TrainingStartRequest.model_validate(config)
# Pass via_api_key explicitly (a direct call leaves it a Depends object).
# MCP drives Studio like the UI session, so it coexists and frees VRAM.
return _dump(await start(request, current_subject = "mcp", via_api_key = False))
@mcp.tool
async def stop_training(save: bool = True) -> dict[str, Any]:
"""Ask the active training process to stop at its next safe checkpoint."""
from routes.training import TrainingStopRequest, stop_training as stop
return _dump(await stop(TrainingStopRequest(save = save), current_subject = "mcp"))
@mcp.tool
async def list_training_runs(limit: int = 50, offset: int = 0) -> dict[str, Any]:
"""List completed and stopped training runs, newest first."""
from routes.training_history import list_training_runs as list_runs
# Clamp here (direct call skips Query bounds); a negative LIMIT = no limit.
limit = _clamp(limit, 1, 200)
offset = max(0, offset)
return _dump(await list_runs(limit = limit, offset = offset, current_subject = "mcp"))
@mcp.tool
def validate_recipe(recipe: dict[str, Any]) -> dict[str, Any]:
"""Validate a Data Recipe with the same validator used by Studio."""
from models.data_recipe import RecipePayload
from routes.data_recipe.validate import validate
return _dump(validate(RecipePayload(recipe = recipe)))
@mcp.tool
def get_recipe_job_status(job_id: str) -> dict[str, Any]:
"""Read the status of a Data Recipe job."""
from routes.data_recipe.jobs import job_status
return _dump(job_status(job_id))
@mcp.tool
def get_recipe_job_dataset(
job_id: str,
limit: int = 20,
offset: int = 0,
) -> dict[str, Any]:
"""Read a bounded page of generated Data Recipe rows."""
from routes.data_recipe.jobs import job_dataset
# Clamp here (direct call skips FastAPI's Query bounds).
limit = _clamp(limit, 1, 500)
offset = max(0, offset)
return _dump(job_dataset(job_id, limit = limit, offset = offset))
@mcp.tool
async def load_checkpoint(
checkpoint_path: str,
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
approved_remote_code_fingerprint: str | None = None,
hf_token: str | None = None,
) -> dict[str, Any]:
"""Load a checkpoint into the export backend.
Export runs in its own subprocess and coexists with training and
inference; it does not unload them, so a load can fail with a clear
out-of-memory error if the GPU is already full. Pass hf_token to load a
gated checkpoint, and approved_remote_code_fingerprint to retry a
trust_remote_code load that was blocked pending review.
"""
from models import LoadCheckpointRequest
from routes.export import load_checkpoint as load
request = LoadCheckpointRequest(
checkpoint_path = checkpoint_path,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
approved_remote_code_fingerprint = approved_remote_code_fingerprint,
hf_token = hf_token,
)
return _dump(await load(request, current_subject = "mcp"))
@mcp.tool
async def export_gguf(
save_directory: str,
quantization_method: str | list[str] = "Q4_K_M",
push_to_hub: bool = False,
repo_id: str | None = None,
hf_token: str | None = None,
imatrix: bool = False,
imatrix_path: str | None = None,
) -> dict[str, Any]:
"""Export the loaded model to GGUF using Studio's existing path validation.
quantization_method may be a single method or a list to produce several
GGUFs from one load. Pass hf_token when push_to_hub is set (the backend
rejects a Hub upload without it). Set imatrix (or imatrix_path) for the
IQ low-bit quants that require an importance matrix.
"""
from models import ExportGGUFRequest
from routes.export import export_gguf as export
request = ExportGGUFRequest(
save_directory = save_directory,
quantization_method = quantization_method,
push_to_hub = push_to_hub,
repo_id = repo_id,
hf_token = hf_token,
imatrix = imatrix,
imatrix_path = imatrix_path,
)
return _dump(await export(request, current_subject = "mcp"))
return mcp
async def _gather_status(*coroutines: Any) -> tuple[Any, ...]:
"""Gather independent status calls without letting one optional backend fail all state."""
import asyncio
results = await asyncio.gather(*coroutines, return_exceptions = True)
return tuple(
{"error": str(result)} if isinstance(result, Exception) else result for result in results
)

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

@ -20,6 +20,7 @@ from loggers import get_logger
import asyncio
import threading
import weakref
from contextlib import ExitStack
import re as _re
@ -28,6 +29,7 @@ import re as _re
from utils.models import extract_model_size_b as _extract_model_size_b
from utils.api_errors import openai_error_body, anthropic_error_body
from core.inference.orchestrator import GenStreamError, GenStreamErrorRaised
from core.inference.llama_admission import (
LlamaAdmissionCancelled,
LlamaAdmissionConfig,
@ -212,6 +214,14 @@ def _friendly_error(exc: Exception) -> str:
return "An internal error occurred"
def _friendly_gen_stream_error(value) -> str:
"""Return a client-safe message for typed local generation errors."""
text = str(value)
if getattr(value, "public", False):
return text
return safe_error_detail(RuntimeError(text), fallback = "An internal error occurred.")
def _friendly_upstream_error(text: str) -> str:
"""Rewrite a raw llama-server error body into an actionable message where we can.
@ -998,6 +1008,7 @@ try:
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
extra_args_disable_mmproj,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
@ -1035,6 +1046,7 @@ except ImportError:
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
extra_args_disable_mmproj,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
@ -1915,16 +1927,57 @@ async def artifact_preview_frame(allow_network: bool = False):
_BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:')
def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
def _detect_safetensors_features(
backend,
chat_template: Optional[str],
tools = None,
) -> dict:
"""Classify reasoning/tool capabilities via the GGUF classifier so flags
match across backends. gpt-oss is overridden: Harmony routes reasoning and
tools through tokenizer channels, not template markup."""
model_id = getattr(backend, "active_model_name", None)
feature_template = chat_template
try:
from core.inference.chat_template_helpers import _selected_template_strings_from_value
selected_templates = _selected_template_strings_from_value(chat_template, tools)
if selected_templates:
feature_template = selected_templates[0]
except Exception:
logger.debug("safetensors_named_template_selection_failed", exc_info = True)
flags = detect_reasoning_flags(
chat_template,
feature_template,
model_identifier = model_id,
log_source = "safetensors",
)
if not flags.get("supports_reasoning"):
try:
from core.inference.chat_template_helpers import (
detect_reasoning_channel_markers_from_template,
)
templates = [chat_template]
models = getattr(backend, "models", None)
model_info = (
models.get(model_id, {})
if isinstance(models, dict) and model_id is not None
else {}
)
if isinstance(model_info, dict):
templates.extend(
(
model_info.get("native_chat_template"),
(model_info.get("chat_template_info") or {}).get("template"),
)
)
if any(
detect_reasoning_channel_markers_from_template(template, tools = tools) is not None
for template in templates
):
flags["supports_reasoning"] = True
flags["reasoning_always_on"] = True
logger.info("safetensors: model always reasons (native channel markers)")
except Exception:
logger.debug("safetensors_native_reasoning_marker_check_failed", exc_info = True)
# Markers any supported parser recognises (template advertises tools but
# uses none -> drop the pill). Reuse the parser's own signal list so this
# gate never drifts (a hand-maintained copy lost the DeepSeek variants);
@ -1938,9 +1991,9 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
)
if (
flags.get("supports_tools")
and chat_template
and not any(m in chat_template for m in _PARSER_MARKERS)
and not _BARE_JSON_NAME_MARKER_RE.search(chat_template)
and isinstance(feature_template, str)
and not any(m in feature_template for m in _PARSER_MARKERS)
and not _BARE_JSON_NAME_MARKER_RE.search(feature_template)
):
logger.info(
"safetensors: template advertises tools but uses an "
@ -3968,6 +4021,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
native_grant_backed = False
model_log_label = request.model_path
gguf_load_stack = ExitStack()
try:
# Validate user pass-through args up front so a managed-flag collision
# returns 400 before any model work.
@ -4187,15 +4241,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
llama_backend = get_llama_cpp_backend()
unsloth_backend = get_inference_backend()
# Unload any active Unsloth model to free VRAM (off the event loop:
# unload takes _gen_lock and can wait on an in-flight stream).
if unsloth_backend.active_model_name:
logger.info(
f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF"
)
await asyncio.to_thread(
unsloth_backend.unload_model, unsloth_backend.active_model_name
)
if config.gguf_hf_repo:
from core.inference.llama_cpp import gguf_load_in_flight
gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo))
# Inherit llama_extra_args from the previous load when the request
# omits the field (the chat-settings Apply path doesn't round-trip
@ -4275,6 +4323,38 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
extra_llama_args,
)
# Block cache writes that would race the download manager. This runs
# after pass-through argument inheritance so a carried --no-mmproj
# changes the companion requirement exactly as it does for the load.
if config.gguf_hf_repo:
from core.inference.llama_cpp import _hub_download_blocks_gguf_load
if await asyncio.to_thread(
_hub_download_blocks_gguf_load,
config.gguf_hf_repo,
config.gguf_variant,
require_mmproj = bool(
config.is_vision and not extra_args_disable_mmproj(extra_llama_args)
),
hf_token = request.hf_token,
):
raise HTTPException(
status_code = 409,
detail = (
f"'{model_log_label}' is currently being downloaded "
"by the download manager. Wait for the download to "
"finish (or cancel it), then load the model."
),
)
# Unload any active Unsloth model only after every hub conflict check.
if unsloth_backend.active_model_name:
logger.info(
f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF"
)
await asyncio.to_thread(
unsloth_backend.unload_model, unsloth_backend.active_model_name
)
# Route to HF or local mode based on config. Run in a thread so the
# event loop stays free for progress polling and other requests
# during the (potentially long) GGUF download + llama-server start.
@ -4639,6 +4719,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
logger.error(f"Error loading model: {e}", exc_info = True)
msg = _maybe_unsupported_message(redacted_msg)
raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}")
finally:
gguf_load_stack.close()
def _requires_trust_remote_code_for_model(
@ -5392,6 +5474,10 @@ async def generate_stream(
if chunk is _DONE:
completed = True
break
if isinstance(chunk, GenStreamError):
yield f"data: {json.dumps({'error': _friendly_gen_stream_error(chunk)})}\n\n"
yield "data: [DONE]\n\n"
return
yield f"data: {json.dumps({'content': chunk})}\n\n"
if completed:
yield "data: [DONE]\n\n"
@ -5405,6 +5491,7 @@ async def generate_stream(
backend.reset_generation_state()
logger.error(f"Error during generation: {e}", exc_info = True)
yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n"
yield "data: [DONE]\n\n"
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
if not completed and not cancel_event.is_set():
@ -7028,6 +7115,13 @@ async def openai_chat_completions(
chunk_text = await asyncio.to_thread(next, gen, _DONE)
if chunk_text is _DONE:
break
if isinstance(chunk_text, GenStreamError):
_msg = _friendly_gen_stream_error(chunk_text)
api_monitor.fail(monitor_id, _msg)
yield _openai_stream_error_sse(
{"error": {"message": _msg, "type": "server_error"}}
)
return
if chunk_text:
api_monitor.append_reply(monitor_id, chunk_text)
yield _chat_content_chunk(
@ -7043,8 +7137,11 @@ async def openai_chat_completions(
raise
except Exception as e:
logger.error(f"Error during audio input streaming: {e}", exc_info = True)
api_monitor.fail(monitor_id, _friendly_error(e))
yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
_msg = _friendly_error(e)
api_monitor.fail(monitor_id, _msg)
yield _openai_stream_error_sse(
{"error": {"message": _msg, "type": "server_error"}}
)
finally:
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
_tracker.__exit__(None, None, None)
@ -7061,7 +7158,15 @@ async def openai_chat_completions(
)
else:
try:
full_text = "".join(audio_input_generate())
full_text = ""
for chunk_text in audio_input_generate():
if isinstance(chunk_text, GenStreamError):
_msg = _friendly_gen_stream_error(chunk_text)
api_monitor.fail(monitor_id, _msg)
raise HTTPException(status_code = 500, detail = _msg)
full_text += chunk_text
except HTTPException:
raise
except Exception as e:
api_monitor.fail(monitor_id, _friendly_error(e))
raise
@ -8605,19 +8710,33 @@ async def openai_chat_completions(
# Classify capability flags from the loaded template.
_sf_model_info = backend.models.get(backend.active_model_name, {})
_sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template")
_sf_features = _detect_safetensors_features(backend, _sf_tpl)
# GGUF parity: enable_thinking templates prefill an unclosed <think>; split into
# reasoning_content deltas so the UI renders the block for safetensors and MLX.
_sf_parse_think = bool(
_sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on")
# Named templates may expose native reasoning only in their ``tool_use``
# branch. Use a truthy placeholder for Studio-managed tools, whose concrete
# schemas are selected below, and the request schemas for client passthrough.
_sf_server_tool_intent = bool(
_effective_enable_tools(payload) or _explicit_studio_tool_loop_requested(payload)
)
# Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode.
_sf_reasoning_prefilled = _sf_reasoning_prefill_mode(
_sf_features,
payload.enable_thinking,
_sf_tpl,
reasoning_effort = payload.reasoning_effort,
_sf_template_tools = payload.tools if payload.tool_choice != "none" else None
if not _sf_template_tools and _sf_server_tool_intent:
_sf_template_tools = ({},)
def _sf_response_protocol(tools = None):
features = _detect_safetensors_features(backend, _sf_tpl, tools = tools)
parse_think = bool(
features.get("supports_reasoning") or features.get("reasoning_always_on")
)
reasoning_prefilled = _sf_reasoning_prefill_mode(
features,
payload.enable_thinking,
_sf_tpl,
reasoning_effort = payload.reasoning_effort,
)
return features, parse_think, reasoning_prefilled
# GGUF parity: split canonical <think> output into reasoning_content. The
# selected template branch must match whether this request renders tools.
_sf_features, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol(
_sf_template_tools
)
def _new_sf_reasoning_extractor():
@ -8767,6 +8886,7 @@ async def openai_chat_completions(
permission_mode = payload.permission_mode,
use_adapter = payload.use_adapter,
stats_holder = _sf_stats_holder,
reasoning_prefilled = _sf_reasoning_prefilled,
)
_sf_tool_sentinel = object()
@ -8826,6 +8946,18 @@ async def openai_chat_completions(
_sf_next_task = None
if event is _sf_tool_sentinel:
break
if isinstance(event, GenStreamError):
backend.reset_generation_state()
_msg = _friendly_gen_stream_error(event)
api_monitor.fail(monitor_id, _msg)
yield _openai_stream_error_sse(
{"error": {"message": _msg, "type": "server_error"}}
)
return
if not isinstance(event, dict):
raise RuntimeError(
f"Invalid safetensors tool event: {type(event).__name__}"
)
if event["type"] == "heartbeat":
# Tool-execution wrapper heartbeat -> SSE keepalive.
@ -8913,6 +9045,11 @@ async def openai_chat_completions(
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
raise
except GenStreamErrorRaised as exc:
backend.reset_generation_state()
_msg = _friendly_gen_stream_error(exc)
api_monitor.fail(monitor_id, _msg)
yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}})
except Exception:
backend.reset_generation_state()
# Generic wire message; full trace stays in the log (CWE-209:
@ -8962,6 +9099,15 @@ async def openai_chat_completions(
for event in gen:
if cancel_event.is_set():
break
if isinstance(event, GenStreamError):
raise HTTPException(
status_code = 500,
detail = _friendly_gen_stream_error(event),
)
if not isinstance(event, dict):
raise RuntimeError(
f"Invalid safetensors tool event: {type(event).__name__}"
)
if event.get("type") == "content":
full_text = _strip_tool_xml_for_display(
event.get("text", ""),
@ -9002,6 +9148,15 @@ async def openai_chat_completions(
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
raise
except GenStreamErrorRaised as exc:
backend.reset_generation_state()
_msg = _friendly_gen_stream_error(exc)
api_monitor.fail(monitor_id, _msg)
raise HTTPException(status_code = 500, detail = _msg)
except HTTPException as exc:
backend.reset_generation_state()
api_monitor.fail(monitor_id, str(exc.detail))
raise
except Exception:
backend.reset_generation_state()
# CWE-209: generic detail; full trace in log.
@ -9088,6 +9243,12 @@ async def openai_chat_completions(
else:
gen_kwargs["tools"] = payload.tools
# The potential tool context above is needed before server/client routing is
# known. This standard path now has the exact schemas that will be rendered,
# so resolve reasoning parsing again to keep empty registries, forced-tool
# misses, and tool_choice="none" on the marker-free template branch.
_, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol(gen_kwargs.get("tools"))
# Request-scoped usage/timings receptacle (filled at gen_done).
stats_holder: dict = {}
@ -9168,6 +9329,14 @@ async def openai_chat_completions(
_next_task = None
if cumulative is _DONE:
break
if isinstance(cumulative, GenStreamError):
backend.reset_generation_state()
_msg = _friendly_gen_stream_error(cumulative)
api_monitor.fail(monitor_id, _msg)
yield _openai_stream_error_sse(
{"error": {"message": _msg, "type": "server_error"}}
)
return
if await request.is_disconnected():
cancel_event.set()
backend.reset_generation_state()
@ -9317,6 +9486,11 @@ async def openai_chat_completions(
try:
full_text = ""
for token in generate():
if isinstance(token, GenStreamError):
backend.reset_generation_state()
_msg = _friendly_gen_stream_error(token)
api_monitor.fail(monitor_id, _msg)
raise HTTPException(status_code = 500, detail = _msg)
full_text = token
# Split prefilled <think> reasoning (GGUF parity); also covers MLX via
@ -9415,6 +9589,8 @@ async def openai_chat_completions(
api_monitor.finish(monitor_id)
return _model_json_response(response)
except HTTPException:
raise
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI completion: {e}", exc_info = True)

20
studio/backend/run.py
View file

@ -1354,11 +1354,23 @@ def run_server(
if secure:
os.environ["UNSLOTH_SECURE"] = "1"
import nest_asyncio
nest_asyncio.apply()
import asyncio
# nest_asyncio is for Colab/IPython, where the main thread already runs a loop
# the blocking waits below would collide with. Apply it only with a loop running
# (a plain CLI start has nothing to nest) and only on Python <= 3.13: on 3.14+
# its global Task patch leaves asyncio.current_task() None (tracking moved into
# C), which also breaks the background uvicorn loop and 500s every request. It
# is archived upstream, so no 3.14 fix is coming; skip it there.
if sys.version_info < (3, 14):
try:
asyncio.get_running_loop()
except RuntimeError:
pass
else:
import nest_asyncio
nest_asyncio.apply()
from threading import Thread, Event
import uvicorn

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

@ -16,7 +16,6 @@ from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
@ -128,15 +127,13 @@ def _kwargs_for(flags: dict, enable_thinking, reasoning_effort):
"""Drive the real backend method with a shim carrying the detected flags."""
from core.inference.llama_cpp import LlamaCppBackend
shim = SimpleNamespace(
_supports_reasoning = flags["supports_reasoning"],
_reasoning_always_on = flags["reasoning_always_on"],
_reasoning_style = flags["reasoning_style"],
_reasoning_effort_levels = flags["reasoning_effort_levels"],
_supports_preserve_thinking = flags["supports_preserve_thinking"],
)
build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim)
return build(enable_thinking, reasoning_effort, None) or {}
shim = object.__new__(LlamaCppBackend)
shim._supports_reasoning = flags["supports_reasoning"]
shim._reasoning_always_on = flags["reasoning_always_on"]
shim._reasoning_style = flags["reasoning_style"]
shim._reasoning_effort_levels = flags["reasoning_effort_levels"]
shim._supports_preserve_thinking = flags["supports_preserve_thinking"]
return shim._request_reasoning_kwargs(enable_thinking, reasoning_effort, None) or {}
def _flags():

740
studio/backend/tests/test_gguf_load_cache_reuse.py Normal file
View file

@ -0,0 +1,740 @@
# 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 cached GGUF reuse and load/download exclusion.
No GPU, network, or subprocesses are required.
"""
from __future__ import annotations
import asyncio
import sys
import threading
import types as _types
from pathlib import Path
from unittest.mock import patch
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Stub optional dependencies before importing the modules under test.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
sys.modules.setdefault("structlog", _structlog_stub)
try:
import httpx # noqa: F401
except ImportError:
_httpx_stub = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
"RequestError",
"HTTPStatusError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
_httpx_stub.Response = type("Response", (), {})
_httpx_stub.Request = type("Request", (), {})
class _FakeTimeout:
def __init__(self, *a, **kw):
pass
_httpx_stub.Timeout = _FakeTimeout
_httpx_stub.Client = type(
"Client",
(),
{
"__init__": lambda self, **kw: None,
"__enter__": lambda self: self,
"__exit__": lambda self, *a: None,
},
)
sys.modules.setdefault("httpx", _httpx_stub)
from huggingface_hub import constants as hf_constants
from core.inference.llama_cpp import (
LlamaCppBackend,
cached_gguf_for_load,
gguf_load_in_flight,
hf_gguf_load_in_flight,
)
REPO = "unsloth/gemma-test-GGUF"
VARIANT = "UD-Q4_K_XL"
MAIN = f"gemma-test-{VARIANT}.gguf"
def _build_cache(
root: Path,
repo_id: str,
files: dict[str, int],
*,
snapshot_sha: str = "a" * 40,
) -> Path:
"""Create ``$root/models--<repo>/snapshots/<sha>/<rel>`` for each entry."""
repo_dir = root / f"models--{repo_id.replace('/', '--')}"
(repo_dir / "blobs").mkdir(parents = True, exist_ok = True)
snap = repo_dir / "snapshots" / snapshot_sha
snap.mkdir(parents = True, exist_ok = True)
for rel, size in files.items():
full = snap / rel
full.parent.mkdir(parents = True, exist_ok = True)
full.write_bytes(b"\0" * size)
return snap
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
return tmp_path
def _fail_download(*_args, **_kwargs):
raise AssertionError("must reuse the cached GGUF instead of downloading")
def _fail_get_paths_info(*_args, **_kwargs):
raise AssertionError("cached reuse must return before the sizing preflight")
class TestLoadReusesCachedCopy:
def test_online_reuse_after_revision_bump(self, hf_cache):
"""A new repo revision does not replace a complete cached model."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / MAIN)
def test_reuse_size_check_uses_cached_snapshot_revision(self, hf_cache):
"""Current-revision size changes do not invalidate an older complete copy."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
revisions: list[str | None] = []
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
revisions.append(revision)
size = 4 if revision == snap.name else 8
return [_types.SimpleNamespace(path = path, size = size) for path in paths]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / MAIN)
assert revisions == [snap.name]
def test_reuse_when_cached_revision_vanished_from_hub(self, hf_cache):
"""The Hub answers an unknown revision with an empty result, not an error."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", lambda *_a, **_k: []),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / MAIN)
def test_truncated_cached_file_is_not_reused(self, hf_cache):
backend = LlamaCppBackend()
_build_cache(hf_cache, REPO, {MAIN: 4})
downloaded: list[str] = []
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
return [_types.SimpleNamespace(path = path, size = 8) for path in paths]
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [MAIN]
assert out == f"/fake/{REPO}/{MAIN}"
def test_truncated_cached_split_shard_is_not_reused(self, hf_cache):
backend = LlamaCppBackend()
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
_build_cache(hf_cache, REPO, {shard1: 8, shard2: 4})
downloaded: list[str] = []
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
return [_types.SimpleNamespace(path = path, size = 8) for path in paths]
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [shard1, shard2]
assert out == f"/fake/{REPO}/{shard1}"
def test_online_reuse_when_reupload_renamed_the_file(self, hf_cache):
"""A renamed variant still reuses its cached file."""
backend = LlamaCppBackend()
old_name = f"gemma-test-old-{VARIANT}.gguf"
snap = _build_cache(hf_cache, REPO, {old_name: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / old_name)
def test_downloads_when_nothing_cached(self, hf_cache):
backend = LlamaCppBackend()
downloaded: list[str] = []
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [MAIN]
assert out == f"/fake/{REPO}/{MAIN}"
def test_force_redownloads_despite_cache(self, hf_cache):
"""A forced download ignores a complete cached copy."""
backend = LlamaCppBackend()
_build_cache(hf_cache, REPO, {MAIN: 4})
downloaded: list[str] = []
def fake_download(
repo_id,
filename,
token = None,
**kwargs,
):
assert kwargs.get("force_download") is True
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT, force = True)
assert downloaded == [MAIN]
assert out == f"/fake/{REPO}/{MAIN}"
def test_split_reused_only_when_colocated(self, hf_cache):
backend = LlamaCppBackend()
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
snap = _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / shard1)
def test_partial_split_set_downloads(self, hf_cache):
"""A partial split set is not reused."""
backend = LlamaCppBackend()
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
_build_cache(hf_cache, REPO, {shard1: 4})
downloaded: list[str] = []
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p is not None]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [shard1, shard2]
assert out == f"/fake/{REPO}/{shard1}"
def test_reuse_prefers_newest_snapshot_after_update(self, hf_cache):
"""Loads prefer the newest complete snapshot."""
import os
backend = LlamaCppBackend()
old_snap = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "a" * 40)
new_snap = _build_cache(hf_cache, REPO, {MAIN: 6}, snapshot_sha = "b" * 40)
os.utime(old_snap, (1_000_000, 1_000_000))
os.utime(new_snap, (2_000_000, 2_000_000))
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(new_snap / MAIN)
def test_low_disk_fallback_reuses_cached_copy(self, hf_cache):
backend = LlamaCppBackend()
fallback = "gemma-test-Q2_K.gguf"
snap = _build_cache(hf_cache, REPO, {fallback: 4})
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
size = 4 if revision == snap.name else 100
return [_types.SimpleNamespace(path = path, size = size) for path in paths]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("shutil.disk_usage", lambda *_a, **_k: _types.SimpleNamespace(free = 10)),
patch.object(
backend,
"_find_smallest_fitting_variant",
lambda *_a, **_k: (fallback, 4, []),
),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / fallback)
def test_companion_prefers_main_snapshot_sibling(self, hf_cache):
"""A cached mmproj is reused from the main model's snapshot."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4, "mmproj-F16.gguf": 2})
def _fail_list(*_args, **_kwargs):
raise AssertionError("snapshot sibling must resolve without a repo listing")
with patch("huggingface_hub.list_repo_files", _fail_list):
out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN))
assert out == str(snap / "mmproj-F16.gguf")
def test_companion_finds_snapshot_through_hf_symlink(self, hf_cache):
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {})
blobs = snap.parent.parent / "blobs"
main_blob = blobs / "main"
mmproj_blob = blobs / "mmproj"
main_blob.write_bytes(b"main")
mmproj_blob.write_bytes(b"mmproj")
try:
(snap / MAIN).symlink_to(main_blob)
(snap / "mmproj-F16.gguf").symlink_to(mmproj_blob)
except OSError as exc:
pytest.skip(f"symlinks unavailable: {exc}")
with patch("huggingface_hub.list_repo_files", _fail_download):
out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN))
assert out == str(snap / "mmproj-F16.gguf")
def test_companion_does_not_download_during_hub_job(self, hf_cache):
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
registry = _types.SimpleNamespace(active_job_refs = lambda _repo: [object()])
with (
patch("huggingface_hub.list_repo_files", _fail_download),
patch("hub.utils.download_registry.get_models_registry", lambda: registry),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN))
assert out is None
class TestCachedGgufForLoadProbe:
def test_complete_copy_found(self, hf_cache):
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN)
def test_absent_copy_is_none(self, hf_cache):
assert cached_gguf_for_load(REPO, VARIANT) is None
def test_partial_split_is_none(self, hf_cache):
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
_build_cache(hf_cache, REPO, {shard1: 4})
assert cached_gguf_for_load(REPO, VARIANT) is None
def test_partial_new_snapshot_does_not_hide_complete_split(self, hf_cache):
import os
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
old = _build_cache(
hf_cache,
REPO,
{shard1: 4, shard2: 4},
snapshot_sha = "a" * 40,
)
new = _build_cache(hf_cache, REPO, {shard1: 4}, snapshot_sha = "b" * 40)
os.utime(old, (1_000_000, 1_000_000))
os.utime(new, (2_000_000, 2_000_000))
assert cached_gguf_for_load(REPO, VARIANT) == str(old / shard1)
def test_split_requires_every_declared_shard(self, hf_cache):
shard1 = f"gemma-test-{VARIANT}-00001-of-00003.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00003.gguf"
_build_cache(hf_cache, REPO, {shard1: 4, shard2: 4})
assert cached_gguf_for_load(REPO, VARIANT) is None
def test_required_mmproj_must_share_main_snapshot(self, hf_cache):
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN)
assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) is None
(snap / "mmproj-F16.gguf").write_bytes(b"mmproj")
assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(snap / MAIN)
def test_required_mmproj_scans_past_newer_main_only_snapshot(self, hf_cache):
import os
old = _build_cache(
hf_cache,
REPO,
{MAIN: 4, "mmproj-F16.gguf": 2},
snapshot_sha = "a" * 40,
)
new = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "b" * 40)
os.utime(old, (1_000_000, 1_000_000))
os.utime(new, (2_000_000, 2_000_000))
assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(old / MAIN)
class TestLoadHubDownloadExclusion:
def test_in_flight_marker_counts_and_normalizes_case(self):
assert not hf_gguf_load_in_flight(REPO)
with gguf_load_in_flight(REPO):
assert hf_gguf_load_in_flight(REPO.upper())
with gguf_load_in_flight(REPO.lower()):
assert hf_gguf_load_in_flight(REPO)
assert hf_gguf_load_in_flight(REPO)
assert not hf_gguf_load_in_flight(REPO)
def test_marker_noops_for_local_loads(self):
with gguf_load_in_flight(None):
assert not hf_gguf_load_in_flight("")
def test_marker_cleared_on_exception(self):
with pytest.raises(RuntimeError):
with gguf_load_in_flight(REPO):
raise RuntimeError("boom")
assert not hf_gguf_load_in_flight(REPO)
def test_hub_download_refused_while_load_in_flight(self):
from fastapi import HTTPException
from hub.schemas.downloads import DownloadModelRequest
from hub.services.models import downloads as dl
body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT)
with (
patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id),
gguf_load_in_flight(REPO),
):
with pytest.raises(HTTPException) as exc_info:
asyncio.run(dl.download_model_response(body))
assert exc_info.value.status_code == 409
assert "load" in exc_info.value.detail.lower()
def test_hub_download_rechecks_marker_before_claim(self):
from fastapi import HTTPException
from hub.schemas.downloads import DownloadModelRequest
from hub.services.models import downloads as dl
scope = None
def mark_load(*_args, **_kwargs):
nonlocal scope
if scope is None:
scope = gguf_load_in_flight(REPO)
scope.__enter__()
return frozenset()
class _Registry:
def claim(self, *_args, admission_check, **_kwargs):
assert admission_check() is False
return False, "admission_blocked"
def current_generation(self, _key):
return 0
registry = _Registry()
body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT)
try:
with (
patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id),
patch.object(dl.gguf_variants, "gguf_variant_blob_hashes", mark_load),
patch.object(dl, "_registry", registry),
):
with pytest.raises(HTTPException) as exc_info:
asyncio.run(dl.download_model_response(body))
finally:
if scope is not None:
scope.__exit__(None, None, None)
assert exc_info.value.status_code == 409
def test_registry_admission_check_prevents_claim(self):
from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP
registry = DownloadRegistry()
claimed, state = registry.claim(
f"{REPO}::{VARIANT}",
TRANSPORT_HTTP,
repo_type = "model",
repo_id = REPO,
variant = VARIANT,
admission_check = lambda: False,
)
assert claimed is False
assert state == "admission_blocked"
assert registry.active_jobs(REPO) == {}
def test_same_variant_job_stays_visible_during_retry_handoff(self):
from hub.utils.download_registry import DownloadRegistry, TRANSPORT_XET
from core.inference.llama_cpp import _hub_download_blocks_gguf_load
registry = DownloadRegistry()
key = f"{REPO}::{VARIANT}"
claimed, _ = registry.claim(
key,
TRANSPORT_XET,
repo_type = "model",
repo_id = REPO,
variant = VARIANT,
)
assert claimed is True
assert registry.has_active_variant(REPO, VARIANT.lower()) is True
registry.release_active_slot(key)
assert registry.active_jobs(REPO) == {}
assert registry.active_job_refs(REPO)
assert registry.has_active_variant(REPO, VARIANT) is True
with (
patch("hub.utils.download_registry.get_models_registry", lambda: registry),
patch(
"core.inference.llama_cpp.cached_gguf_for_load",
side_effect = AssertionError("same-variant jobs must block before cache reuse"),
),
):
assert _hub_download_blocks_gguf_load(REPO, VARIANT) is True
registry.set_job(key, "complete")
assert registry.has_active_variant(REPO, VARIANT) is False
def test_other_variant_job_still_allows_complete_cached_load(self):
from core.inference.llama_cpp import _hub_download_blocks_gguf_load
from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP
registry = DownloadRegistry()
registry.claim(
f"{REPO}::Q8_0",
TRANSPORT_HTTP,
repo_type = "model",
repo_id = REPO,
variant = "Q8_0",
)
with (
patch("hub.utils.download_registry.get_models_registry", lambda: registry),
patch(
"core.inference.llama_cpp.cached_gguf_for_load",
return_value = "/cached/model.gguf",
) as cached_probe,
):
assert _hub_download_blocks_gguf_load(REPO, VARIANT) is False
cached_probe.assert_called_once_with(
REPO,
VARIANT,
require_mmproj = False,
verify_sizes = True,
hf_token = None,
)
def test_cancelled_request_keeps_marker_until_load_thread_finishes(self):
from core.inference.llama_cpp import _with_gguf_load_marker
started = threading.Event()
release = threading.Event()
finished = threading.Event()
class FakeBackend:
@_with_gguf_load_marker
def load_model(self, *, hf_repo):
started.set()
release.wait(timeout = 2)
finished.set()
return True
async def scenario():
with patch(
"core.inference.llama_cpp._hub_download_blocks_gguf_load",
return_value = False,
):
task = asyncio.create_task(
asyncio.to_thread(FakeBackend().load_model, hf_repo = REPO)
)
assert await asyncio.to_thread(started.wait, 1)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert hf_gguf_load_in_flight(REPO)
release.set()
assert await asyncio.to_thread(finished.wait, 1)
for _ in range(100):
if not hf_gguf_load_in_flight(REPO):
break
await asyncio.sleep(0.001)
assert not hf_gguf_load_in_flight(REPO)
asyncio.run(scenario())
def test_load_marker_precedes_hub_guard_and_unload(self):
source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text()
gguf_branch = source[source.index("if config.is_gguf:") :]
assert (
gguf_branch.index("enter_context(gguf_load_in_flight")
< gguf_branch.index("if request.llama_extra_args is None")
< gguf_branch.index("_hub_download_blocks_gguf_load")
< gguf_branch.index("unsloth_backend.unload_model")
)
llama_source = (
Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py"
).read_text()
assert "@_with_gguf_load_marker\n def load_model(" in llama_source

74
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -1061,6 +1061,80 @@ def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch):
]
def test_same_turn_duplicate_does_not_drop_later_parallel_call(monkeypatch):
# One batch: search(a), search(a) [duplicate], search(b). The duplicate is an
# internal no-op, but the distinct search(b) after it must still run, and the
# no-op nudge must land after the tool results rather than splitting them.
batch = [
_sse(
{
"tool_calls": [
{
"index": 0,
"id": "call_a1",
"type": "function",
"function": {"name": "web_search", "arguments": json.dumps({"query": "a"})},
},
{
"index": 1,
"id": "call_a2",
"type": "function",
"function": {"name": "web_search", "arguments": json.dumps({"query": "a"})},
},
{
"index": 2,
"id": "call_b",
"type": "function",
"function": {"name": "web_search", "arguments": json.dumps({"query": "b"})},
},
]
}
),
_done(),
]
final_stream = [_sse({"content": "Final answer."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [batch, final_stream], payloads)
calls: list[dict] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append(arguments)
return "search-result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 3,
)
)
# Both distinct calls ran; the duplicate did not (old `break` dropped search(b)).
assert calls == [{"query": "a"}, {"query": "b"}]
assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [
"call_a1",
"call_b",
]
# The next generation's conversation must be well-formed: the assistant lists
# only the executed calls (no orphan for the duplicate), the two tool results
# follow contiguously, and the no-op nudge lands after them, never between.
conv = payloads[1]["messages"]
asst = next(m for m in conv if m["role"] == "assistant" and m.get("tool_calls"))
assert [tc.get("id") for tc in asst["tool_calls"]] == ["call_a1", "call_b"]
after = conv[conv.index(asst) + 1 :]
assert [m["role"] for m in after[:2]] == ["tool", "tool"]
assert [m.get("tool_call_id") for m in after[:2]] == ["call_a1", "call_b"]
assert after[2]["role"] == "user" # deferred duplicate nudge, after the results
assert after[2]["content"].startswith(
"One earlier request to call tool 'web_search' in this batch was not executed"
)
assert "previous tool request" not in after[2]["content"].lower()
def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch):
same_turn_render_calls = [
_sse(

43
studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
View file

@ -373,6 +373,7 @@ def test_kill_orphaned_servers_returns_count():
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
):
n = LlamaCppBackend._kill_orphaned_servers()
assert n == 1, "only the Studio-owned orphan should be counted"
@ -384,11 +385,53 @@ def test_kill_orphaned_servers_returns_count():
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
):
assert LlamaCppBackend._kill_orphaned_servers() == 0
assert killed == []
def test_kill_orphaned_servers_spares_live_parent():
"""A Studio-owned llama-server whose parent is still running is not an
orphan (a live Studio or the user's shell owns it) and must never be
killed; only the true orphan (parent gone) is reaped."""
import os
mypid = os.getpid()
fake_path = "/tmp/unsloth-test-llama/llama-server"
killed: list[int] = []
class _FakeProc:
def __init__(self, pid, name, exe):
self.info = {"pid": pid, "name": name, "exe": exe}
def kill(self):
killed.append(self.info["pid"])
live_parent = _FakeProc(mypid + 1, "llama-server", fake_path)
true_orphan = _FakeProc(mypid + 2, "llama-server", fake_path)
fake_psutil = _types.ModuleType("psutil")
fake_psutil.NoSuchProcess = type("NoSuchProcess", (Exception,), {})
fake_psutil.AccessDenied = type("AccessDenied", (Exception,), {})
fake_psutil.ZombieProcess = type("ZombieProcess", (Exception,), {})
fake_psutil.process_iter = lambda attrs = None: [live_parent, true_orphan]
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
patch.object(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)),
patch.object(
LlamaCppBackend,
"_pid_parent_is_alive",
staticmethod(lambda pid: pid == mypid + 1),
),
):
n = LlamaCppBackend._kill_orphaned_servers()
assert n == 1, "only the true orphan should be reaped"
assert killed == [mypid + 2], "the live-parent server must be spared"
def test_startup_reaper_arms_settle_timestamp():
"""__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an
orphan (so the first load_model waits for VRAM to settle), and leaves the

290
studio/backend/tests/test_mcp_server.py Normal file
View file

@ -0,0 +1,290 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import sys
import types
import pytest
from mcp_server import BearerTokenMiddleware, _clamp, _dump, create_studio_mcp
def _get_tool(name):
tools = asyncio.run(create_studio_mcp().list_tools())
return {tool.name: tool for tool in tools}[name]
def test_studio_mcp_registers_control_plane_tools():
tools = asyncio.run(create_studio_mcp().list_tools())
assert {tool.name for tool in tools} == {
"studio_status",
"list_local_models",
"get_training_status",
"start_training",
"stop_training",
"list_training_runs",
"validate_recipe",
"get_recipe_job_status",
"get_recipe_job_dataset",
"load_checkpoint",
"export_gguf",
}
def test_dump_serializes_pydantic_values():
class Response:
def model_dump(self, *, mode):
assert mode == "json"
return {"ok": True}
assert _dump(Response()) == {"ok": True}
assert _dump({"already": "json"}) == {"already": "json"}
def test_bearer_token_middleware_rejects_wrong_token():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "http", "headers": [(b"authorization", b"Bearer wrong")]},
None,
send,
)
)
assert events[0]["status"] == 401
assert "app" not in events
def test_bearer_token_middleware_closes_unauthorized_websocket():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "websocket", "headers": []},
None,
send,
)
)
assert events == [{"type": "websocket.close", "code": 4401}]
def test_bearer_token_middleware_rejects_non_ascii_authorization():
# A non-ASCII bearer value must produce a clean 401, not a 500. Comparing on
# bytes avoids the str hmac.compare_digest TypeError on non-ASCII input.
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "http", "headers": [(b"authorization", b"Bearer \xff\xff")]},
None,
send,
)
)
assert events[0]["status"] == 401
assert "app" not in events
def test_bearer_token_middleware_accepts_correct_token():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "http", "headers": [(b"authorization", b"Bearer secret")]},
None,
send,
)
)
assert events == ["app"]
def test_bearer_token_middleware_requires_non_empty_token():
async def app(scope, receive, send):
pass
for bad in ("", " "):
with pytest.raises(ValueError):
BearerTokenMiddleware(app, bad)
def test_bearer_token_middleware_rejects_non_ascii_token():
async def app(scope, receive, send):
pass
# non-ASCII tokens cannot be transmitted in an HTTP header by a standard
# client, so they are rejected at construction instead of locking out.
for bad in ("töken", "\U0001f600"):
with pytest.raises(ValueError):
BearerTokenMiddleware(app, bad)
def test_bearer_token_middleware_passes_through_non_http_scopes():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(middleware({"type": "lifespan"}, None, send))
assert events == ["app"]
def test_clamp_restricts_to_inclusive_bounds():
assert _clamp(5, 1, 200) == 5
assert _clamp(-10, 1, 200) == 1
assert _clamp(10_000, 1, 200) == 200
assert _clamp(0, 1, 500) == 1
assert _clamp(1_000, 1, 500) == 500
def test_export_and_checkpoint_tools_expose_forwarded_fields():
export_props = set(_get_tool("export_gguf").parameters["properties"])
assert {"hf_token", "imatrix", "imatrix_path"} <= export_props
checkpoint_props = set(_get_tool("load_checkpoint").parameters["properties"])
assert {"hf_token", "approved_remote_code_fingerprint"} <= checkpoint_props
def _stub_module(monkeypatch, name, **attrs):
module = types.ModuleType(name)
for key, value in attrs.items():
setattr(module, key, value)
if "." in name:
module.__path__ = [] # mark package-like so submodule imports resolve
monkeypatch.setitem(sys.modules, name, module)
return module
def test_export_gguf_forwards_hf_token_and_imatrix(monkeypatch):
captured = {}
class FakeExportGGUFRequest:
def __init__(self, **kwargs):
captured.update(kwargs)
async def fake_export(request, current_subject):
return {"current_subject": current_subject}
_stub_module(monkeypatch, "models", ExportGGUFRequest = FakeExportGGUFRequest)
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.export", export_gguf = fake_export)
tool = _get_tool("export_gguf")
result = asyncio.run(
tool.fn(
save_directory = "/tmp/out",
quantization_method = ["Q4_K_M", "Q8_0"],
push_to_hub = True,
repo_id = "me/model",
hf_token = "hf_secret",
imatrix = True,
imatrix_path = "/tmp/imatrix.dat",
)
)
assert captured["hf_token"] == "hf_secret"
assert captured["imatrix"] is True
assert captured["imatrix_path"] == "/tmp/imatrix.dat"
assert captured["quantization_method"] == ["Q4_K_M", "Q8_0"]
assert result["current_subject"] == "mcp"
def test_load_checkpoint_forwards_token_and_fingerprint(monkeypatch):
captured = {}
class FakeLoadCheckpointRequest:
def __init__(self, **kwargs):
captured.update(kwargs)
async def fake_load(request, current_subject):
return {"current_subject": current_subject}
_stub_module(monkeypatch, "models", LoadCheckpointRequest = FakeLoadCheckpointRequest)
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.export", load_checkpoint = fake_load)
tool = _get_tool("load_checkpoint")
asyncio.run(
tool.fn(
checkpoint_path = "/tmp/ckpt",
approved_remote_code_fingerprint = "sha256:abc",
hf_token = "hf_secret",
)
)
assert captured["hf_token"] == "hf_secret"
assert captured["approved_remote_code_fingerprint"] == "sha256:abc"
def test_list_training_runs_clamps_pagination(monkeypatch):
captured = {}
async def fake_list_runs(limit, offset, current_subject):
captured["limit"] = limit
captured["offset"] = offset
return {"ok": True}
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.training_history", list_training_runs = fake_list_runs)
tool = _get_tool("list_training_runs")
asyncio.run(tool.fn(limit = 10_000, offset = -5))
assert captured["limit"] == 200
assert captured["offset"] == 0
def test_get_recipe_job_dataset_clamps_pagination(monkeypatch):
captured = {}
def fake_job_dataset(job_id, limit, offset):
captured["limit"] = limit
captured["offset"] = offset
return {"ok": True}
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.data_recipe")
_stub_module(monkeypatch, "routes.data_recipe.jobs", job_dataset = fake_job_dataset)
tool = _get_tool("get_recipe_job_dataset") # this tool is synchronous
tool.fn(job_id = "job-1", limit = -1, offset = -9)
assert captured["limit"] == 1
assert captured["offset"] == 0

166
studio/backend/tests/test_mlx_inference_backend.py
View file

@ -523,3 +523,169 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
assert render["kwargs"]["enable_thinking"] is True
assert render["kwargs"]["reasoning_effort"] == "medium"
assert render["kwargs"]["preserve_thinking"] is True
def test_mlx_text_normalizes_native_reasoning_and_close_releases_lock(monkeypatch):
_install_fake_mlx(monkeypatch)
from core.inference.mlx_inference import MLXInferenceBackend
monkeypatch.setattr(
"core.inference.chat_template_helpers.apply_chat_template_for_generation",
lambda *_args, **_kwargs: "prompt",
raising = True,
)
monkeypatch.setattr(
"core.inference.chat_template_helpers.render_with_native_template_fallback",
lambda formatted_prompt, **_kwargs: SimpleNamespace(
prompt = formatted_prompt,
reasoning_channel_markers = ("<|channel>thought\n", "<channel|>"),
),
raising = True,
)
mlx_lm_pkg = types.ModuleType("mlx_lm")
mlx_lm_sample = types.ModuleType("mlx_lm.sample_utils")
mlx_lm_sample.make_sampler = lambda **_kw: object()
mlx_lm_sample.make_logits_processors = lambda **_kw: None
class _Resp:
def __init__(self, text, tok):
self.text = text
self.token = tok
def _stream_generate(_model, _tokenizer, **_kw):
yield _Resp("<|channel>thought\n", 10)
yield _Resp("r", 11)
yield _Resp("<channel|>", 12)
yield _Resp("a", 13)
mlx_lm_pkg.stream_generate = _stream_generate
monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg)
monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample)
backend = MLXInferenceBackend()
backend._model = object()
backend._tokenizer = SimpleNamespace(all_special_tokens = [])
backend._is_vlm = False
assert list(
backend.generate_chat_response(
messages = [{"role": "user", "content": "ping"}],
max_new_tokens = 4,
)
) == ["<think>", "<think>r", "<think>r</think>", "<think>r</think>a"]
gen = backend.generate_chat_response(
messages = [{"role": "user", "content": "ping"}],
max_new_tokens = 4,
)
assert next(gen) == "<think>"
assert backend._generation_lock.locked()
gen.close()
assert not backend._generation_lock.locked()
def test_mlx_text_native_metadata_preserves_prefilled_think_snapshots(monkeypatch):
_install_fake_mlx(monkeypatch)
from core.inference.mlx_inference import MLXInferenceBackend
monkeypatch.setattr(
"core.inference.chat_template_helpers.apply_chat_template_for_generation",
lambda *_args, **_kwargs: "prompt<think>\n",
raising = True,
)
monkeypatch.setattr(
"core.inference.chat_template_helpers.render_with_native_template_fallback",
lambda formatted_prompt, **_kwargs: SimpleNamespace(
prompt = formatted_prompt,
reasoning_channel_markers = ("<|channel>thought", "<channel|>"),
),
raising = True,
)
mlx_lm_pkg = types.ModuleType("mlx_lm")
mlx_lm_sample = types.ModuleType("mlx_lm.sample_utils")
mlx_lm_sample.make_sampler = lambda **_kw: object()
mlx_lm_sample.make_logits_processors = lambda **_kw: None
class _Resp:
def __init__(self, text, tok):
self.text = text
self.token = tok
def _stream_generate(_model, _tokenizer, **_kw):
yield _Resp("reason", 10)
yield _Resp("</think>", 11)
yield _Resp("answer", 12)
mlx_lm_pkg.stream_generate = _stream_generate
monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg)
monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample)
backend = MLXInferenceBackend()
backend._model = object()
backend._tokenizer = SimpleNamespace(all_special_tokens = [])
backend._is_vlm = False
snapshots = list(
backend.generate_chat_response(
messages = [{"role": "user", "content": "ping"}],
max_new_tokens = 3,
)
)
assert snapshots == [
"<think>\n",
"<think>\nreason",
"<think>\nreason</think>",
"<think>\nreason</think>answer",
]
assert all(current.startswith(previous) for previous, current in zip(snapshots, snapshots[1:]))
def test_mlx_vlm_normalizes_native_reasoning_channels(monkeypatch):
_install_fake_mlx(monkeypatch)
from core.inference.mlx_inference import MLXInferenceBackend
monkeypatch.setattr(
"core.inference.chat_template_helpers.apply_chat_template_for_generation",
lambda *_args, **_kwargs: "prompt",
raising = True,
)
mlx_vlm_pkg = types.ModuleType("mlx_vlm")
class _Resp:
def __init__(self, text, tok):
self.text = text
self.token = tok
def _stream_generate(_model, _processor, _prompt, _images, **_kw):
yield _Resp("<|channel>thought\n", 10)
yield _Resp("vision", 11)
yield _Resp("<channel|>", 12)
yield _Resp(" answer", 13)
mlx_vlm_pkg.stream_generate = _stream_generate
monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm_pkg)
backend = MLXInferenceBackend()
backend._model = SimpleNamespace(config = SimpleNamespace())
backend._processor = SimpleNamespace(
chat_template = "<|channel>thought\n...<channel|>",
all_special_tokens = [],
apply_chat_template = lambda *_args, **_kwargs: "prompt",
)
backend._is_vlm = True
assert list(
backend.generate_chat_response(
messages = [{"role": "user", "content": "describe"}],
image = object(),
max_new_tokens = 4,
)
) == [
"<think>",
"<think>vision",
"<think>vision</think>",
"<think>vision</think> answer",
]

28
studio/backend/tests/test_mtp_drafter_companion.py
View file

@ -328,6 +328,7 @@ def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch):
pick,
label,
cancel_event = None,
near_path = None,
):
captured["pick"] = pick
return None
@ -428,6 +429,32 @@ def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch
assert got is not None and Path(got).parent.parent.name == "newest"
def test_download_mtp_prefers_main_snapshot_offline(tmp_path, monkeypatch):
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
snapshots = tmp_path / "models--unsloth--gemma" / "snapshots"
old = snapshots / "old"
new = snapshots / "new"
old.mkdir(parents = True)
new.mkdir(parents = True)
main = old / "gemma-UD-Q4_K_XL.gguf"
old_drafter = old / "mtp-gemma.gguf"
new_drafter = new / "mtp-gemma.gguf"
main.write_bytes(b"main")
old_drafter.write_bytes(b"old")
new_drafter.write_bytes(b"new")
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda _repo: [new, old])
got = LlamaCppBackend()._download_mtp(
hf_repo = "unsloth/gemma-GGUF",
near_path = str(main),
)
assert got == str(old_drafter)
def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch):
# Online, do not reuse a cached copy: go to the download path so a changed
# drafter is refetched (hf_hub_download checks the current revision).
@ -447,6 +474,7 @@ def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch):
pick,
label,
cancel_event = None,
near_path = None,
):
reached["hit"] = True
return None

40
studio/backend/tests/test_offline_gguf_cache_fallback.py
View file

@ -244,9 +244,7 @@ class TestGgufVariantFileResolution:
def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial(
self, monkeypatch, hf_cache
):
# Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download
# resumes the partial current-ref download and revalidates the revision instead
# of serving an older snapshot's same-name blob.
# Keep coverage for offline reuse; online reuse is tested separately.
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
backend = LlamaCppBackend()
repo = "unsloth/vision-GGUF"
@ -292,8 +290,7 @@ class TestGgufVariantFileResolution:
def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it(
self, monkeypatch, hf_cache
):
# Case-variant cross-dir reuse is offline-only; online the canonical repo id
# resolves up front and hf_hub_download fetches the current revision.
# Keep coverage for case-insensitive offline cache lookup.
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
backend = LlamaCppBackend()
canonical_repo = "unsloth/gemma-4-E2B-it-GGUF"
@ -348,45 +345,26 @@ class TestGgufVariantFileResolution:
assert out == str(snap / gguf_file)
assert seen_repos
def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache):
# Online, an older same-name snapshot must not be served (it may be a stale
# revision); hf_hub_download is called so the current revision is fetched and
# its etag revalidated.
def test_download_online_reuses_complete_cached_snapshot(self, monkeypatch, hf_cache):
# Loads reuse complete cached models across repo revisions.
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
backend = LlamaCppBackend()
repo = "unsloth/vision-GGUF"
_build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40)
downloaded: list[str] = []
snap = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40)
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p]
def fake_download(
repo_id,
filename,
token = None,
**kwargs,
):
downloaded.append(filename)
return f"/fresh/{filename}"
def fail_download(*_args, **_kwargs):
raise AssertionError("must reuse the cached GGUF instead of downloading")
with (
patch(
"huggingface_hub.list_repo_files",
lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"],
),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download),
):
out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL")
assert downloaded == ["model-UD-Q4_K_XL.gguf"]
assert out == "/fresh/model-UD-Q4_K_XL.gguf"
assert out == str(snap / "model-UD-Q4_K_XL.gguf")
def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache):
# HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline

25
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -1054,7 +1054,7 @@ class TestChatCompletionRequestToolFields:
monkeypatch.setattr(
inference_route,
"_detect_safetensors_features",
lambda backend, chat_template: {"supports_tools": True},
lambda backend, chat_template, tools = None: {"supports_tools": True},
)
monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inference_route, "api_monitor", monitor)
@ -6620,6 +6620,29 @@ class TestApiMonitorAudioInput:
assert entry["reply"] == "hello world"
assert monitor.active_count() == 0
def failing_chunks():
yield "partial"
raise RuntimeError("generation failed")
self._patch_audio_backend(monkeypatch, failing_chunks())
error_monitor = ApiMonitor(max_entries = 3)
monkeypatch.setattr(inf_mod, "api_monitor", error_monitor)
error_response = await openai_chat_completions(
payload,
request = request,
current_subject = "test",
)
error_chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in error_response.body_iterator
]
assert '"type": "server_error"' in error_chunks[-1]
assert error_chunks[-1].endswith("data: [DONE]\n\n")
[error_entry] = error_monitor.snapshot()
assert error_entry["status"] == "error"
assert error_monitor.active_count() == 0
asyncio.run(_run())
def test_non_gguf_tts_auto_route_records_monitor(self, monkeypatch):

50
studio/backend/tests/test_rag_parsing.py
View file

@ -54,6 +54,56 @@ def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch):
assert "#" not in text and "|" not in text # plain text path emits no Markdown markup
def test_pdf_bytes_use_same_extraction_path(tmp_path, monkeypatch):
from core.rag import config, parsers
monkeypatch.setattr(config, "PDF_MARKDOWN", False)
pdf = tmp_path / "table.pdf"
_table_pdf(pdf)
from_file = parsers.parse(str(pdf))
from_bytes, total_pages = parsers.parse_pdf_bytes(pdf.read_bytes())
assert [page.text for page in from_bytes] == [page.text for page in from_file]
assert total_pages == len(from_file)
def test_pdf_bytes_limit_pages_before_extraction(monkeypatch):
import pymupdf
from core.rag import config, parsers
monkeypatch.setattr(config, "PDF_MARKDOWN", False)
doc = pymupdf.open()
for marker in ("page one", "page two", "page three"):
page = doc.new_page()
page.insert_text((40, 40), marker)
data = doc.tobytes()
doc.close()
pages, total_pages = parsers.parse_pdf_bytes(data, max_pages = 2)
assert len(pages) == 2
assert "page two" in pages[-1].text
assert total_pages == 3 # full count, not the 2 extracted
def test_pdf_markdown_receives_page_limit(monkeypatch):
from core.rag import parsers
captured = {}
class _FakePymupdf4llm:
@staticmethod
def to_markdown(doc, **kwargs):
captured.update(kwargs)
return [{"text": "page"} for _ in kwargs["pages"]]
class _Doc:
page_count = 100
monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm)
assert parsers._pdf_markdown(_Doc(), range(2)) == ["page", "page"]
assert captured == {"page_chunks": True, "show_progress": False, "pages": [0, 1]}
def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch):
# The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the
# newer layout-only OCR knobs or Markdown extraction silently loses policy control.

53
studio/backend/tests/test_safetensors_capability_advertise.py
View file

@ -417,6 +417,59 @@ def test_detect_safetensors_features_gemma_native_tool_call_keeps_tools_on():
assert flags["supports_tools"] is True
def test_detect_safetensors_features_gemma_native_reasoning_is_parseable_not_prefilled():
"""Native Gemma channels are normalized to <think>, then split by the route."""
from routes.inference import _detect_safetensors_features, _sf_reasoning_prefill_mode
tpl_with_gemma_native = "{% if add_generation_prompt %}<|channel>thought\n<channel|>{% endif %}"
backend = SimpleNamespace(
active_model_name = "unsloth/gemma-4-E2B-it",
models = {
"unsloth/gemma-4-E2B-it": {
"native_chat_template": tpl_with_gemma_native,
"chat_template_info": {"template": "override has no native markers"},
}
},
)
flags = _detect_safetensors_features(backend, "override has no native markers")
missing_arg_flags = _detect_safetensors_features(backend, None)
assert flags["supports_reasoning"] is True
assert flags["reasoning_always_on"] is True
assert missing_arg_flags["supports_reasoning"] is True
assert _sf_reasoning_prefill_mode(flags, None, tpl_with_gemma_native) is False
def test_detect_safetensors_features_selects_native_reasoning_from_tool_template():
"""Request tools select a marker-bearing named template without affecting default chat."""
from routes.inference import _detect_safetensors_features
named_template = {
"default": "plain default template",
"tool_use": "{% if tools %}<|channel>thought\n<channel|>{% endif %}",
}
backend = SimpleNamespace(
active_model_name = "custom/named-native-reasoning",
models = {
"custom/named-native-reasoning": {
"native_chat_template": named_template,
"chat_template_info": {"template": "{% if tools %}<tool_call>{% endif %}"},
}
},
)
default_flags = _detect_safetensors_features(backend, "plain override")
tool_flags = _detect_safetensors_features(
backend,
"plain override",
tools = [{"type": "function"}],
)
assert default_flags["supports_reasoning"] is False
assert tool_flags["supports_reasoning"] is True
assert tool_flags["reasoning_always_on"] is True
# Qwen3.5 family pin: the live GGUF + safetensors templates both wrap tool
# calls as ``<tool_call>\n<function=name>...``. Faithful slice so the
# classifier never silently regresses for this family.

132
studio/backend/tests/test_safetensors_reasoning_stream.py
View file

@ -215,3 +215,135 @@ def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort():
swallowed = _replay_sf_reasoning_stream(events, prefilled = True)
assert swallowed["visible"] == ""
assert swallowed["reasoning"] == "The capital of France is Paris."
def test_native_reasoning_streamer_selected_and_errors_raise():
import threading
import pytest
torch = pytest.importorskip("torch")
inf = pytest.importorskip("core.inference.inference")
class Batch(dict):
def to(self, _device):
return self
class Tok:
chat_template = "<|channel>thought\n...<channel|>"
all_special_tokens = []
eos_token_id = 1
pad_token_id = None
pieces = {10: "<|channel>thought\n", 11: "r", 12: "<channel|>", 13: "a"}
def __call__(self, *_args, **_kwargs):
return Batch({"input_ids": torch.zeros((1, 1), dtype = torch.long)})
def decode(self, ids, **_kwargs):
return "".join(self.pieces.get(int(token_id), "") for token_id in ids)
class Model:
device = "cpu"
generation_config = type("Cfg", (), {"eos_token_id": 1})()
config = generation_config
def __init__(self, fail = False):
self.fail = fail
self.kwargs = None
def generate(self, **kwargs):
self.kwargs = kwargs
streamer = kwargs["streamer"]
streamer.put(torch.zeros((1, 1), dtype = torch.long))
for token_id in [10, 11, 12, 13]:
streamer.put(torch.tensor([token_id]))
if self.fail:
raise RuntimeError("boom")
backend = inf.InferenceBackend.__new__(inf.InferenceBackend)
backend.active_model_name = "gemma-test"
backend._generation_lock = threading.Lock()
backend.models = {"gemma-test": {"model": Model(), "tokenizer": Tok()}}
assert list(backend.generate_stream("prompt", max_new_tokens = 4))[-1] == "<think>r</think>a"
backend.models["gemma-test"]["model"] = Model(fail = True)
with pytest.raises(inf._GenerationThreadError, match = "boom"):
list(backend.generate_stream("prompt", max_new_tokens = 4))
def test_text_only_vlm_fallback_resolves_native_markers_off():
import threading
import pytest
torch = pytest.importorskip("torch")
inf = pytest.importorskip("core.inference.inference")
class Batch(dict):
def to(self, _device):
return self
class Tokenizer:
all_special_tokens = []
eos_token_id = 1
pad_token_id = None
def __call__(self, *_args, **_kwargs):
return Batch({"input_ids": torch.zeros((1, 1), dtype = torch.long)})
class Processor:
chat_template = "<|channel>thought\n...<channel|>"
tokenizer = Tokenizer()
class Model:
device = "cpu"
generation_config = type("Cfg", (), {"eos_token_id": 1})()
config = generation_config
def generate(self, **_kwargs):
return None
class EmptyStreamer:
def __next__(self):
raise StopIteration
def end(self):
return None
captured = {}
backend = inf.InferenceBackend.__new__(inf.InferenceBackend)
backend.active_model_name = "vision-test"
backend._generation_lock = threading.Lock()
backend.models = {
"vision-test": {
"model": Model(),
"processor": Processor(),
"tokenizer": Processor(),
}
}
backend.format_chat_prompt = lambda *_args, **_kwargs: "manual text-only prompt"
def make_streamer(*_args, **kwargs):
captured.update(kwargs)
return EmptyStreamer()
backend._make_text_streamer = make_streamer
assert (
list(
backend._generate_vision_response(
messages = [{"role": "user", "content": "hello"}],
system_prompt = "",
image = None,
temperature = 0.7,
top_p = 0.9,
top_k = 40,
min_p = 0.0,
max_new_tokens = 1,
repetition_penalty = 1.0,
)
)
== []
)
assert captured["reasoning_channel_markers"] is None
assert captured["reasoning_channel_markers_resolved"] is True

245
studio/backend/tests/test_safetensors_tool_loop.py
View file

@ -2843,6 +2843,61 @@ class TestLoopBehaviour:
]
assert len(duplicate_nudges) == 1
def test_same_turn_duplicate_does_not_drop_later_parallel_call(self):
# Turn 1 runs search(x). Turn 2's batch is [search(x) duplicate, python]:
# the duplicate is a no-op, but python after it must still run, and the
# no-op nudge must land after python's result rather than splitting it.
captured_messages: list[list[dict]] = []
turns = iter(
[
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
[
'<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
'<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>'
],
["final"],
]
)
def fake_single_turn(messages, active_tools = None):
captured_messages.append([dict(m) for m in messages])
chunks = next(turns)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
exec_fn = FakeExecuteTool(["search-x", "py-result"])
_collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "hi"}],
tools = [
{"type": "function", "function": {"name": "web_search"}},
{"type": "function", "function": {"name": "python"}},
],
execute_tool = exec_fn,
max_tool_iterations = 4,
)
)
# Turn-1 search and turn-2 python both ran; the turn-2 duplicate search did not.
assert exec_fn.calls == [
("web_search", {"query": "x"}),
("python", {"code": "print(1)"}),
]
conv = captured_messages[-1]
turn2 = [m for m in conv if m.get("role") == "assistant" and m.get("tool_calls")][-1]
assert [tc["function"]["name"] for tc in turn2["tool_calls"]] == ["python"]
after = conv[conv.index(turn2) + 1 :]
assert after[0]["role"] == "tool" and after[0]["content"] == "py-result"
assert after[1]["role"] == "user" # deferred duplicate nudge, after the result
assert after[1]["content"].startswith(
"One earlier request to call tool 'web_search' in this batch was not executed"
)
assert "previous tool request" not in after[1]["content"].lower()
def test_duplicate_tool_call_internal_noop_allows_distinct_followup_tool(self):
captured_messages: list[list[dict]] = []
captured_tool_names: list[list[str]] = []
@ -3150,6 +3205,196 @@ class TestLoopBehaviour:
class TestLoopRePrompt:
"""Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``."""
def test_reasoning_intent_does_not_reprompt_a_visible_answer(self):
generations = 0
def _gen(_messages, active_tools = None):
nonlocal generations
generations += 1
yield (
"<think>Let me prepare the requested summary carefully.</think>"
"This is the final visible answer."
)
exec_fn = FakeExecuteTool([])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "summarize this"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
nudge_tool_calls = True,
)
)
assert generations == 1
assert exec_fn.calls == []
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1].endswith("This is the final visible answer.")
def test_prefilled_reasoning_intent_does_not_reprompt_a_visible_answer(self):
generations = 0
def _gen(_messages, active_tools = None):
nonlocal generations
generations += 1
yield "Let me prepare the requested summary carefully.</think>This is the final visible answer."
exec_fn = FakeExecuteTool([])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "summarize this"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
nudge_tool_calls = True,
reasoning_prefilled = True,
)
)
assert generations == 1
assert exec_fn.calls == []
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1].endswith("This is the final visible answer.")
def test_prefilled_reasoning_with_reemitted_think_does_not_reprompt(self):
generations = 0
def _gen(_messages, active_tools = None):
nonlocal generations
generations += 1
yield (
"Let me prepare the requested summary carefully."
"<think>more private planning</think>This is the final visible answer."
)
exec_fn = FakeExecuteTool([])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "summarize this"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
nudge_tool_calls = True,
reasoning_prefilled = True,
)
)
assert generations == 1
assert exec_fn.calls == []
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1].endswith("This is the final visible answer.")
def test_prefilled_reasoning_with_later_think_does_not_reprompt(self):
generations = 0
def _gen(_messages, active_tools = None):
nonlocal generations
generations += 1
yield (
"private prefilled planning</think>"
"<think>Let me prepare the requested summary carefully.</think>"
"This is the final visible answer."
)
exec_fn = FakeExecuteTool([])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "summarize this"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
nudge_tool_calls = True,
reasoning_prefilled = True,
)
)
assert generations == 1
assert exec_fn.calls == []
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1].endswith("This is the final visible answer.")
def test_reasoning_only_intent_still_reprompts_and_uses_a_tool(self):
loop, exec_fn = _make_loop(
turns = [
["<think>Let me search for that.</think>"],
['<tool_call>{"name":"web_search","arguments":{"query":"cats"}}</tool_call>'],
["Here is the answer."],
],
exec_results = ["result"],
nudge_tool_calls = True,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})]
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1] == "Here is the answer."
def test_prefilled_no_close_reasoning_intent_still_reprompts(self):
loop, exec_fn = _make_loop(
turns = [
["I need more context.<think>Let me search for that."],
['<tool_call>{"name":"web_search","arguments":{"query":"cats"}}</tool_call>'],
["Here is the answer."],
],
exec_results = ["result"],
nudge_tool_calls = True,
reasoning_prefilled = True,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})]
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1] == "Here is the answer."
def test_prefilled_reasoning_prefix_is_kept_for_reasoning_only_reprompt(self):
loop, exec_fn = _make_loop(
turns = [
["Let me search for that.</think><think>checking details</think>"],
['<tool_call>{"name":"web_search","arguments":{"query":"cats"}}</tool_call>'],
["Here is the answer."],
],
exec_results = ["result"],
nudge_tool_calls = True,
reasoning_prefilled = True,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})]
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1] == "Here is the answer."
def test_reprompt_history_uses_visible_intent_text(self):
captured: list[list[dict]] = []
def _gen(messages, active_tools = None):
captured.append([dict(message) for message in messages])
if len(captured) == 1:
yield "<think>private planning details</think>Let me search for that."
elif len(captured) == 2:
yield '<tool_call>{"name":"web_search","arguments":{"query":"cats"}}</tool_call>'
else:
yield "Here is the answer."
exec_fn = FakeExecuteTool(["result"])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "find cats"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
nudge_tool_calls = True,
)
)
assert exec_fn.calls == [("web_search", {"query": "cats"})]
assert captured[1][1] == {"role": "assistant", "content": "Let me search for that."}
contents = [e["text"] for e in events if e["type"] == "content"]
assert contents[-1] == "Here is the answer."
def test_intent_signal_triggers_reprompt(self):
# Turn 1: intent signal, no tool call.
# Turn 2 (re-prompt): proper tool call -> executes.

55
studio/backend/tests/test_sf_client_tools_passthrough.py
View file

@ -177,6 +177,15 @@ def _sse_objects(chunks):
# ── Non-streaming ─────────────────────────────────────────────────
def test_non_reasoning_backend_keeps_literal_think_tags(monkeypatch):
backend = _ScriptedBackend(_fixed("show <think>example</think> tags"))
response = _call(_request(stream = False), monkeypatch, backend, supports_tools = False)
message = _json_body(response)["choices"][0]["message"]
assert message["content"] == "show <think>example</think> tags"
assert message["reasoning_content"] is None
def test_xml_healed_to_tool_calls_non_streaming(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
@ -485,6 +494,52 @@ def test_streaming_no_tools_verbatim(monkeypatch):
assert finishes == ["stop"]
def test_streaming_gen_stream_error_is_not_model_text(monkeypatch):
from core.inference.orchestrator import GenStreamError
class _ErrorAfterPartial(_ScriptedBackend):
def __init__(self):
super().__init__(_fixed())
def generate_chat_response(self, **_kwargs):
yield "<think>partial"
yield GenStreamError("Error: /tmp/secret traceback")
backend = _ErrorAfterPartial()
payload = _request(stream = True)
response = _call(payload, monkeypatch, backend, supports_tools = False)
chunks = _collect_sse(response)
objs = _sse_objects(chunks)
deltas = [o.get("choices", [{}])[0].get("delta", {}) for o in objs if o.get("choices")]
assert any("partial" in json.dumps(delta) for delta in deltas)
assert not any("/tmp/secret" in json.dumps(delta) for delta in deltas)
errors = [o["error"]["message"] for o in objs if "error" in o]
assert errors == ["An internal error occurred."]
assert any(
"data: [DONE]" in (chunk.decode() if isinstance(chunk, bytes) else chunk)
for chunk in chunks
)
def test_server_tool_streaming_invalid_event_is_error(monkeypatch):
class _InvalidEventBackend(_ScriptedBackend):
def __init__(self):
super().__init__(_fixed())
def generate_chat_completion_with_tools(self, **_kwargs):
yield {"type": "content", "text": "partial"}
yield "not-an-event"
backend = _InvalidEventBackend()
payload = _request(tools = [LOOKUP_TOOL], enable_tools = True, stream = True)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
errors = [o["error"]["message"] for o in objs if "error" in o]
assert errors == ["An internal error occurred."]
def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch):
# Repeated then shrunk cumulative snapshots must not double-heal.
backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML))

177
studio/backend/tests/test_think_prefill_reemit.py
View file

@ -2,7 +2,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Unit tests for detect_think_prefill.
Unit tests for local reasoning-stream helpers.
Reasoning templates (Qwen3.6-style) end the generation prompt with an open
``<think>\\n`` so the model starts reasoning immediately. skip_prompt
@ -16,7 +16,13 @@ import sys
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference.chat_template_helpers import detect_think_prefill
from core.inference.chat_template_helpers import (
ReasoningChannelNormalizer,
detect_reasoning_channel_markers,
detect_reasoning_channel_markers_from_model_info,
detect_think_prefill,
render_with_native_template_fallback,
)
QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n"
@ -87,3 +93,170 @@ def test_guard_emits_when_think_not_special():
def test_guard_default_and_empty_keep_emitting():
assert detect_think_prefill(QWEN_PROMPT + "<think>\n", None) == "<think>\n"
assert detect_think_prefill(QWEN_PROMPT + "<think>\n", []) == "<think>\n"
def test_gemma_channel_detection_uses_active_template_not_token_metadata():
class TemplateTokenizer:
chat_template = {"default": "...<|channel>thought\\n{{ eoc_token }}"}
class NamedTemplateTokenizer:
chat_template = {
"default": "plain assistant template",
"tool_use": "...<|channel>thought\\n{{ eoc_token }}",
}
class TokenMetadataOnly:
chat_template = None
soc_token = "<|channel>"
eoc_token = "<channel|>"
class NamedTemplateProcessor:
chat_template = {
"default": "plain processor default",
"tool_use": "<|channel>thought\nprocessor tool template<channel|>",
}
tokenizer = TokenMetadataOnly()
def apply_chat_template(self, *_args, **_kwargs):
raise NotImplementedError
expected = ("<|channel>thought", "<channel|>")
assert detect_reasoning_channel_markers(TemplateTokenizer()) == expected
assert detect_reasoning_channel_markers(NamedTemplateTokenizer()) is None
assert (
detect_reasoning_channel_markers(
NamedTemplateTokenizer(), tools = [{"function": {"name": "web_search"}}]
)
== expected
)
assert detect_reasoning_channel_markers(NamedTemplateTokenizer(), tools = []) is None
assert (
detect_reasoning_channel_markers(
NamedTemplateProcessor(), tools = [{"function": {"name": "web_search"}}]
)
is None
)
assert detect_reasoning_channel_markers(TokenMetadataOnly()) is None
def test_gemma_channel_detection_tries_no_argument_getter_fallback():
class FallbackTokenizer:
chat_template = "plain fallback template"
def get_chat_template(self, **kwargs):
if kwargs:
raise ValueError("tools are not supported")
return "...<|channel>thought\n<channel|>"
assert detect_reasoning_channel_markers(
FallbackTokenizer(), tools = [{"function": {"name": "web_search"}}]
) == ("<|channel>thought", "<channel|>")
def test_native_template_fallback_returns_selected_reasoning_metadata():
from types import SimpleNamespace
messages = [{"role": "user", "content": "hi"}]
tools = [{"type": "function", "function": {"name": "web_search"}}]
def render(tokenizer, msgs, *, tools, **_kw):
body = "".join(message["content"] for message in msgs)
suffix = "|TOOLS" if tools else ""
return body + suffix if tokenizer.chat_template == "NATIVE <|channel>thought\n" else body
result = render_with_native_template_fallback(
formatted_prompt = "hi",
tokenizer = SimpleNamespace(chat_template = "OVERRIDE"),
model_info = {
"native_chat_template": "NATIVE <|channel>thought\n",
"tokenizer": SimpleNamespace(chat_template = "OVERRIDE"),
},
active_model_name = "gemma-test",
messages = messages,
tools = tools,
apply_fn = render,
return_metadata = True,
)
assert result.prompt == "hi|TOOLS"
assert result.reasoning_channel_markers == ("<|channel>thought", "<channel|>")
def test_cached_native_template_metadata_recovers_reasoning_markers_without_tools():
from types import SimpleNamespace
model_info = {"chat_template_info": {"template": "native <|channel>thought\n<channel|>"}}
assert detect_reasoning_channel_markers_from_model_info(
SimpleNamespace(chat_template = "override has no native markers"),
model_info,
tools = None,
) == ("<|channel>thought", "<channel|>")
result = render_with_native_template_fallback(
formatted_prompt = "prompt from override",
tokenizer = SimpleNamespace(chat_template = "override has no native markers"),
model_info = model_info,
active_model_name = "gemma-test",
messages = [{"role": "user", "content": "hi"}],
tools = None,
return_metadata = True,
)
assert result.prompt == "prompt from override"
assert result.reasoning_channel_markers == ("<|channel>thought", "<channel|>")
def test_cached_native_markers_do_not_describe_live_tool_template():
from types import SimpleNamespace
tools = [{"type": "function", "function": {"name": "web_search"}}]
class LiveTokenizer:
chat_template = "live tool template without native markers"
def render(_tokenizer, _messages, *, tools, **_kwargs):
return "prompt with tools" if tools else "prompt without tools"
result = render_with_native_template_fallback(
formatted_prompt = "prompt with tools",
tokenizer = LiveTokenizer(),
model_info = {
"chat_template_info": {"template": "native <|channel>thought\n<channel|>"},
"tokenizer": SimpleNamespace(),
},
active_model_name = "gemma-test",
messages = [{"role": "user", "content": "hi"}],
tools = tools,
apply_fn = render,
return_metadata = True,
)
assert result.prompt == "prompt with tools"
assert result.reasoning_channel_markers is None
def test_gemma_channel_normalization_is_prefix_monotonic_and_preserves_tools():
parser = ReasoningChannelNormalizer("<|channel>thought", "<channel|>")
output = ""
snapshots = []
for chunk in (
"<|chan",
"nel>thought",
"\nReason",
"<chan",
"nel|><|tool_call>web_search<tool_call|>",
):
delta = parser.feed(chunk)
if delta:
output += delta
snapshots.append(output)
assert snapshots == [
"<think>",
"<think>Reason",
"<think>Reason</think><|tool_call>web_search<tool_call|>",
]
assert snapshots[1].startswith(snapshots[0])
compact = ReasoningChannelNormalizer("<|channel>thought", "<channel|>")
assert compact.feed("<|channel>thought<channel|>answer") + compact.finish() == (
"<think></think>answer"
)

28
studio/backend/tests/test_tool_loop_controller.py
View file

@ -13,6 +13,7 @@ if _BACKEND_DIR not in sys.path:
from core.inference.tool_loop_controller import (
ToolLoopController,
append_deferred_nudges,
canonical_tool_call_key,
coerce_tool_arguments,
status_for_tool,
@ -21,6 +22,22 @@ from core.inference.tool_loop_controller import (
)
def test_append_deferred_nudges_merges_deduped_into_one_message():
conversation = [{"role": "assistant", "tool_calls": [1]}, {"role": "tool", "content": "r"}]
nudges = [
{"role": "user", "content": "duplicate"},
{"role": "user", "content": "duplicate"}, # dropped: same content
{"role": "user", "content": "disabled foo"},
]
append_deferred_nudges(conversation, nudges)
# One user message, after the results, with distinct contents joined.
assert conversation[2:] == [{"role": "user", "content": "duplicate\n\ndisabled foo"}]
# Empty is a no-op.
before = list(conversation)
append_deferred_nudges(conversation, [])
assert conversation == before
def _tool(name: str) -> dict:
return {"type": "function", "function": {"name": name}}
@ -111,6 +128,10 @@ def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools():
assert not duplicate.should_execute
assert not duplicate.emit_visible_events
duplicate_nudge = completion.model_message()["content"]
assert duplicate_nudge.startswith(
"One earlier request to call tool 'web_search' in this batch was not executed"
)
assert "previous tool request" not in duplicate_nudge.lower()
assert "already completed successfully" in duplicate_nudge
assert "different enabled tool" in duplicate_nudge
assert completion.model_message()["role"] == "user"
@ -165,7 +186,12 @@ def test_empty_enabled_tool_list_blocks_all_tool_calls():
assert decision.action == "disabled"
assert not decision.emit_visible_events
assert completion.model_message()["role"] == "user"
assert "not enabled" in completion.model_message()["content"]
disabled_nudge = completion.model_message()["content"]
assert disabled_nudge.startswith(
"One earlier request to call tool 'web_search' in this batch was not executed"
)
assert "previous tool request" not in disabled_nudge.lower()
assert "not enabled" in disabled_nudge
assert controller.force_final_answer
assert controller.active_tools() == []

138
studio/backend/tests/test_web_fetch_binary_guard.py
View file

@ -59,6 +59,18 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
return tools._fetch_page_text("https://example.com/thing", timeout = 5)
def _pdf_bytes(*page_texts: str) -> bytes:
pymupdf = pytest.importorskip("pymupdf")
doc = pymupdf.open()
for text in page_texts:
page = doc.new_page()
if text:
page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), text, fontsize = 11)
data = doc.tobytes()
doc.close()
return data
@pytest.mark.parametrize(
"content_type,expected",
[
@ -90,10 +102,119 @@ def test_is_text_candidate_content_type(content_type, expected):
assert tools._is_text_candidate_content_type(content_type) is expected
def test_pdf_rejected_by_content_type(monkeypatch):
out = _fetch_with(monkeypatch, b"%PDF-1.7\n\xff\xd8\xff\x00\x89PNG" * 200, "application/pdf")
assert "<EFBFBD>" not in out
assert "non-text content" in out and "application/pdf" in out
@pytest.mark.parametrize(
"content_type",
["application/pdf", "application/octet-stream", "text/html", "text/plain", None],
)
def test_pdf_text_extracted(monkeypatch, content_type):
out = _fetch_with(
monkeypatch,
_pdf_bytes("First page marker", "Second page marker"),
content_type,
)
assert "## Page 1\n\nFirst page marker" in out
assert "## Page 2" in out and "Second page marker" in out
assert "binary content" not in out and "non-text content" not in out
@pytest.mark.parametrize("content_type", ["application/pdf", "text/plain"])
def test_malformed_pdf_returns_safe_placeholder(monkeypatch, content_type):
out = _fetch_with(monkeypatch, b"%PDF-1.7\nnot a complete PDF", content_type)
assert out == "(PDF content could not be read as text)"
def test_pdf_without_text_layer_reported(monkeypatch):
out = _fetch_with(monkeypatch, _pdf_bytes(""), "application/pdf")
assert out == "(PDF contains no extractable text)"
def test_encrypted_pdf_returns_safe_placeholder(monkeypatch):
pymupdf = pytest.importorskip("pymupdf")
doc = pymupdf.open()
doc.new_page().insert_text((40, 40), "private text")
data = doc.tobytes(
encryption = pymupdf.PDF_ENCRYPT_AES_256,
owner_pw = "owner",
user_pw = "secret",
)
doc.close()
out = _fetch_with(monkeypatch, data, "application/pdf")
assert out == "(PDF content could not be read as text)"
def test_pdf_download_limit_enforced(monkeypatch):
monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", 256)
out = _fetch_with(monkeypatch, _pdf_bytes("Readable but oversized"), "application/pdf")
assert out == "(PDF content exceeds the download limit; not readable as text)"
def test_mislabeled_pdf_is_read_past_text_download_cap(monkeypatch):
body = _pdf_bytes("Cross-reference data was fetched")
monkeypatch.setattr(tools, "_MAX_FETCH_BYTES", 128)
monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", len(body) + 100)
out = _fetch_with(monkeypatch, body, "text/plain")
assert "Cross-reference data was fetched" in out
def test_pdf_extraction_caps_pages_and_intermediate_text(monkeypatch):
from core.rag.parsers import Page
seen = {}
def fake_parse(data, *, max_pages = None):
seen["max_pages"] = max_pages
pages = [Page(text = "x" * 1000, page_number = i, char_count = 1000) for i in range(1, 51)]
return pages, 60 # document actually has more pages than the cap
monkeypatch.setattr("core.rag.parsers.parse_pdf_bytes", fake_parse)
text = tools._extract_pdf_text(b"unused")
assert seen["max_pages"] == tools._MAX_WEB_PDF_PAGES
assert len(text) <= tools._MAX_PAGE_CHARS
assert "text limited to 16,000 characters" in text
assert "page processing capped at 50 pages" in text
def test_pdf_exactly_at_page_cap_not_marked_capped(monkeypatch):
from core.rag.parsers import Page
# Exactly _MAX_WEB_PDF_PAGES pages are fully read, so no "capped" marker.
monkeypatch.setattr(
"core.rag.parsers.parse_pdf_bytes",
lambda data, *, max_pages = None: (
[Page(text = "short", page_number = i, char_count = 5) for i in range(1, 51)],
50,
),
)
text = tools._extract_pdf_text(b"unused")
assert "page processing capped" not in text
assert "## Page 50\n\nshort" in text
def test_pdf_page_cap_does_not_claim_later_pages_are_textless(monkeypatch):
from core.rag.parsers import Page
monkeypatch.setattr(
"core.rag.parsers.parse_pdf_bytes",
lambda data, *, max_pages = None: (
[Page(text = "", page_number = i, char_count = 0) for i in range(1, 51)],
60,
),
)
assert tools._extract_pdf_text(b"unused") == (
"(PDF contains no extractable text in the first 50 pages)"
)
def test_pdf_result_discarded_after_fetch_deadline(monkeypatch):
clock = {"time": 1000.0}
monkeypatch.setattr(tools.time, "monotonic", lambda: clock["time"])
def slow_extract(data):
clock["time"] += 10.0
return "late PDF text"
monkeypatch.setattr(tools, "_extract_pdf_text", slow_extract)
out = _fetch_with(monkeypatch, _pdf_bytes("Readable text"), "application/pdf")
assert out == "Failed to fetch URL: timed out."
def test_text_octet_stream_kept_after_sniffing(monkeypatch):
@ -154,7 +275,6 @@ def test_valid_utf8_binary_caught_by_control_chars(monkeypatch):
@pytest.mark.parametrize(
"magic",
[
b"%PDF-",
b"PK\x03\x04",
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",
b"\x1f\x8b",
@ -180,8 +300,8 @@ def test_text_labeled_binary_caught_by_magic(monkeypatch, magic):
b"\t\xef\xbb\xbf ",
],
)
def test_pdf_magic_after_harmless_prefix(monkeypatch, prefix):
body = prefix + b"%PDF-1.7\n" + b"1 0 obj<</Type/Catalog>>endobj\n" * 100
def test_binary_magic_after_harmless_prefix(monkeypatch, prefix):
body = prefix + b"\x1f\x8b" + b" printable text-heavy body" * 100
out = _fetch_with(monkeypatch, body, "text/plain")
assert "binary content" in out
@ -243,10 +363,10 @@ def test_html_page_unaffected(monkeypatch):
def test_content_type_sanitized_in_message(monkeypatch):
# Do not echo obs-folded header content into the model response.
out = _fetch_with(monkeypatch, b"\x00\x01\x02" * 500, "application/pdf\r\n data: injected")
out = _fetch_with(monkeypatch, b"PK\x03\x04" * 500, "application/zip\r\n data: injected")
assert "\n" not in out and "\r" not in out
assert "injected" not in out
assert "application/pdf" in out
assert "application/zip" in out
@pytest.mark.parametrize(

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

@ -49,6 +49,11 @@ logger = get_logger(__name__)
# and spawn workers copy os.environ. setdefault so an explicit user override wins.
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
# Studio workers can import MLX without importing unsloth first, so mirror the
# package bootstrap here. Keep an explicit user value authoritative.
if platform.system() == "Darwin" and platform.machine() == "arm64":
os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1")
# ========== Device Enum ==========

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

@ -1551,7 +1551,7 @@ export function AppSidebar() {
<SidebarMenuButton
size="lg"
aria-label={t("shell.accountMenu", { name: displayTitle })}
className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-[14px] group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] pl-2 pr-[45px] py-[3px] rounded-[14px] group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
>
<div className="flex shrink-0 items-center">
<UserAvatar
@ -1561,21 +1561,12 @@ export function AppSidebar() {
className="!size-[32px] group-data-[collapsible=icon]:!rounded-full"
/>
</div>
<div className="flex flex-col gap-px leading-tight group-data-[collapsible=icon]:hidden">
{/* min-w-0 so long names truncate instead of overflowing;
pr on the button reserves room for the settings cog */}
<div className="flex min-w-0 flex-1 flex-col gap-px leading-tight group-data-[collapsible=icon]:hidden">
<span className="truncate font-heading text-[13.5px] tracking-[0.025em] dark:tracking-[0.04em] font-semibold text-nav-fg">{displayTitle}</span>
<span className="truncate text-[11.5px] tracking-nav text-muted-foreground">Unsloth</span>
</div>
{/* settings cog (replaces the up/down chevron) */}
<span
aria-hidden="true"
className="ml-auto flex size-[32px] shrink-0 items-center justify-center text-muted-foreground group-data-[collapsible=icon]:hidden"
>
<HugeiconsIcon
icon={Settings02Icon}
strokeWidth={1.5}
className="!size-[18px]"
/>
</span>
</SidebarMenuButton>
</DropdownMenuTrigger>
<DropdownMenuContent
@ -1690,6 +1681,20 @@ export function AppSidebar() {
)}
</DropdownMenuContent>
</DropdownMenu>
{/* settings cog; sibling of the trigger (buttons cannot nest),
overlaid on the row's right edge, opens settings directly */}
<button
type="button"
aria-label={t("shell.navigation.settings")}
onClick={() => useSettingsDialogStore.getState().openDialog()}
className="absolute right-2 top-1/2 flex size-[32px] -translate-y-1/2 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-black/10 hover:text-foreground dark:hover:bg-white/10 focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring group-data-[collapsible=icon]:hidden"
>
<HugeiconsIcon
icon={Settings02Icon}
strokeWidth={1.5}
className="!size-[18px]"
/>
</button>
</SidebarMenuItem>
</SidebarMenu>
</SidebarFooter>

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

@ -1326,7 +1326,9 @@ const ThreadWelcome: FC<{
useEffect(() => {
// Prefer the nickname; otherwise first name only. Blank falls back to none.
const name = nickname.trim() || (displayName.trim().split(/\s+/)[0] ?? "");
const raw = nickname.trim() || (displayName.trim().split(/\s+/)[0] ?? "");
// Cap very long names so the greeting stays on one line.
const name = raw.length > 20 ? `${raw.slice(0, 20)}…` : raw;
setWelcome(buildWelcome(new Date().getHours(), name));
}, [displayName, nickname]);
@ -2707,6 +2709,7 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
const setCodeToolsEnabled = useChatRuntimeStore((s) => s.setCodeToolsEnabled);
const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled);
const showCanvasMenuItem = useChatRuntimeStore((s) => s.showCanvasMenuItem);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const setMcpEnabledForChat = useChatRuntimeStore(
(s) => s.setMcpEnabledForChat,
@ -2957,7 +2960,8 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
</DropdownMenuSubContent>
</DropdownMenuSub>
),
canvas: (
// Hidden by default; enabled from Settings > Chat > Canvas.
canvas: showCanvasMenuItem ? (
<DropdownMenuItem
className={artifactsEnabled ? "text-primary font-medium" : undefined}
onSelect={() => setArtifactsEnabled(!artifactsEnabled)}
@ -2968,7 +2972,7 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
) : null}
</DropdownMenuItem>
),
) : null,
bypassPermissions: <BypassPermissionsMenuItem />,
projects: (
<DropdownMenuSub>

59
studio/frontend/src/components/ui/dropdown-menu.tsx
View file

@ -2,10 +2,11 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { DropdownMenu as DropdownMenuPrimitive } from "radix-ui";
import type * as React from "react";
import * as React from "react";
import { Tick02Icon } from "@/lib/tick-icon";
import { ChevronRightStandardIcon } from "@/lib/chevron-icons";
import { useIsMobile } from "@/hooks/use-mobile";
import { cn } from "@/lib/utils";
import { HugeiconsIcon } from "@hugeicons/react";
@ -206,6 +207,14 @@ function DropdownMenuShortcut({
);
}
function assignRef<T>(ref: React.Ref<T> | undefined, value: T | null) {
if (typeof ref === "function") {
ref(value);
} else if (ref) {
ref.current = value;
}
}
function DropdownMenuSub({
...props
}: React.ComponentProps<typeof DropdownMenuPrimitive.Sub>) {
@ -242,17 +251,63 @@ function DropdownMenuSubTrigger({
function DropdownMenuSubContent({
className,
sideOffset,
style,
ref,
...props
}: React.ComponentProps<typeof DropdownMenuPrimitive.SubContent>) {
const isMobile = useIsMobile();
const [contentWidth, setContentWidth] = React.useState(0);
const resizeObserverRef = React.useRef<ResizeObserver | null>(null);
const composedRef = React.useCallback(
(
element: React.ComponentRef<
typeof DropdownMenuPrimitive.SubContent
> | null,
) => {
resizeObserverRef.current?.disconnect();
resizeObserverRef.current = null;
assignRef(ref, element);
if (!element) return;
const updateContentWidth = () => {
setContentWidth(element.offsetWidth);
};
updateContentWidth();
if (typeof ResizeObserver !== "undefined") {
resizeObserverRef.current = new ResizeObserver(updateContentWidth);
resizeObserverRef.current.observe(element);
}
},
[ref],
);
React.useEffect(
() => () => {
resizeObserverRef.current?.disconnect();
},
[],
);
const compactSideOffset =
isMobile && contentWidth > 0 ? -contentWidth : sideOffset;
return (
// Portaled like DropdownMenuContent: rendered inline, the fixed popper
// wrapper is a descendant of the parent menu's scroll container, so any
// transform there turns on overflow clipping and hides the submenu.
<DropdownMenuPrimitive.Portal>
<DropdownMenuPrimitive.SubContent
ref={composedRef}
data-slot="dropdown-menu-sub-content"
sideOffset={compactSideOffset}
style={{
...style,
visibility:
isMobile && contentWidth === 0 ? "hidden" : style?.visibility,
}}
className={cn(
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 bg-popover text-popover-foreground min-w-36 rounded-lg p-1 duration-100 z-50 origin-(--radix-dropdown-menu-content-transform-origin) overflow-hidden",
"data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 bg-popover text-popover-foreground min-w-36 max-w-[calc(100vw-2rem)] rounded-lg p-1 duration-100 z-50 origin-(--radix-dropdown-menu-content-transform-origin) overflow-hidden",
className,
)}
{...props}

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

@ -613,6 +613,7 @@ export function SharedComposer({
);
const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled);
const showCanvasMenuItem = useChatRuntimeStore((s) => s.showCanvasMenuItem);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const setMcpEnabledForChat = useChatRuntimeStore(
@ -1411,7 +1412,8 @@ export function SharedComposer({
</DropdownMenuSubContent>
</DropdownMenuSub>
),
canvas: (
// Hidden by default; enabled from Settings > Chat > Canvas.
canvas: showCanvasMenuItem ? (
<DropdownMenuItem
className={artifactsEnabled ? "text-primary font-medium" : undefined}
onSelect={() => setArtifactsEnabled(!artifactsEnabled)}
@ -1422,7 +1424,7 @@ export function SharedComposer({
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
) : null}
</DropdownMenuItem>
),
) : null,
bypassPermissions: <BypassPermissionsMenuItem />,
projects: (
<DropdownMenuSub>

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

@ -2,7 +2,11 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { RememberedLoadSettings } from "@/components/assistant-ui/model-selector/remembered-load-settings";
import { cancelStagedModelDownload } from "@/features/hub";
import {
cancelStagedModelDownload,
mirrorHfTokenInto,
useHfTokenStore,
} from "@/features/hub";
import { toast } from "@/lib/toast";
import { create } from "zustand";
import { isExternalModelId, parseExternalModelId } from "../external-providers";
@ -23,14 +27,15 @@ import {
savePersistedChatSettingsPatch,
} from "../utils/chat-settings-storage";
import { useExternalProvidersStore } from "./external-providers-store";
import { PLUS_MENU_PINS_STORAGE_KEY } from "./plus-menu-prefs-store";
const HF_TOKEN_KEY = "unsloth_hf_token";
const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed";
export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled";
export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled";
export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled";
export const CHAT_ARTIFACTS_ENABLED_KEY = "unsloth_chat_artifacts_enabled";
export const CHAT_SHOW_CANVAS_MENU_ITEM_KEY =
"unsloth_chat_show_canvas_menu_item";
export const CHAT_COLLAPSE_HTML_ARTIFACTS_KEY =
"unsloth_chat_collapse_html_artifacts";
export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY =
@ -357,6 +362,24 @@ function saveBool(key: string, value: boolean): void {
}
}
// The visibility flag shipped after the menu pins, so when it is absent,
// profiles that had explicitly pinned Canvas keep it visible.
function loadShowCanvasMenuItem(): boolean {
const stored = loadOptionalBool(CHAT_SHOW_CANVAS_MENU_ITEM_KEY);
if (stored !== null) return stored;
if (!canUseStorage()) return false;
try {
const raw = localStorage.getItem(PLUS_MENU_PINS_STORAGE_KEY);
if (raw === null) return false;
const parsed = JSON.parse(raw) as {
state?: { pins?: { canvas?: boolean } };
};
return parsed.state?.pins?.canvas === true;
} catch {
return false;
}
}
/**
* "full" is intentionally not restorable: it disables the sandbox and every
* confirmation gate, so it must be re-enabled (through the warning dialog)
@ -474,17 +497,6 @@ export function saveSpeculativeType(value: string | null): void {
}
}
function notifyHfTokenChanged(value: string): void {
if (!canUseStorage()) return;
try {
window.dispatchEvent(
new CustomEvent(HF_TOKEN_CHANGED_EVENT, { detail: value }),
);
} catch {
// ignore
}
}
/** A local model staged for a deferred load (see `pendingSelection`). Shape is
* a subset of the load hook's `SelectedModelInput`, structurally assignable. */
export type PendingModelSelection = {
@ -642,6 +654,8 @@ type ChatRuntimeStore = {
codeToolsEnabled: boolean;
imageToolsEnabled: boolean;
artifactsEnabled: boolean;
// Whether the Canvas toggle is offered in the composer + menu (hidden by default).
showCanvasMenuItem: boolean;
collapseHtmlArtifacts: boolean;
allowArtifactNetworkAccess: boolean;
mcpEnabledForChat: boolean;
@ -818,6 +832,7 @@ type ChatRuntimeStore = {
enabled: boolean,
options?: { persist?: boolean },
) => void;
setShowCanvasMenuItem: (enabled: boolean) => void;
setCollapseHtmlArtifacts: (enabled: boolean) => void;
setAllowArtifactNetworkAccess: (enabled: boolean) => void;
setMcpEnabledForChat: (enabled: boolean) => void;
@ -1120,7 +1135,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
runningByThreadId: {},
cancelByThreadId: {},
autoTitle: false,
hfToken: loadString(HF_TOKEN_KEY, ""),
hfToken: useHfTokenStore.getState().token,
modelsError: null,
lastModelLoadError: null,
activeGgufVariant: null,
@ -1147,6 +1162,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false),
imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false),
artifactsEnabled: loadBool(CHAT_ARTIFACTS_ENABLED_KEY, false),
showCanvasMenuItem: loadShowCanvasMenuItem(),
collapseHtmlArtifacts: loadBool(CHAT_COLLAPSE_HTML_ARTIFACTS_KEY, false),
allowArtifactNetworkAccess: loadBool(
CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY,
@ -1324,11 +1340,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle);
return { autoTitle };
}),
setHfToken: (hfToken) => {
saveString(HF_TOKEN_KEY, hfToken);
set({ hfToken });
notifyHfTokenChanged(hfToken);
},
setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken),
setModelsError: (modelsError) => set({ modelsError }),
setLastModelLoadError: (lastModelLoadError) => set({ lastModelLoadError }),
setCheckpoint: (modelId, ggufVariant) =>
@ -1504,6 +1516,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
}
return { artifactsEnabled };
}),
setShowCanvasMenuItem: (showCanvasMenuItem) =>
set(() => {
saveBool(CHAT_SHOW_CANVAS_MENU_ITEM_KEY, showCanvasMenuItem);
return { showCanvasMenuItem };
}),
setCollapseHtmlArtifacts: (collapseHtmlArtifacts) =>
set((state) => {
saveBool(CHAT_COLLAPSE_HTML_ARTIFACTS_KEY, collapseHtmlArtifacts);
@ -1807,6 +1824,12 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
setContextUsage: (contextUsage) => set({ contextUsage }),
}));
// Mirror token edits made through the shared store (e.g. Studio's field).
const unsubscribeHfTokenMirror = mirrorHfTokenInto(useChatRuntimeStore);
if (import.meta.hot) {
import.meta.hot.dispose(unsubscribeHfTokenMirror);
}
export function resolveSpeculativeSettingsForLoad({
usePersistedPreference = false,
}: {

4
studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts
View file

@ -44,6 +44,8 @@ const DEFAULT_PINS: Record<PlusMenuItemId, boolean> = {
bypassPermissions: false,
};
export const PLUS_MENU_PINS_STORAGE_KEY = "unsloth_plus_menu_pins";
export interface PlusMenuPrefsState {
pins: Record<PlusMenuItemId, boolean>;
setPin: (id: PlusMenuItemId, value: boolean) => void;
@ -72,7 +74,7 @@ export const usePlusMenuPrefsStore = create<PlusMenuPrefsState>()(
})),
}),
{
name: "unsloth_plus_menu_pins",
name: PLUS_MENU_PINS_STORAGE_KEY,
// Backfill any ids added in a later release so persisted state from an
// older version still resolves every menu item.
merge: (persisted, current) => {

5
studio/frontend/src/features/hub/index.ts
View file

@ -2,3 +2,8 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { cancelStagedModelDownload } from "./download-manager";
export {
getHfToken,
mirrorHfTokenInto,
useHfTokenStore,
} from "./stores/hf-token-store";

25
studio/frontend/src/features/hub/stores/hf-token-store.ts
View file

@ -5,11 +5,9 @@ import { create } from "zustand";
import { bumpInventoryVersion } from "./inventory-events";
const HF_TOKEN_KEY = "unsloth_hf_token";
const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed";
const LEGACY_TRAINING_KEY = "unsloth_training_config_v1";
let storageSyncStarted = false;
let storageSyncListener: ((event: StorageEvent) => void) | null = null;
let tokenChangedListener: ((event: Event) => void) | null = null;
function canUseStorage(): boolean {
return typeof window !== "undefined";
@ -63,10 +61,6 @@ function stopStorageSync(): void {
window.removeEventListener("storage", storageSyncListener);
storageSyncListener = null;
}
if (tokenChangedListener !== null) {
window.removeEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener);
tokenChangedListener = null;
}
storageSyncStarted = false;
}
@ -100,10 +94,6 @@ export const useHfTokenStore = create<HfTokenStore>((set) => {
applyToken(event.newValue ?? "", false);
};
window.addEventListener("storage", storageSyncListener);
tokenChangedListener = (event) => {
applyToken((event as CustomEvent<string>).detail ?? "", false);
};
window.addEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener);
}
return {
@ -117,6 +107,21 @@ export function getHfToken(): string {
return useHfTokenStore.getState().token;
}
// Keep a plain zustand store's `hfToken` field in sync with the shared token:
// seed the current value, then mirror later edits. Returns the unsubscribe so
// callers can wire it to HMR disposal.
export function mirrorHfTokenInto<T extends { hfToken: string }>(store: {
getState: () => T;
setState: (partial: Partial<T>) => void;
}): () => void {
store.setState({ hfToken: getHfToken() } as Partial<T>);
return useHfTokenStore.subscribe((state) => {
if (store.getState().hfToken !== state.token) {
store.setState({ hfToken: state.token } as Partial<T>);
}
});
}
// HF's JS client throws on a non-empty token that isn't `hf_...` instead of
// browsing anonymously, so treat anything malformed as no token.
export function hfApiToken(

2
studio/frontend/src/features/settings/components/api-monitor-console.tsx
View file

@ -111,7 +111,7 @@ function MonitorEntry({
const reply = replyText || (entry.status === "running" ? "Waiting..." : "No reply");
return (
<article className="rounded-lg border border-border/70 bg-background">
<article className="min-w-0 rounded-lg border border-border/70 bg-background">
<button
type="button"
onClick={onToggle}

3
studio/frontend/src/features/settings/settings-dialog.tsx
View file

@ -253,7 +253,8 @@ export function SettingsDialog() {
<DialogDescription className="sr-only">
{t("settings.dialog.description")}
</DialogDescription>
<div className="flex h-full min-h-0 max-sm:flex-col">
{/* Keep tab content from expanding the dialog grid. */}
<div className="flex h-full min-h-0 min-w-0 w-full max-sm:flex-col">
<aside className="font-heading flex w-[248px] shrink-0 flex-col border-r border-sidebar-border bg-muted/20 p-2 dark:border-r-0 max-sm:w-full max-sm:border-r-0 max-sm:border-b max-sm:border-sidebar-border">
<div className="relative mx-1 mt-3 mb-2 max-sm:hidden">
<HugeiconsIcon

17
studio/frontend/src/features/settings/tabs/chat-tab.tsx
View file

@ -175,6 +175,12 @@ export function ChatTab() {
const [clearing, setClearing] = useState(false);
const autoTitle = useChatRuntimeStore((state) => state.autoTitle);
const setAutoTitle = useChatRuntimeStore((state) => state.setAutoTitle);
const showCanvasMenuItem = useChatRuntimeStore(
(state) => state.showCanvasMenuItem,
);
const setShowCanvasMenuItem = useChatRuntimeStore(
(state) => state.setShowCanvasMenuItem,
);
const collapseHtmlArtifacts = useChatRuntimeStore(
(state) => state.collapseHtmlArtifacts,
);
@ -408,9 +414,16 @@ export function ChatTab() {
>
{PLUS_MENU_SETTINGS.map((item) => (
<SettingsRow key={item.id} label={item.label} icon={item.icon}>
{/* Canvas toggles menu visibility; the rest toggle pin placement. */}
<Switch
checked={plusPins[item.id]}
onCheckedChange={() => togglePlusPin(item.id)}
checked={
item.id === "canvas" ? showCanvasMenuItem : plusPins[item.id]
}
onCheckedChange={
item.id === "canvas"
? setShowCanvasMenuItem
: () => togglePlusPin(item.id)
}
/>
</SettingsRow>
))}

7
studio/frontend/src/features/settings/tabs/general-tab.tsx
View file

@ -654,9 +654,10 @@ export function GeneralTab() {
description={t("settings.general.rag.embeddingModelDescription", {
defaultModel: embeddingModel?.defaultEmbeddingModel ?? "",
})}
className="max-[360px]:flex-col max-[360px]:items-stretch max-[360px]:gap-3"
>
<div className="flex flex-col items-end gap-1">
<div className="flex items-center gap-2">
<div className="flex flex-col items-end gap-1 max-[360px]:w-full">
<div className="flex items-center gap-2 max-[360px]:w-full">
<EmbeddingModelCombobox
value={draftEmbeddingModel}
onChange={(next) => {
@ -668,7 +669,7 @@ export function GeneralTab() {
disabled={!embeddingModel}
placeholder={embeddingModel?.defaultEmbeddingModel ?? ""}
ariaLabel={t("settings.general.rag.embeddingModel")}
className="w-[220px]"
className="w-[220px] max-[360px]:min-w-0 max-[360px]:flex-1"
/>
<Button
variant="outline"

20
studio/frontend/src/features/studio/sections/dataset-section.tsx
View file

@ -54,7 +54,7 @@ import {
// Imported directly from the store module rather than the "@/features/training"
// barrel to avoid an import cycle (the barrel re-exports this section's siblings).
import { hasSeparateStreamingEvalSplit } from "@/features/training/stores/training-config-store";
import { useDebouncedValue, useHfTokenValidation } from "@/hooks";
import { useDebouncedValue } from "@/hooks";
import { translate, useT } from "@/i18n";
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
import { toast } from "@/lib/toast";
@ -398,9 +398,6 @@ export function DatasetSection() {
enabled: pickerTab === "huggingface",
});
const { error: tokenValidationError, isChecking: isCheckingToken } =
useHfTokenValidation(hfToken);
const hfResultIds = useMemo(() => {
const ids = hfResults.map((r) => r.id);
if (dataset && !ids.includes(dataset)) {
@ -683,11 +680,7 @@ export function DatasetSection() {
title={t("studio.dataset.title")}
description={t("studio.dataset.description")}
accent="indigo"
className={`dark:shadow-border ${
advancedOpen || (datasetSource === "upload" && uploadedFile)
? "min-h-studio-config-column"
: "h-studio-config-column"
}`}
className="dark:shadow-border min-h-studio-config-column"
>
<div className="flex min-w-0 flex-col gap-4">
{(() => {
@ -1009,9 +1002,9 @@ export function DatasetSection() {
</ComboboxContent>
</Combobox>
</div>
{(tokenValidationError ?? hfSearchError) && (
{hfSearchError && (
<p className="text-xs text-destructive">
{tokenValidationError ?? hfSearchError}
{hfSearchError}
{" — "}
<a
href="https://huggingface.co/settings/tokens"
@ -1023,11 +1016,6 @@ export function DatasetSection() {
</a>
</p>
)}
{isCheckingToken && (
<p className="text-xs text-muted-foreground">
{t("studio.dataset.checkingToken")}
</p>
)}
{pickerTab !== activeSourceTab && (
<p className="text-[11px] text-muted-foreground">
{t("studio.dataset.browsingSource", {

7
studio/frontend/src/features/studio/sections/params-section.tsx
View file

@ -192,7 +192,6 @@ export function ParamsSection(): ReactElement {
const showVisionImageSize = showVisionLora && !isDeepseekOcr;
const [loraOpen, setLoraOpen] = useState(false);
const [hyperOpen, setHyperOpen] = useState(false);
const needsExpandedHeight = isCpt || (isLora && loraOpen) || hyperOpen;
const [ctxInput, setCtxInput] = useState(String(store.contextLength));
const ctxAnchorRef = useRef<HTMLDivElement>(null);
const ctxItems = CONTEXT_LENGTHS.map(String);
@ -233,11 +232,7 @@ export function ParamsSection(): ReactElement {
title={t("studio.params.title")}
description={t("studio.params.description")}
accent="orange"
className={`${
needsExpandedHeight
? "min-h-studio-config-column"
: "h-studio-config-column"
} duration-150`}
className="min-h-studio-config-column"
>
<div className="flex flex-col gap-4">
<div className="flex flex-col gap-2">

3
studio/frontend/src/features/studio/sections/training-section.tsx
View file

@ -53,7 +53,6 @@ export function TrainingSection() {
((!store.isVisionModel && store.isDatasetImage === true) ||
(!store.isAudioModel && store.isDatasetAudio === true));
const configValidation = validateTrainingConfig(store);
const hasMessage = !!(startError || isIncompatible || (!configValidation.ok && configValidation.message));
const fileInputRef = useRef<HTMLInputElement>(null);
const handleFileUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
@ -125,7 +124,7 @@ export function TrainingSection() {
title={t("studio.training.title")}
description={t("studio.training.description")}
accent="blue"
className={hasMessage ? "min-h-studio-config-column" : "h-studio-config-column"}
className="min-h-studio-config-column"
>
<div className="flex flex-col gap-4">
{/* Loss chart */}

24
studio/frontend/src/features/training/stores/training-config-store.ts
View file

@ -3,6 +3,7 @@
import { CPT_TARGET_MODULES, DEFAULT_HYPERPARAMS, LR_DEFAULT_CPT, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS, TARGET_MODULES } from "@/config/training";
import { authFetch } from "@/features/auth";
import { getHfToken, mirrorHfTokenInto, useHfTokenStore } from "@/features/hub";
import { isAdapterMethod } from "@/types/training";
import type { DatasetFormat } from "@/types/training";
import type { ModelType, StepNumber, TrainingMethod } from "@/types/training";
@ -117,7 +118,9 @@ let _datasetFormatAutoForcedByCpt = false;
// modelType / isVisionModel / isAudioModel persist so multimodal-only UI
// paints right on reload; the model-config fetch still re-derives them.
// hfToken mirrors the shared hf-token-store and is persisted there instead.
const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set([
"hfToken",
"isCheckingVision",
"isEmbeddingModel",
"isLoadingModelDefaults",
@ -632,8 +635,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
),
);
},
setHfToken: (hfToken) =>
set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }),
setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken),
setDatasetSource: (datasetSource) => set({ datasetSource }),
selectHfDataset: (dataset) => {
_datasetCheckController?.abort();
@ -923,7 +925,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
_learningRateManuallySet = false;
_yamlLearningRate = undefined;
clearCptDatasetFormatTracking();
set(initialState);
set({ ...initialState, hfToken: getHfToken() });
},
resetToModelDefaults: () => {
const { selectedModel } = get();
@ -947,7 +949,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
},
{
name: "unsloth_training_config_v1",
version: 11,
version: 12,
migrate: (persisted, version) => {
const s = persisted as Record<string, unknown>;
if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@ -1000,6 +1002,15 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
// own version guard.
s.datasetStreaming ??= false;
}
if (version < 12) {
// hfToken moved to the shared hf-token-store; seed it once so an
// existing Studio-only token isn't lost.
const legacyToken = typeof s.hfToken === "string" ? s.hfToken.trim() : "";
if (legacyToken && !getHfToken()) {
useHfTokenStore.getState().setToken(legacyToken);
}
delete s.hfToken;
}
return s as unknown as TrainingConfigStore;
},
partialize: partializePersistedState,
@ -1022,3 +1033,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
},
),
);
const unsubscribeHfTokenMirror = mirrorHfTokenInto(useTrainingConfigStore);
if (import.meta.hot) {
import.meta.hot.dispose(unsubscribeHfTokenMirror);
}

6
studio/frontend/src/index.css
View file

@ -2048,15 +2048,11 @@ html[data-chat-font] .aui-root {
}
}
/* Fine-tuning Studio: equal default height, expandable when needed (md+) */
/* Fine-tuning Studio: equal minimum card height, grows with content (md+) */
.min-h-studio-config-column {
@apply md:min-h-[520px];
}
.h-studio-config-column {
@apply md:h-[520px];
}
[data-streamdown="unordered-list"] {
list-style-type: disc;
list-style-position: outside;

82
tests/python/test_get_lora_parameters_bias_fp8_block_size.py Normal file
View file

@ -0,0 +1,82 @@
import ast
from pathlib import Path
def _load_function(name):
# Extract a function from kernels/utils.py without importing unsloth (which
# needs a GPU / torch / bitsandbytes). The functions exercised here only use
# getattr and the _FP8_WEIGHT_DTYPES name on the paths under test.
source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
funcs = [
node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name
]
assert len(funcs) == 1, (name, funcs)
namespace = {"getattr": getattr, "_FP8_WEIGHT_DTYPES": ()}
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
exec(compile(module, str(source), "exec"), namespace)
return namespace[name]
class _Obj:
pass
def _make_disabled_block_fp8_proj(block_size):
# A merged/disabled projection whose base layer is a block-fp8 weight that
# ships a non-default block size on its checkpoint.
weight = _Obj()
weight.quant_state = _Obj()
base_layer = _Obj()
base_layer.weight = weight
base_layer.quant_method = "fp8"
base_layer.block_size = block_size
base_layer.bias = None
proj = _Obj()
proj.base_layer = base_layer
proj.merged = True
proj.disable_adapters = True
return proj, weight.quant_state
def test_bias_variant_propagates_fp8_block_size_on_disabled_path():
# Downstream fp8 kernels read getattr(weight_scale, "block_size", [128, 128]),
# so the checkpoint's real block size must survive the merged/disabled path,
# exactly as it does for the non-bias sibling get_lora_parameters.
get_lora_parameters_bias = _load_function("get_lora_parameters_bias")
proj, weight_scale = _make_disabled_block_fp8_proj([64, 128])
get_lora_parameters_bias(proj)
assert getattr(weight_scale, "block_size", [128, 128]) == [64, 128]
def _make_decompressed_merged_proj():
# A merged compressed-tensors layer that was decompressed back to bf16. It keeps
# quant_method == "fp8" from the checkpoint metadata, but the live weight is bf16
# so there is no quant state to attach a block size to.
weight = _Obj()
weight.dtype = "bfloat16"
base_layer = _Obj()
base_layer.weight = weight
base_layer.quant_method = "fp8"
base_layer.block_size = [128, 128]
base_layer.bias = None
proj = _Obj()
proj.base_layer = base_layer
proj.merged = True
proj.disable_adapters = True
return proj
def test_bias_variant_keeps_none_quant_state_for_decompressed_layer():
# Such a layer has no quant state, and fast_linear_forward relies on getting
# W_quant None back so it can fall back to a plain matmul, so setting the block
# size must not assume a quant state is present.
get_lora_parameters_bias = _load_function("get_lora_parameters_bias")
W, W_quant = get_lora_parameters_bias(_make_decompressed_merged_proj())[:2]
assert W_quant is None
assert getattr(W, "block_size", None) == [128, 128]

81
tests/python/test_get_lora_parameters_fp8_block_size.py Normal file
View file

@ -0,0 +1,81 @@
import ast
from pathlib import Path
def _load_function(name):
# Extract a function from kernels/utils.py without importing unsloth (which
# needs a GPU / torch / bitsandbytes). get_lora_parameters only uses getattr,
# hasattr and the _FP8_WEIGHT_DTYPES name on the paths under test.
source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
funcs = [
node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name
]
assert len(funcs) == 1, (name, funcs)
namespace = {"getattr": getattr, "hasattr": hasattr, "_FP8_WEIGHT_DTYPES": ()}
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
exec(compile(module, str(source), "exec"), namespace)
return namespace[name]
class _Obj:
pass
def _make_disabled_block_fp8_proj(block_size):
# A merged/disabled projection whose base layer is a block-fp8 weight that
# ships a non-default block size on its checkpoint.
weight = _Obj()
weight.quant_state = _Obj()
base_layer = _Obj()
base_layer.weight = weight
base_layer.quant_method = "fp8"
base_layer.block_size = block_size
proj = _Obj()
proj.base_layer = base_layer
proj.merged = True
proj.disable_adapters = True
return proj, weight.quant_state
def test_propagates_fp8_block_size_on_disabled_path():
# get_lora_parameters already sets block_size before its early return; downstream
# fp8 kernels read getattr(weight_scale, "block_size", [128, 128]), so the
# checkpoint's real block size must survive the merged/disabled path.
get_lora_parameters = _load_function("get_lora_parameters")
proj, weight_scale = _make_disabled_block_fp8_proj([64, 128])
get_lora_parameters(proj)
assert getattr(weight_scale, "block_size", [128, 128]) == [64, 128]
def _make_decompressed_merged_proj():
# A merged compressed-tensors layer that was decompressed back to bf16. It keeps
# quant_method == "fp8" from the checkpoint metadata, but the live weight is bf16
# so there is no quant state to attach a block size to.
weight = _Obj()
weight.dtype = "bfloat16"
base_layer = _Obj()
base_layer.weight = weight
base_layer.quant_method = "fp8"
base_layer.block_size = [128, 128]
proj = _Obj()
proj.base_layer = base_layer
proj.merged = True
proj.disable_adapters = True
return proj
def test_keeps_none_quant_state_for_decompressed_layer():
# Mirrors the get_lora_parameters_bias guard: with no quant state, assigning
# W_quant.block_size must not assume one is present, or it raises AttributeError
# on None. fast_lora relies on getting W_quant None back to fall back to a plain
# matmul, so this path must stay crash-free.
get_lora_parameters = _load_function("get_lora_parameters")
W, W_quant = get_lora_parameters(_make_decompressed_merged_proj())[:2]
assert W_quant is None
assert getattr(W, "block_size", None) == [128, 128]

76
tests/saving/test_export_dispatch.py
View file

@ -8,6 +8,8 @@ regressions that pure AST checks cannot (e.g. wrong scheme/suffix/outtype passed
from __future__ import annotations
import inspect
import pytest
import unsloth.save as save_mod
@ -126,6 +128,80 @@ def test_gguf_lora_push_to_hub_is_rejected(tmp_path):
)
# The above rejection points users at push_to_hub_gguf(save_method='lora'), so that path
# has to work; it is only ever exercised here.
def test_push_to_hub_gguf_lora_dispatches(monkeypatch):
seen = {}
monkeypatch.setattr(
save_mod,
"_unsloth_save_lora_gguf",
lambda model, tok, sd, **kw: seen.update(kw),
)
save_mod.unsloth_push_to_hub_gguf(
_FakeModel(),
"repo/id",
tokenizer = object(),
save_method = "lora",
quantization_method = "q8_0",
)
assert seen.get("outtype") == "q8_0"
assert seen.get("push_to_hub") is True
def test_push_to_hub_gguf_lora_skips_non_main_process(monkeypatch):
calls = []
monkeypatch.setattr(
save_mod,
"_unsloth_save_lora_gguf",
lambda *a, **kw: calls.append(kw),
)
result = save_mod.unsloth_push_to_hub_gguf(
_FakeModel(),
"repo/id",
tokenizer = object(),
save_method = "lora",
is_main_process = False,
)
assert result is None
assert calls == []
def test_push_to_hub_gguf_skips_non_main_process_before_merged_conversion(monkeypatch):
calls = []
monkeypatch.setattr(
save_mod,
"unsloth_save_pretrained_gguf",
lambda **kw: calls.append(kw),
)
result = save_mod.unsloth_push_to_hub_gguf(
_FakeModel(),
"repo/id",
tokenizer = object(),
is_main_process = False,
)
assert result is None
assert calls == []
def test_push_to_hub_gguf_preserves_positional_max_shard_size():
bound = inspect.signature(save_mod.unsloth_push_to_hub_gguf).bind(
_FakeModel(),
"repo/id",
object(),
"q4_k_m",
None,
None,
None,
None,
"token",
"50GB",
)
assert bound.arguments["max_shard_size"] == "50GB"
assert "is_main_process" not in bound.arguments
# -- torchao PTQ / QAT dispatch ------------------------------------------------------------

58
tests/security/test_release_desktop_permissions.py Normal file
View file

@ -0,0 +1,58 @@
"""Permission-boundary checks for the desktop release workflow."""
from pathlib import Path
import yaml
REPO_ROOT = Path(__file__).resolve().parents[2]
WORKFLOW = REPO_ROOT / ".github" / "workflows" / "release-desktop.yml"
def _workflow():
return yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8"))
def test_only_publish_job_can_write_repository_contents():
workflow = _workflow()
assert workflow["permissions"] == {"contents": "read"}
write_jobs = [
name
for name, job in workflow["jobs"].items()
if job.get("permissions", {}).get("contents") == "write"
]
assert write_jobs == ["publish-release"]
def test_build_matrix_hands_off_assets_without_release_credentials():
jobs = _workflow()["jobs"]
build = jobs["build"]
publish = jobs["publish-release"]
assert "permissions" not in build
tauri_steps = [
step
for step in build["steps"]
if step.get("uses", "").startswith("tauri-apps/tauri-action@")
]
assert len(tauri_steps) == 3
for step in tauri_steps:
assert "GITHUB_TOKEN" not in step.get("env", {})
assert not {"releaseId", "tagName", "releaseName"} & step.get("with", {}).keys()
assert any(
step.get("uses", "").startswith("actions/upload-artifact@") for step in build["steps"]
)
assert any(
step.get("uses", "").startswith("actions/download-artifact@") for step in publish["steps"]
)
assert "build" in publish["needs"]
release_step = next(
step
for step in publish["steps"]
if step.get("name") == "Create or validate versioned release"
)
assert "gh release view" in release_step["run"]
assert "--json tagName,isDraft,isPrerelease" in release_step["run"]

41
tests/studio/test_compact_dropdown_submenus.py Normal file
View file

@ -0,0 +1,41 @@
"""Compact viewport contracts for nested dropdown menus."""
from pathlib import Path
REPO = Path(__file__).resolve().parents[2]
DROPDOWN_MENU = REPO / "studio/frontend/src/components/ui/dropdown-menu.tsx"
FRONTEND_SRC = REPO / "studio/frontend/src"
def test_shared_submenu_uses_its_layout_width_on_mobile():
source = DROPDOWN_MENU.read_text(encoding = "utf-8")
assert 'import { useIsMobile } from "@/hooks/use-mobile";' in source
assert "element.offsetWidth" in source
assert "element.getBoundingClientRect().width" not in source
assert "new ResizeObserver(updateContentWidth)" in source
assert "isMobile && contentWidth > 0 ? -contentWidth : sideOffset" in source
assert "sideOffset={compactSideOffset}" in source
assert 'isMobile && contentWidth === 0 ? "hidden"' in source
assert "-248" not in source
def test_shared_submenu_never_exceeds_the_compact_viewport():
source = DROPDOWN_MENU.read_text(encoding = "utf-8")
assert "max-w-[calc(100vw-2rem)]" in source
def test_consumers_do_not_duplicate_compact_offset_logic():
for path in FRONTEND_SRC.rglob("*.tsx"):
if path == DROPDOWN_MENU:
continue
source = path.read_text(encoding = "utf-8")
assert "compactSubmenuOffset" not in source, path
def test_all_submenu_consumers_use_the_shared_primitive():
for path in FRONTEND_SRC.rglob("*.tsx"):
if path == DROPDOWN_MENU:
continue
source = path.read_text(encoding = "utf-8")
assert "DropdownMenuPrimitive.SubContent" not in source, path

35
tests/studio/test_settings_compact_overflow_contract.py Normal file
View file

@ -0,0 +1,35 @@
"""Responsive overflow contracts for the settings dialog."""
from pathlib import Path
REPO = Path(__file__).resolve().parents[2]
SETTINGS_DIALOG = REPO / "studio/frontend/src/features/settings/settings-dialog.tsx"
API_MONITOR = REPO / "studio/frontend/src/features/settings/components/api-monitor-console.tsx"
GENERAL_TAB = REPO / "studio/frontend/src/features/settings/tabs/general-tab.tsx"
def test_dialog_content_can_shrink_inside_the_dialog_grid():
source = SETTINGS_DIALOG.read_text(encoding = "utf-8")
assert "flex h-full min-h-0 min-w-0 w-full max-sm:flex-col" in source
assert "relative flex min-h-0 min-w-0 flex-1 flex-col" in source
def test_api_monitor_entries_and_expanded_text_can_shrink():
source = API_MONITOR.read_text(encoding = "utf-8")
assert (
'<article className="min-w-0 rounded-lg border border-border/70 bg-background">' in source
)
assert (
'<section className="flex min-w-0 flex-col rounded-lg border border-border/70 bg-background">'
in source
)
assert source.count('className="max-h-44 overflow-auto whitespace-pre-wrap break-words') == 2
def test_embedding_model_controls_stack_on_the_narrowest_viewports():
source = GENERAL_TAB.read_text(encoding = "utf-8")
assert (
'className="max-[360px]:flex-col max-[360px]:items-stretch max-[360px]:gap-3"'
) in source
assert 'className="w-[220px] max-[360px]:min-w-0 max-[360px]:flex-1"' in source

55
tests/test_raw_text.py
View file

@ -241,6 +241,61 @@ def test_raw_text_loader():
os.unlink(test_file)
def test_smart_chunk_text_single_chunk_no_eos_returns_plain_list():
"""smart_chunk_text's single-chunk branch must return a plain list for
input_ids even when the tokenizer has no eos_token_id, matching the
multi-chunk branch's unconditional tolist()/list() conversion."""
class MockTensor:
def __init__(self, data):
self.data = data
def __getitem__(self, idx):
return self.data
def __len__(self):
return len(self.data)
def tolist(self):
return self.data
class MockTokenizerNoEos:
def __init__(self):
self.eos_token = None
self.eos_token_id = None
def __call__(
self,
text,
return_tensors = None,
add_special_tokens = False,
):
token_ids = list(range(len(text.split())))
if return_tensors == "pt":
return {"input_ids": [MockTensor(token_ids)]}
return {"input_ids": token_ids}
def decode(
self,
token_ids,
skip_special_tokens = False,
):
return " ".join(f"word_{i}" for i in token_ids)
loader = RawTextDataLoader(MockTokenizerNoEos(), chunk_size = 2048, stride = 512)
result = loader.smart_chunk_text(
"hello world short text", chunk_size = 2048, stride = 512, return_tokenized = True
)
input_ids = result[0]["input_ids"]
assert isinstance(
input_ids, list
), f"input_ids should be a plain list even without an eos_token_id, got {type(input_ids)}"
assert input_ids == [0, 1, 2, 3], f"unexpected input_ids: {input_ids}"
print("✅ test_smart_chunk_text_single_chunk_no_eos_returns_plain_list passed!")
return True
if __name__ == "__main__":
success = test_raw_text_loader()
success = test_smart_chunk_text_single_chunk_no_eos_returns_plain_list() and success
sys.exit(0 if success else 1)

5
unsloth/__init__.py
View file

@ -16,6 +16,11 @@ import os, importlib.util, platform
os.environ["UNSLOTH_IS_PRESENT"] = "1"
# Relax Metal's context-store timeout before MLX modules can initialize Metal.
# Keep an explicit user value authoritative.
if platform.system() == "Darwin" and platform.machine() == "arm64":
os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1")
# ── Windows console UTF-8 safety ─────────────────────────────────────────────
# Legacy Windows consoles (cp1252) can't encode Unsloth's emoji/box-drawing
# glyphs and crash with UnicodeEncodeError. Force stdout/stderr to UTF-8 only on

2
unsloth/dataprep/raw_text.py
View file

@ -154,9 +154,9 @@ class RawTextDataLoader:
if len(tokens) <= chunk_size:
# Fits in a single chunk
if return_tokenized:
tokens = tokens.tolist() if hasattr(tokens, "tolist") else list(tokens)
eos_token_id = getattr(self.tokenizer, "eos_token_id", None)
if eos_token_id is not None:
tokens = tokens.tolist() if hasattr(tokens, "tolist") else list(tokens)
tokens.append(eos_token_id)
attention_mask = [1] * len(tokens)

18
unsloth/kernels/utils.py
View file

@ -325,7 +325,10 @@ def get_lora_parameters(proj):
if getattr(base_layer, "quant_method", None) == "fp8":
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
# A decompressed compressed-tensors layer keeps quant_method == "fp8" while its
# weight is back to bf16, so it has no quant state to carry the block size.
if W_quant is not None:
W_quant.block_size = W.block_size
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
if getattr(proj, "disable_adapters", True) or proj.merged:
@ -375,14 +378,17 @@ def get_lora_parameters_bias(proj):
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
if getattr(proj, "disable_adapters", True) or proj.merged:
return W, W_quant, None, None, None, base_layer.bias
if getattr(base_layer, "quant_method", None) == "fp8":
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
# A decompressed compressed-tensors layer keeps quant_method == "fp8" while its
# weight is back to bf16, so it has no quant state to carry the block size.
if W_quant is not None:
W_quant.block_size = W.block_size
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
if getattr(proj, "disable_adapters", True) or proj.merged:
return W, W_quant, None, None, None, base_layer.bias
adapter = getattr(proj, "active_adapters", None)
if adapter is None:

6
unsloth/save.py
View file

@ -3143,6 +3143,7 @@ def unsloth_push_to_hub_gguf(
datasets: Optional[List[str]] = None,
save_method: str = None,
imatrix_file = None,
is_main_process: bool = True,
):
"""
Same as .push_to_hub(...) except 4bit weights are auto
@ -3175,11 +3176,11 @@ def unsloth_push_to_hub_gguf(
"""
if tokenizer is None:
raise ValueError("Unsloth: Saving to GGUF must have a tokenizer.")
if not is_main_process:
return None
# save_method="lora" exports the adapter itself as a GGUF LoRA (not a merged model).
if save_method is not None and str(save_method).lower() == "lora":
if not is_main_process:
return None # only the main rank converts and uploads, like the local lora branch
_qm = quantization_method
if isinstance(_qm, (list, tuple)) and len(_qm) == 1:
_qm = _qm[0] # the gguf API allows a list; unwrap a single outtype
@ -3233,6 +3234,7 @@ def unsloth_push_to_hub_gguf(
first_conversion = first_conversion,
push_to_hub = False, # Never push from here
token = token, # forwarded so imatrix_file=True can read a gated/private upstream
is_main_process = is_main_process,
max_shard_size = max_shard_size,
safe_serialization = safe_serialization,
temporary_location = temporary_location,