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

Merge branch 'main' into feature/chat-api

Browse source
This commit is contained in:
Roland Tannous 2026-05-05 20:35:27 +04:00 • committed by GitHub
commit a64dd983d5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
443 changed files with 76715 additions and 5838 deletions
Download patch file Download diff file Expand all files Collapse all files

13
.github/dependabot.yml vendored
View file

@ -9,19 +9,6 @@ updates:
actions:
patterns: ["*"]
- package-ecosystem: "pip"
directories:
- "/"
- "/studio/backend/plugins/data-designer-unstructured-seed"
- "/studio/backend/requirements"
- "/unsloth/kernels/moe"
schedule:
interval: "weekly"
open-pull-requests-limit: 10
groups:
pip:
patterns: ["*"]
- package-ecosystem: "bun"
directory: "/studio/frontend"
schedule:

226
.github/workflows/release-desktop.yml vendored Normal file
View file

@ -0,0 +1,226 @@
name: Release Desktop App
on:
workflow_dispatch:
inputs:
draft:
description: 'Create as draft release'
type: boolean
default: true
permissions:
contents: write
jobs:
build:
strategy:
fail-fast: false
max-parallel: 1
matrix:
include:
- platform: macos-latest
args: '--target aarch64-apple-darwin'
label: macOS (Apple Silicon)
# - platform: macos-latest
# args: '--target x86_64-apple-darwin'
# label: macOS (Intel)
- platform: ubuntu-22.04
args: ''
label: Linux (x64)
- platform: windows-latest
args: ''
label: Windows (x64)
name: Build ${{ matrix.label }}
runs-on: ${{ matrix.platform }}
env:
FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true
steps:
- uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5
# ── Linux dependencies ──
- name: Install Linux dependencies
if: matrix.platform == 'ubuntu-22.04'
run: |
sudo apt-get update
sudo apt-get install -y libwebkit2gtk-4.1-dev libayatana-appindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf
# ── Node.js ──
- name: Setup Node.js
uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020
with:
node-version: 24
- name: Install pinned Tauri CLI
run: npm install --save-dev --prefix studio @tauri-apps/cli@2.10.1
- name: Verify pinned Tauri CLI
shell: bash
run: |
out="$(npx --prefix studio tauri --version)"
echo "$out"
if [ "$out" != "tauri-cli 2.10.1" ]; then
echo "Expected tauri-cli 2.10.1, got $out" >&2
exit 1
fi
- name: Install frontend dependencies
working-directory: studio/frontend
run: npm install
- name: Verify backend package is published
shell: bash
run: |
node <<'JS'
const { readFileSync } = require('node:fs');
(async () => {
const cargo = readFileSync('studio/src-tauri/Cargo.toml', 'utf8');
const match = cargo.match(/^version\s*=\s*"([^"]+)"/m);
if (!match) throw new Error('Could not read desktop app version');
const appVersion = match[1];
const response = await fetch(`https://pypi.org/pypi/unsloth/${appVersion}/json`);
if (!response.ok) {
const message = 'Publish unsloth=={app_version} to PyPI before the desktop release';
throw new Error(`${message.replace('{app_version}', appVersion)} (HTTP ${response.status})`);
}
})();
JS
# ── Rust ──
- name: Install Rust stable
uses: dtolnay/rust-toolchain@stable
with:
targets: ${{ matrix.platform == 'macos-latest' && 'aarch64-apple-darwin,x86_64-apple-darwin' || '' }}
- name: Rust cache
uses: swatinem/rust-cache@42dc69e1aa15d09112580998cf2ef0119e2e91ae
with:
workspaces: 'studio/src-tauri -> target'
# ── macOS: import signing certificate ──
- name: Import Apple certificate
if: matrix.platform == 'macos-latest'
env:
APPLE_CERTIFICATE: ${{ secrets.APPLE_CERTIFICATE }}
APPLE_CERTIFICATE_PASSWORD: ${{ secrets.APPLE_CERTIFICATE_PASSWORD }}
KEYCHAIN_PASSWORD: ${{ secrets.KEYCHAIN_PASSWORD }}
run: |
echo $APPLE_CERTIFICATE | base64 --decode > certificate.p12
security create-keychain -p "$KEYCHAIN_PASSWORD" build.keychain
security default-keychain -s build.keychain
security unlock-keychain -p "$KEYCHAIN_PASSWORD" build.keychain
security set-keychain-settings -t 3600 -u build.keychain
security import certificate.p12 -k build.keychain -P "$APPLE_CERTIFICATE_PASSWORD" -T /usr/bin/codesign
security set-key-partition-list -S apple-tool:,apple:,codesign: -s -k "$KEYCHAIN_PASSWORD" build.keychain
security find-identity -v -p codesigning build.keychain
rm -f certificate.p12
# ── Windows: install Azure Trusted Signing CLI ──
- name: Install trusted-signing-cli
if: matrix.platform == 'windows-latest'
run: |
cargo install trusted-signing-cli --version 0.9.0 --locked
echo "$env:USERPROFILE\.cargo\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
# ── Windows: verify signing CLI is accessible ──
- name: Verify trusted-signing-cli
if: matrix.platform == 'windows-latest'
run: |
Write-Output "PATH: $env:PATH"
Get-Command trusted-signing-cli -ErrorAction SilentlyContinue || Write-Output "trusted-signing-cli NOT in PATH"
trusted-signing-cli --version || Write-Output "trusted-signing-cli failed to run"
# ── Linux: build + sign + upload ──
- name: Build Linux app
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 }}
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: desktop-v__VERSION__
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal).
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> 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: false
args: -v ${{ matrix.args }}
# ── macOS: build + sign + notarize + upload ──
- name: Build macOS app
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 }}
APPLE_ID: ${{ secrets.APPLE_ID }}
APPLE_PASSWORD: ${{ secrets.APPLE_PASSWORD }}
APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }}
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: desktop-v__VERSION__
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal).
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> 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: false
args: -v ${{ matrix.args }}
# ── Windows: build + sign + upload ──
- name: Build Windows app
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 }}
AZURE_CLIENT_SECRET: ${{ secrets.AZURE_CLIENT_SECRET }}
AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }}
AZURE_TRUSTED_SIGNING_ACCOUNT_NAME: ${{ secrets.AZURE_TRUSTED_SIGNING_ACCOUNT_NAME }}
AZURE_CERTIFICATE_PROFILE_NAME: ${{ secrets.AZURE_CERTIFICATE_PROFILE_NAME }}
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: desktop-v__VERSION__
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal).
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> 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: false
args: -v ${{ matrix.args }}

12
.gitignore vendored
View file

@ -204,6 +204,18 @@ tmp/
**/node_modules/
auth.db
# Tauri local build/generated output
studio/src-tauri/target/
studio/src-tauri/gen/
studio/src-tauri/artifacts/
studio/src-tauri/icons/android/
studio/src-tauri/icons/ios/
studio/src-tauri/icons/128x128@2x.png
studio/src-tauri/icons/64x64.png
studio/src-tauri/icons/Square*Logo.png
studio/src-tauri/icons/StoreLogo.png
studio/src-tauri/icons/squarehq.png
# Local working docs
**/CLAUDE.md
**/claude.md

2
.pre-commit-config.yaml
View file

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

68
README.md
View file

@ -1,34 +1,50 @@
<h1 align="center" style="margin:0;">
<a href="https://unsloth.ai/docs"><picture>
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/unslothai/unsloth/main/images/STUDIO%20WHITE%20LOGO.png">
<source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/unslothai/unsloth/main/images/STUDIO%20BLACK%20LOGO.png">
<img alt="Unsloth logo" src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/STUDIO%20BLACK%20LOGO.png" height="60" style="max-width:100%;">
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20white%20text.png">
<source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20black%20text.png">
<img alt="Unsloth logo" src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20black%20text.png" height="80" style="max-width:100%;">
</picture></a>
</h1>
<h3 align="center" style="margin: 0; margin-top: 0;">
Run and train AI models with a unified local interface.
Unsloth Studio lets you run and train models locally.
</h3>
<p align="center">
<a href="#-features">Features</a> •
<a href="#-quickstart">Quickstart</a> •
<a href="#-install">Quickstart</a> •
<a href="#-free-notebooks">Notebooks</a> •
<a href="https://unsloth.ai/docs">Documentation</a> •
<a href="https://www.reddit.com/r/unsloth/">Reddit</a>
<a href="https://unsloth.ai/docs">Documentation</a>
</p>
<a href="https://unsloth.ai/docs/new/studio">
<img alt="unsloth studio ui homepage" src="https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png" style="max-width: 100%; margin-bottom: 0;"></a>
<br>
<a href="https://unsloth.ai/docs/new/studio">
<img alt="unsloth studio ui homepage" src="https://github.com/user-attachments/assets/53ae17a9-d975-44ef-9686-efb4ebd0454d" style="max-width: 100%; margin-bottom: 0;"></a>
Unsloth Studio (Beta) lets you run and train text, [audio](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning), [embedding](https://unsloth.ai/docs/new/embedding-finetuning), [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) models on Windows, Linux and macOS.
## ⚡ Get started
#### macOS, Linux, WSL:
```bash
curl -fsSL https://unsloth.ai/install.sh | sh
```
#### Windows:
```powershell
irm https://unsloth.ai/install.ps1 | iex
```
#### Community:
- [Discord](https://discord.gg/unsloth)
- [𝕏 (Twitter)](https://x.com/UnslothAI)
- [Reddit](https://reddit.com/r/unsloth)
## ⭐ Features
Unsloth provides several key features for both inference and training:
Unsloth Studio (Beta) lets you run and train text, [audio](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning), [embedding](https://unsloth.ai/docs/new/embedding-finetuning), [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) models on Windows, Linux and macOS.
### Inference
* **Search + download + run models** including GGUF, LoRA adapters, safetensors
* **Export models**: [Save or export](https://unsloth.ai/docs/new/studio/export) models to GGUF, 16-bit safetensors and other formats.
* **Tool calling**: Support for [self-healing tool calling](https://unsloth.ai/docs/new/studio/chat#auto-healing-tool-calling) and web search
* **[Code execution](https://unsloth.ai/docs/new/studio/chat#code-execution)**: lets LLMs test code in Claude artifacts and sandbox environments
* [Auto-tune inference parameters](https://unsloth.ai/docs/new/studio/chat#auto-parameter-tuning) and customize chat templates.
* **[API inference endpoint](https://unsloth.ai/docs/basics/api)**: Deploy and run local LLMs in Claude Code, Codex tools with Unsloth
* [Auto set inference settings](https://unsloth.ai/docs/new/studio/chat#auto-parameter-tuning) and customize chat templates.
* We work directly with teams behind [gpt-oss](https://docs.unsloth.ai/new/gpt-oss-how-to-run-and-fine-tune#unsloth-fixes-for-gpt-oss), [Qwen3](https://www.reddit.com/r/LocalLLaMA/comments/1kaodxu/qwen3_unsloth_dynamic_ggufs_128k_context_bug_fixes/), [Llama 4](https://github.com/ggml-org/llama.cpp/pull/12889), [Mistral](models/tutorials/devstral-how-to-run-and-fine-tune.md), [Gemma 1-3](https://news.ycombinator.com/item?id=39671146), and [Phi-4](https://unsloth.ai/blog/phi4), where we’ve fixed bugs that improve model accuracy.
* Upload images, audio, PDFs, code, DOCX and more file types to chat with.
### Training
@ -40,7 +56,7 @@ Unsloth provides several key features for both inference and training:
* **Observability**: Monitor training live, track loss and GPU usage and customize graphs.
* [Multi-GPU](https://unsloth.ai/docs/basics/multi-gpu-training-with-unsloth) training is supported, with major improvements coming soon.
## ⚡ Quickstart
## 📥 Install
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.
### Unsloth Studio (web UI)
@ -64,8 +80,9 @@ irm https://unsloth.ai/install.ps1 | iex
#### Launch
```bash
unsloth studio -H 0.0.0.0 -p 8888
unsloth studio -p 8888
```
> For cloud VMs or LAN access, add `-H 0.0.0.0` to bind on all interfaces.
#### Update
To update, use the same install commands as above. Or run (does not work on Windows):
@ -109,18 +126,19 @@ You can use the same Docker image as Unsloth Studio.
For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth). <br>
To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
## ✨ Free Notebooks
## 📒 Free Notebooks
Train for free with our notebooks. Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Add dataset, run, then deploy your trained model.
Train for free with our notebooks. You can use our new [free Unsloth Studio notebook](https://colab.research.google.com/github/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb) to run and train models for free in a web UI.
Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Add dataset, run, then deploy your trained model.
| Model | Free Notebooks | Performance | Memory use |
|-----------|---------|--------|----------|
| **Gemma 4 (E2B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma4_(E2B)-Vision.ipynb) | 1.5x faster | 50% less |
| **Qwen3.5 (4B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_5_(4B)_Vision.ipynb) | 1.5x faster | 60% less |
| **gpt-oss (20B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-Fine-tuning.ipynb) | 2x faster | 70% less |
| **Qwen3.5 GSPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_5_(4B)_Vision_GRPO.ipynb) | 2x faster | 70% less |
| **gpt-oss (20B): GRPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb) | 2x faster | 80% less |
| **Qwen3: Advanced GRPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_(4B)-GRPO.ipynb) | 2x faster | 70% less |
| **Gemma 3 (4B) Vision** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3_(4B)-Vision.ipynb) | 1.7x faster | 60% less |
| **embeddinggemma (300M)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/EmbeddingGemma_(300M).ipynb) | 2x faster | 20% less |
| **Mistral Ministral 3 (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Ministral_3_VL_(3B)_Vision.ipynb) | 1.5x faster | 60% less |
| **Llama 3.1 (8B) Alpaca** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-Alpaca.ipynb) | 2x faster | 70% less |
@ -132,6 +150,9 @@ Train for free with our notebooks. Read our [guide](https://unsloth.ai/docs/get-
- See detailed documentation for Unsloth [here](https://unsloth.ai/docs)
## 🦥 Unsloth News
- **API inference endpoint**: Deploy and run local LLMs in Claude Code, Codex tools. [Guide](https://unsloth.ai/docs/basics/api)
- **Qwen3.6**: Qwen3.6-35B-A3B can now be trained and run in Unsloth Studio. [Blog](https://unsloth.ai/docs/models/qwen3.6)
- **Gemma 4**: Run and train Google’s new models directly in Unsloth. [Blog](https://unsloth.ai/docs/models/gemma-4)
- **Introducing Unsloth Studio**: our new web UI for running and training LLMs. [Blog](https://unsloth.ai/docs/new/studio)
- **Qwen3.5** - 0.8B, 2B, 4B, 9B, 27B, 35-A3B, 112B-A10B are now supported. [Guide + notebooks](https://unsloth.ai/docs/models/qwen3.5/fine-tune)
- Train **MoE LLMs 12x faster** with 35% less VRAM - DeepSeek, GLM, Qwen and gpt-oss. [Blog](https://unsloth.ai/docs/new/faster-moe)
@ -140,7 +161,6 @@ Train for free with our notebooks. Read our [guide](https://unsloth.ai/docs/get-
- New RoPE & MLP **Triton Kernels** & **Padding Free + Packing**: 3x faster training & 30% less VRAM. [Blog](https://unsloth.ai/docs/new/3x-faster-training-packing)
- **500K Context**: Training a 20B model with >500K context is now possible on an 80GB GPU. [Blog](https://unsloth.ai/docs/blog/500k-context-length-fine-tuning)
- **FP8 & Vision RL**: You can now do FP8 & VLM GRPO on consumer GPUs. [FP8 Blog](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) • [Vision RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/vision-reinforcement-learning-vlm-rl)
- **gpt-oss** by OpenAI: Read our [RL blog](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/gpt-oss-reinforcement-learning), [Flex Attention](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/long-context-gpt-oss-training) blog and [Guide](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune).
## 📥 Advanced Installation
The below advanced instructions are for Unsloth Studio. For Unsloth Core advanced installation, [view our docs](https://unsloth.ai/docs/get-started/install/pip-install#advanced-pip-installation).
@ -149,7 +169,7 @@ The below advanced instructions are for Unsloth Studio. For Unsloth Core advance
git clone https://github.com/unslothai/unsloth
cd unsloth
./install.sh --local
unsloth studio -H 0.0.0.0 -p 8888
unsloth studio -p 8888
```
Then to update :
```bash
@ -162,7 +182,7 @@ git clone https://github.com/unslothai/unsloth.git
cd unsloth
Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass
.\install.ps1 --local
unsloth studio -H 0.0.0.0 -p 8888
unsloth studio -p 8888
```
Then to update :
```bash
@ -175,11 +195,11 @@ git clone https://github.com/unslothai/unsloth
cd unsloth
git checkout nightly
./install.sh --local
unsloth studio -H 0.0.0.0 -p 8888
unsloth studio -p 8888
```
Then to launch every time:
```bash
unsloth studio -H 0.0.0.0 -p 8888
unsloth studio -p 8888
```
#### Nightly: Windows:
@ -190,11 +210,11 @@ cd unsloth
git checkout nightly
Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass
.\install.ps1 --local
unsloth studio -H 0.0.0.0 -p 8888
unsloth studio -p 8888
```
Then to launch every time:
```bash
unsloth studio -H 0.0.0.0 -p 8888
unsloth studio -p 8888
```
#### Uninstall

BIN
images/unsloth logo only.png
View file

Binary file not shown.
Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 56 KiB

After

Width:  |  Height:  |  Size: 59 KiB

Before After
Before After

BIN
images/unsloth new logo.png
View file

Binary file not shown.
Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 59 KiB

After

Width:  |  Height:  |  Size: 59 KiB

Before After
Before After

512
install.ps1
View file

@ -8,15 +8,90 @@ function Install-UnslothStudio {
$ErrorActionPreference = "Stop"
$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq "1")
# ── Tauri structured output ──
function Write-TauriLog {
param([string]$Tag, [string]$Message)
if ($TauriMode) {
Write-Host "[TAURI:$Tag] $Message"
}
}
function Format-TauriDiagBool {
param([bool]$Value)
if ($Value) { return "true" }
return "false"
}
function Get-TauriDiagArch {
$arch = [string]$env:PROCESSOR_ARCHITECTURE
if ([string]::IsNullOrWhiteSpace($arch)) {
try { $arch = [System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture.ToString() } catch { $arch = "unknown" }
}
$arch = $arch.ToLowerInvariant()
switch ($arch) {
"amd64" { return "x86_64" }
"x64" { return "x86_64" }
"arm64" { return "arm64" }
"x86" { return "x86" }
default { return ($arch -replace '[^a-z0-9_.-]', '_') }
}
}
function Get-TauriTorchIndexFamily {
param([string]$TorchIndexUrl)
if ($SkipTorch) { return "none" }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf }
if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf }
return "auto"
}
function Get-TauriGpuBranch {
param([string]$TorchIndexFamily)
if ($SkipTorch) { return "no_torch" }
if ($TorchIndexFamily -like "cu*") { return "cuda" }
if ($TorchIndexFamily -like "rocm*") { return "rocm" }
if ($TorchIndexFamily -eq "cpu") { return "cpu" }
return "unknown"
}
function Write-TauriDiag {
param(
[string]$GpuBranch = "unknown",
[string]$TorchIndexFamily = "none",
[string]$PythonVersionForDiag = $PythonVersion
)
if ([string]::IsNullOrWhiteSpace($PythonVersionForDiag)) { $PythonVersionForDiag = "unknown" }
Write-TauriLog "DIAG" "diag_schema=1 platform=windows arch=$(Get-TauriDiagArch) python_version=$($PythonVersionForDiag.ToLowerInvariant()) skip_torch=$(Format-TauriDiagBool $SkipTorch) mac_intel=false gpu_branch=$GpuBranch torch_index_family=$TorchIndexFamily"
}
function Exit-InstallFailure {
param(
[Parameter(Mandatory = $true)][string]$Message,
[int]$Code = 1
)
if ($Code -eq 0) { $Code = 1 }
Write-TauriLog "ERROR" $Message
if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) {
Restore-StudioVenvRollback
}
if ($TauriMode) {
exit $Code
}
}
# ── Parse flags ──
$StudioLocalInstall = $false
$PackageName = "unsloth"
$RepoRoot = ""
$TauriMode = $false
$SkipTorch = $false
$argList = $args
for ($i = 0; $i -lt $argList.Count; $i++) {
switch ($argList[$i]) {
"--local" { $StudioLocalInstall = $true }
"--tauri" { $TauriMode = $true }
"--no-torch" { $SkipTorch = $true }
"--verbose" { $script:UnslothVerbose = $true }
"-v" { $script:UnslothVerbose = $true }
@ -24,7 +99,7 @@ function Install-UnslothStudio {
$i++
if ($i -ge $argList.Count) {
Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red
return
return (Exit-InstallFailure "--package requires an argument.")
}
$PackageName = $argList[$i]
}
@ -40,10 +115,16 @@ function Install-UnslothStudio {
$RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path
if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) {
Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red
return
return (Exit-InstallFailure "--local must be run from the unsloth repo root")
}
}
# Validate --package to prevent injection into shell/Python commands
if ($PackageName -notmatch '^[a-zA-Z0-9][a-zA-Z0-9._-]*$') {
Write-Host "[ERROR] --package name contains invalid characters (allowed: a-z A-Z 0-9 . _ -)" -ForegroundColor Red
return (Exit-InstallFailure "--package name contains invalid characters")
}
$PythonVersion = "3.13"
$StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
$VenvDir = Join-Path $StudioHome "unsloth_studio"
@ -100,22 +181,115 @@ function Install-UnslothStudio {
Write-Host ""
# ── Helper: refresh PATH from registry (deduplicating entries) ──
# Merge order: venv Scripts (if active) > Machine > User > current $env:Path.
# Dedup compares both raw and expanded forms (%VAR% vs literal).
function Refresh-SessionPath {
$machine = [System.Environment]::GetEnvironmentVariable("Path", "Machine")
$user = [System.Environment]::GetEnvironmentVariable("Path", "User")
$merged = "$machine;$user;$env:Path"
$venvScripts = if ($env:VIRTUAL_ENV) { Join-Path $env:VIRTUAL_ENV "Scripts" } else { $null }
$sources = @()
if ($venvScripts) { $sources += $venvScripts }
$sources += @($machine, $user, $env:Path)
$merged = ($sources | Where-Object { $_ }) -join ";"
$seen = @{}
$unique = @()
$unique = New-Object System.Collections.Generic.List[string]
foreach ($p in $merged -split ";") {
$key = $p.TrimEnd("\").ToLowerInvariant()
if ($key -and -not $seen.ContainsKey($key)) {
$seen[$key] = $true
$unique += $p
$rawKey = $p.Trim().Trim('"').TrimEnd("\").ToLowerInvariant()
$expKey = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd("\").ToLowerInvariant()
if ($rawKey -and -not $seen.ContainsKey($rawKey) -and -not $seen.ContainsKey($expKey)) {
$seen[$rawKey] = $true
if ($expKey -and $expKey -ne $rawKey) { $seen[$expKey] = $true }
$unique.Add($p)
}
}
$env:Path = $unique -join ";"
}
# ── Helper: safely add a directory to the persistent User PATH ──
# Direct registry access preserves REG_EXPAND_SZ (avoids dotnet/runtime#1442).
# Append (default) keeps existing tools first; Prepend for must-win entries.
function Add-ToUserPath {
param(
[Parameter(Mandatory = $true)][string]$Directory,
[ValidateSet('Append','Prepend')]
[string]$Position = 'Append'
)
try {
$regKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Environment')
try {
$rawPath = $regKey.GetValue('Path', '', [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames)
[string[]]$entries = if ($rawPath) { $rawPath -split ';' } else { @() } # string[] prevents scalar collapse
$normalDir = $Directory.Trim().Trim('"').TrimEnd('\').ToLowerInvariant()
$expNormalDir = [Environment]::ExpandEnvironmentVariables($Directory).Trim().Trim('"').TrimEnd('\').ToLowerInvariant()
$kept = New-Object System.Collections.Generic.List[string]
$matchIndices = New-Object System.Collections.Generic.List[int]
for ($i = 0; $i -lt $entries.Count; $i++) {
$stripped = $entries[$i].Trim().Trim('"')
$rawNorm = $stripped.TrimEnd('\').ToLowerInvariant()
$expNorm = [Environment]::ExpandEnvironmentVariables($stripped).TrimEnd('\').ToLowerInvariant()
$isMatch = ($rawNorm -and ($rawNorm -eq $normalDir -or $rawNorm -eq $expNormalDir)) -or
($expNorm -and ($expNorm -eq $normalDir -or $expNorm -eq $expNormalDir))
if ($isMatch) {
$matchIndices.Add($i)
continue
}
$kept.Add($entries[$i])
}
$alreadyPresent = $matchIndices.Count -gt 0
if ($alreadyPresent -and $Position -eq 'Append') { # Append: idempotent no-op
return $false
}
if ($alreadyPresent -and $Position -eq 'Prepend' -and # Prepend: no-op if already at front
$matchIndices.Count -eq 1 -and $matchIndices[0] -eq 0) {
return $false
}
# One-time backup under HKCU\Software\Unsloth\PathBackup
if ($rawPath) {
try {
$backupKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Software\Unsloth')
try {
$existingBackup = $backupKey.GetValue('PathBackup', $null)
if (-not $existingBackup) {
$backupKey.SetValue('PathBackup', $rawPath, [Microsoft.Win32.RegistryValueKind]::ExpandString)
}
} finally {
$backupKey.Close()
}
} catch { }
}
if (-not $rawPath) {
Write-Host "[WARN] User PATH is empty - initializing with $Directory" -ForegroundColor Yellow
}
$newPath = if ($rawPath) {
if ($Position -eq 'Prepend') {
(@($Directory) + $kept) -join ';'
} else {
($kept + @($Directory)) -join ';'
}
} else {
$Directory
}
if ($newPath -ceq $rawPath) { # no actual change
return $false
}
$regKey.SetValue('Path', $newPath, [Microsoft.Win32.RegistryValueKind]::ExpandString)
# Broadcast WM_SETTINGCHANGE via dummy env-var roundtrip.
# [NullString]::Value avoids PS 7.5+/.NET 9 $null-to-"" coercion.
try {
$d = "UnslothPathRefresh_$([guid]::NewGuid().ToString('N').Substring(0,8))"
[Environment]::SetEnvironmentVariable($d, '1', 'User')
[Environment]::SetEnvironmentVariable($d, [NullString]::Value, 'User')
} catch { }
return $true
} finally {
$regKey.Close()
}
} catch {
Write-Host "[WARN] Could not update User PATH: $($_.Exception.Message)" -ForegroundColor Yellow
return $false
}
}
function step {
param(
[Parameter(Mandatory = $true)][string]$Label,
@ -378,7 +552,7 @@ try {
} catch {}
exit 1
}
`$studioCommand = '& "' + `$studioExe + '" studio -H 0.0.0.0 -p ' + `$launchPort
`$studioCommand = '& "' + `$studioExe + '" studio -p ' + `$launchPort
`$launchArgs = @(
'-NoExit',
'-NoProfile',
@ -516,11 +690,12 @@ shell.Run cmd, 0, False
}
# ── Check winget ──
Write-TauriLog "STEP" "Checking system dependencies"
if (-not (Get-Command winget -ErrorAction SilentlyContinue)) {
step "winget" "not available" "Red"
substep "Install it from https://aka.ms/getwinget" "Yellow"
substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow"
return
return (Exit-InstallFailure "winget is not available")
}
# ── Helper: detect a working Python 3.11-3.13 on the system ──
@ -595,6 +770,7 @@ shell.Run cmd, 0, False
# ── Install Python if no compatible version (3.11-3.13) found ──
# Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
Write-TauriLog "STEP" "Installing Python"
$DetectedPython = Find-CompatiblePython
if ($DetectedPython) {
step "python" "Python $($DetectedPython.Version) already installed"
@ -638,11 +814,17 @@ shell.Run cmd, 0, False
Write-Host " Please install Python $PythonVersion manually from https://www.python.org/downloads/" -ForegroundColor Yellow
Write-Host " Make sure to check 'Add Python to PATH' during installation." -ForegroundColor Yellow
Write-Host " Then re-run this installer." -ForegroundColor Yellow
return
return (Exit-InstallFailure "Python installation failed")
}
}
$DiagPythonVersion = $PythonVersion
if ($DetectedPython) { $DiagPythonVersion = $DetectedPython.Version }
$InitialGpuBranch = "unknown"
if ($SkipTorch) { $InitialGpuBranch = "no_torch" }
Write-TauriDiag -GpuBranch $InitialGpuBranch -TorchIndexFamily "none" -PythonVersionForDiag $DiagPythonVersion
# ── Install uv if not present ──
Write-TauriLog "STEP" "Installing uv package manager"
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
substep "installing uv package manager..."
$prevEAP = $ErrorActionPreference
@ -653,7 +835,7 @@ shell.Run cmd, 0, False
# Fallback: if winget didn't put uv on PATH, try the PowerShell installer
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
substep "trying alternative uv installer..." "Yellow"
powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1")
Refresh-SessionPath
}
}
@ -661,23 +843,81 @@ shell.Run cmd, 0, False
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
step "uv" "could not be installed" "Red"
substep "Install it from https://docs.astral.sh/uv/" "Yellow"
return
return (Exit-InstallFailure "uv could not be installed")
}
# ── Create venv (migrate old layout if possible, otherwise fresh) ──
# Pass the resolved executable path to uv so it does not re-resolve
# a version string back to a conda interpreter.
Write-TauriLog "STEP" "Creating virtual environment"
if (-not (Test-Path $StudioHome)) {
New-Item -ItemType Directory -Path $StudioHome -Force | Out-Null
}
$VenvPython = Join-Path $VenvDir "Scripts\python.exe"
$_Migrated = $false
$script:StudioVenvRollbackDir = $null
$script:StudioVenvRollbackTarget = $VenvDir
$script:StudioVenvRollbackActive = $false
function Start-StudioVenvRollback {
param([Parameter(Mandatory = $true)][string]$ExistingDir)
$stamp = Get-Date -Format "yyyyMMddHHmmss"
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID"
$suffix = 0
while (Test-Path $candidate) {
$suffix++
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix"
}
Move-Item -Path $ExistingDir -Destination $candidate -ErrorAction Stop
$script:StudioVenvRollbackDir = $candidate
$script:StudioVenvRollbackTarget = $ExistingDir
$script:StudioVenvRollbackActive = $true
substep "previous environment preserved for rollback"
}
function Restore-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
$target = $script:StudioVenvRollbackTarget
if (-not $backup -or -not (Test-Path $backup)) {
$script:StudioVenvRollbackActive = $false
return
}
substep "restoring previous environment after failed install..." "Yellow"
try {
if (Test-Path $target) {
Remove-Item -Recurse -Force $target -ErrorAction SilentlyContinue
}
Move-Item -Path $backup -Destination $target -Force -ErrorAction Stop
substep "restored previous environment"
$script:StudioVenvRollbackActive = $false
$script:StudioVenvRollbackDir = $null
} catch {
Write-Host "[WARN] Could not restore previous environment from $backup to $target" -ForegroundColor Yellow
Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow
}
}
function Complete-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
if ($backup -and (Test-Path $backup)) {
Remove-Item -Recurse -Force $backup -ErrorAction SilentlyContinue
}
$script:StudioVenvRollbackActive = $false
$script:StudioVenvRollbackDir = $null
}
if (Test-Path $VenvPython) {
# New layout already exists -- nuke for fresh install
substep "removing existing environment for fresh install..."
Remove-Item -Recurse -Force $VenvDir
# New layout already exists -- replace only after preserving rollback copy.
substep "preserving existing environment for rollback..."
try {
Start-StudioVenvRollback -ExistingDir $VenvDir
} catch {
Write-Host "[ERROR] Could not prepare existing environment for reinstall: $($_.Exception.Message)" -ForegroundColor Red
return (Exit-InstallFailure "Could not prepare existing environment for reinstall")
}
} elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) {
# Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating
$OldVenv = Join-Path $StudioHome ".venv"
@ -686,18 +926,23 @@ shell.Run cmd, 0, False
$prevEAP2 = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
& $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null
$torchOk = ($LASTEXITCODE -eq 0)
} catch { $torchOk = $false }
if ($SkipTorch) {
& $OldPy -c "import sys; print(sys.executable)" 2>$null | Out-Null
} else {
& $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null
}
$legacyOk = ($LASTEXITCODE -eq 0)
} catch { $legacyOk = $false }
$ErrorActionPreference = $prevEAP2
if ($torchOk) {
if ($legacyOk) {
substep "legacy environment is healthy -- migrating..."
Move-Item -Path $OldVenv -Destination $VenvDir -Force
substep "moved .venv -> unsloth_studio"
$_Migrated = $true
} else {
substep "legacy environment failed validation -- creating fresh environment" "Yellow"
Remove-Item -Recurse -Force $OldVenv -ErrorAction SilentlyContinue
$invalidVenv = Join-Path $StudioHome (".venv.invalid.{0}.{1}" -f (Get-Date -Format "yyyyMMddHHmmss"), $PID)
Move-Item -Path $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue
}
} elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) {
# CWD-relative venv from old install.ps1 -- migrate to absolute path
@ -714,7 +959,7 @@ shell.Run cmd, 0, False
$venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" }
if ($venvExit -ne 0) {
Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to create virtual environment (exit code $venvExit)" $venvExit)
}
} else {
step "venv" "using migrated environment"
@ -754,7 +999,7 @@ shell.Run cmd, 0, False
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
# Mirrors Get-PytorchCudaTag in setup.ps1.
function Get-TorchIndexUrl {
$baseUrl = "https://download.pytorch.org/whl"
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = & $NvidiaSmiExe 2>&1 | Out-String
@ -772,6 +1017,9 @@ shell.Run cmd, 0, False
return "$baseUrl/cu126"
}
$TorchIndexUrl = Get-TorchIndexUrl
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
$GpuBranch = Get-TauriGpuBranch $TorchIndexFamily
Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version
# ── Print CPU-only hint when no GPU detected ──
if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") {
@ -815,11 +1063,12 @@ shell.Run cmd, 0, False
if ($_Migrated) {
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
Write-TauriLog "STEP" "Installing unsloth"
substep "upgrading unsloth in migrated environment..."
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo }
if ($baseInstallExit -eq 0) {
$NoTorchReq = Find-NoTorchRuntimeFile
if ($NoTorchReq) {
@ -827,37 +1076,45 @@ shell.Run cmd, 0, False
}
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
}
if ($StudioLocalInstall) {
substep "overlaying local repo (editable)..."
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
if ($overlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
}
substep "overlaying unsloth-zoo from git main..."
$zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
if ($zooOverlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
}
}
} elseif ($TorchIndexUrl) {
if ($SkipTorch) {
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
}
}
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo }
if ($baseInstallExit -eq 0) {
$NoTorchReq = Find-NoTorchRuntimeFile
if ($NoTorchReq) {
@ -865,13 +1122,13 @@ shell.Run cmd, 0, False
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.18" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.1" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
}
if ($StudioLocalInstall) {
@ -879,29 +1136,85 @@ shell.Run cmd, 0, False
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
if ($overlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
}
substep "overlaying unsloth-zoo from git main..."
$zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
if ($zooOverlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
}
}
} else {
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.18" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.1" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
}
substep "overlaying local repo (editable)..."
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
if ($overlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
}
substep "overlaying unsloth-zoo from git main..."
$zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
if ($zooOverlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "$PackageName" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
}
}
}
# Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped
# for --local: the editable install above already makes _PACKAGE_ROOT in
# unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__).
# Source paths match the Tauri bundle layout in studio/src-tauri/tauri.conf.json,
# which bundles install_python_stack.py at the bundle root next to install.ps1.
if ($TauriMode) {
$rawPath = if ($PSCommandPath) { $PSCommandPath } else { $MyInvocation.ScriptName }
if ($rawPath) {
# Strip leading \\?\ extended-length prefix if the launcher passed one.
$scriptDir = Split-Path -Parent ($rawPath -replace '^\\\\\?\\', '')
$overlayMap = [ordered]@{
"install_python_stack.py" = "Lib\site-packages\studio\install_python_stack.py"
}
foreach ($rel in $overlayMap.Keys) {
$src = Join-Path $scriptDir $rel
$dst = Join-Path $VenvDir $overlayMap[$rel]
if (-not (Test-Path $src)) { continue }
$dstParent = Split-Path -Parent $dst
if (-not (Test-Path $dstParent)) {
Write-Host "[WARN] Overlay target dir missing: $dstParent; studio setup may use stale bundled file" -ForegroundColor Yellow
continue
}
try {
if (-not (Test-Path $dst)) {
# Backfill: target file missing but parent dir exists.
Copy-Item $src $dst -Force
substep ("backfilled bundled " + (Split-Path -Leaf $rel))
} else {
# Hash-compare so re-runs are no-ops when files already match.
$srcHash = (Get-FileHash $src -Algorithm SHA256).Hash
$dstHash = (Get-FileHash $dst -Algorithm SHA256).Hash
if ($srcHash -ne $dstHash) {
Copy-Item $src $dst -Force
substep ("applied bundled " + (Split-Path -Leaf $rel))
}
}
} catch {
Write-Host "[WARN] Could not overlay $($rel): $($_.Exception.Message); studio setup may use stale bundled file" -ForegroundColor Yellow
}
}
}
}
@ -909,6 +1222,7 @@ shell.Run cmd, 0, False
# ── Run studio setup ──
# setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools,
# CUDA Toolkit, Node.js, and other dependencies automatically via winget.
Write-TauriLog "STEP" "Running studio setup"
step "setup" "running unsloth studio setup..."
$UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe"
if (-not (Test-Path $UnslothExe)) {
@ -916,12 +1230,14 @@ shell.Run cmd, 0, False
Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow
Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow
Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow
return
return (Exit-InstallFailure "unsloth CLI was not installed correctly")
}
# Tell setup.ps1 to skip base package installation (install.ps1 already did it)
$env:SKIP_STUDIO_BASE = "1"
$env:STUDIO_PACKAGE_NAME = $PackageName
$env:UNSLOTH_NO_TORCH = if ($SkipTorch) { "true" } else { "false" }
# Tauri desktop app bundles its own frontend — skip Node/npm/frontend build
$env:SKIP_STUDIO_FRONTEND = if ($TauriMode) { "1" } else { "0" }
# Always set STUDIO_LOCAL_INSTALL explicitly to avoid stale values from
# a previous --local run in the same PowerShell session.
if ($StudioLocalInstall) {
@ -936,37 +1252,117 @@ shell.Run cmd, 0, False
# and bypass the fast-path version check from PR #4667.
$studioArgs = @('studio', 'setup')
if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
& $UnslothExe @studioArgs
$setupExit = $LASTEXITCODE
$env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1"
try {
& $UnslothExe @studioArgs
$setupExit = $LASTEXITCODE
} finally {
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
}
if ($setupExit -ne 0) {
Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit)
}
# ── Expose `unsloth` via a shim dir containing only unsloth.exe ──
# We do NOT add the venv Scripts dir to PATH (it also holds python.exe
# and pip.exe, which would hijack the user's system interpreter).
# Hardlink preferred; falls back to copy if cross-volume or non-NTFS.
#
# Remove the legacy venv Scripts PATH entry that older installers wrote.
$LegacyScriptsDir = Join-Path $VenvDir "Scripts"
try {
$legacyKey = [Microsoft.Win32.Registry]::CurrentUser.CreateSubKey('Environment')
try {
$rawPath = $legacyKey.GetValue('Path', '', [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames)
if ($rawPath) {
[string[]]$pathEntries = $rawPath -split ';'
$normalLegacy = $LegacyScriptsDir.Trim().Trim('"').TrimEnd('\').ToLowerInvariant()
$expNormalLegacy = [Environment]::ExpandEnvironmentVariables($LegacyScriptsDir).Trim().Trim('"').TrimEnd('\').ToLowerInvariant()
$filtered = @($pathEntries | Where-Object {
$stripped = $_.Trim().Trim('"')
$rawNorm = $stripped.TrimEnd('\').ToLowerInvariant()
$expNorm = [Environment]::ExpandEnvironmentVariables($stripped).TrimEnd('\').ToLowerInvariant()
($rawNorm -ne $normalLegacy -and $rawNorm -ne $expNormalLegacy) -and
($expNorm -ne $normalLegacy -and $expNorm -ne $expNormalLegacy)
})
$cleanedPath = $filtered -join ';'
if ($cleanedPath -ne $rawPath) {
$legacyKey.SetValue('Path', $cleanedPath, [Microsoft.Win32.RegistryValueKind]::ExpandString)
try {
$d = "UnslothPathRefresh_$([guid]::NewGuid().ToString('N').Substring(0,8))"
[Environment]::SetEnvironmentVariable($d, '1', 'User')
[Environment]::SetEnvironmentVariable($d, [NullString]::Value, 'User')
} catch { }
}
}
} finally {
$legacyKey.Close()
}
} catch { }
$ShimDir = Join-Path $StudioHome "bin"
New-Item -ItemType Directory -Force -Path $ShimDir | Out-Null
$ShimExe = Join-Path $ShimDir "unsloth.exe"
# try/catch: if unsloth.exe is locked (Studio running), keep the old shim.
$shimUpdated = $false
try {
if (Test-Path $ShimExe) { Remove-Item $ShimExe -Force -ErrorAction Stop }
try {
New-Item -ItemType HardLink -Path $ShimExe -Target $UnslothExe -ErrorAction Stop | Out-Null
} catch {
Copy-Item -Path $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy
}
$shimUpdated = $true
} catch {
if (Test-Path $ShimExe) {
Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow
Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow
Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow
Write-Host " Continuing with the existing launcher." -ForegroundColor Yellow
} else {
Write-Host "[WARN] Could not create unsloth launcher at $ShimExe" -ForegroundColor Yellow
Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow
Write-Host " Launch unsloth studio directly via '$UnslothExe' until the next successful install." -ForegroundColor Yellow
}
}
# Only add to PATH when the launcher actually exists on disk.
$pathAdded = $false
if (Test-Path $ShimExe) {
$pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend'
}
if ($shimUpdated -and $pathAdded) {
step "path" "added unsloth launcher to PATH"
}
Refresh-SessionPath # sync current session with registry
Complete-StudioVenvRollback
# ── Tauri mode: done, skip shortcuts and auto-launch ──
if ($TauriMode) {
Write-TauriLog "DONE" ""
return
}
New-StudioShortcuts -UnslothExePath $UnslothExe
# ── Add venv Scripts dir to User PATH so `unsloth studio` works from any terminal ──
$ScriptsDir = Join-Path $VenvDir "Scripts"
$UserPath = [System.Environment]::GetEnvironmentVariable("Path", "User")
if (-not $UserPath -or $UserPath -notlike "*$ScriptsDir*") {
if ($UserPath) {
[System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir;$UserPath", "User")
} else {
[System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User")
}
Refresh-SessionPath
step "path" "added unsloth to PATH"
}
# Launch studio automatically in interactive terminals;
# in non-interactive environments (CI, Docker) just print instructions.
# In interactive terminals, ask the user before starting Studio.
# In non-interactive environments (CI, Docker) just print instructions.
$IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
if ($IsInteractive) {
& $UnslothExe studio -H 0.0.0.0 -p 8888
Write-Host ""
$reply = Read-Host " Start Unsloth Studio now? [Y/n]"
if ([string]::IsNullOrWhiteSpace($reply) -or $reply -match '^[Yy]') {
& $UnslothExe studio -p 8888
} else {
step "launch" "to start later, run:"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
Write-Host ""
}
} else {
step "launch" "manual commands:"
substep "& `"$VenvDir\Scripts\Activate.ps1`""
substep "unsloth studio -H 0.0.0.0 -p 8888"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
Write-Host ""
}
}

768
install.sh
View file

File diff suppressed because it is too large Load diff

9
pyproject.toml
View file

@ -53,6 +53,7 @@ studio = [
"frontend/*.yaml",
"frontend/.git*",
"backend/requirements/**/*",
"backend/plugins/**/*",
"backend/core/data_recipe/oxc-validator/*.json",
"backend/core/data_recipe/oxc-validator/*.mjs",
]
@ -82,13 +83,13 @@ huggingfacenotorch = [
"huggingface_hub>=0.34.0",
"hf_transfer",
"diffusers",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.3.0",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
"trl>=0.18.2,!=0.19.0,<=0.24.0",
"sentence-transformers",
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.3.7",
"unsloth_zoo>=2026.5.1",
"torchvision",
"unsloth[triton]",
]
@ -578,10 +579,10 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.3.7",
"unsloth_zoo>=2026.5.1",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.3.0",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
"datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0",
"sentencepiece>=0.2.0",
"tqdm",

169
scripts/install_gemma4_mlx.sh Executable file
View file

@ -0,0 +1,169 @@
#!/bin/bash
set -e
# ============================================================
# Gemma 4 MLX — One-command setup + inference
#
# Usage:
# bash install_gemma4_mlx.sh [--venv-dir DIR]
#
# This script:
# 1. Creates a Python virtual environment
# 2. Installs uv, mlx-vlm, transformers
# ============================================================
# ── Output style (inspired by unsloth/install.sh) ─────────────
RULE=""
_rule_i=0
while [ "$_rule_i" -lt 52 ]; do
RULE="${RULE}─"
_rule_i=$((_rule_i + 1))
done
if [ -n "${NO_COLOR:-}" ]; then
C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then
_ESC="$(printf '\033')"
C_TITLE="${_ESC}[38;5;117m"
C_DIM="${_ESC}[38;5;245m"
C_OK="${_ESC}[38;5;108m"
C_WARN="${_ESC}[38;5;136m"
C_ERR="${_ESC}[91m"
C_RST="${_ESC}[0m"
else
C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
fi
step() { printf " ${C_DIM}%-18.18s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
substep() { printf " ${C_DIM}%-18s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; }
fail() { step "error" "$1" "$C_ERR"; exit 1; }
# ── Parse flags ───────────────────────────────────────────────
VENV_DIR=""
_next_is_venv=false
for arg in "$@"; do
if [ "$_next_is_venv" = true ]; then
VENV_DIR="$arg"
_next_is_venv=false
continue
fi
case "$arg" in
--venv-dir) _next_is_venv=true ;;
esac
done
# Default venv location
if [ -z "$VENV_DIR" ]; then
VENV_DIR="$HOME/.unsloth/unsloth_gemma4_mlx"
fi
# ── Banner ────────────────────────────────────────────────────
echo ""
printf " ${C_TITLE}%s${C_RST}\n" "💎 Gemma 4 MLX Installer"
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""
# ── Platform check ────────────────────────────────────────────
if [ "$(uname)" != "Darwin" ]; then
fail "MLX requires macOS with Apple Silicon. Detected: $(uname)"
fi
_ARCH=$(uname -m)
if [ "$_ARCH" != "arm64" ]; then
step "warning" "Apple Silicon recommended (detected: $_ARCH)" "$C_WARN"
fi
step "platform" "macOS ($_ARCH)"
# ── Detect Python ─────────────────────────────────────────────
PYTHON=""
for _candidate in python3.12 python3.11 python3.13 python3; do
if command -v "$_candidate" >/dev/null 2>&1; then
PYTHON="$_candidate"
break
fi
done
if [ -z "$PYTHON" ]; then
fail "Python 3 not found. Install via: brew install python@3.12"
fi
_PY_VERSION=$("$PYTHON" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}')")
step "python" "$PYTHON ($_PY_VERSION)"
# ── Create virtual environment ────────────────────────────────
if [ -x "$VENV_DIR/bin/python" ]; then
step "venv" "using existing environment"
substep "$VENV_DIR"
else
step "venv" "creating virtual environment"
substep "$VENV_DIR"
mkdir -p "$(dirname "$VENV_DIR")"
"$PYTHON" -m venv "$VENV_DIR"
fi
# ── Install uv ───────────────────────────────────────────────
if ! command -v uv >/dev/null 2>&1; then
step "uv" "installing uv package manager..."
_uv_tmp=$(mktemp)
curl -LsSf "https://astral.sh/uv/install.sh" -o "$_uv_tmp"
sh "$_uv_tmp" </dev/null >/dev/null 2>&1
rm -f "$_uv_tmp"
if [ -f "$HOME/.local/bin/env" ]; then
. "$HOME/.local/bin/env"
fi
export PATH="$HOME/.local/bin:$PATH"
substep "done"
else
step "uv" "found $(uv --version 2>/dev/null || echo 'uv')"
fi
_VENV_PY="$VENV_DIR/bin/python"
# ── Install dependencies ──────────────────────────────────────
step "install" "installing mlx-vlm..."
uv pip install --python "$_VENV_PY" -q mlx-vlm
substep "done"
step "install" "installing transformers>=5.5.0..."
if uv pip install --python "$_VENV_PY" -q "transformers>=5.5.0" 2>/dev/null; then
substep "installed from PyPI"
else
substep "PyPI install failed (Python <3.10?), trying GitHub..."
if uv pip install --python "$_VENV_PY" -q "git+https://github.com/huggingface/transformers.git@v5.5-release" 2>/dev/null; then
substep "installed from huggingface/transformers v5.5-release"
else
step "warning" "could not install transformers>=5.5.0" "$C_WARN"
substep "tried: PyPI, huggingface/transformers v5.5-release"
fi
fi
# ── Verify installation ──────────────────────────────────────
if "$_VENV_PY" -c "import mlx_vlm"; then
substep "mlx-vlm verified"
else
fail "Installation verification failed."
fi
# ── Done ──────────────────────────────────────────────────────
echo ""
printf " ${C_TITLE}%s${C_RST}\n" "Gemma 4 MLX installed!"
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""
step "available models" "unsloth/gemma-4-E2B-it-UD-MLX-4bit"
substep "unsloth/gemma-4-E4B-it-UD-MLX-4bit"
substep "unsloth/gemma-4-26b-a4b-it-UD-MLX-4bit"
substep "unsloth/gemma-4-31b-it-UD-MLX-4bit"
echo ""
step "venv activate" "source ${VENV_DIR}/bin/activate"
echo ""
step "text chat" "python -m mlx_vlm.chat --model unsloth/gemma-4-E2B-it-UD-MLX-4bit"
echo ""
step "vision chat" "python -m mlx_vlm.chat --model unsloth/gemma-4-31b-it-UD-MLX-4bit"
substep "Use /image path/to/image.jpg to load an image"
echo ""
step "gradio UI" "python -m mlx_vlm.chat_ui --model unsloth/gemma-4-31b-it-UD-MLX-4bit"
echo ""
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""

191
scripts/install_qwen3_6_mlx.sh Normal file
View file

@ -0,0 +1,191 @@
#!/bin/bash
set -e
# ============================================================
# Qwen3.6 MLX — One-command setup + inference
#
# Usage:
# bash install_qwen3_6_mlx.sh [--venv-dir DIR]
#
# This script:
# 1. Creates a Python virtual environment
# 2. Installs uv, mlx-vlm, transformers, torch, torchvision
# ============================================================
# ── Output style (inspired by unsloth/install.sh) ─────────────
RULE=""
_rule_i=0
while [ "$_rule_i" -lt 52 ]; do
RULE="${RULE}─"
_rule_i=$((_rule_i + 1))
done
if [ -n "${NO_COLOR:-}" ]; then
C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then
_ESC="$(printf '\033')"
C_TITLE="${_ESC}[38;5;117m"
C_DIM="${_ESC}[38;5;245m"
C_OK="${_ESC}[38;5;108m"
C_WARN="${_ESC}[38;5;136m"
C_ERR="${_ESC}[91m"
C_RST="${_ESC}[0m"
else
C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
fi
step() { printf " ${C_DIM}%-18.18s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
substep() { printf " ${C_DIM}%-18s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; }
fail() { step "error" "$1" "$C_ERR"; exit 1; }
# ── Parse flags ───────────────────────────────────────────────
VENV_DIR=""
_next_is_venv=false
for arg in "$@"; do
if [ "$_next_is_venv" = true ]; then
VENV_DIR="$arg"
_next_is_venv=false
continue
fi
case "$arg" in
--venv-dir) _next_is_venv=true ;;
esac
done
# Default venv location
if [ -z "$VENV_DIR" ]; then
VENV_DIR="$HOME/.unsloth/unsloth_qwen3_6_mlx"
fi
# ── Banner ────────────────────────────────────────────────────
echo ""
printf " ${C_TITLE}%s${C_RST}\n" "Qwen3.6 MLX Installer"
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""
# ── Platform check ────────────────────────────────────────────
if [ "$(uname)" != "Darwin" ]; then
fail "MLX requires macOS with Apple Silicon. Detected: $(uname)"
fi
_ARCH=$(uname -m)
if [ "$_ARCH" != "arm64" ]; then
step "warning" "Apple Silicon recommended (detected: $_ARCH)" "$C_WARN"
fi
step "platform" "macOS ($_ARCH)"
# ── Detect Python ─────────────────────────────────────────────
PYTHON=""
for _candidate in python3.12 python3.11 python3.13 python3; do
if command -v "$_candidate" >/dev/null 2>&1; then
PYTHON="$_candidate"
break
fi
done
if [ -z "$PYTHON" ]; then
fail "Python 3 not found. Install via: brew install python@3.12"
fi
_PY_VERSION=$("$PYTHON" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}')")
step "python" "$PYTHON ($_PY_VERSION)"
# ── Create virtual environment ────────────────────────────────
if [ -x "$VENV_DIR/bin/python" ]; then
step "venv" "using existing environment"
substep "$VENV_DIR"
else
step "venv" "creating virtual environment"
substep "$VENV_DIR"
mkdir -p "$(dirname "$VENV_DIR")"
"$PYTHON" -m venv "$VENV_DIR"
fi
# ── Install uv ───────────────────────────────────────────────
if ! command -v uv >/dev/null 2>&1; then
step "uv" "installing uv package manager..."
_uv_tmp=$(mktemp)
curl -LsSf "https://astral.sh/uv/install.sh" -o "$_uv_tmp"
sh "$_uv_tmp" </dev/null
rm -f "$_uv_tmp"
if [ -f "$HOME/.local/bin/env" ]; then
. "$HOME/.local/bin/env"
fi
export PATH="$HOME/.local/bin:$PATH"
substep "done"
else
step "uv" "found $(uv --version 2>/dev/null || echo 'uv')"
fi
_VENV_PY="$VENV_DIR/bin/python"
# ── Install dependencies ──────────────────────────────────────
step "install" "installing mlx-vlm..."
uv pip install --python "$_VENV_PY" -q mlx-vlm
substep "done"
step "install" "installing transformers>=5.2.0..."
if uv pip install --python "$_VENV_PY" -q "transformers>=5.2.0"; then
substep "installed from PyPI"
else
substep "PyPI install failed, trying GitHub..."
if uv pip install --python "$_VENV_PY" -q "git+https://github.com/huggingface/transformers.git"; then
substep "installed from huggingface/transformers main"
else
fail "Could not install transformers>=5.2.0 (required for Qwen3.5/3.6 model support). Please check your Python version (>=3.10 required) and network connection, then try again."
fi
fi
step "install" "installing torch + torchvision (needed for Qwen3 VL processor)..."
uv pip install --python "$_VENV_PY" -q torch torchvision
substep "done"
# ── Verify installation ──────────────────────────────────────
if "$_VENV_PY" -c "import mlx_vlm; import torch; import torchvision; import transformers"; then
substep "mlx-vlm + torch + transformers verified"
else
fail "Installation verification failed. Please ensure Python >=3.10 and try again."
fi
# ── Apply patches for multi-turn image chat ──────────────────
_PATCH_BASE="https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/fix/ui-fix/unsloth/models/patches/mlx_vlm_qwen3_5"
_SITE_PKGS=$("$_VENV_PY" -c "import site; print(site.getsitepackages()[0])")
step "patch" "fixing multi-turn image chat..."
if curl -sSLf "${_PATCH_BASE}/qwen3_5.py" -o "${_SITE_PKGS}/mlx_vlm/models/qwen3_5/qwen3_5.py"; then
substep "patched qwen3_5.py (MRoPE position reset)"
else
step "warning" "failed to download qwen3_5.py patch — multi-turn image chat may not work" "$C_WARN"
fi
if curl -sSLf "${_PATCH_BASE}/generate.py" -o "${_SITE_PKGS}/mlx_vlm/generate.py"; then
substep "patched generate.py (mask trim on cache reuse)"
else
step "warning" "failed to download generate.py patch — multi-turn image chat may not work" "$C_WARN"
fi
# Clear pycache so patches take effect
find "${_SITE_PKGS}/mlx_vlm" -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true
substep "cleared bytecode cache"
# ── Done ──────────────────────────────────────────────────────
echo ""
printf " ${C_TITLE}%s${C_RST}\n" "Qwen3.6 MLX installed!"
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""
step "available models" "unsloth/Qwen3.6-35B-A3B-UD-MLX-3bit"
substep "unsloth/Qwen3.6-35B-A3B-UD-MLX-4bit"
substep "unsloth/Qwen3.6-35B-A3B-MLX-8bit"
echo ""
step "venv activate" "source ${VENV_DIR}/bin/activate"
echo ""
step "vision chat" "python -m mlx_vlm.chat --model unsloth/Qwen3.6-35B-A3B-UD-MLX-4bit"
substep "Use /image path/to/image.jpg to load an image"
echo ""
step "gradio UI" "python -m mlx_vlm.chat_ui --model unsloth/Qwen3.6-35B-A3B-UD-MLX-4bit"
echo ""
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""

2
studio/Unsloth_Studio_Colab.ipynb
View file

@ -64,7 +64,7 @@
"id": "27e68f91"
},
"outputs": [],
"source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh"
"source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local"
},
{
"cell_type": "markdown",

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

@ -1,6 +1,14 @@
{
"_comment": "Per-model-family inference parameter defaults. Sources: (1) Ollama params blobs, (2) Existing Unsloth Studio YAML configs. Patterns ordered longest-match-first.",
"families": {
"qwen3.6": {
"temperature": 0.7,
"top_p": 0.8,
"top_k": 20,
"min_p": 0.0,
"repetition_penalty": 1.0,
"presence_penalty": 1.5
},
"qwen3.5": {
"temperature": 0.7,
"top_p": 0.8,
@ -93,6 +101,14 @@
"min_p": 0.0,
"repetition_penalty": 1.0
},
"gemma-4": {
"temperature": 1.0,
"top_p": 0.95,
"top_k": 64,
"min_p": 0.0,
"repetition_penalty": 1.0,
"presence_penalty": 0.0
},
"gemma-3n": {
"temperature": 1.0,
"top_p": 0.95,
@ -361,12 +377,12 @@
}
},
"patterns": [
"qwen3.5",
"qwen3.6", "qwen3.5",
"qwen3-coder", "qwen3-next", "qwen3-vl", "qwen3",
"qwen2.5-coder", "qwen2.5-vl", "qwen2.5-omni", "qwen2.5-math", "qwen2.5",
"qwen2-vl", "qwen2",
"qwq",
"gemma-3n", "gemma-3", "medgemma", "gemma-2",
"gemma-4", "gemma-3n", "gemma-3", "medgemma", "gemma-2",
"llama-4", "llama-3.3", "llama-3.2", "llama-3.1", "llama-3",
"phi-4", "phi-3",
"mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral",

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-26B-A4B-it
# Also applies to: google/gemma-4-26B-A4B-it, unsloth/gemma-4-26B-A4B-it-GGUF
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-26B-A4B (base/pretrained)
# Also applies to: google/gemma-4-26B-A4B
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-31B-it
# Also applies to: google/gemma-4-31B-it, unsloth/gemma-4-31B-it-GGUF
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-31B (base/pretrained)
# Also applies to: google/gemma-4-31B
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-E2B-it
# Also applies to: google/gemma-4-E2B-it, unsloth/gemma-4-E2B-it-GGUF
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-E2B (base/pretrained)
# Also applies to: google/gemma-4-E2B
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-E4B-it
# Also applies to: google/gemma-4-E4B-it, unsloth/gemma-4-E4B-it-GGUF
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

47
studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml Normal file
View file

@ -0,0 +1,47 @@
# Model defaults for unsloth/gemma-4-E4B (base/pretrained)
# Also applies to: google/gemma-4-E4B
training:
trust_remote_code: false
max_seq_length: 2048
num_epochs: 0
learning_rate: 2e-4
batch_size: 2
gradient_accumulation_steps: 4
warmup_steps: 5
max_steps: 30
save_steps: 30
weight_decay: 0.001
random_seed: 3407
packing: false
train_on_completions: true
gradient_checkpointing: "unsloth"
optim: "adamw_8bit"
lr_scheduler_type: "linear"
lora:
lora_r: 8
lora_alpha: 8
lora_dropout: 0.0
target_modules:
- "all-linear"
use_rslora: false
use_loftq: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
finetune_mlp_modules: true
logging:
enable_wandb: false
wandb_project: "llm-finetuning"
enable_tensorboard: false
tensorboard_dir: "runs"
log_frequency: 10
inference:
trust_remote_code: false
temperature: 1.0
top_p: 0.95
top_k: 64
min_p: 0.0

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

@ -10,10 +10,12 @@ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
import jwt
from .storage import (
API_KEY_PREFIX,
get_jwt_secret,
get_user_and_secret,
load_jwt_secret,
save_refresh_token,
validate_api_key,
verify_refresh_token,
)
@ -50,6 +52,8 @@ def _decode_subject_without_verification(token: str) -> Optional[str]:
def create_access_token(
subject: str,
expires_delta: Optional[timedelta] = None,
*,
desktop: bool = False,
) -> str:
"""
Create a signed JWT for the given subject (e.g. username).
@ -57,6 +61,8 @@ def create_access_token(
Tokens are valid across restarts because the signing secret is stored in SQLite.
"""
to_encode = {"sub": subject}
if desktop:
to_encode["desktop"] = True
expire = datetime.now(timezone.utc) + (
expires_delta or timedelta(minutes = ACCESS_TOKEN_EXPIRE_MINUTES)
)
@ -68,7 +74,29 @@ def create_access_token(
)
def create_refresh_token(subject: str) -> str:
def is_desktop_access_token(token: str) -> bool:
"""Return true only for a valid desktop-issued JWT access token."""
if token.startswith(API_KEY_PREFIX):
return False
subject = _decode_subject_without_verification(token)
if subject is None:
return False
record = get_user_and_secret(subject)
if record is None:
return False
_salt, _pwd_hash, jwt_secret, _must_change_password = record
try:
payload = jwt.decode(token, jwt_secret, algorithms = [ALGORITHM])
except jwt.InvalidTokenError:
return False
return payload.get("sub") == subject and payload.get("desktop") is True
def create_refresh_token(subject: str, *, desktop: bool = False) -> str:
"""
Create a random refresh token, store its hash in SQLite, and return it.
@ -76,21 +104,28 @@ def create_refresh_token(subject: str) -> str:
"""
token = secrets.token_urlsafe(48)
expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS)
save_refresh_token(token, subject, expires_at.isoformat())
save_refresh_token(token, subject, expires_at.isoformat(), is_desktop = desktop)
return token
def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[str]]:
def refresh_access_token(
refresh_token: str,
) -> Tuple[Optional[str], Optional[str], bool]:
"""
Validate a refresh token and issue a new access token.
The refresh token itself is NOT consumed — it stays valid until expiry.
Returns a new access_token or None if the refresh token is invalid/expired.
"""
username = verify_refresh_token(refresh_token)
if username is None:
return None, None
return create_access_token(subject = username), username
verified = verify_refresh_token(refresh_token)
if verified is None:
return None, None, False
username, is_desktop = verified
return (
create_access_token(subject = username, desktop = is_desktop),
username,
is_desktop,
)
def reload_secret() -> None:
@ -137,6 +172,18 @@ async def _get_current_subject(
...
"""
token = credentials.credentials
# --- API key path (sk-unsloth-...) ---
if token.startswith(API_KEY_PREFIX):
username = validate_api_key(token)
if username is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid or expired API key",
)
return username
# --- JWT path ---
subject = _decode_subject_without_verification(token)
if subject is None:
raise HTTPException(
@ -159,7 +206,8 @@ async def _get_current_subject(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid token payload",
)
if must_change_password and not allow_password_change:
is_desktop = payload.get("desktop") is True
if must_change_password and not allow_password_change and not is_desktop:
raise HTTPException(
status_code = status.HTTP_403_FORBIDDEN,
detail = "Password change required",

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

@ -6,6 +6,7 @@ SQLite storage for authentication data (user credentials + JWT secret).
"""
import hashlib
import os
import secrets
import sqlite3
from datetime import datetime, timezone
@ -54,6 +55,10 @@ def generate_bootstrap_password() -> str:
# before the user changes the password.
ensure_dir(_BOOTSTRAP_PW_PATH.parent)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password)
try:
os.chmod(_BOOTSTRAP_PW_PATH, 0o600)
except OSError:
pass
return _bootstrap_password
@ -63,6 +68,17 @@ def get_bootstrap_password() -> Optional[str]:
return _bootstrap_password
def _load_bootstrap_password() -> Optional[str]:
"""Load an existing bootstrap password without creating one."""
global _bootstrap_password
_bootstrap_password = None
if _BOOTSTRAP_PW_PATH.is_file():
bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
if bootstrap_password:
_bootstrap_password = bootstrap_password
return _bootstrap_password
def clear_bootstrap_password() -> None:
"""Delete the persisted bootstrap password file (called after password change)."""
global _bootstrap_password
@ -72,7 +88,22 @@ def clear_bootstrap_password() -> None:
def _hash_token(token: str) -> str:
"""SHA-256 hash helper used for refresh token storage."""
"""SHA-256 hash helper used for refresh token storage.
Plain SHA-256 is intentional here: refresh tokens are high-entropy
random strings from ``secrets.token_urlsafe(48)`` (384 bits of
entropy), so a slow KDF (Argon2 / bcrypt / PBKDF2) provides zero
additional security — no attacker can brute-force 2^384 regardless
of hash speed — while adding tens of ms of CPU to every refresh.
See the OWASP Password Storage Cheat Sheet on fast-vs-slow hashing
of high-entropy inputs.
API keys use the separate ``_pbkdf2_api_key`` helper below, which
runs PBKDF2-HMAC-SHA256 with a persistent server-side salt — not
for cryptographic reasons (128-bit random tokens don't need slow
hashing), but because CodeQL's ``py/weak-sensitive-data-hashing``
query mislabels API keys as passwords and demands a KDF.
"""
return hashlib.sha256(token.encode("utf-8")).hexdigest()
@ -99,7 +130,39 @@ def get_connection() -> sqlite3.Connection:
id INTEGER PRIMARY KEY,
token_hash TEXT NOT NULL,
username TEXT NOT NULL,
expires_at TEXT NOT NULL
expires_at TEXT NOT NULL,
is_desktop INTEGER NOT NULL DEFAULT 0
);
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS api_keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL,
key_prefix TEXT NOT NULL,
key_hash TEXT NOT NULL UNIQUE,
name TEXT NOT NULL DEFAULT '',
created_at TEXT NOT NULL,
last_used_at TEXT,
expires_at TEXT,
is_active INTEGER NOT NULL DEFAULT 1,
is_internal INTEGER NOT NULL DEFAULT 0
);
"""
)
api_key_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")
}
if "is_internal" not in api_key_columns:
conn.execute(
"ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS app_secrets (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
"""
)
@ -108,10 +171,107 @@ def get_connection() -> sqlite3.Connection:
conn.execute(
"ALTER TABLE auth_user ADD COLUMN must_change_password INTEGER NOT NULL DEFAULT 0"
)
refresh_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")
}
if "is_desktop" not in refresh_columns:
conn.execute(
"ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0"
)
conn.commit()
return conn
# ── API-key PBKDF2 salt ────────────────────────────────────────────────
#
# Module-level cache for the persistent API-key PBKDF2 salt. Populated
# lazily on first use via ``_get_or_create_api_key_pbkdf2_salt``. Not
# protected by a lock because (a) the ``INSERT OR IGNORE`` provides
# atomicity at the SQLite layer and (b) concurrent populations converge
# on the same value, so the worst case is a harmless duplicate read on
# startup.
_api_key_pbkdf2_salt_cache: Optional[bytes] = None
def _get_or_create_api_key_pbkdf2_salt() -> bytes:
"""Return the persistent API-key PBKDF2 salt, generating it once if missing.
Stored as a hex-encoded 32-byte random value in the ``app_secrets``
table under key ``"api_key_pbkdf2_salt"``. Regenerated only if the row
is missing (i.e. fresh install, or operator manually deleted the row
and accepts invalidating existing API keys).
"""
global _api_key_pbkdf2_salt_cache
if _api_key_pbkdf2_salt_cache is not None:
return _api_key_pbkdf2_salt_cache
conn = get_connection()
try:
cur = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
("api_key_pbkdf2_salt",),
)
row = cur.fetchone()
if row is None:
new_value = secrets.token_hex(32) # 32 bytes -> 64 hex chars
conn.execute(
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
("api_key_pbkdf2_salt", new_value),
)
conn.commit()
cur = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
("api_key_pbkdf2_salt",),
)
row = cur.fetchone()
salt = bytes.fromhex(row["value"])
finally:
conn.close()
_api_key_pbkdf2_salt_cache = salt
return salt
_API_KEY_PBKDF2_ITERATIONS = 100_000
DESKTOP_SECRET_PREFIX = "desktop-"
_DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"
_DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at"
def _pbkdf2_api_key(raw_key: str) -> str:
"""PBKDF2-HMAC-SHA256 an API key with a persistent server-side salt.
Used for API-key storage ONLY, not refresh tokens. Matches the
PBKDF2 algorithm + iteration count used by the password hasher in
``auth/hashing.py`` so the codebase is consistent on which KDF it
uses for credential storage.
Notes on why a slow KDF here is *only* a CodeQL appeasement and
*not* a cryptographic requirement: API keys are cryptographically
random 128-bit tokens (via ``secrets.token_hex``), so brute force
against 2^128 is infeasible regardless of hash speed. CodeQL's
``py/weak-sensitive-data-hashing`` query mislabels these tokens as
"password" sensitive data and then demands a KDF from its
allowlist (Argon2 / scrypt / bcrypt / PBKDF2). Per the query's
own recommendation page we use PBKDF2. The persistent salt is
still loaded from ``app_secrets`` so an attacker dumping the
``api_keys`` table alone cannot derive hashes for candidate
tokens without also obtaining the salt row.
"""
salt = _get_or_create_api_key_pbkdf2_salt()
dk = hashlib.pbkdf2_hmac(
"sha256",
raw_key.encode("utf-8"),
salt,
_API_KEY_PBKDF2_ITERATIONS,
)
return dk.hex()
def _pbkdf2_desktop_secret(raw_secret: str) -> str:
return _pbkdf2_api_key(raw_secret)
def is_initialized() -> bool:
"""Check if auth is ready for login (at least one user exists in DB)."""
conn = get_connection()
@ -253,6 +413,10 @@ def ensure_default_admin() -> bool:
Uses a randomly generated diceware passphrase as the bootstrap password.
Returns True when the default admin was created in this call.
"""
if get_user_and_secret(DEFAULT_ADMIN_USERNAME) is not None:
_load_bootstrap_password()
return False
bootstrap_pw = generate_bootstrap_password()
try:
create_initial_user(
@ -285,12 +449,19 @@ def update_password(username: str, new_password: str) -> bool:
conn.commit()
if cursor.rowcount > 0:
clear_bootstrap_password()
clear_desktop_secret()
return cursor.rowcount > 0
finally:
conn.close()
def save_refresh_token(token: str, username: str, expires_at: str) -> None:
def save_refresh_token(
token: str,
username: str,
expires_at: str,
*,
is_desktop: bool = False,
) -> None:
"""
Store a hashed refresh token with its associated username and expiry.
"""
@ -299,21 +470,21 @@ def save_refresh_token(token: str, username: str, expires_at: str) -> None:
try:
conn.execute(
"""
INSERT INTO refresh_tokens (token_hash, username, expires_at)
VALUES (?, ?, ?)
INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop)
VALUES (?, ?, ?, ?)
""",
(token_hash, username, expires_at),
(token_hash, username, expires_at, int(is_desktop)),
)
conn.commit()
finally:
conn.close()
def verify_refresh_token(token: str) -> Optional[str]:
def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
"""
Verify a refresh token and return the username.
Verify a refresh token and return the username plus desktop marker.
Returns the username if valid and not expired, None otherwise.
Returns the username and desktop marker if valid and not expired, None otherwise.
The token is NOT consumed — it stays valid until it expires.
"""
token_hash = _hash_token(token)
@ -328,7 +499,7 @@ def verify_refresh_token(token: str) -> Optional[str]:
cur = conn.execute(
"""
SELECT id, username, expires_at FROM refresh_tokens
SELECT id, username, expires_at, is_desktop FROM refresh_tokens
WHERE token_hash = ?
""",
(token_hash,),
@ -344,7 +515,7 @@ def verify_refresh_token(token: str) -> Optional[str]:
conn.commit()
return None
return row["username"]
return row["username"], bool(row["is_desktop"])
finally:
conn.close()
@ -357,3 +528,208 @@ def revoke_user_refresh_tokens(username: str) -> None:
conn.commit()
finally:
conn.close()
def create_desktop_secret() -> str:
"""Create/rotate the local desktop credential and return it once."""
ensure_default_admin()
raw_secret = DESKTOP_SECRET_PREFIX + secrets.token_urlsafe(48)
secret_hash = _pbkdf2_desktop_secret(raw_secret)
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
conn.execute(
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
(_DESKTOP_SECRET_HASH_KEY, secret_hash),
)
conn.execute(
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
(_DESKTOP_SECRET_CREATED_AT_KEY, now),
)
conn.commit()
return raw_secret
finally:
conn.close()
def validate_desktop_secret(raw_secret: str) -> Optional[str]:
"""Return the real admin username when the desktop secret matches."""
if not raw_secret.startswith(DESKTOP_SECRET_PREFIX):
return None
if get_user_and_secret(DEFAULT_ADMIN_USERNAME) is None:
return None
secret_hash = _pbkdf2_desktop_secret(raw_secret)
conn = get_connection()
try:
cur = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_DESKTOP_SECRET_HASH_KEY,),
)
row = cur.fetchone()
if row is None:
return None
if not secrets.compare_digest(row["value"], secret_hash):
return None
return DEFAULT_ADMIN_USERNAME
finally:
conn.close()
def clear_desktop_secret() -> None:
"""Remove backend-side desktop auth state."""
conn = get_connection()
try:
conn.execute(
"DELETE FROM app_secrets WHERE key IN (?, ?)",
(_DESKTOP_SECRET_HASH_KEY, _DESKTOP_SECRET_CREATED_AT_KEY),
)
conn.commit()
finally:
conn.close()
# ---------------------------------------------------------------------------
# API key management
# ---------------------------------------------------------------------------
API_KEY_PREFIX = "sk-unsloth-"
def create_api_key(
username: str,
name: str,
expires_at: Optional[str] = None,
internal: bool = False,
) -> Tuple[str, dict]:
"""Create a new API key for *username*.
Returns ``(raw_key, row_dict)`` where *raw_key* is shown to the user
exactly once. The database only stores the PBKDF2 hash.
Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe
runs) that should not appear in user-facing key listings.
"""
raw_key = API_KEY_PREFIX + secrets.token_hex(16)
key_hash = _pbkdf2_api_key(raw_key)
key_prefix = raw_key[len(API_KEY_PREFIX) : len(API_KEY_PREFIX) + 8]
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(
username,
key_prefix,
key_hash,
name,
now,
expires_at,
1 if internal else 0,
),
)
conn.commit()
cur = conn.execute("SELECT * FROM api_keys WHERE key_hash = ?", (key_hash,))
row = cur.fetchone()
return raw_key, dict(row)
finally:
conn.close()
def list_api_keys(username: str, include_internal: bool = False) -> list:
"""Return API keys for *username*. Internal workflow keys are hidden
by default so they do not clutter user-facing UIs."""
conn = get_connection()
try:
if include_internal:
cur = conn.execute(
"""
SELECT id, username, key_prefix, name, created_at, last_used_at,
expires_at, is_active, is_internal
FROM api_keys
WHERE username = ?
ORDER BY created_at DESC
""",
(username,),
)
else:
cur = conn.execute(
"""
SELECT id, username, key_prefix, name, created_at, last_used_at,
expires_at, is_active, is_internal
FROM api_keys
WHERE username = ? AND is_internal = 0
ORDER BY created_at DESC
""",
(username,),
)
return [dict(row) for row in cur.fetchall()]
finally:
conn.close()
def revoke_api_key(username: str, key_id: int) -> bool:
"""Soft-delete an API key. Returns True if a matching row was found."""
conn = get_connection()
try:
cursor = conn.execute(
"UPDATE api_keys SET is_active = 0 WHERE id = ? AND username = ?",
(key_id, username),
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def revoke_internal_api_key(key_id: int) -> bool:
"""Revoke an internal workflow-minted key without requiring a username.
Used by the recipe runner to retire its sk-unsloth-* key once the job
terminates, shrinking the window a leaked key could be abused.
"""
conn = get_connection()
try:
cursor = conn.execute(
"UPDATE api_keys SET is_active = 0 WHERE id = ? AND is_internal = 1",
(key_id,),
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def validate_api_key(raw_key: str) -> Optional[str]:
"""Validate *raw_key* and return the owning username, or ``None``.
Also updates ``last_used_at`` on success.
"""
key_hash = _pbkdf2_api_key(raw_key)
conn = get_connection()
try:
cur = conn.execute(
"SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?",
(key_hash,),
)
row = cur.fetchone()
if row is None:
return None
if not row["is_active"]:
return None
if row["expires_at"] is not None:
expires = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires:
return None
conn.execute(
"UPDATE api_keys SET last_used_at = ? WHERE id = ?",
(datetime.now(timezone.utc).isoformat(), row["id"]),
)
conn.commit()
return row["username"]
finally:
conn.close()

5
studio/backend/colab.py
View file

@ -66,7 +66,10 @@ def show_link(port: int = 8888):
<svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="white"><polygon points="5,3 19,12 5,21"/></svg>
Open Unsloth Studio
</a>
<p style="color: #333333; margin: 16px 0 0 0; font-size: 13px; font-family: monospace;">
<p style="color: #333333; margin: 12px 0 0 0; font-size: 14px; font-weight: bold;">
If the link doesn't work, you can scroll down to view the UI generated directly in Colab.
</p>
<p style="color: #333333; margin: 16px 0 0 0; font-size: 13px; font-family: monospace; font-weight: bold;">
{short_url}
</p>
</div>

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

@ -31,6 +31,7 @@ __all__ = [
# Config
"ModelConfig",
"is_vision_model",
"scan_trained_models",
"scan_trained_loras",
"load_model_defaults",
"get_base_model_from_lora",
@ -72,6 +73,7 @@ def __getattr__(name):
if name in (
"is_vision_model",
"ModelConfig",
"scan_trained_models",
"scan_trained_loras",
"load_model_defaults",
"get_base_model_from_lora",
@ -79,14 +81,15 @@ def __getattr__(name):
from utils.models import (
is_vision_model,
ModelConfig,
scan_trained_loras,
scan_trained_models,
load_model_defaults,
get_base_model_from_lora,
)
globals()["is_vision_model"] = is_vision_model
globals()["ModelConfig"] = ModelConfig
globals()["scan_trained_loras"] = scan_trained_loras
globals()["scan_trained_models"] = scan_trained_models
globals()["scan_trained_loras"] = scan_trained_models
globals()["load_model_defaults"] = load_model_defaults
globals()["get_base_model_from_lora"] = get_base_model_from_lora
return globals()[name]

1
studio/backend/core/data_recipe/jobs/constants.py
View file

@ -9,6 +9,7 @@ STAGE_PREVIEW = "preview"
STAGE_DAG = "dag"
STAGE_HEALTHCHECK = "healthcheck"
STAGE_SAMPLING = "sampling"
STAGE_SOURCE = "source"
STAGE_COLUMN_CONFIG = "column_config"
STAGE_GENERATING = "generating"
STAGE_BATCH = "batch"

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

@ -33,6 +33,60 @@ from .worker import run_job_process
_CTX = mp.get_context("spawn")
def _github_source_estimated_total(recipe: dict) -> int | None:
seed_config = recipe.get("seed_config")
if not isinstance(seed_config, dict):
return None
source = seed_config.get("source")
if not isinstance(source, dict) or source.get("seed_type") != "github_repo":
return None
repos_raw = source.get("repos")
repos = (
[repo for repo in repos_raw if isinstance(repo, str) and repo.strip()]
if isinstance(repos_raw, list)
else []
)
item_types_raw = source.get("item_types")
item_types = (
[
item
for item in item_types_raw
if isinstance(item, str) and item in {"issues", "pulls", "commits"}
]
if isinstance(item_types_raw, list)
else []
)
try:
limit = int(source.get("limit") or 0)
except (TypeError, ValueError):
return None
if not repos or not item_types or limit <= 0:
return None
return len(repos) * len(item_types) * limit
def _source_progress_status(job: Job) -> dict[str, Any] | None:
progress = job.source_progress
if progress is None:
return None
return {
"source": progress.source,
"status": progress.status,
"repo": progress.repo,
"resource": progress.resource,
"page": progress.page,
"page_items": progress.page_items,
"fetched_items": progress.fetched_items,
"estimated_total": progress.estimated_total,
"percent": progress.percent,
"rate_remaining": progress.rate_remaining,
"retry_after_sec": progress.retry_after_sec,
"message": progress.message,
"updated_at": progress.updated_at,
}
@dataclass
class Subscription:
replay: list[dict]
@ -71,8 +125,20 @@ class JobManager:
self._pump_thread: threading.Thread | None = None
self._seq: int = 0
def start(self, *, recipe: dict, run: dict) -> str:
"""Spawn the job subprocess (one at a time, no cap)."""
def start(
self,
*,
recipe: dict,
run: dict,
internal_api_key_id: int | None = None,
) -> str:
"""Spawn the job subprocess (one at a time, no cap).
``internal_api_key_id`` is the row id of a workflow-scoped
sk-unsloth-* key minted by the route layer for local providers.
JobManager revokes it when the job reaches a terminal state so the
key's live window is no longer than the run.
"""
llm_columns = recipe.get("columns") or []
llm_column_count = 0
if isinstance(llm_columns, list):
@ -92,18 +158,29 @@ class JobManager:
job_id = uuid.uuid4().hex
self._job = Job(job_id = job_id, status = "pending", started_at = time.time())
self._job.progress_columns_total = llm_column_count
self._job.source_progress_estimated_total = _github_source_estimated_total(
recipe
)
self._job.internal_api_key_id = internal_api_key_id
self._events.clear()
self._seq = 0
run_payload = dict(run)
run_payload["_job_id"] = job_id
mp_q = _CTX.Queue()
proc = _CTX.Process(
target = run_job_process,
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
daemon = True,
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
proc.start()
with native_path_secret_removed_for_child_start():
mp_q = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_job_process,),
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
daemon = True,
)
proc.start()
self._mp_q = mp_q
self._proc = proc
@ -163,6 +240,7 @@ class JobManager:
"ok": job.column_progress.ok,
"failed": job.column_progress.failed,
},
"source_progress": _source_progress_status(job),
"model_usage": {
name: {
"model": usage.model,
@ -405,6 +483,7 @@ class JobManager:
for e in self._drain_queue(mp_q):
self._handle_event(job, e)
retired_job: Job | None = None
with self._lock:
if self._job and self._job.status in {
"pending",
@ -429,6 +508,9 @@ class JobManager:
"job_id": self._job.job_id,
}
)
retired_job = self._job
if retired_job is not None:
self._retire_workflow_key(retired_job)
return
def _handle_event(self, job: Job, event: dict) -> None:
@ -436,6 +518,7 @@ class JobManager:
et = event.get("type")
msg = event.get("message") if et == "log" else None
terminal = False
with self._lock:
if self._job is None or self._job.job_id != job.job_id:
return
@ -452,18 +535,43 @@ class JobManager:
if self._job.progress.total and self._job.progress.total > 0:
self._job.progress.done = self._job.progress.total
self._job.progress.percent = 100.0
terminal = True
if et == EVENT_JOB_ERROR:
self._job.status = "error"
self._job.finished_at = time.time()
self._job.error = event.get("error") or "error"
terminal = True
if et == EVENT_JOB_CANCELLED:
terminal = True
if msg:
upd = parse_log_message(msg)
if upd:
apply_update(self._job, upd)
if terminal:
self._retire_workflow_key(job)
self._emit(event)
def _retire_workflow_key(self, job: Job) -> None:
"""Revoke the workflow-scoped sk-unsloth-* key, if one was minted.
Best-effort: revocation failures are swallowed. The key would
expire on its own after 24h, so a missed revoke is a latency
concern, not a correctness one.
"""
key_id = getattr(job, "internal_api_key_id", None)
if not key_id:
return
try:
from auth import storage # deferred: avoids circular import
storage.revoke_internal_api_key(int(key_id))
except Exception:
pass
job.internal_api_key_id = None
_JOB_MANAGER: JobManager | None = None

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

@ -4,6 +4,7 @@
from __future__ import annotations
import re
import time
from dataclasses import dataclass
from typing import Any
@ -17,9 +18,10 @@ from .constants import (
STAGE_PREVIEW,
STAGE_PROFILING,
STAGE_SAMPLING,
STAGE_SOURCE,
USAGE_RESET_STAGES,
)
from .types import Job, ModelUsage, Progress
from .types import Job, ModelUsage, Progress, SourceProgress
@dataclass(frozen = True)
@ -41,6 +43,7 @@ class ParsedUpdate:
usage_requests_total: int | None = None
usage_rpm: float | None = None
usage_section_start: bool | None = None
source_progress: SourceProgress | None = None
# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI.
@ -61,9 +64,165 @@ _RE_USAGE_TOKENS = re.compile(
_RE_USAGE_REQUESTS = re.compile(
r"requests:\s*success=(?P<success>\d+),\s*failed=(?P<failed>\d+),\s*total=(?P<total>\d+),\s*rpm=(?P<rpm>[0-9.]+)"
)
_RE_GITHUB_PAGE = re.compile(
r"^\[(?P<repo>[^\]\s]+/[^\]\s]+)\]\s+"
r"(?P<resource>issues|PRs|commits)\s+page\s+(?P<page>\d+)\s+"
r"\(\+(?P<items>\d+)\).*?\bremaining=(?P<remaining>\d+)",
re.IGNORECASE,
)
_RE_GITHUB_RATE_LIMIT = re.compile(
r"Rate limit hit\. Sleeping (?P<seconds>\d+)s until reset\.",
re.IGNORECASE,
)
_RE_GITHUB_SECONDARY_RATE_LIMIT = re.compile(
r"Secondary rate limit(?: on REST)?\. Sleep (?P<seconds>\d+)s\.",
re.IGNORECASE,
)
_RE_GITHUB_REST_RATE_LIMIT = re.compile(
r"REST 403/429, sleep (?P<seconds>\d+)",
re.IGNORECASE,
)
_RE_GITHUB_TRANSIENT = re.compile(
r"^(?P<api>GraphQL|REST) (?P<code>\d{3}) transient, retrying",
re.IGNORECASE,
)
_RE_GITHUB_NETWORK_RETRY = re.compile(
r"^(?P<api>GraphQL|REST) network error: .* Retry\.",
re.IGNORECASE,
)
_RE_GITHUB_TRIAL_LIMIT = re.compile(
r"Trial limit reached for (?P<resource>issues|PRs|commits) \((?P<items>\d+)\)",
re.IGNORECASE,
)
_RE_GITHUB_COMPLETE = re.compile(
r"Scraper complete\. GraphQL calls=\d+ REST calls=\d+",
re.IGNORECASE,
)
def parse_log_message(msg: str) -> ParsedUpdate | None:
m = _RE_GITHUB_PAGE.search(msg)
if m:
resource_raw = m.group("resource")
resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower()
repo = m.group("repo")
page = int(m.group("page"))
page_items = int(m.group("items"))
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "fetching",
repo = repo,
resource = resource,
page = page,
page_items = page_items,
rate_remaining = int(m.group("remaining")),
message = (
f"Scraping GitHub source: {repo} "
f"{resource} page {page} (+{page_items})"
),
),
)
m = _RE_GITHUB_RATE_LIMIT.search(msg)
if m:
seconds = int(m.group("seconds"))
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub rate limit. "
"Studio will resume automatically."
),
),
)
m = _RE_GITHUB_SECONDARY_RATE_LIMIT.search(msg)
if m:
seconds = int(m.group("seconds"))
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub secondary rate limit. "
"Studio will resume automatically."
),
),
)
m = _RE_GITHUB_REST_RATE_LIMIT.search(msg)
if m:
seconds = int(m.group("seconds"))
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub rate limit. "
"Studio will resume automatically."
),
),
)
m = _RE_GITHUB_TRIAL_LIMIT.search(msg)
if m:
resource_raw = m.group("resource")
resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower()
items = int(m.group("items"))
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "fetching",
resource = resource,
message = f"GitHub {resource} trial limit reached ({items}).",
),
)
m = _RE_GITHUB_TRANSIENT.search(msg)
if m:
api = m.group("api")
code = m.group("code")
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "retrying",
message = f"GitHub {api} returned {code}; retrying automatically.",
),
)
m = _RE_GITHUB_NETWORK_RETRY.search(msg)
if m:
api = m.group("api")
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "retrying",
message = f"GitHub {api} request failed; retrying automatically.",
),
)
if _RE_GITHUB_COMPLETE.search(msg):
return ParsedUpdate(
stage = STAGE_SOURCE,
source_progress = SourceProgress(
source = "github",
status = "completed",
message = "GitHub source scrape complete.",
),
)
m = _RE_SAMPLERS.search(msg)
if m:
return ParsedUpdate(
@ -172,6 +331,8 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
job.batch.idx = update.batch_idx
if update.batch_total is not None:
job.batch.total = update.batch_total
if update.source_progress is not None:
_apply_source_progress(job, update.source_progress)
if update.stage in USAGE_RESET_STAGES:
# usage summary is a short block so we reset once we move into the next stage.
@ -216,6 +377,67 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
usage.rpm = update.usage_rpm
def _apply_source_progress(job: Job, progress: SourceProgress) -> None:
previous = job.source_progress
now = time.time()
page_items = progress.page_items
if progress.repo and progress.resource and progress.page is not None:
page_key = f"{progress.repo}:{progress.resource}:{progress.page}"
count_key = f"{progress.repo}:{progress.resource}"
if page_key not in job._source_seen_pages:
job._source_seen_pages.add(page_key)
job._source_counts[count_key] = int(
job._source_counts.get(count_key, 0)
) + int(page_items or 0)
fetched_items = sum(job._source_counts.values())
if fetched_items <= 0:
fetched_items = progress.fetched_items or (
previous.fetched_items if previous else None
)
estimated_total = (
progress.estimated_total
or job.source_progress_estimated_total
or (previous.estimated_total if previous else None)
)
percent: float | None = progress.percent
if percent is None and estimated_total and fetched_items is not None:
raw_percent = (float(fetched_items) / float(max(1, estimated_total))) * 100.0
percent = 100.0 if progress.status == "completed" else min(99.0, raw_percent)
if percent is None and previous is not None:
percent = previous.percent
job.source_progress = SourceProgress(
source = "github",
status = progress.status or (previous.status if previous else None),
repo = progress.repo or (previous.repo if previous else None),
resource = progress.resource or (previous.resource if previous else None),
page = (
progress.page
if progress.page is not None
else (previous.page if previous else None)
),
page_items = (
page_items
if page_items is not None
else (previous.page_items if previous else None)
),
fetched_items = fetched_items,
estimated_total = estimated_total,
percent = percent,
rate_remaining = (
progress.rate_remaining
if progress.rate_remaining is not None
else (previous.rate_remaining if previous else None)
),
retry_after_sec = progress.retry_after_sec,
message = progress.message or (previous.message if previous else None),
updated_at = now,
)
def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
if not job.rows:
return column_progress

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

@ -35,6 +35,23 @@ class BatchProgress:
total: int | None = None
@dataclass
class SourceProgress:
source: str = "github"
status: str | None = None
repo: str | None = None
resource: str | None = None
page: int | None = None
page_items: int | None = None
fetched_items: int | None = None
estimated_total: int | None = None
percent: float | None = None
rate_remaining: int | None = None
retry_after_sec: int | None = None
message: str | None = None
updated_at: float | None = None
@dataclass
class ModelUsage:
model: str
@ -57,6 +74,7 @@ class Job:
progress: Progress = field(default_factory = Progress)
column_progress: Progress = field(default_factory = Progress)
batch: BatchProgress = field(default_factory = BatchProgress)
source_progress: SourceProgress | None = None
rows: int | None = None
cols: int | None = None
error: str | None = None
@ -70,8 +88,15 @@ class Job:
processor_artifacts: dict[str, Any] | None = None
model_usage: dict[str, ModelUsage] = field(default_factory = dict)
progress_columns_total: int | None = None
source_progress_estimated_total: int | None = None
completed_columns: list[str] = field(default_factory = list)
# Id of the internal sk-unsloth-* API key minted for a local-model
# workflow. Revoked when the job terminates so the key's live window
# matches the run rather than its 24h TTL.
internal_api_key_id: int | None = None
_current_usage_model: str | None = None
_in_usage_summary: bool = False
_seen_generation_columns: list[str] = field(default_factory = list)
_column_done: dict[str, int] = field(default_factory = dict)
_source_counts: dict[str, int] = field(default_factory = dict)
_source_seen_pages: set[str] = field(default_factory = set)

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

@ -21,6 +21,15 @@ from ..service import build_config_builder, create_data_designer
from utils.paths import ensure_dir, recipe_datasets_root
_ARTIFACT_ROOT = recipe_datasets_root()
_RE_GITHUB_CURSOR = re.compile(r"\bcursor=[^\s,]+")
_RE_SECRET_TOKEN = re.compile(
r"\b(?:(?:ghp|gho|ghu|ghs|ghr|github_pat)_[A-Za-z0-9_]+|sk-unsloth-[A-Za-z0-9]+)"
)
def _sanitize_log_message(message: str) -> str:
message = _RE_GITHUB_CURSOR.sub("cursor=<redacted>", message)
return _RE_SECRET_TOKEN.sub("<redacted-token>", message)
class _QueueLogHandler(logging.Handler):
@ -35,7 +44,7 @@ class _QueueLogHandler(logging.Handler):
"ts": record.created,
"level": record.levelname,
"logger": record.name,
"message": record.getMessage(),
"message": _sanitize_log_message(record.getMessage()),
}
self._q.put(event)
except (OSError, RuntimeError, ValueError):
@ -119,10 +128,16 @@ def run_job_process(
# Attach queue logger directly to `data_designer` so parser events survive root resets.
handler = _QueueLogHandler(event_queue)
handler.setLevel(logging.INFO)
data_designer_logger = logging.getLogger("data_designer")
data_designer_logger.addHandler(handler)
data_designer_logger.setLevel(logging.INFO)
data_designer_logger.propagate = True
for logger_name in (
"data_designer",
"scraper",
"gh_client",
"data_designer_github_repo_seed",
):
logger = logging.getLogger(logger_name)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
logger.propagate = True
if run_config_raw:
designer.set_run_config(RunConfig.model_validate(run_config_raw))
@ -180,8 +195,8 @@ def run_job_process(
{
"type": EVENT_JOB_ERROR,
"ts": time.time(),
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"error": _sanitize_log_message(str(exc)),
"stack": _sanitize_log_message(traceback.format_exc(limit = 20)),
}
)

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

@ -33,6 +33,12 @@ _OXC_TOOL_DIR = Path(__file__).resolve().parent / "oxc-validator"
_OXC_RUNNER_PATH = _OXC_TOOL_DIR / "validate.mjs"
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
@dataclass(frozen = True)
class OxcLocalCallableValidatorSpec:
name: str
@ -243,7 +249,7 @@ def _run_oxc_batch(
}
try:
tmp_dir = ensure_dir(oxc_validator_tmp_root())
env = dict(os.environ)
env = child_env_without_native_path_secret()
tmp_dir_str = str(tmp_dir)
env["TMPDIR"] = tmp_dir_str
env["TMP"] = tmp_dir_str
@ -256,6 +262,7 @@ def _run_oxc_batch(
capture_output = True,
check = False,
env = env,
**_windows_hidden_subprocess_kwargs(),
)
except (OSError, ValueError) as exc:
logger.warning("OXC subprocess launch failed: %s", exc)

2
studio/backend/core/data_recipe/oxc-validator/package.json
View file

@ -4,7 +4,7 @@
"version": "0.0.1",
"type": "module",
"dependencies": {
"oxc-parser": "^0.121.0",
"oxc-parser": "^0.123.0",
"oxlint": "^1.51.0"
}
}

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

@ -167,12 +167,7 @@ def _validate_recipe_runtime_support(
recipe: dict[str, Any],
model_providers: list[Any],
) -> None:
if not _recipe_has_llm_columns(recipe):
raise ValueError(
"Recipe Studio currently requires at least one AI generation step."
)
if not model_providers:
if _recipe_has_llm_columns(recipe) and not model_providers:
raise ValueError("Add a Provider connection block before running this recipe.")
@ -266,6 +261,21 @@ def create_data_designer(
model_providers = build_model_providers(recipe)
_validate_recipe_runtime_support(recipe, model_providers)
# DataDesigner requires at least one model provider in its registry even
# when the pipeline contains no LLM columns. Supply a lightweight stub
# so sampler/expression-only recipes can run without a real provider.
if not model_providers:
from data_designer.config.models import ModelProvider
model_providers = [
ModelProvider(
name = "_unused",
endpoint = "http://localhost",
provider_type = "openai",
api_key = None,
)
]
return DataDesigner(
artifact_path = artifact_path,
model_providers = model_providers,

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

@ -28,6 +28,8 @@ from core.inference import get_inference_backend
logger = get_logger(__name__)
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
def _is_wsl():
"""Detect if running under Windows Subsystem for Linux."""
@ -310,7 +312,7 @@ class ExportBackend:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""
Export merged model (for PEFT models).
@ -323,14 +325,21 @@ class ExportBackend:
private: Whether to make the repo private
Returns:
Tuple of (success: bool, message: str)
Tuple of (success, message, output_path). output_path is the
resolved absolute on-disk directory of the saved model when
``save_directory`` was set, else None.
"""
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first."
return False, "No model loaded. Please select a checkpoint first.", None
if not self.is_peft:
return False, "This is not a PEFT model. Use 'Export Base Model' instead."
return (
False,
"This is not a PEFT model. Use 'Export Base Model' instead.",
None,
)
output_path: Optional[str] = None
try:
# Determine save method
if format_type == "4-bit (FP4)":
@ -354,6 +363,7 @@ class ExportBackend:
# Write export metadata so the Chat page can identify the base model
self._write_export_metadata(save_directory)
logger.info(f"Model saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Push to hub if requested
if push_to_hub:
@ -361,6 +371,7 @@ class ExportBackend:
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
None,
)
logger.info(f"Pushing merged model to Hub: {repo_id}")
@ -378,14 +389,14 @@ class ExportBackend:
)
logger.info(f"Model pushed successfully to {repo_id}")
return True, "Model exported successfully"
return True, "Model exported successfully", output_path
except Exception as e:
logger.error(f"Error exporting merged model: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Export failed: {str(e)}"
return False, f"Export failed: {str(e)}", None
def export_base_model(
self,
@ -395,22 +406,26 @@ class ExportBackend:
hf_token: Optional[str] = None,
private: bool = False,
base_model_id: Optional[str] = None,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""
Export base model (for non-PEFT models).
Returns:
Tuple of (success: bool, message: str)
Tuple of (success, message, output_path). output_path is the
resolved absolute on-disk directory of the saved model when
``save_directory`` was set, else None.
"""
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first."
return False, "No model loaded. Please select a checkpoint first.", None
if self.is_peft:
return (
False,
"This is a PEFT model. Use 'Merged Model' export type instead.",
None,
)
output_path: Optional[str] = None
try:
# Save locally if requested
if save_directory:
@ -424,6 +439,7 @@ class ExportBackend:
# Write export metadata so the Chat page can identify the base model
self._write_export_metadata(save_directory)
logger.info(f"Model saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Push to hub if requested
if push_to_hub:
@ -431,6 +447,7 @@ class ExportBackend:
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
None,
)
logger.info(f"Pushing base model to Hub: {repo_id}")
@ -472,16 +489,16 @@ class ExportBackend:
)
logger.info(f"Model pushed successfully to {repo_id}")
else:
return False, "Local save directory required for Hub upload"
return False, "Local save directory required for Hub upload", None
return True, "Model exported successfully"
return True, "Model exported successfully", output_path
except Exception as e:
logger.error(f"Error exporting base model: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Export failed: {str(e)}"
return False, f"Export failed: {str(e)}", None
def export_gguf(
self,
@ -490,7 +507,7 @@ class ExportBackend:
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""
Export model in GGUF format.
@ -502,15 +519,43 @@ class ExportBackend:
hf_token: Hugging Face token
Returns:
Tuple of (success: bool, message: str)
Tuple of (success, message, output_path). output_path is the
resolved absolute on-disk directory containing the .gguf
files when ``save_directory`` was set, else None.
"""
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first."
return False, "No model loaded. Please select a checkpoint first.", None
output_path: Optional[str] = None
try:
# Convert quantization method to lowercase for unsloth
quant_method = quantization_method.lower()
# Pin convert_hf_to_gguf.py to the same llama.cpp ref as the
# llama-quantize binary (Studio installs at a tagged ref via
# setup.sh) so it can't drift past the pinned binary's gguf API.
# Set before both branches; hub-only export has save_directory == "".
global _LLAMA_CPP_SCRIPTS_WARNING_EMITTED
try:
from unsloth_zoo.llama_cpp import (
LLAMA_CPP_DEFAULT_DIR,
_resolve_local_convert_script, # noqa: F401
)
os.environ.setdefault(
"UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR
)
except ImportError:
if not _LLAMA_CPP_SCRIPTS_WARNING_EMITTED:
logger.warning(
"Unsloth: installed unsloth_zoo does not honor "
"UNSLOTH_LLAMA_CPP_SCRIPTS_DIR; convert_hf_to_gguf.py will "
"still be downloaded from llama.cpp master and may drift "
"past the pinned llama-quantize binary. Upgrade unsloth_zoo "
"to activate the local script pin."
)
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
@ -601,6 +646,7 @@ class ExportBackend:
abs_save_dir,
"\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)",
)
output_path = str(Path(abs_save_dir).resolve())
# Push to hub if requested
if push_to_hub:
@ -608,6 +654,7 @@ class ExportBackend:
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
None,
)
logger.info(f"Pushing GGUF model to Hub: {repo_id}")
@ -620,14 +667,18 @@ class ExportBackend:
)
logger.info(f"GGUF model pushed successfully to {repo_id}")
return True, f"GGUF model exported successfully ({quantization_method})"
return (
True,
f"GGUF model exported successfully ({quantization_method})",
output_path,
)
except Exception as e:
logger.error(f"Error exporting GGUF model: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"GGUF export failed: {str(e)}"
return False, f"GGUF export failed: {str(e)}", None
def export_lora_adapter(
self,
@ -636,19 +687,22 @@ class ExportBackend:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""
Export LoRA adapter only (not merged).
Returns:
Tuple of (success: bool, message: str)
Tuple of (success, message, output_path). output_path is the
resolved absolute on-disk directory of the saved adapter
when ``save_directory`` was set, else None.
"""
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first."
return False, "No model loaded. Please select a checkpoint first.", None
if not self.is_peft:
return False, "This is not a PEFT model. No adapter to export."
return False, "This is not a PEFT model. No adapter to export.", None
output_path: Optional[str] = None
try:
# Save locally if requested
if save_directory:
@ -659,6 +713,7 @@ class ExportBackend:
self.current_model.save_pretrained(save_directory)
self.current_tokenizer.save_pretrained(save_directory)
logger.info(f"Adapter saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Push to hub if requested
if push_to_hub:
@ -666,6 +721,7 @@ class ExportBackend:
return (
False,
"Repository ID and Hugging Face token required for Hub upload",
None,
)
logger.info(f"Pushing LoRA adapter to Hub: {repo_id}")
@ -676,14 +732,14 @@ class ExportBackend:
)
logger.info(f"Adapter pushed successfully to {repo_id}")
return True, "LoRA adapter exported successfully"
return True, "LoRA adapter exported successfully", output_path
except Exception as e:
logger.error(f"Error exporting LoRA adapter: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Adapter export failed: {str(e)}"
return False, f"Adapter export failed: {str(e)}", None
# Global export backend instance

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

@ -16,19 +16,25 @@ Pattern follows core/inference/orchestrator.py.
import atexit
import structlog
from collections import deque
from loggers import get_logger
import multiprocessing as mp
import queue
import threading
import time
from pathlib import Path
from typing import Any, List, Optional, Tuple
from typing import Any, Deque, Dict, List, Optional, Tuple
from utils.paths import outputs_root
logger = get_logger(__name__)
_CTX = mp.get_context("spawn")
# Maximum number of captured log lines kept in memory per export
# orchestrator. Acts as scrollback for the live export log panel in the
# UI. 4000 lines is ~1 MB worst-case at 256 chars/line.
_LOG_BUFFER_MAXLEN = 4000
class ExportOrchestrator:
"""
@ -44,6 +50,9 @@ class ExportOrchestrator:
self._proc: Optional[mp.Process] = None
self._cmd_queue: Any = None
self._resp_queue: Any = None
# Serializes export operations (load_checkpoint, export_*,
# cleanup) so concurrent HTTP requests can never interleave
# commands on the subprocess queue. Previously unused.
self._lock = threading.Lock()
# Local state mirrors (updated from subprocess responses)
@ -51,30 +60,131 @@ class ExportOrchestrator:
self.is_vision: bool = False
self.is_peft: bool = False
# ── Live log capture ─────────────────────────────────────
# Thread-safe ring buffer of log lines forwarded from the
# worker subprocess. Powers the GET /api/export/logs/stream
# SSE endpoint that the export dialog consumes.
self._log_buffer: Deque[Dict[str, Any]] = deque(maxlen = _LOG_BUFFER_MAXLEN)
self._log_lock = threading.Lock()
# Monotonically increasing sequence number. Never reset across
# operations, so SSE clients can use it as a stable cursor even
# if clear_logs() is called mid-session.
self._log_seq: int = 0
# Snapshot of _log_seq captured at the start of the current run
# (updated by clear_logs()). The SSE endpoint defaults its
# cursor to this value so a client that connects AFTER the
# worker has already emitted its first lines still sees the
# full run. Every line appended during the current run has seq
# strictly greater than _run_start_seq, and every line from
# prior runs has seq less than or equal to it.
self._run_start_seq: int = 0
# True while an export operation (load/export/cleanup) is
# running. The SSE endpoint ends the stream 1 second after
# this flips back to False to drain any trailing log lines.
self._export_active: bool = False
atexit.register(self._cleanup)
logger.info("ExportOrchestrator initialized (subprocess mode)")
# ------------------------------------------------------------------
# Live log capture helpers
# ------------------------------------------------------------------
def _append_log(self, entry: Dict[str, Any]) -> None:
"""Append a log line from the worker subprocess to the buffer.
Entries look like {"type": "log", "stream": "stdout"|"stderr",
"line": "...", "ts": ...}. Each is stamped with a monotonic
seq number before it lands in the buffer so SSE clients can
cursor through new lines.
"""
line = entry.get("line")
if not line:
return
with self._log_lock:
self._log_seq += 1
self._log_buffer.append(
{
"seq": self._log_seq,
"stream": entry.get("stream", "stdout"),
"line": line,
"ts": entry.get("ts", time.time()),
}
)
def clear_logs(self) -> None:
"""Drop any buffered log lines from a previous operation.
Called at the start of each export op so the UI shows only the
output of the current run. The seq counter is NOT reset, so an
SSE client that captured the cursor before clear_logs() will
still see new lines (with strictly greater seq numbers).
Also snapshots the current seq into ``_run_start_seq`` so the
SSE endpoint can anchor its default cursor at the start of
this run. Anything appended after this call has seq strictly
greater than the snapshot and is reachable via
``get_logs_since(get_run_start_seq())``.
"""
with self._log_lock:
self._log_buffer.clear()
self._run_start_seq = self._log_seq
def get_logs_since(self, cursor: int) -> Tuple[List[Dict[str, Any]], int]:
"""Return log entries with seq > cursor, plus the new cursor."""
with self._log_lock:
new_entries = [entry for entry in self._log_buffer if entry["seq"] > cursor]
if new_entries:
return new_entries, new_entries[-1]["seq"]
return [], cursor
def get_current_log_seq(self) -> int:
"""Return the current seq counter without reading any entries."""
with self._log_lock:
return self._log_seq
def get_run_start_seq(self) -> int:
"""Return the seq value captured at the start of the current run.
The SSE endpoint uses this as the default cursor so a client
that connects AFTER the worker has already started emitting
output still sees every line from the current run.
"""
with self._log_lock:
return self._run_start_seq
def is_export_active(self) -> bool:
"""True while an export / load / cleanup command is running."""
return self._export_active
# ------------------------------------------------------------------
# Subprocess lifecycle
# ------------------------------------------------------------------
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new export subprocess."""
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
from .worker import run_export_process
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
with native_path_secret_removed_for_child_start():
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._proc = _CTX.Process(
target = run_export_process,
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"config": config,
},
daemon = True,
)
self._proc.start()
self._proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_export_process,),
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"config": config,
},
daemon = True,
)
self._proc.start()
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
@ -179,8 +289,26 @@ class ExportOrchestrator:
error_msg = resp.get("error", "Unknown error")
raise RuntimeError(f"Subprocess error: {error_msg}")
if rtype == "log":
# Forwarded stdout/stderr line from the worker process.
self._append_log(resp)
continue
if rtype == "status":
logger.info("Export subprocess status: %s", resp.get("message", ""))
message = resp.get("message", "")
logger.info("Export subprocess status: %s", message)
# Surface status messages in the live log panel too so
# users see high level progress (e.g. "Importing
# Unsloth...", "Loading checkpoint: ...") alongside
# subprocess output.
if message:
self._append_log(
{
"stream": "status",
"line": message,
"ts": resp.get("ts", time.time()),
}
)
continue
# Other response types during wait — skip
@ -231,37 +359,47 @@ class ExportOrchestrator:
"hf_token": hf_token,
}
# Always kill existing subprocess and spawn fresh.
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
with self._lock:
# Start a fresh log buffer for this operation so the UI
# sees only the current run's output.
self.clear_logs()
self._export_active = True
try:
# Always kill existing subprocess and spawn fresh.
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
logger.info(
"Spawning fresh export subprocess for '%s'", checkpoint_path
)
self._spawn_subprocess(sub_config)
try:
resp = self._wait_response("loaded", timeout = 300)
except RuntimeError as exc:
self._shutdown_subprocess(timeout = 5)
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return False, str(exc)
try:
resp = self._wait_response("loaded")
except RuntimeError as exc:
self._shutdown_subprocess(timeout = 5)
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return False, str(exc)
if resp.get("success"):
self.current_checkpoint = resp.get("checkpoint")
self.is_vision = resp.get("is_vision", False)
self.is_peft = resp.get("is_peft", False)
logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path)
return True, resp.get("message", "Loaded successfully")
else:
error = resp.get("message", "Failed to load checkpoint")
logger.error("Failed to load checkpoint: %s", error)
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return False, error
if resp.get("success"):
self.current_checkpoint = resp.get("checkpoint")
self.is_vision = resp.get("is_vision", False)
self.is_peft = resp.get("is_peft", False)
logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path)
return True, resp.get("message", "Loaded successfully")
else:
error = resp.get("message", "Failed to load checkpoint")
logger.error("Failed to load checkpoint: %s", error)
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return False, error
finally:
self._export_active = False
def export_merged_model(
self,
@ -271,7 +409,7 @@ class ExportOrchestrator:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""Export merged PEFT model."""
return self._run_export(
"merged",
@ -293,7 +431,7 @@ class ExportOrchestrator:
hf_token: Optional[str] = None,
private: bool = False,
base_model_id: Optional[str] = None,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""Export base model (non-PEFT)."""
return self._run_export(
"base",
@ -314,7 +452,7 @@ class ExportOrchestrator:
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""Export model in GGUF format."""
return self._run_export(
"gguf",
@ -334,7 +472,7 @@ class ExportOrchestrator:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
) -> Tuple[bool, str]:
) -> Tuple[bool, str, Optional[str]]:
"""Export LoRA adapter only."""
return self._run_export(
"lora",
@ -347,46 +485,74 @@ class ExportOrchestrator:
},
)
def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str]:
"""Send an export command to the subprocess and wait for result."""
if not self._ensure_subprocess_alive():
return False, "No export subprocess running. Load a checkpoint first."
def _run_export(
self, export_type: str, params: dict
) -> Tuple[bool, str, Optional[str]]:
"""Send an export command to the subprocess and wait for result.
cmd = {"type": "export", "export_type": export_type, **params}
Returns ``(success, message, output_path)``. ``output_path`` is the
resolved on-disk directory the worker actually wrote to (None when
the export only pushed to Hub or failed before any file was
written). Surfaced via the export route's ``details.output_path``
so the dialog's success screen can show the user where the model
landed.
"""
with self._lock:
if not self._ensure_subprocess_alive():
return (
False,
"No export subprocess running. Load a checkpoint first.",
None,
)
try:
self._send_cmd(cmd)
resp = self._wait_response(
f"export_{export_type}_done",
timeout = 3600, # GGUF for 30B+ models can take 30+ min
)
return resp.get("success", False), resp.get("message", "")
except RuntimeError as exc:
return False, str(exc)
self.clear_logs()
self._export_active = True
try:
cmd = {"type": "export", "export_type": export_type, **params}
try:
self._send_cmd(cmd)
resp = self._wait_response(
f"export_{export_type}_done",
timeout = 3600, # GGUF for 30B+ models can take 30+ min
)
return (
resp.get("success", False),
resp.get("message", ""),
resp.get("output_path"),
)
except RuntimeError as exc:
return False, str(exc), None
finally:
self._export_active = False
def cleanup_memory(self) -> bool:
"""Cleanup export-related models from memory."""
if not self._ensure_subprocess_alive():
# No subprocess — just clear local state
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return True
with self._lock:
if not self._ensure_subprocess_alive():
# No subprocess — just clear local state
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return True
try:
self._send_cmd({"type": "cleanup"})
resp = self._wait_response("cleanup_done", timeout = 30)
success = resp.get("success", False)
except RuntimeError:
success = False
self._export_active = True
try:
try:
self._send_cmd({"type": "cleanup"})
resp = self._wait_response("cleanup_done", timeout = 30)
success = resp.get("success", False)
except RuntimeError:
success = False
# Shut down subprocess after cleanup — no model loaded
self._shutdown_subprocess()
# Shut down subprocess after cleanup — no model loaded
self._shutdown_subprocess()
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return success
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
return success
finally:
self._export_active = False
def scan_checkpoints(
self, outputs_dir: str = str(outputs_root())

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

@ -17,10 +17,12 @@ Pattern follows core/inference/worker.py and core/training/worker.py.
from __future__ import annotations
import errno
import structlog
from loggers import get_logger
import os
import sys
import threading
import time
import traceback
from pathlib import Path
@ -29,38 +31,164 @@ from typing import Any
logger = get_logger(__name__)
def _activate_transformers_version(model_name: str) -> None:
"""Activate the correct transformers version BEFORE any ML imports.
# Gate that controls whether captured stdout/stderr lines are forwarded
# to the parent's resp_queue (and from there to the export-dialog SSE
# stream). Closed by default so the noisy bootstrap phase -- transformers
# venv activation, Unsloth/torch imports, base-model resolution, "Top
# GGUF/hub models" lists, vision detection, weight loading bars -- is
# suppressed in the UI. _handle_export() opens the gate at the start of
# the actual export work and leaves it open; the orchestrator always
# spawns a fresh subprocess for the next checkpoint load (see
# orchestrator._spawn_subprocess) which resets this state.
#
# Lines dropped while the gate is closed are still echoed to the saved
# original stdout/stderr fds so the server console / log file keeps the
# full output for debugging.
_log_forward_gate = threading.Event()
If the model needs transformers 5.x, prepend the pre-installed .venv_t5/
directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/).
def _setup_log_capture(resp_queue: Any) -> None:
"""Redirect fds 1 and 2 through pipes so every line printed by this
worker process and any child process it spawns is forwarded to the
parent process via resp_queue as {"type": "log", ...} messages.
Must be called BEFORE LogConfig.setup_logging and BEFORE any ML
imports, otherwise library handlers may capture the original stderr
reference and bypass the pipe.
Lines are also echoed back to the original stdout/stderr so the
server console keeps receiving the full subprocess output, even
while ``_log_forward_gate`` is closed.
"""
try:
saved_out_fd = os.dup(1)
saved_err_fd = os.dup(2)
except OSError:
# dup failed (exotic platforms) - give up quietly, export still
# works, just no live log streaming.
return
try:
r_out, w_out = os.pipe()
r_err, w_err = os.pipe()
except OSError:
os.close(saved_out_fd)
os.close(saved_err_fd)
return
try:
os.dup2(w_out, 1)
os.dup2(w_err, 2)
except OSError:
for fd in (saved_out_fd, saved_err_fd, r_out, w_out, r_err, w_err):
try:
os.close(fd)
except OSError:
pass
return
# Close the write ends we just dup2'd (fds 1 and 2 are the real
# write ends now).
os.close(w_out)
os.close(w_err)
# Replace Python's sys.stdout/sys.stderr with line-buffered writers
# bound to the (now-redirected) fds 1 and 2.
try:
sys.stdout = os.fdopen(1, "w", buffering = 1, encoding = "utf-8", errors = "replace")
sys.stderr = os.fdopen(2, "w", buffering = 1, encoding = "utf-8", errors = "replace")
except Exception:
pass
def _reader(read_fd: int, stream_name: str, echo_fd: int) -> None:
buf = bytearray()
while True:
try:
chunk = os.read(read_fd, 4096)
except OSError as exc:
if exc.errno == errno.EBADF:
break
continue
if not chunk:
break
# Echo to the original fd so the server console still sees
# the full output.
try:
os.write(echo_fd, chunk)
except OSError:
pass
buf.extend(chunk)
# Split on \n OR \r so tqdm-style progress bars update.
while True:
nl = -1
for i, b in enumerate(buf):
if b == 0x0A or b == 0x0D:
nl = i
break
if nl < 0:
break
line = bytes(buf[:nl]).decode("utf-8", errors = "replace")
del buf[: nl + 1]
if not line:
continue
if not _log_forward_gate.is_set():
# Gate closed (bootstrap phase) -- already echoed to
# the saved console fd above; drop the line so the
# export dialog doesn't see import / vendoring noise.
continue
try:
resp_queue.put_nowait(
{
"type": "log",
"stream": stream_name,
"line": line,
"ts": time.time(),
}
)
except Exception:
# Queue put failed (full, closed, etc.) - drop the
# line rather than crash the reader thread.
pass
if buf and _log_forward_gate.is_set():
try:
resp_queue.put_nowait(
{
"type": "log",
"stream": stream_name,
"line": bytes(buf).decode("utf-8", errors = "replace"),
"ts": time.time(),
}
)
except Exception:
pass
t_out = threading.Thread(
target = _reader,
args = (r_out, "stdout", saved_out_fd),
daemon = True,
name = "export-log-stdout",
)
t_err = threading.Thread(
target = _reader,
args = (r_err, "stderr", saved_err_fd),
daemon = True,
name = "export-log-stderr",
)
t_out.start()
t_err.start()
def _activate_transformers_version(model_name: str) -> None:
"""Activate the correct transformers version BEFORE any ML imports."""
# Ensure backend is on path for utils imports
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from utils.transformers_version import (
needs_transformers_5,
_resolve_base_model,
_ensure_venv_t5_exists,
_VENV_T5_DIR,
)
from utils.transformers_version import activate_transformers_for_subprocess
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
if not _ensure_venv_t5_exists():
raise RuntimeError(
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
)
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
# Propagate to child subprocesses (e.g. GGUF converter)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
else:
logger.info("Using default transformers (4.57.x) for %s", model_name)
activate_transformers_for_subprocess(model_name)
def _send_response(resp_queue: Any, response: dict) -> None:
@ -78,6 +206,19 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
load_in_4bit = cmd.get("load_in_4bit", True)
trust_remote_code = cmd.get("trust_remote_code", False)
# Auto-enable trust_remote_code for NemotronH/Nano models.
if not trust_remote_code:
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_cp_lower = checkpoint_path.lower()
if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")
):
trust_remote_code = True
logger.info(
"Auto-enabled trust_remote_code for Nemotron model: %s",
checkpoint_path,
)
try:
_send_response(
resp_queue,
@ -126,9 +267,17 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
export_type = cmd["export_type"] # "merged", "base", "gguf", "lora"
response_type = f"export_{export_type}_done"
# Open the log forwarding gate so the user sees the actual export
# progress (Unsloth merge bars, file copies, GGUF conversion, etc.)
# in the live log panel. The gate stays open for the rest of this
# subprocess's life; the orchestrator spawns a fresh subprocess for
# the next checkpoint load, which resets the gate to closed.
_log_forward_gate.set()
output_path: Any = None
try:
if export_type == "merged":
success, message = backend.export_merged_model(
success, message, output_path = backend.export_merged_model(
save_directory = cmd.get("save_directory", ""),
format_type = cmd.get("format_type", "16-bit (FP16)"),
push_to_hub = cmd.get("push_to_hub", False),
@ -137,7 +286,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
private = cmd.get("private", False),
)
elif export_type == "base":
success, message = backend.export_base_model(
success, message, output_path = backend.export_base_model(
save_directory = cmd.get("save_directory", ""),
push_to_hub = cmd.get("push_to_hub", False),
repo_id = cmd.get("repo_id"),
@ -146,7 +295,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
base_model_id = cmd.get("base_model_id"),
)
elif export_type == "gguf":
success, message = backend.export_gguf(
success, message, output_path = backend.export_gguf(
save_directory = cmd.get("save_directory", ""),
quantization_method = cmd.get("quantization_method", "Q4_K_M"),
push_to_hub = cmd.get("push_to_hub", False),
@ -154,7 +303,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
hf_token = cmd.get("hf_token"),
)
elif export_type == "lora":
success, message = backend.export_lora_adapter(
success, message, output_path = backend.export_lora_adapter(
save_directory = cmd.get("save_directory", ""),
push_to_hub = cmd.get("push_to_hub", False),
repo_id = cmd.get("repo_id"),
@ -170,6 +319,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
"type": response_type,
"success": success,
"message": message,
"output_path": output_path,
"ts": time.time(),
},
)
@ -181,6 +331,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
"type": response_type,
"success": False,
"message": str(exc),
"output_path": None,
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
@ -226,10 +377,26 @@ def run_export_process(
"""
import queue as _queue
# Install fd-level stdout/stderr capture FIRST so every subsequent
# print and every child process inherits the redirected fds. This
# is what powers the live export log stream in the UI.
_setup_log_capture(resp_queue)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
# Force unbuffered output from any child Python process (e.g. the
# GGUF converter) so their prints surface in the log stream as they
# happen rather than at the end.
os.environ["PYTHONUNBUFFERED"] = "1"
# tqdm defaults to a 10-second mininterval when stdout is not a tty
# (which it isn't here -- we redirected fd 1/2 to a pipe). That makes
# multi-step progress bars look frozen in the export log panel. Force
# frequent flushes so the user sees movement during merge / GGUF
# conversion. Has no effect on single-step bars (e.g. "Copying 1
# files") which only emit start/end events regardless.
os.environ.setdefault("TQDM_MININTERVAL", "0.5")
import warnings
from loggers.config import LogConfig

576
studio/backend/core/inference/anthropic_compat.py Normal file
View file

@ -0,0 +1,576 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""
Anthropic Messages API ↔ OpenAI format translation utilities.
Pure functions and a stateful stream emitter — no FastAPI, no I/O.
"""
from __future__ import annotations
import json
from typing import Any, Optional, Union
def _anthropic_image_block_to_openai_part(block: dict) -> Optional[dict]:
"""Translate one Anthropic ``image`` block to an OpenAI ``image_url`` part.
Accepts both source shapes:
- ``{"type": "base64", "media_type": "image/jpeg", "data": "..."}``
- ``{"type": "url", "url": "https://..."}``
Returns ``None`` when the source is malformed so the caller can skip it.
"""
source = block.get("source") or {}
stype = source.get("type")
if stype == "base64":
data = source.get("data")
if not data:
return None
media_type = source.get("media_type") or "image/jpeg"
return {
"type": "image_url",
"image_url": {"url": f"data:{media_type};base64,{data}"},
}
if stype == "url":
url = source.get("url")
if not url:
return None
return {"type": "image_url", "image_url": {"url": url}}
return None
def anthropic_messages_to_openai(
messages: list[dict],
system: Optional[Union[str, list]] = None,
) -> list[dict]:
"""Convert Anthropic messages + system to OpenAI-format message dicts.
User messages that carry ``image`` blocks are emitted as OpenAI
multimodal content arrays (``[{type: "text", ...}, {type: "image_url", ...}]``)
so they flow through llama-server's native vision pathway.
"""
result: list[dict] = []
# System prompt
if system:
if isinstance(system, str):
result.append({"role": "system", "content": system})
elif isinstance(system, list):
parts = []
for block in system:
if isinstance(block, dict) and block.get("type") == "text":
parts.append(block["text"])
elif isinstance(block, str):
parts.append(block)
if parts:
result.append({"role": "system", "content": "\n".join(parts)})
for msg in messages:
role = msg["role"] if isinstance(msg, dict) else msg.role
content = msg["content"] if isinstance(msg, dict) else msg.content
if isinstance(content, str):
result.append({"role": role, "content": content})
continue
if role == "assistant":
# Assistant content carries text + tool_use; images aren't
# part of Anthropic's assistant content model.
text_parts: list[str] = []
tool_calls: list[dict] = []
for block in content:
b = block if isinstance(block, dict) else block.model_dump()
btype = b.get("type", "")
if btype == "text":
text_parts.append(b["text"])
elif btype == "tool_use":
tool_calls.append(
{
"id": b["id"],
"type": "function",
"function": {
"name": b["name"],
"arguments": json.dumps(b["input"]),
},
}
)
msg_dict: dict[str, Any] = {"role": "assistant"}
if text_parts:
msg_dict["content"] = "\n".join(text_parts)
if tool_calls:
msg_dict["tool_calls"] = tool_calls
result.append(msg_dict)
continue
if role == "user":
# Build an ordered part list so text/image interleaving is
# preserved (e.g. [text, image, text, image]). tool_result
# blocks become their own OpenAI "tool" role messages.
user_parts: list[dict] = []
has_image = False
tool_results: list[dict] = []
for block in content:
b = block if isinstance(block, dict) else block.model_dump()
btype = b.get("type", "")
if btype == "text":
user_parts.append({"type": "text", "text": b["text"]})
elif btype == "image":
part = _anthropic_image_block_to_openai_part(b)
if part is not None:
user_parts.append(part)
has_image = True
elif btype == "tool_result":
tc = b.get("content", "")
if isinstance(tc, list):
tc = " ".join(
p["text"]
for p in tc
if isinstance(p, dict) and p.get("type") == "text"
)
tool_results.append(
{
"role": "tool",
"tool_call_id": b["tool_use_id"],
"content": str(tc),
}
)
if has_image:
result.append({"role": "user", "content": user_parts})
else:
# No images — collapse text parts to a plain string so
# existing text-only callers keep their simple shape.
text = "\n".join(p["text"] for p in user_parts)
if text:
result.append({"role": "user", "content": text})
for tr in tool_results:
result.append(tr)
return result
def anthropic_tools_to_openai(tools: list) -> list[dict]:
"""Convert Anthropic tool definitions to OpenAI function-tool format."""
result = []
for t in tools:
td = t if isinstance(t, dict) else t.model_dump()
result.append(
{
"type": "function",
"function": {
"name": td["name"],
"description": td.get("description", ""),
"parameters": td.get("input_schema", {}),
},
}
)
return result
def anthropic_tool_choice_to_openai(tc: Any) -> Any:
"""Translate Anthropic `tool_choice` into OpenAI `tool_choice`.
Anthropic formats (all dict shapes with a ``type`` discriminator):
- ``{"type": "auto"}`` → ``"auto"``
- ``{"type": "any"}`` → ``"required"``
- ``{"type": "none"}`` → ``"none"``
- ``{"type": "tool", "name": "get_weather"}``
→ ``{"type": "function", "function": {"name": "get_weather"}}``
Returns ``None`` for ``None`` or any unrecognized shape (caller may
then fall back to its own default, typically ``"auto"``).
"""
if tc is None:
return None
if not isinstance(tc, dict):
return None
t = tc.get("type")
if t == "auto":
return "auto"
if t == "any":
return "required"
if t == "none":
return "none"
if t == "tool":
name = tc.get("name")
if not name:
return None
return {"type": "function", "function": {"name": name}}
return None
def build_anthropic_sse_event(event_type: str, data: dict) -> str:
"""Format a single Anthropic SSE event."""
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
class AnthropicStreamEmitter:
"""Converts generator events from generate_chat_completion_with_tools()
into Anthropic Messages SSE strings."""
def __init__(self) -> None:
self.block_index: int = 0
self._text_block_open: bool = False
self._prev_text: str = ""
self._usage: dict = {}
def start(self, message_id: str, model: str) -> list[str]:
"""Emit message_start and open the first text content block."""
events = []
events.append(
build_anthropic_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": message_id,
"type": "message",
"role": "assistant",
"content": [],
"model": model,
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
},
},
)
)
events.extend(self._open_text_block())
return events
def feed(self, event: dict) -> list[str]:
"""Process one generator event, return SSE strings."""
etype = event.get("type", "")
if etype == "content":
return self._handle_content(event)
elif etype == "tool_start":
return self._handle_tool_start(event)
elif etype == "tool_end":
return self._handle_tool_end(event)
elif etype == "metadata":
self._usage = event.get("usage", {})
return []
# status events — no Anthropic equivalent
return []
def finish(self, stop_reason: str = "end_turn") -> list[str]:
"""Close any open block and emit message_delta + message_stop."""
events = []
if self._text_block_open:
events.append(self._close_block())
events.append(
build_anthropic_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
"usage": {
"output_tokens": self._usage.get("completion_tokens", 0),
},
},
)
)
events.append(
build_anthropic_sse_event(
"message_stop",
{
"type": "message_stop",
},
)
)
return events
def _handle_content(self, event: dict) -> list[str]:
cumulative = event.get("text", "")
new_text = cumulative[len(self._prev_text) :]
self._prev_text = cumulative
if not new_text:
return []
if not self._text_block_open:
events = self._open_text_block()
else:
events = []
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {"type": "text_delta", "text": new_text},
},
)
)
return events
def _handle_tool_start(self, event: dict) -> list[str]:
events = []
# Close current text block if open
if self._text_block_open:
events.append(self._close_block())
# Open a tool_use block
self.block_index += 1
events.append(
build_anthropic_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": self.block_index,
"content_block": {
"type": "tool_use",
"id": event.get("tool_call_id", ""),
"name": event.get("tool_name", ""),
"input": {},
},
},
)
)
# Emit the arguments as input_json_delta
args = event.get("arguments", {})
if args:
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {
"type": "input_json_delta",
"partial_json": json.dumps(args),
},
},
)
)
return events
def _handle_tool_end(self, event: dict) -> list[str]:
events = []
# Close the tool_use block
events.append(self._close_block())
# Emit custom tool_result event (non-standard, ignored by SDKs)
events.append(
build_anthropic_sse_event(
"tool_result",
{
"type": "tool_result",
"tool_use_id": event.get("tool_call_id", ""),
"content": event.get("result", ""),
},
)
)
# Open a new text block for the model's next response
self.block_index += 1
events.extend(self._open_text_block())
# Reset text tracking for the next synthesis turn
self._prev_text = ""
return events
def _open_text_block(self) -> list[str]:
self._text_block_open = True
return [
build_anthropic_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": self.block_index,
"content_block": {"type": "text", "text": ""},
},
)
]
def _close_block(self) -> str:
self._text_block_open = False
return build_anthropic_sse_event(
"content_block_stop",
{
"type": "content_block_stop",
"index": self.block_index,
},
)
class AnthropicPassthroughEmitter:
"""Converts llama-server's OpenAI-format streaming chunks into Anthropic SSE.
Used for the client-side tool-use pass-through path: the client (e.g. Claude
Code) sends its own tool definitions in the ``tools`` field and expects to
execute them itself. We forward them to llama-server and translate the
streaming response back to Anthropic format without executing anything.
"""
def __init__(self) -> None:
self.block_index: int = -1
self._current_block_type: Optional[str] = None # "text" | "tool_use" | None
self._tool_call_states: dict = {} # delta index -> {block_index, id, name}
self._usage: dict = {}
self._stop_reason: str = "end_turn"
def start(self, message_id: str, model: str) -> list[str]:
return [
build_anthropic_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": message_id,
"type": "message",
"role": "assistant",
"content": [],
"model": model,
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
},
},
)
]
def feed_chunk(self, chunk: dict) -> list[str]:
"""Process one OpenAI streaming chat.completion.chunk."""
events: list[str] = []
# usage-only chunks carry token totals
usage = chunk.get("usage")
if usage:
self._usage = usage
choices = chunk.get("choices") or []
if not choices:
return events
choice = choices[0]
delta = choice.get("delta") or {}
finish_reason = choice.get("finish_reason")
# ── Text content ──
content = delta.get("content")
if content:
if self._current_block_type != "text":
if self._current_block_type is not None:
events.append(self._close_current_block())
events.extend(self._open_text_block())
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {"type": "text_delta", "text": content},
},
)
)
# ── Tool calls (streaming deltas) ──
tool_calls = delta.get("tool_calls") or []
for tc in tool_calls:
tc_idx = tc.get("index", 0)
fn = tc.get("function") or {}
if tc_idx not in self._tool_call_states:
# New tool call — close prior block, open tool_use block
if self._current_block_type is not None:
events.append(self._close_current_block())
tc_id = tc.get("id", "")
tc_name = fn.get("name", "")
self.block_index += 1
self._current_block_type = "tool_use"
self._tool_call_states[tc_idx] = {
"block_index": self.block_index,
"id": tc_id,
"name": tc_name,
}
events.append(
build_anthropic_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": self.block_index,
"content_block": {
"type": "tool_use",
"id": tc_id,
"name": tc_name,
"input": {},
},
},
)
)
args_delta = fn.get("arguments", "")
if args_delta:
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self._tool_call_states[tc_idx]["block_index"],
"delta": {
"type": "input_json_delta",
"partial_json": args_delta,
},
},
)
)
# ── Finish reason ──
if finish_reason:
if finish_reason == "tool_calls":
self._stop_reason = "tool_use"
elif finish_reason == "length":
self._stop_reason = "max_tokens"
else:
self._stop_reason = "end_turn"
return events
def finish(self) -> list[str]:
events: list[str] = []
if self._current_block_type is not None:
events.append(self._close_current_block())
events.append(
build_anthropic_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {
"stop_reason": self._stop_reason,
"stop_sequence": None,
},
"usage": {
"output_tokens": self._usage.get("completion_tokens", 0),
},
},
)
)
events.append(
build_anthropic_sse_event(
"message_stop",
{"type": "message_stop"},
)
)
return events
def _open_text_block(self) -> list[str]:
self.block_index += 1
self._current_block_type = "text"
return [
build_anthropic_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": self.block_index,
"content_block": {"type": "text", "text": ""},
},
)
]
def _close_current_block(self) -> str:
idx = self.block_index
self._current_block_type = None
return build_anthropic_sse_event(
"content_block_stop",
{
"type": "content_block_stop",
"index": idx,
},
)

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

@ -8,6 +8,7 @@ Supports: SNAC (Orpheus), CSM (Sesame), BiCodec (Spark), DAC (OuteTTS)
import io
import re
import subprocess
import wave
import structlog
from loggers import get_logger
@ -16,6 +17,11 @@ from typing import Optional, Tuple
import numpy as np
import torch
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
logger = get_logger(__name__)
@ -81,7 +87,6 @@ class AudioCodecManager:
return
import os
import sys
import subprocess
# Clone SparkAudio/Spark-TTS GitHub repo for the sparktts Python package
# (same approach as training — the HF model repos don't contain the package)
@ -101,6 +106,8 @@ class AudioCodecManager:
spark_code_dir,
],
check = True,
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
)
if spark_code_dir not in sys.path:
@ -119,7 +126,6 @@ class AudioCodecManager:
return
import os
import sys
import subprocess
# Clone OuteTTS repo (same pattern as Spark-TTS / BiCodec)
# The pip package has problematic dependencies; the notebook clones and
@ -139,6 +145,8 @@ class AudioCodecManager:
outetts_code_dir,
],
check = True,
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
)
# Remove files that pull in heavy / incompatible dependencies
# (matches notebook: gguf_model.py is under models/, others under outetts/)

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

@ -6,6 +6,15 @@
import utils.hardware.hardware as hw
DEFAULT_MODELS_GGUF = [
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",
"unsloth/gemma-4-26B-A4B-it-GGUF",
"unsloth/Qwen3.6-35B-A3B-GGUF",
"unsloth/Qwen3.5-4B-GGUF",
"unsloth/Qwen3.5-9B-GGUF",
"unsloth/Qwen3.5-35B-A3B-GGUF",
"unsloth/Qwen3.5-0.8B-GGUF",
"unsloth/Llama-3.2-1B-Instruct-GGUF",
"unsloth/Llama-3.2-3B-Instruct-GGUF",
"unsloth/Llama-3.1-8B-Instruct-GGUF",
@ -15,6 +24,19 @@ DEFAULT_MODELS_GGUF = [
]
DEFAULT_MODELS_STANDARD = [
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",
"unsloth/gemma-4-26B-A4B-it-GGUF",
"unsloth/Qwen3.6-35B-A3B-GGUF",
"unsloth/Qwen3.5-4B-GGUF",
"unsloth/Qwen3.5-9B-GGUF",
"unsloth/Qwen3.5-35B-A3B-GGUF",
"unsloth/Qwen3.5-0.8B-GGUF",
"unsloth/gemma-4-E2B-it",
"unsloth/gemma-4-E4B-it",
"unsloth/gemma-4-31B-it",
"unsloth/gemma-4-26B-A4B-it",
"unsloth/Qwen3-4B-Instruct-2507",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
"unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",

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

@ -253,6 +253,10 @@ class InferenceBackend:
"""
Load any model: base, LoRA adapter, text, or vision.
"""
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
if max_seq_length <= 0:
max_seq_length = 2048
try:
model_name = config.identifier

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

File diff suppressed because it is too large Load diff

120
studio/backend/core/inference/llama_server_args.py Normal file
View file

@ -0,0 +1,120 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Validator for user-supplied llama-server pass-through args.
Studio runs llama-server as a managed subprocess and lets callers pass
extra flags directly (CLI: ``unsloth run ... --top-k 20``; HTTP:
``LoadRequest.llama_extra_args``). This module is the boundary that
rejects only flags Studio fundamentally cannot share with the user --
model identity, the auth key, and the network endpoint Studio's HTTP
proxy targets. Anything else passes through.
User-supplied args are appended to ``cmd`` after Studio's auto-set
flags, so llama.cpp's last-wins CLI parsing makes the user's value
override the auto-set one. That covers tunable knobs the user might
reasonably want to override -- ``-c``/``--ctx-size``,
``-np``/``--parallel``, ``-fa``/``--flash-attn``,
``-ngl``/``--gpu-layers``, ``-t``/``--threads``, ``-fit``/``--fit*``,
``--cache-type-k/v``, ``--chat-template-file/-kwargs``,
``--spec-*``, ``--jinja``/``--no-jinja``,
``--no-context-shift``/``--context-shift``, sampling params, etc.
Reference: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
"""
from __future__ import annotations
from typing import Iterable, Optional
# Each group is the full set of aliases (short + long) for one
# hard-denied flag, taken from the llama-server README. If llama.cpp
# adds a new alias for an existing denied flag, extend the relevant
# group.
#
# Flags NOT in this list (e.g. -c, --parallel, --flash-attn, -ngl,
# -t/--threads, --jinja, --no-context-shift, --fit*, --cache-type-*,
# --chat-template-*, --spec-*) pass through and override Studio's
# auto-set version via llama.cpp's last-wins CLI parsing.
_DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Model identity -- Studio resolves the model from LoadRequest and
# passes -m / mmproj after downloading from HF if needed. A second
# -m would point at a different model than the one Studio thinks
# is loaded.
frozenset({"-m", "--model"}),
frozenset({"-mu", "--model-url"}),
frozenset({"-dr", "--docker-repo"}),
frozenset({"-hf", "-hfr", "--hf-repo"}),
frozenset({"-hff", "--hf-file"}),
frozenset({"-hfv", "-hfrv", "--hf-repo-v"}),
frozenset({"-hffv", "--hf-file-v"}),
frozenset({"-hft", "--hf-token"}),
frozenset({"-mm", "--mmproj"}),
frozenset({"-mmu", "--mmproj-url"}),
# Networking -- Studio binds llama-server's port and reverse-proxies
# HTTP traffic to it. Retargeting host/port/path/prefix would
# orphan Studio's proxy and the UI would lose the server.
frozenset({"--host"}),
frozenset({"--port"}),
frozenset({"--path"}),
frozenset({"--api-prefix"}),
frozenset({"--reuse-port"}),
# Auth / TLS -- Studio terminates auth at its own layer; an
# upstream --api-key would shadow Studio's UNSLOTH_DIRECT_STREAM
# key, and TLS on llama-server would break the local proxy hop.
frozenset({"--api-key"}),
frozenset({"--api-key-file"}),
frozenset({"--ssl-key-file"}),
frozenset({"--ssl-cert-file"}),
# Single-model server -- Studio runs one model per llama-server
# process and serves its own UI. Enabling multi-model loading or
# llama-server's built-in web UI changes the surface clients see.
frozenset({"--webui", "--no-webui"}),
frozenset({"--models-dir"}),
frozenset({"--models-preset"}),
frozenset({"--models-max"}),
frozenset({"--models-autoload", "--no-models-autoload"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
def _flag_name(token: str) -> Optional[str]:
"""Return the flag name for a token, or None if it isn't a flag.
Peels ``--key=value`` to the bare ``--key``. Plain numeric values
like ``-1`` or ``-0.5`` (e.g. ``--seed -1``) are values, not flags;
llama-server short-form flags always start with a letter.
"""
if not token.startswith("-") or token in {"-", "--"}:
return None
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
return None
return token.split("=", 1)[0]
def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
"""Validate user-supplied llama-server args.
Returns the args as a flat list ready to extend the llama-server
command. Raises ``ValueError`` (with the offending flag in the
message) the moment a token resolves to a Studio-managed flag.
"""
if not args:
return []
out: list[str] = []
for raw in args:
token = str(raw)
flag = _flag_name(token)
if flag is not None and flag in _DENYLIST:
raise ValueError(
f"llama-server flag '{flag}' is managed by Unsloth Studio "
f"and cannot be passed as an extra arg"
)
out.append(token)
return out
def is_managed_flag(flag: str) -> bool:
"""True if ``flag`` is a Studio-managed llama-server flag."""
return flag in _DENYLIST

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

@ -109,12 +109,13 @@ class InferenceOrchestrator:
self._top_models_ready.wait(timeout = 5)
top_gguf = self._top_gguf_cache or []
top_hub = self._top_hub_cache or []
# GGUFs first, then hub models, then static fallbacks.
# Curated static defaults first (editorial picks like new models),
# then HF download-ranked models to backfill.
# Send extras so the frontend still has 4 per category
# after removing already-downloaded models.
result: list[str] = []
seen: set[str] = set()
for m in top_gguf + top_hub + self._static_models:
for m in self._static_models + top_gguf + top_hub:
if m not in seen:
result.append(m)
seen.add(m)
@ -165,23 +166,30 @@ class InferenceOrchestrator:
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new inference subprocess."""
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
from .worker import run_inference_process
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._cancel_event = _CTX.Event()
with native_path_secret_removed_for_child_start():
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._cancel_event = _CTX.Event()
self._proc = _CTX.Process(
target = run_inference_process,
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"cancel_event": self._cancel_event,
"config": config,
},
daemon = True,
)
self._proc.start()
self._proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_inference_process,),
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"cancel_event": self._cancel_event,
"config": config,
},
daemon = True,
)
self._proc.start()
logger.info("Inference subprocess started (pid=%s)", self._proc.pid)
def _cancel_generation(self) -> None:
@ -707,6 +715,17 @@ class InferenceOrchestrator:
def unload_model(self, model_name: str) -> bool:
"""Unload a model from the subprocess."""
if model_name in self.loading_models:
logger.info(
"Cancelling in-flight load for model '%s' by terminating subprocess",
model_name,
)
self._shutdown_subprocess(timeout = 0.5)
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()
return True
if not self._ensure_subprocess_alive():
# No subprocess — just clear local state
self.models.pop(model_name, None)

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

@ -14,6 +14,8 @@ import os
os.environ["UNSLOTH_IS_PRESENT"] = "1"
import random
import re
import shlex
import ssl
import subprocess
import sys
@ -26,11 +28,240 @@ from loggers import get_logger
logger = get_logger(__name__)
_EXEC_TIMEOUT = 300 # 5 minutes
# Pre-import modules used in _sandbox_preexec at module level so that
# the preexec_fn closure does not trigger the import machinery in the
# forked child (which can deadlock in multi-threaded servers).
_libc = None
if sys.platform == "linux":
try:
import ctypes
import ctypes.util
_libc_name = ctypes.util.find_library("c")
if _libc_name:
_libc = ctypes.CDLL(_libc_name, use_errno = True)
except (OSError, AttributeError):
pass
_resource = None
if sys.platform != "win32":
try:
import resource as _resource
except ImportError:
pass
# Strict raster-image allowlist for sandbox file serving.
# No .svg (XSS risk via embedded scripts), no .html, no .pdf.
_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
_MAX_OUTPUT_CHARS = 8000 # truncate long output
_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
_BLOCKED_COMMANDS_COMMON = frozenset(
{
"rm",
"sudo",
"su",
"dd",
"chmod",
"chown",
"mkfs",
"shutdown",
"reboot",
"passwd",
"mount",
"umount",
"fdisk",
"kill",
"killall",
"pkill",
}
)
_BLOCKED_COMMANDS_WIN = frozenset(
{
"rmdir",
"takeown",
"icacls",
"runas",
"powershell",
"pwsh",
}
)
_BLOCKED_COMMANDS = (
_BLOCKED_COMMANDS_COMMON | _BLOCKED_COMMANDS_WIN
if sys.platform == "win32"
else _BLOCKED_COMMANDS_COMMON
)
def _find_blocked_commands(command: str) -> set[str]:
"""Detect blocked commands using shlex tokenization and regex scanning.
Catches: full paths (/usr/bin/sudo), quoted strings ("sudo"),
split-quotes (su""do), backslash escapes (\\rm), and command-position
words after ;, |, &&, $().
"""
blocked = set()
# 1. shlex tokenization (handles quotes, escapes, concatenation)
try:
tokens = (
shlex.split(command)
if sys.platform != "win32"
else shlex.split(command, posix = False)
)
except ValueError:
tokens = command.split()
for token in tokens:
base = os.path.basename(token).lower()
# Strip common Windows executable extensions so that
# runas.exe, shutdown.bat, etc. match the blocklist.
stem, ext = os.path.splitext(base)
if ext in {".exe", ".com", ".bat", ".cmd"}:
base = stem
if base in _BLOCKED_COMMANDS:
blocked.add(base)
# 2. Regex: catch blocked words at shell command boundaries
# (semicolons, pipes, &&, ||, backticks, $(), <(), subshells, newlines)
# Uses a single combined pattern for all blocked words.
# Handles optional Unix path prefix (/usr/bin/) and Windows drive
# letter prefix (C:\Windows\...\).
lowered = command.lower()
if _BLOCKED_COMMANDS:
words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS))
pattern = (
rf"(?:^|[;&|`\n(]\s*|[$]\(\s*|<\(\s*)"
rf"(?:[\w./\\-]*/|[a-zA-Z]:[/\\][\w./\\-]*)?"
rf"({words_alt})(?:\.(?:exe|com|bat|cmd))?\b"
)
blocked.update(re.findall(pattern, lowered))
# 3. Check for nested shell invocations (bash -c 'sudo whoami',
# bash -lc '...', bash --login -c '...', cmd /c '...').
# When a -c or /c flag is found, look backwards for a shell name
# (skipping intermediate flags like --login, -l, -x) and recursively
# scan the nested command string.
_SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"}
_SHELLS_WIN = {"cmd", "cmd.exe"}
for i, token in enumerate(tokens):
tok_lower = token.lower()
# Match -c exactly, or combined flags ending in c (e.g. -lc, -xc)
is_unix_c = tok_lower == "-c" or (
tok_lower.startswith("-")
and tok_lower.endswith("c")
and not tok_lower.startswith("--")
)
is_win_c = tok_lower == "/c"
if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens):
continue
# Look backwards past any flags to find the shell binary.
# On Unix, flags start with - (skip those). On Windows, flags
# start with / but so do absolute paths, so only skip short
# single-char /X flags (not /bin/bash style paths).
for j in range(i - 1, -1, -1):
prev = tokens[j]
if prev.startswith("-"):
continue # skip Unix flags like --login, -l
if is_win_c and prev.startswith("/") and len(prev) <= 3:
continue # skip Windows flags like /s, /q (not /bin/bash)
prev_base = os.path.basename(prev).lower()
if is_unix_c and prev_base in _SHELLS:
blocked |= _find_blocked_commands(tokens[i + 1])
elif is_win_c and prev_base in _SHELLS_WIN:
blocked |= _find_blocked_commands(tokens[i + 1])
break # stop at first non-flag token
return blocked
def _build_safe_env(workdir: str) -> dict[str, str]:
"""Build a minimal, credential-free environment for sandboxed subprocesses.
Strips HF_TOKEN, WANDB_API_KEY, AWS_*, GH_TOKEN, LD_PRELOAD, DYLD_*, etc.
Preserves the active Python interpreter and virtualenv directories in PATH
so that pip, uv, and packages installed in the Studio runtime remain
accessible.
"""
# Start with the directory containing the running Python interpreter
# so that subprocess calls to 'python', 'pip', etc. resolve to the
# same environment the Studio server is running in.
exe_dir = os.path.dirname(sys.executable)
path_entries = [exe_dir] if exe_dir else []
# If a virtualenv is active, include its bin/Scripts directory.
venv = os.environ.get("VIRTUAL_ENV")
if venv:
venv_bin = os.path.join(venv, "Scripts" if sys.platform == "win32" else "bin")
if venv_bin not in path_entries:
path_entries.append(venv_bin)
if sys.platform == "win32":
sysroot = os.environ.get("SystemRoot", r"C:\Windows")
path_entries.extend([os.path.join(sysroot, "System32"), sysroot])
else:
path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"])
# Deduplicate while preserving order
deduped = list(dict.fromkeys(p for p in path_entries if p))
env = {
"PATH": os.pathsep.join(deduped),
"HOME": workdir,
"TMPDIR": workdir,
"LANG": os.environ.get("LANG", "C.UTF-8"),
"TERM": "dumb",
"PYTHONIOENCODING": "utf-8",
}
if venv:
env["VIRTUAL_ENV"] = venv
# Windows needs SystemRoot for Python/subprocess to work
if sys.platform == "win32":
env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows")
return env
def _sandbox_preexec():
"""Pre-exec hook: drop privilege escalation ability and set resource limits.
On Linux, applies PR_SET_NO_NEW_PRIVS so sudo/su/pkexec fail at the
kernel level. On Linux and macOS, sets RLIMIT_FSIZE.
No-op on Windows (use creationflags instead).
Note: RLIMIT_NPROC is intentionally NOT set because Linux enforces it
per real UID, not per process tree, so it would starve the Studio
server and other sessions sharing the same user account.
All modules and handles are resolved at import time (module level) so
this function does not trigger Python imports in the forked child,
avoiding potential deadlocks in multi-threaded servers.
"""
if _libc is not None:
try:
# PR_SET_NO_NEW_PRIVS = 38, arg2 = 1 (enable)
_libc.prctl(38, 1, 0, 0, 0)
except (OSError, AttributeError):
pass # Not available (container, old kernel, etc.)
if _resource is not None:
try:
# Limit file size to 100MB (prevents disk filling)
_resource.setrlimit(
_resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024)
)
except (ValueError, OSError):
pass
def _get_shell_cmd(command: str) -> list[str]:
"""Return the platform-appropriate shell invocation for a command string."""
if sys.platform == "win32":
return ["cmd", "/c", command]
return ["bash", "-c", command]
# Per-session working directories so each chat thread gets its own sandbox.
# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id.
# Falls back to a shared ~/studio_sandbox/_default for API callers without a
# session_id.
_workdirs: dict[str, str] = {}
@ -51,7 +282,7 @@ def _get_workdir(session_id: str | None = None) -> str:
if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)):
workdir = os.path.join(sandbox_root, "_invalid")
else:
workdir = sandbox_root
workdir = os.path.join(sandbox_root, "_default")
os.makedirs(workdir, exist_ok = True)
_workdirs[key] = workdir
return _workdirs[key]
@ -424,6 +655,7 @@ def _check_signal_escape_patterns(code: str):
signal_tampering = []
exception_catching = []
shell_escapes = []
warnings = []
def _ast_name_matches(node, names):
@ -441,10 +673,84 @@ def _check_signal_escape_patterns(code: str):
return full_name in names
return False
# Dangerous os/subprocess functions that can execute shell commands
_SHELL_EXEC_FUNCS = frozenset(
{
"os.system",
"os.popen",
"os.popen2",
"os.popen3",
"os.popen4",
"os.execl",
"os.execle",
"os.execlp",
"os.execlpe",
"os.execv",
"os.execve",
"os.execvp",
"os.execvpe",
"os.spawnl",
"os.spawnle",
"os.spawnlp",
"os.spawnlpe",
"os.spawnv",
"os.spawnve",
"os.spawnvp",
"os.spawnvpe",
"os.posix_spawn",
"os.posix_spawnp",
"subprocess.run",
"subprocess.call",
"subprocess.check_call",
"subprocess.check_output",
"subprocess.Popen",
"subprocess.getoutput",
"subprocess.getstatusoutput",
}
)
def _extract_string_from_node(node):
"""Extract a plain string value from an AST node, if it is a constant."""
if isinstance(node, ast.Constant) and isinstance(node.value, str):
return node.value
return None
def _extract_strings_from_list(node):
"""Extract string elements from an AST List or Tuple node."""
if isinstance(node, (ast.List, ast.Tuple)):
parts = []
for elt in node.elts:
s = _extract_string_from_node(elt)
if s is not None:
parts.append(s)
return parts
return []
# Keyword argument names that carry command content (as opposed to
# control flags like check=True, text=True, capture_output=True).
_CMD_KWARGS = frozenset({"args", "command", "executable", "path", "file"})
def _check_args_for_blocked(args_nodes):
"""Check if any call arguments contain blocked commands."""
found = set()
for arg in args_nodes:
s = _extract_string_from_node(arg)
if s is not None:
found |= _find_blocked_commands(s)
strs = _extract_strings_from_list(arg)
for s in strs:
found |= _find_blocked_commands(s)
return found
class SignalEscapeVisitor(ast.NodeVisitor):
def __init__(self):
self.imports_signal = False
self.signal_aliases = {"signal"}
self.os_aliases = {"os"}
self.subprocess_aliases = {"subprocess"}
# Maps bare function names to their fully-qualified form
# for from-import tracking (e.g. "system" -> "os.system")
self.shell_exec_aliases: dict[str, str] = {}
self.loop_depth = 0
def visit_Import(self, node):
@ -453,6 +759,10 @@ def _check_signal_escape_patterns(code: str):
self.imports_signal = True
if alias.asname:
self.signal_aliases.add(alias.asname)
elif alias.name == "os":
self.os_aliases.add(alias.asname or "os")
elif alias.name == "subprocess":
self.subprocess_aliases.add(alias.asname or "subprocess")
self.generic_visit(node)
def visit_ImportFrom(self, node):
@ -470,6 +780,16 @@ def _check_signal_escape_patterns(code: str):
"alarm",
):
self.signal_aliases.add(alias.asname or alias.name)
elif node.module in ("os", "subprocess"):
if node.module == "os":
self.os_aliases.add("os")
else:
self.subprocess_aliases.add("subprocess")
# Track from-imports of dangerous functions
for alias in node.names:
fq = f"{node.module}.{alias.name}"
if fq in _SHELL_EXEC_FUNCS:
self.shell_exec_aliases[alias.asname or alias.name] = fq
self.generic_visit(node)
def visit_While(self, node):
@ -534,6 +854,111 @@ def _check_signal_escape_patterns(code: str):
"description": "Modifies signal mask (may block SIGALRM)",
}
)
# --- Shell escape detection ---
# Resolve the fully qualified function name for os.*/subprocess.*
shell_func = None
if isinstance(func, ast.Attribute):
if isinstance(func.value, ast.Name):
if func.value.id in self.os_aliases:
shell_func = f"os.{func.attr}"
elif func.value.id in self.subprocess_aliases:
shell_func = f"subprocess.{func.attr}"
elif isinstance(func, ast.Name):
# Check from-import aliases: from os import system; system(...)
shell_func = self.shell_exec_aliases.get(func.id)
if shell_func and shell_func in _SHELL_EXEC_FUNCS:
# Expand **kwargs dicts to inspect their keys
expanded_kwargs: dict[str, ast.AST] = {}
has_opaque_kwargs = False
for kw in node.keywords:
if kw.arg is not None:
expanded_kwargs[kw.arg] = kw.value
elif isinstance(kw.value, ast.Dict):
for k, v in zip(kw.value.keys, kw.value.values):
key = _extract_string_from_node(k) if k else None
if key is not None:
expanded_kwargs[key] = v
else:
has_opaque_kwargs = True
cmd_kw_values = [
v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS
]
all_call_args = list(node.args) + cmd_kw_values
blocked_in_args = _check_args_for_blocked(all_call_args)
if has_opaque_kwargs:
# Can't inspect dynamic **kwargs -- flag as unsafe
shell_escapes.append(
{
"type": "shell_escape_dynamic",
"line": node.lineno,
"description": (
f"{shell_func}() called with dynamic **kwargs"
),
}
)
elif blocked_in_args:
shell_escapes.append(
{
"type": "shell_escape",
"line": node.lineno,
"description": (
f"{shell_func}() invokes blocked command(s): "
f"{', '.join(sorted(blocked_in_args))}"
),
}
)
else:
# Only flag dynamic args for functions that interpret
# strings as shell commands, or when shell= might be
# enabled. Treat any non-literal-False shell= value
# as potentially True (conservative).
_STRING_SHELL_FUNCS = frozenset(
{
"os.system",
"os.popen",
"os.popen2",
"os.popen3",
"os.popen4",
"subprocess.getoutput",
"subprocess.getstatusoutput",
}
)
shell_node = expanded_kwargs.get("shell")
shell_safe = shell_node is None or (
isinstance(shell_node, ast.Constant)
and shell_node.value is False
)
if shell_func in _STRING_SHELL_FUNCS or not shell_safe:
def _is_safe_literal(n):
if _extract_string_from_node(n) is not None:
return True
if isinstance(n, (ast.List, ast.Tuple)):
return all(
_extract_string_from_node(e) is not None
for e in n.elts
)
return False
has_non_literal = any(
not _is_safe_literal(a) for a in all_call_args
)
if has_non_literal:
shell_escapes.append(
{
"type": "shell_escape_dynamic",
"line": node.lineno,
"description": (
f"{shell_func}() called with non-literal "
f"shell command (potential shell escape)"
),
}
)
self.generic_visit(node)
def visit_ExceptHandler(self, node):
@ -549,7 +974,12 @@ def _check_signal_escape_patterns(code: str):
}
)
elif isinstance(node.type, ast.Name):
if node.type.id in ("TimeoutError", "BaseException", "Exception"):
# Only flag BaseException and TimeoutError, NOT Exception.
# except Exception does not catch SystemExit or
# KeyboardInterrupt, so it cannot suppress timeout
# enforcement. Flagging Exception causes false positives
# on normal error-handling patterns.
if node.type.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
"type": f"catches_{node.type.id}_in_loop",
@ -560,7 +990,7 @@ def _check_signal_escape_patterns(code: str):
elif isinstance(node.type, ast.Tuple):
for elt in node.type.elts:
if isinstance(elt, ast.Name):
if elt.id in ("TimeoutError", "BaseException", "Exception"):
if elt.id in ("TimeoutError", "BaseException"):
exception_catching.append(
{
"type": f"catches_{elt.id}_in_loop",
@ -576,10 +1006,15 @@ def _check_signal_escape_patterns(code: str):
if visitor.imports_signal and not signal_tampering:
warnings.append("Code imports 'signal' module - review manually for safety")
is_safe = len(signal_tampering) == 0 and len(exception_catching) == 0
is_safe = (
len(signal_tampering) == 0
and len(exception_catching) == 0
and len(shell_escapes) == 0
)
return is_safe, {
"signal_tampering": signal_tampering,
"exception_catching": exception_catching,
"shell_escapes": shell_escapes,
"warnings": warnings,
}
@ -591,13 +1026,27 @@ def _check_code_safety(code: str) -> str | None:
"""
safe, info = _check_signal_escape_patterns(code)
if not safe:
# SyntaxError from ast.parse -- let these through so the subprocess
# produces a normal Python traceback instead of a misleading
# "unsafe code detected" message.
if info.get("error"):
return None
reasons = [
item.get("description", "") for item in info.get("signal_tampering", [])
]
return (
f"Error: unsafe code detected ({'; '.join(reasons)}). "
f"Please remove signal manipulation from your code."
)
shell_reasons = [
item.get("description", "") for item in info.get("shell_escapes", [])
]
exception_reasons = [
item.get("description", "") for item in info.get("exception_catching", [])
]
all_reasons = [r for r in reasons + shell_reasons + exception_reasons if r]
if all_reasons:
return (
f"Error: unsafe code detected ({'; '.join(all_reasons)}). "
f"Please remove unsafe patterns from your code."
)
return None
@ -634,6 +1083,17 @@ def _python_exec(
tmp_path = None
workdir = _get_workdir(session_id)
# Snapshot image mtimes so we detect both new and overwritten files.
_before: dict[str, int] = {}
if os.path.isdir(workdir):
for _name in os.listdir(workdir):
if os.path.splitext(_name)[1].lower() in _IMAGE_EXTS:
_p = os.path.join(workdir, _name)
if os.path.isfile(_p):
try:
_before[_name] = os.stat(_p).st_mtime_ns
except OSError:
pass
try:
fd, tmp_path = tempfile.mkstemp(
suffix = ".py", prefix = "studio_exec_", dir = workdir
@ -641,13 +1101,20 @@ def _python_exec(
with os.fdopen(fd, "w") as f:
f.write(code)
proc = subprocess.Popen(
[sys.executable, tmp_path],
safe_env = _build_safe_env(workdir)
popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
cwd = workdir,
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _sandbox_preexec
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
proc = subprocess.Popen([sys.executable, tmp_path], **popen_kwargs)
# Spawn cancel watcher if we have a cancel event
if cancel_event is not None:
@ -669,7 +1136,29 @@ def _python_exec(
result = output or ""
if proc.returncode != 0:
result = f"Exit code {proc.returncode}:\n{result}"
return _truncate(result) if result.strip() else "(no output)"
result = _truncate(result) if result.strip() else "(no output)"
# Detect new or overwritten image files and append sentinel for frontend
if session_id and os.path.isdir(workdir):
new_images = []
for _name in os.listdir(workdir):
if os.path.splitext(_name)[1].lower() not in _IMAGE_EXTS:
continue
_p = os.path.join(workdir, _name)
if not os.path.isfile(_p):
continue
try:
_mtime = os.stat(_p).st_mtime_ns
except OSError:
continue
if _name not in _before or _mtime != _before[_name]:
new_images.append(_name)
if new_images:
import json as _json
result += f"\n__IMAGES__:{_json.dumps(sorted(new_images))}"
return result
except Exception as e:
return f"Execution error: {e}"
@ -691,21 +1180,27 @@ def _bash_exec(
if not command or not command.strip():
return "No command provided."
# Block dangerous commands
tokens = set(command.lower().split())
blocked = tokens & _BASH_BLOCKED_WORDS
# Block dangerous commands (shlex + regex based)
blocked = _find_blocked_commands(command)
if blocked:
return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
try:
workdir = _get_workdir(session_id)
proc = subprocess.Popen(
["bash", "-c", command],
safe_env = _build_safe_env(workdir)
popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
cwd = workdir,
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _sandbox_preexec
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
proc = subprocess.Popen(_get_shell_cmd(command), **popen_kwargs)
if cancel_event is not None:
watcher = threading.Thread(

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

@ -34,37 +34,15 @@ from utils.hardware import apply_gpu_ids
def _activate_transformers_version(model_name: str) -> None:
"""Activate the correct transformers version BEFORE any ML imports.
If the model needs transformers 5.x, prepend the pre-installed .venv_t5/
directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/).
"""
"""Activate the correct transformers version BEFORE any ML imports."""
# Ensure backend is on path for utils imports
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from utils.transformers_version import (
needs_transformers_5,
_resolve_base_model,
_ensure_venv_t5_exists,
_VENV_T5_DIR,
)
from utils.transformers_version import activate_transformers_for_subprocess
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
if not _ensure_venv_t5_exists():
raise RuntimeError(
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
)
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
# Propagate to child subprocesses (e.g. GGUF converter)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
else:
logger.info("Using default transformers (4.57.x) for %s", model_name)
activate_transformers_for_subprocess(model_name)
def _decode_image(image_base64: str):
@ -145,6 +123,8 @@ def _get_hf_download_state(
blobs_dirs: list[Path] = []
if model_names:
from utils.paths import resolve_cached_repo_id_case
for name in model_names:
if not name:
continue
@ -154,6 +134,7 @@ def _get_hf_download_state(
# relative paths, and Windows paths.
if name.startswith(("/", ".", "~")) or "\\" in name:
continue
name = resolve_cached_repo_id_case(name)
# HF cache dir format: models--org--name (slashes -> --)
cache_dir_name = "models--" + name.replace("/", "--")
blobs_dir = cache / cache_dir_name / "blobs"
@ -306,19 +287,21 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
except Exception as e:
logger.warning("Could not read adapter_config.json: %s", e)
# Auto-enable trust_remote_code for unsloth/* transformers 5.x models
# (matches the training worker logic in core/training/worker.py)
# Auto-enable trust_remote_code for NemotronH/Nano models only.
# NemotronH has config parsing bugs requiring trust_remote_code=True.
# Other transformers 5.x models are native and do NOT need it.
# NOTE: Must NOT match Llama-Nemotron (standard Llama architecture).
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code:
from utils.transformers_version import needs_transformers_5
model_name = config["model_name"]
if needs_transformers_5(model_name) and model_name.lower().startswith(
"unsloth/"
_mn_lower = model_name.lower()
if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/")
):
trust_remote_code = True
logger.info(
"Auto-enabled trust_remote_code for unsloth/* transformers 5.x model: %s",
"Auto-enabled trust_remote_code for Nemotron model: %s",
model_name,
)

75
studio/backend/core/training/resume.py Normal file
View file

@ -0,0 +1,75 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Helpers for validating resumable training outputs."""
from pathlib import Path
from typing import Optional
from utils.paths import outputs_root, resolve_output_dir
def _is_under_outputs(path: Path) -> bool:
resolved = path.resolve(strict = False)
root = outputs_root().resolve(strict = False)
try:
resolved.relative_to(root)
return True
except ValueError:
return False
def has_resume_state(path_value: Optional[str]) -> bool:
if not path_value:
return False
return get_resume_checkpoint_path(path_value) is not None
def _checkpoint_step(path: Path) -> int:
try:
return int(path.name.removeprefix("checkpoint-"))
except ValueError:
return -1
def get_resume_checkpoint_path(path_value: str) -> Optional[str]:
path = resolve_output_dir(path_value)
if not _is_under_outputs(path) or not path.is_dir():
return None
if (path / "trainer_state.json").is_file():
return str(path)
checkpoints = [
child
for child in path.glob("checkpoint-*")
if child.is_dir() and (child / "trainer_state.json").is_file()
]
if not checkpoints:
return None
return str(max(checkpoints, key = _checkpoint_step))
def normalize_resume_output_dir(path_value: str) -> str:
path = resolve_output_dir(path_value)
if not _is_under_outputs(path):
raise ValueError("Resume checkpoint must be inside Studio outputs.")
return str(path)
def can_resume_run(run: dict) -> bool:
if run.get("resumed_later"):
return False
final_step = run.get("final_step")
total_steps = run.get("total_steps")
has_remaining_steps = (
not isinstance(final_step, int)
or not isinstance(total_steps, int)
or total_steps <= 0
or final_step < total_steps
)
return (
run.get("status") == "stopped"
and has_remaining_steps
and has_resume_state(run.get("output_dir"))
)

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

@ -49,6 +49,7 @@ from unsloth.chat_templates import get_chat_template
import json
import threading
import math
import subprocess
import structlog
from loggers import get_logger
import time
@ -69,6 +70,11 @@ from utils.paths import (
)
from trl import SFTTrainer, SFTConfig
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
logger = get_logger(__name__)
@ -190,7 +196,11 @@ class UnslothTrainer:
self._cuda_audio_used = False
# --- Detect VLM ---
vision = is_vision_model(model_name) if not self.is_audio else False
vision = (
is_vision_model(model_name, hf_token = hf_token)
if not self.is_audio
else False
)
self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image
logger.info(
@ -367,6 +377,7 @@ class UnslothTrainer:
def _finalize_training(self, output_dir, label = ""):
"""Save model after training and update progress. Used by all training branches."""
if self.should_stop and self.save_on_stop:
self.trainer._save_checkpoint(self.trainer.model, trial = None)
self.trainer.save_model()
self.tokenizer.save_pretrained(output_dir)
self._patch_adapter_config(output_dir)
@ -558,7 +569,11 @@ class UnslothTrainer:
self._cuda_audio_used = False
# VLM: vision model with image dataset (mutually exclusive with audio paths)
vision = is_vision_model(model_name) if not self.is_audio else False
vision = (
is_vision_model(model_name, hf_token = hf_token)
if not self.is_audio
else False
)
self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image
self.model_name = model_name
self.max_seq_length = max_seq_length
@ -1757,6 +1772,8 @@ class UnslothTrainer:
spark_code_dir,
],
check = True,
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
)
if spark_code_dir not in sys.path:
@ -1974,8 +1991,6 @@ class UnslothTrainer:
device = "cuda" if torch.cuda.is_available() else "cpu"
# Clone OuteTTS repo (same as audio_codecs._load_dac)
import subprocess
base_dir = os.path.dirname(os.path.abspath(__file__))
outetts_code_dir = os.path.join(base_dir, "inference", "OuteTTS")
outetts_pkg = os.path.join(outetts_code_dir, "outetts")
@ -1992,6 +2007,8 @@ class UnslothTrainer:
outetts_code_dir,
],
check = True,
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
)
for fpath in [
os.path.join(outetts_pkg, "models", "gguf_model.py"),
@ -2815,7 +2832,9 @@ class UnslothTrainer:
total_steps = total, status_message = "Starting CSM training..."
)
logger.info(f"CSM training config: {config}\n")
self.trainer.train()
self.trainer.train(
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
)
self._finalize_training(output_dir, "CSM")
return
@ -2854,7 +2873,9 @@ class UnslothTrainer:
total_steps = total, status_message = "Starting SNAC training..."
)
logger.info(f"SNAC training config: {config}\n")
self.trainer.train()
self.trainer.train(
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
)
self._finalize_training(output_dir, "SNAC")
return
@ -2900,7 +2921,9 @@ class UnslothTrainer:
total_steps = total, status_message = "Starting Whisper training..."
)
logger.info(f"Whisper training config: {config}\n")
self.trainer.train()
self.trainer.train(
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
)
self._finalize_training(output_dir, "Whisper")
return
@ -3395,7 +3418,9 @@ class UnslothTrainer:
# ========== START TRAINING ==========
self._update_progress(status_message = "Starting training...")
logger.info("Starting training...\n")
self.trainer.train()
self.trainer.train(
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
)
# ========== SAVE MODEL ==========
self._finalize_training(output_dir)

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

@ -29,6 +29,10 @@ from typing import Optional, Tuple, Any
import matplotlib.pyplot as plt
from utils.hardware import prepare_gpu_selection
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
logger = get_logger(__name__)
@ -185,6 +189,7 @@ class TrainingBackend:
"wandb_project": kwargs.get("wandb_project", "unsloth-training"),
"enable_tensorboard": kwargs.get("enable_tensorboard", False),
"tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
"resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
"trust_remote_code": kwargs.get("trust_remote_code", False),
"gpu_ids": kwargs.get("gpu_ids"),
}
@ -212,20 +217,22 @@ class TrainingBackend:
from .worker import run_training_process
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
proc = _CTX.Process(
target = run_training_process,
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
try:
proc.start()
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
proc.start()
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
return False

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

@ -16,26 +16,40 @@ from __future__ import annotations
import structlog
from loggers import get_logger
import os
import platform
import shutil
import sys
import time
import traceback
import json
import subprocess as _sp
from pathlib import Path
from typing import Any
import urllib.error
import urllib.request
from typing import Any, Callable
logger = get_logger(__name__)
from utils.hardware import apply_gpu_ids
from utils.wheel_utils import (
direct_wheel_url,
flash_attn_wheel_url,
install_wheel,
probe_torch_wheel_env,
url_exists,
)
def _output_dir_from_resume_checkpoint(
resume_from_checkpoint: str | None,
) -> str | None:
if not resume_from_checkpoint:
return None
path = Path(resume_from_checkpoint)
return str(path.parent if path.name.startswith("checkpoint-") else path)
_CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4"
_CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1"
_MAMBA_SSM_RELEASE_TAG = "v2.3.1"
_MAMBA_SSM_PACKAGE_VERSION = "2.3.1"
_FLASH_ATTN_RUNTIME_MIN_SEQ_LEN = 32768
_FLASH_ATTN_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL"
def _model_wants_causal_conv1d(model_name: str) -> bool:
@ -45,6 +59,8 @@ def _model_wants_causal_conv1d(model_name: str) -> bool:
for key in (
"qwen3.5",
"qwen3_5",
"qwen3.6",
"qwen3_6",
"qwen3-next",
"qwen3_next",
"nemotron_h",
@ -59,206 +75,186 @@ def _model_wants_causal_conv1d(model_name: str) -> bool:
)
def _causal_conv1d_platform_tag() -> str | None:
machine = platform.machine().lower()
if sys.platform.startswith("linux"):
if machine in {"x86_64", "amd64"}:
return "linux_x86_64"
if machine in {"aarch64", "arm64"}:
return "linux_aarch64"
return None
# No prebuilt wheels published for macOS or Windows
return None
def _probe_causal_conv1d_env() -> dict[str, str] | None:
try:
probe = _sp.run(
[
sys.executable,
"-c",
(
"import json, sys, re, torch; "
"parts = torch.__version__.split('+', 1)[0].split('.')[:2]; "
"minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; "
"torch_mm = parts[0] + '.' + minor; "
"print(json.dumps({"
"'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', "
"'torch_mm': torch_mm, "
"'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', "
"'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()"
"}))"
),
],
stdout = _sp.PIPE,
stderr = _sp.PIPE,
text = True,
timeout = 30,
)
except _sp.TimeoutExpired:
logger.warning("Torch environment probe timed out after 30s")
return None
if probe.returncode != 0:
logger.warning(
"Failed to probe torch environment for causal-conv1d wheel:\n%s",
probe.stdout,
)
return None
try:
return json.loads(probe.stdout.strip())
except json.JSONDecodeError:
logger.warning(
"Failed to parse torch environment probe output: %s", probe.stdout
)
return None
def _direct_wheel_url(
*,
filename_prefix: str,
package_version: str,
release_tag: str,
release_base_url: str,
env: dict[str, str] | None = None,
) -> str | None:
env = env or _probe_causal_conv1d_env()
platform_tag = _causal_conv1d_platform_tag()
if env is None or platform_tag is None or not env.get("cuda_major"):
return None
filename = (
f"{filename_prefix}-{package_version}"
f"+cu{env['cuda_major']}torch{env['torch_mm']}"
f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}-{platform_tag}.whl"
)
return f"{release_base_url}/{release_tag}/{filename}"
def _url_exists(url: str) -> bool:
try:
request = urllib.request.Request(url, method = "HEAD")
with urllib.request.urlopen(request, timeout = 10):
return True
except urllib.error.HTTPError as exc:
if exc.code == 404:
return False
logger.warning("Unexpected HTTP error while probing %s: %s", url, exc)
return False
except Exception as exc:
logger.warning("Failed to probe %s: %s", url, exc)
return False
def _install_package_wheel_first(
*,
event_queue: Any,
import_name: str,
display_name: str,
pypi_name: str,
pypi_version: str,
filename_prefix: str,
release_tag: str,
release_base_url: str,
) -> None:
pypi_version: str | None = None,
filename_prefix: str | None = None,
release_tag: str | None = None,
release_base_url: str | None = None,
wheel_url_builder: Callable[[dict[str, str] | None], str | None] | None = None,
pypi_spec: str | None = None,
pypi_status_message: str | None = None,
) -> bool:
try:
__import__(import_name)
logger.info("%s already installed", display_name)
return
return True
except ImportError:
pass
env = _probe_causal_conv1d_env()
wheel_url = _direct_wheel_url(
filename_prefix = filename_prefix,
package_version = pypi_version,
release_tag = release_tag,
release_base_url = release_base_url,
env = env,
)
env = probe_torch_wheel_env(timeout = 30)
if wheel_url_builder is not None:
wheel_url = wheel_url_builder(env)
else:
wheel_url = direct_wheel_url(
filename_prefix = filename_prefix,
package_version = pypi_version,
release_tag = release_tag,
release_base_url = release_base_url,
env = env,
)
if wheel_url is None:
logger.info("No compatible %s wheel candidate", display_name)
else:
if _url_exists(wheel_url):
_send_status(event_queue, f"Installing prebuilt {display_name} wheel...")
installed = False
# Try uv first if available, then fall back to pip
if shutil.which("uv"):
uv_cmd = [
"uv",
"pip",
"install",
"--python",
sys.executable,
"--no-deps",
wheel_url,
]
result = _sp.run(
uv_cmd,
stdout = _sp.PIPE,
stderr = _sp.STDOUT,
text = True,
)
if result.returncode == 0:
installed = True
else:
logger.warning(
"uv failed to install %s wheel:\n%s",
display_name,
result.stdout,
)
if not installed:
pip_cmd = [
sys.executable,
"-m",
"pip",
"install",
"--no-deps",
wheel_url,
]
result = _sp.run(
pip_cmd,
stdout = _sp.PIPE,
stderr = _sp.STDOUT,
text = True,
)
if result.returncode == 0:
installed = True
else:
logger.warning(
"pip failed to install %s wheel:\n%s",
display_name,
result.stdout,
)
if installed:
elif url_exists(wheel_url):
_send_status(event_queue, f"Installing prebuilt {display_name} wheel...")
for installer, result in install_wheel(
wheel_url,
python_executable = sys.executable,
use_uv = bool(shutil.which("uv")),
run = _sp.run,
):
if result.returncode == 0:
logger.info("Installed prebuilt %s wheel successfully", display_name)
return
return True
logger.warning(
"%s failed to install %s wheel:\n%s",
installer,
display_name,
result.stdout,
)
else:
logger.info("No published %s wheel found: %s", display_name, wheel_url)
is_hip = env and env.get("hip_version")
if is_hip and not shutil.which("hipcc"):
logger.error(
"%s requires hipcc for source compilation on ROCm. "
"Install the ROCm HIP SDK: https://rocm.docs.amd.com",
display_name,
)
_send_status(
event_queue,
f"{display_name}: hipcc not found (ROCm HIP SDK required)",
)
return False
if pypi_spec is None:
pypi_spec = f"{pypi_name}=={pypi_version}"
if pypi_status_message is None:
if is_hip:
pypi_status_message = (
f"Compiling {display_name} from source for ROCm "
"(this may take several minutes)..."
)
else:
logger.info("No published %s wheel found: %s", display_name, wheel_url)
pypi_status_message = f"Installing {display_name} from PyPI..."
_send_status(event_queue, pypi_status_message)
# Prefer uv for faster dependency resolution when available
plain_pypi_install = pypi_version is None
if plain_pypi_install:
if shutil.which("uv"):
pypi_cmd = [
"uv",
"pip",
"install",
"--python",
sys.executable,
pypi_spec,
]
else:
pypi_cmd = [sys.executable, "-m", "pip", "install", pypi_spec]
else:
if shutil.which("uv"):
pypi_cmd = [
"uv",
"pip",
"install",
"--python",
sys.executable,
"--no-build-isolation",
"--no-deps",
]
# Avoid stale cache artifacts from partial HIP source builds
if is_hip:
pypi_cmd.append("--no-cache")
pypi_cmd.append(pypi_spec)
else:
pypi_cmd = [
sys.executable,
"-m",
"pip",
"install",
"--no-build-isolation",
"--no-deps",
"--no-cache-dir",
pypi_spec,
]
# Source compilation on ROCm can take 10-30 minutes; use a generous
# timeout. Non-HIP installs preserve the pre-existing "no timeout"
# behaviour so unrelated slow installs (e.g. causal-conv1d source
# build on Linux aarch64 or unsupported torch/CUDA combinations)
# are not aborted at 5 minutes by this PR.
_run_kwargs: dict[str, Any] = {
"stdout": _sp.PIPE,
"stderr": _sp.STDOUT,
"text": True,
}
if is_hip:
_run_kwargs["timeout"] = 1800
try:
result = _sp.run(pypi_cmd, **_run_kwargs)
except _sp.TimeoutExpired:
logger.error(
"%s installation timed out after %ds",
display_name,
_run_kwargs.get("timeout"),
)
_send_status(
event_queue,
f"{display_name} installation timed out after "
f"{_run_kwargs.get('timeout')}s",
)
return False
_send_status(event_queue, f"Installing {display_name} from PyPI...")
pypi_cmd = [
sys.executable,
"-m",
"pip",
"install",
"--no-build-isolation",
"--no-deps",
"--no-cache-dir",
f"{pypi_name}=={pypi_version}",
]
result = _sp.run(
pypi_cmd,
stdout = _sp.PIPE,
stderr = _sp.STDOUT,
text = True,
)
if result.returncode != 0:
logger.error("Failed to install %s from PyPI:\n%s", display_name, result.stdout)
return
if is_hip:
# Surface a clear error for ROCm source build failures
error_lines = (result.stdout or "").strip().splitlines()
snippet = "\n".join(error_lines[-5:]) if error_lines else "(no output)"
logger.error(
"Failed to compile %s for ROCm:\n%s",
display_name,
result.stdout,
)
_send_status(
event_queue,
f"Failed to compile {display_name} for ROCm. "
"Check that hipcc and ROCm development headers are installed.\n"
f"{snippet}",
)
else:
logger.error(
"Failed to install %s from PyPI:\n%s",
display_name,
result.stdout,
)
return False
logger.info("Installed %s from PyPI", display_name)
if is_hip:
logger.info("Compiled and installed %s from source for ROCm", display_name)
else:
logger.info("Installed %s from PyPI", display_name)
return True
def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:
@ -305,38 +301,41 @@ def _ensure_mamba_ssm(event_queue: Any, model_name: str) -> None:
)
def _activate_transformers_version(model_name: str) -> None:
"""Activate the correct transformers version BEFORE any ML imports.
def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool:
if os.getenv(_FLASH_ATTN_SKIP_ENV) == "1":
return False
if max_seq_length < _FLASH_ATTN_RUNTIME_MIN_SEQ_LEN:
return False
return sys.platform.startswith("linux")
If the model needs transformers 5.x, prepend the pre-installed .venv_t5/
directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/).
"""
def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None:
if not _should_try_runtime_flash_attn_install(max_seq_length):
return
installed = _install_package_wheel_first(
event_queue = event_queue,
import_name = "flash_attn",
display_name = "flash-attn",
pypi_name = "flash-attn",
wheel_url_builder = flash_attn_wheel_url,
pypi_spec = "flash-attn",
pypi_status_message = "Installing flash-attn from PyPI for long-context training...",
)
if not installed:
_send_status(event_queue, "Continuing without flash-attn")
def _activate_transformers_version(model_name: str) -> None:
"""Activate the correct transformers version BEFORE any ML imports."""
# Ensure backend is on path for utils imports
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from utils.transformers_version import (
needs_transformers_5,
_resolve_base_model,
_ensure_venv_t5_exists,
_VENV_T5_DIR,
)
from utils.transformers_version import activate_transformers_for_subprocess
resolved = _resolve_base_model(model_name)
if needs_transformers_5(resolved):
if not _ensure_venv_t5_exists():
raise RuntimeError(
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
)
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
# Propagate to child subprocesses (e.g. GGUF converter)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
else:
logger.info("Using default transformers (4.57.x) for %s", model_name)
activate_transformers_for_subprocess(model_name)
def run_training_process(
@ -386,20 +385,22 @@ def run_training_process(
)
return
# ── 1a. Auto-enable trust_remote_code for unsloth/* transformers 5.x models ──
# Some newer architectures (e.g. NemotronH) have config parsing bugs in
# transformers that require trust_remote_code=True as a workaround.
# Only auto-enable for unsloth/* prefixed models (trusted source).
from utils.transformers_version import needs_transformers_5
# ── 1a. Auto-enable trust_remote_code for NemotronH/Nano models ──
# NemotronH has config parsing bugs in transformers that require
# trust_remote_code=True as a workaround. Other transformers 5.x models
# (Qwen3.5, Gemma 4, etc.) are native and do NOT need it — enabling it
# bypasses the compiler (disabling fused CE).
# NOTE: Must NOT match Llama-Nemotron (standard Llama architecture).
_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
_lowered = model_name.lower()
if (
needs_transformers_5(model_name)
and model_name.lower().startswith("unsloth/")
any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS)
and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/"))
and not config.get("trust_remote_code", False)
):
config["trust_remote_code"] = True
logger.info(
"Auto-enabled trust_remote_code for unsloth/* transformers 5.x model: %s",
"Auto-enabled trust_remote_code for Nemotron model: %s",
model_name,
)
@ -407,6 +408,10 @@ def run_training_process(
try:
_ensure_causal_conv1d_fast_path(event_queue, model_name)
_ensure_mamba_ssm(event_queue, model_name)
_ensure_flash_attn_for_long_context(
event_queue,
int(config.get("max_seq_length", 2048)),
)
except Exception as exc:
event_queue.put(
{
@ -761,7 +766,10 @@ def run_training_process(
return
# Generate output dir
output_dir = config.get("output_dir")
resume_from_checkpoint = config.get("resume_from_checkpoint")
output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint(
resume_from_checkpoint
)
if not output_dir:
output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}"
output_dir = str(resolve_output_dir(output_dir))
@ -809,6 +817,7 @@ def run_training_process(
max_seq_length = config.get("max_seq_length", 2048),
optim = config.get("optim", "adamw_8bit"),
lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
resume_from_checkpoint = resume_from_checkpoint,
)
_tqdm_stop.set()
@ -825,10 +834,13 @@ def run_training_process(
}
)
else:
saved_output_dir = (
None if trainer.should_stop and not trainer.save_on_stop else output_dir
)
event_queue.put(
{
"type": "complete",
"output_dir": output_dir,
"output_dir": saved_output_dir,
"status_message": progress.status_message or "Training completed",
"ts": time.time(),
}
@ -1113,11 +1125,15 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
)
return
output_dir = config.get("output_dir")
resume_from_checkpoint = config.get("resume_from_checkpoint")
output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint(
resume_from_checkpoint
)
if not output_dir:
output_dir = str(
resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}")
)
output_dir = str(resolve_output_dir(output_dir))
num_epochs = config.get("num_epochs", 2)
batch_size = config.get("batch_size", 256)
@ -1225,7 +1241,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
callbacks = [_EmbeddingProgressCallback()],
)
trainer.train()
trainer.train(resume_from_checkpoint = resume_from_checkpoint)
except Exception as e:
event_queue.put(
{
@ -1251,6 +1267,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
_send_status(event_queue, "Saving model...")
try:
if _should_stop and _save_on_stop:
trainer._save_checkpoint(trainer.model, trial = None)
model.save_pretrained(output_dir)
model.tokenizer.save_pretrained(output_dir)
logger.info("Embedding model saved to %s", output_dir)

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

@ -22,6 +22,8 @@ from typing import Optional
import structlog
from loggers.handlers import filter_sensitive_data
class LogConfig:
"""Structured logging configuration for the application.
@ -44,12 +46,22 @@ class LogConfig:
# Fallback to INFO if an invalid level is provided
log_level = getattr(logging, log_level_name, logging.INFO)
if sys.platform == "win32":
for stream in (sys.stdout, sys.stderr):
if hasattr(stream, "reconfigure"):
try:
stream.reconfigure(encoding = "utf-8", errors = "replace")
except Exception:
pass
structlog.configure(
processors = [
# Reorder processors to control field order
structlog.processors.TimeStamper(fmt = "iso"), # timestamp first
structlog.processors.add_log_level, # level second
structlog.contextvars.merge_contextvars,
structlog.processors.format_exc_info,
filter_sensitive_data,
# Custom processor to flatten the extra field
lambda logger, method_name, event_dict: {
"timestamp": event_dict.get("timestamp"),

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

@ -15,6 +15,7 @@ Key Components:
- get_logger: Factory function for structured loggers
"""
import re
import time
from typing import Callable
@ -22,7 +23,12 @@ import structlog
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from utils.native_path_leases import redact_native_paths
logger = structlog.get_logger(__name__)
_NATIVE_PATH_LEASE_RE = re.compile(
r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"
)
class LoggingMiddleware(BaseHTTPMiddleware):
@ -75,6 +81,12 @@ def filter_sensitive_data(logger, method_name, event_dict):
"""Structlog processor to filter out base64 data from logs."""
def filter_value(value):
if isinstance(value, str):
try:
value = redact_native_paths(value)
except Exception:
pass
value = _NATIVE_PATH_LEASE_RE.sub(r"\1<redacted native path lease>", value)
if (
isinstance(value, str)
and len(value) > 100
@ -83,12 +95,22 @@ def filter_sensitive_data(logger, method_name, event_dict):
# Likely base64 data, truncate it
return value[:20] + "..."
elif isinstance(value, dict):
return {k: filter_value(v) for k, v in value.items()}
return {
k: "<redacted native path lease>"
if str(k).replace("_", "").lower() == "nativepathlease"
else filter_value(v)
for k, v in value.items()
}
elif isinstance(value, list):
return [filter_value(item) for item in value]
return value
return {k: filter_value(v) for k, v in event_dict.items()}
return {
k: "<redacted native path lease>"
if str(k).replace("_", "").lower() == "nativepathlease"
else filter_value(v)
for k, v in event_dict.items()
}
def get_logger(name: str) -> structlog.BoundLogger:

58
studio/backend/main.py
View file

@ -27,6 +27,7 @@ import mimetypes
import shutil
import warnings
from contextlib import asynccontextmanager
from importlib.metadata import PackageNotFoundError, version as package_version
# Fix broken Windows registry MIME types. Some Windows installs map .js to
# "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes
@ -61,6 +62,7 @@ from routes import (
datasets_router,
export_router,
inference_router,
inference_studio_router,
models_router,
providers_router,
training_history_router,
@ -77,6 +79,28 @@ from utils.hardware import (
import utils.hardware.hardware as _hw_module
from utils.cache_cleanup import clear_unsloth_compiled_cache
from utils.native_path_leases import native_path_leases_supported
def get_unsloth_version() -> str:
try:
return package_version("unsloth")
except PackageNotFoundError:
pass
version_file = (
_Path(__file__).resolve().parents[2] / "unsloth" / "models" / "_utils.py"
)
try:
for line in version_file.read_text(encoding = "utf-8").splitlines():
if line.startswith("__version__ = "):
return line.split("=", 1)[1].strip().strip('"').strip("'")
except OSError:
pass
return "dev"
UNSLOTH_VERSION = get_unsloth_version()
@asynccontextmanager
@ -146,7 +170,7 @@ async def lifespan(app: FastAPI):
# Create FastAPI app
app = FastAPI(
title = "Unsloth UI Backend",
version = "1.0.0",
version = UNSLOTH_VERSION,
description = "Backend API for Unsloth UI - Training and Model Management",
lifespan = lifespan,
)
@ -163,9 +187,24 @@ logger = LogConfig.setup_logging(
app.add_middleware(LoggingMiddleware)
# CORS middleware
_api_only = os.environ.get("UNSLOTH_API_ONLY") == "1"
_cors_origins = ["*"]
if _api_only:
_cors_origins = [
"tauri://localhost", # Linux/macOS Tauri webview
"http://tauri.localhost", # Windows Tauri webview
"http://localhost", # dev fallback
"http://localhost:5173", # Tauri dev/Vite
"http://127.0.0.1:5173", # Tauri dev/Vite fallback
]
_cors_origin_regex = None
else:
_cors_origin_regex = None
app.add_middleware(
CORSMiddleware,
allow_origins = ["*"], # In production, specify allowed origins
allow_origins = _cors_origins,
allow_origin_regex = _cors_origin_regex,
allow_credentials = True,
allow_methods = ["*"],
allow_headers = ["*"],
@ -178,6 +217,9 @@ app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"])
app.include_router(training_router, prefix = "/api/train", tags = ["training"])
app.include_router(models_router, prefix = "/api/models", tags = ["models"])
app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
# Studio-only inference endpoints (cancel, etc.) are intentionally NOT
# exposed on the /v1 OpenAI-compat prefix below.
app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"])
# OpenAI-compatible endpoints: mount the same inference router at /v1
# so external tools (Open WebUI, SillyTavern, etc.) can use the
@ -205,8 +247,12 @@ async def health_check():
"status": "healthy",
"timestamp": datetime.now().isoformat(),
"service": "Unsloth UI Backend",
"version": UNSLOTH_VERSION,
"device_type": device_type,
"chat_only": _hw_module.CHAT_ONLY,
"desktop_protocol_version": 1,
"supports_desktop_auth": True,
"native_path_leases_supported": native_path_leases_supported(),
}
@ -244,6 +290,7 @@ async def get_system_info():
import platform
import psutil
from utils.hardware import get_device
from utils.hardware.hardware import _backend_label
visibility_info = get_backend_visible_gpu_info()
gpu_info = {
@ -257,7 +304,10 @@ async def get_system_info():
return {
"platform": platform.platform(),
"python_version": platform.python_version(),
"device_backend": get_device().value,
# Use the centralized _backend_label helper so the /api/system
# endpoint reports "rocm" on AMD hosts instead of "cuda", matching
# the /api/hardware and /api/gpu-visibility endpoints.
"device_backend": _backend_label(get_device()),
"cpu_count": psutil.cpu_count(),
"memory": {
"total_gb": round(memory.total / 1e9, 2),
@ -356,7 +406,7 @@ def setup_frontend(app: FastAPI, build_path: Path):
@app.get("/{full_path:path}")
async def serve_frontend(full_path: str):
if full_path.startswith("api"):
if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")):
return {"error": "API endpoint not found"}
file_path = (build_path / full_path).resolve()

49
studio/backend/models/auth.py
View file

@ -5,6 +5,8 @@
Pydantic schemas for Authentication API
"""
from typing import Optional
from pydantic import BaseModel, Field
@ -15,6 +17,12 @@ class AuthLoginRequest(BaseModel):
password: str = Field(..., description = "Password")
class DesktopLoginRequest(BaseModel):
"""Desktop-only local secret exchange payload."""
secret: str = Field(..., description = "Desktop local auth secret")
class RefreshTokenRequest(BaseModel):
"""Refresh token payload to obtain new access + refresh tokens."""
@ -45,3 +53,44 @@ class ChangePasswordRequest(BaseModel):
new_password: str = Field(
..., min_length = 8, description = "Replacement password (minimum 8 characters)"
)
# ---------------------------------------------------------------------------
# API key schemas
# ---------------------------------------------------------------------------
class CreateApiKeyRequest(BaseModel):
"""Request body to create a new API key."""
name: str = Field(..., description = "Human-readable label for this key")
expires_in_days: Optional[int] = Field(
None, description = "Number of days until the key expires (None = never)"
)
class ApiKeyResponse(BaseModel):
"""Public representation of an API key (never contains the raw key)."""
id: int
name: str
key_prefix: str = Field(
..., description = "First 8 characters after sk-unsloth- for display"
)
created_at: str
last_used_at: Optional[str] = None
expires_at: Optional[str] = None
is_active: bool
class CreateApiKeyResponse(BaseModel):
"""Returned once when a key is created -- ``key`` is never shown again."""
key: str = Field(..., description = "Full API key (shown once)")
api_key: ApiKeyResponse
class ApiKeyListResponse(BaseModel):
"""List of API keys for the authenticated user."""
api_keys: list[ApiKeyResponse]

635
studio/backend/models/inference.py
View file

@ -11,13 +11,16 @@ import time
import uuid
from typing import Annotated, Any, Dict, Literal, Optional, List, Union
from pydantic import BaseModel, Discriminator, Field, Tag
from pydantic import BaseModel, Discriminator, Field, Tag, model_validator
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
model_path: str = Field(..., description = "Model identifier or local path")
native_path_lease: Optional[str] = Field(
None, description = "Frontend-visible signed native path grant"
)
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
@ -48,6 +51,20 @@ class LoadRequest(BaseModel):
None,
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
)
speculative_type: Optional[str] = Field(
None,
description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.",
)
llama_extra_args: Optional[List[str]] = Field(
None,
description = (
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
"Studio-managed flags (model identity, port, context length, GPU placement, "
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
"non-GGUF models."
),
)
class UnloadRequest(BaseModel):
@ -65,6 +82,9 @@ class ValidateModelRequest(BaseModel):
"""
model_path: str = Field(..., description = "Model identifier or local path")
native_path_lease: Optional[str] = Field(
None, description = "Frontend-visible signed native path grant"
)
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
@ -90,6 +110,10 @@ class ValidateModelResponse(BaseModel):
is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)")
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
is_vision: bool = Field(False, description = "Whether this is a vision-capable model")
requires_trust_remote_code: bool = Field(
False,
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
class GenerateRequest(BaseModel):
@ -133,6 +157,10 @@ class LoadResponse(BaseModel):
inference: dict = Field(
..., description = "Inference parameters (temperature, top_p, top_k, min_p)"
)
requires_trust_remote_code: bool = Field(
False,
description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
)
context_length: Optional[int] = Field(
None, description = "Model's native context length (from GGUF metadata)"
)
@ -145,12 +173,20 @@ class LoadResponse(BaseModel):
)
supports_reasoning: bool = Field(
False,
description = "Whether model supports thinking/reasoning mode (enable_thinking)",
description = "Whether model supports thinking/reasoning mode (enable_thinking or reasoning_effort)",
)
reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
)
reasoning_always_on: bool = Field(
False,
description = "Whether reasoning is always on (hardcoded <think> tags, not toggleable)",
)
supports_preserve_thinking: bool = Field(
False,
description = "Whether the template understands the optional preserve_thinking kwarg (Qwen3.6-style)",
)
supports_tools: bool = Field(
False,
description = "Whether model supports tool calling (web search, etc.)",
@ -163,6 +199,10 @@ class LoadResponse(BaseModel):
None,
description = "Jinja2 chat template string (from GGUF metadata or tokenizer)",
)
speculative_type: Optional[str] = Field(
None,
description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
)
class UnloadResponse(BaseModel):
@ -172,6 +212,39 @@ class UnloadResponse(BaseModel):
model: str = Field(..., description = "Model identifier that was unloaded")
class LoadProgressResponse(BaseModel):
"""Progress of the active GGUF load, sampled on demand.
Used by the UI to show a real progress bar during the
post-download warmup window (mmap + CUDA upload), rather than a
generic "Starting model..." spinner that freezes for minutes on
large MoE models.
"""
phase: Optional[str] = Field(
None,
description = (
"Load phase: 'mmap' (weights paging into RAM via mmap), "
"'ready' (llama-server reported healthy), or null when no "
"load is in flight."
),
)
bytes_loaded: int = Field(
0,
description = (
"Bytes of the model already resident in the llama-server "
"process (VmRSS on Linux)."
),
)
bytes_total: int = Field(
0,
description = "Total bytes across all GGUF shards for the active model.",
)
fraction: float = Field(
0.0, description = "bytes_loaded / bytes_total, clamped to 0..1."
)
class InferenceStatusResponse(BaseModel):
"""Current inference backend status"""
@ -205,15 +278,31 @@ class InferenceStatusResponse(BaseModel):
inference: Optional[Dict[str, Any]] = Field(
None, description = "Recommended inference parameters for the active model"
)
requires_trust_remote_code: bool = Field(
False,
description = "Whether the active model requires trust_remote_code to be enabled for loading.",
)
supports_reasoning: bool = Field(
False, description = "Whether the active model supports reasoning/thinking mode"
)
reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
"enable_thinking",
description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
)
reasoning_always_on: bool = Field(
False, description = "Whether reasoning is always on (not toggleable)"
)
supports_preserve_thinking: bool = Field(
False,
description = "Whether the active model's template understands the optional preserve_thinking kwarg",
)
supports_tools: bool = Field(
False, description = "Whether the active model supports tool calling"
)
chat_template: Optional[str] = Field(
None,
description = "Jinja2 chat template string for the active model",
)
context_length: Optional[int] = Field(
None, description = "Context length of the active model"
)
@ -225,6 +314,10 @@ class InferenceStatusResponse(BaseModel):
None,
description = "Model's native context length from GGUF metadata (not capped by VRAM)",
)
speculative_type: Optional[str] = Field(
None,
description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
)
# =====================================================================
@ -281,14 +374,69 @@ class ChatMessage(BaseModel):
``content`` may be a plain string (text-only) or a list of
content parts for multimodal messages (OpenAI vision format).
Assistant messages that only contain tool calls may set ``content``
to ``None`` with ``tool_calls`` populated. ``role="tool"`` messages
carry the result of a client-executed tool call and require
``tool_call_id`` per the OpenAI spec.
"""
role: Literal["system", "user", "assistant"] = Field(
role: Literal["system", "user", "assistant", "tool"] = Field(
..., description = "Message role"
)
content: Union[str, list[ContentPart]] = Field(
..., description = "Message content (string or multimodal parts)"
content: Optional[Union[str, list[ContentPart]]] = Field(
None, description = "Message content (string or multimodal parts)"
)
tool_call_id: Optional[str] = Field(
None,
description = "OpenAI tool-result messages: id of the tool call this result belongs to.",
)
tool_calls: Optional[list[dict]] = Field(
None,
description = "OpenAI assistant messages: structured tool calls the model decided to make.",
)
name: Optional[str] = Field(
None,
description = "OpenAI tool-result messages: name of the tool whose result this is.",
)
@model_validator(mode = "after")
def _validate_role_shape(self) -> "ChatMessage":
# Enforce the per-role OpenAI spec shape at the request boundary.
# Without this, malformed messages (e.g. user entries with no
# content, tool_calls on a user/system role, role="tool" without
# tool_call_id) would be silently forwarded to llama-server via
# the passthrough path, surfacing as opaque upstream errors or
# broken tool-call reconciliation downstream.
# Tool-call metadata must appear only on the appropriate role.
if self.tool_calls is not None and self.role != "assistant":
raise ValueError('"tool_calls" is only valid on role="assistant" messages.')
if self.tool_call_id is not None and self.role != "tool":
raise ValueError('"tool_call_id" is only valid on role="tool" messages.')
if self.name is not None and self.role != "tool":
raise ValueError('"name" is only valid on role="tool" messages.')
# Per-role content requirements. OpenAI-compatible clients may send
# ``content=""`` for image-only turns when the image travels in a
# companion field such as Studio's ``image_base64`` extension, so treat
# empty strings as present content for user/system messages.
if self.role == "tool":
if not self.tool_call_id:
raise ValueError(
'role="tool" messages require "tool_call_id" per the OpenAI spec.'
)
if not self.content:
raise ValueError('role="tool" messages require non-empty "content".')
elif self.role == "assistant":
# Assistant messages may omit content when tool_calls is set.
if not self.content and not self.tool_calls:
raise ValueError(
'role="assistant" messages require either "content" or "tool_calls".'
)
else: # "user" | "system"
if self.content is None or self.content == []:
raise ValueError(f'role="{self.role}" messages require "content".')
return self
class ChatCompletionRequest(BaseModel):
@ -298,18 +446,49 @@ class ChatCompletionRequest(BaseModel):
Extensions (non-OpenAI fields) are marked with 'x-unsloth'.
"""
# Accept unknown fields defensively so future OpenAI fields (seed,
# response_format, logprobs, frequency_penalty, etc.) don't get
# silently dropped by Pydantic before route code runs. Mirrors
# AnthropicMessagesRequest and ResponsesRequest.
model_config = {"extra": "allow"}
model: str = Field(
"default",
description = "Model identifier (informational; the active model is used)",
)
messages: list[ChatMessage] = Field(..., description = "Conversation messages")
stream: bool = Field(True, description = "Whether to stream the response via SSE")
stream: bool = Field(
False,
description = (
"Whether to stream the response via SSE. Default matches OpenAI's "
"spec (`false`); opt into streaming by sending `stream: true`."
),
)
temperature: float = Field(0.6, ge = 0.0, le = 2.0)
top_p: float = Field(0.95, ge = 0.0, le = 1.0)
max_tokens: Optional[int] = Field(
None, ge = 1, description = "Maximum tokens to generate (None = until EOS)"
)
presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty")
stop: Optional[Union[str, list[str]]] = Field(
None,
description = "OpenAI stop sequences: a single string or list of strings at which generation halts.",
)
tools: Optional[list[dict]] = Field(
None,
description = (
"OpenAI function-tool definitions. When provided without `enable_tools=true`, "
"Studio forwards the tools to the backend so the model returns structured "
"tool_calls for the client to execute (standard OpenAI function calling)."
),
)
tool_choice: Optional[Union[str, dict]] = Field(
None,
description = (
"OpenAI tool choice: 'auto' | 'required' | 'none' | "
"{'type': 'function', 'function': {'name': ...}}"
),
)
# ── Unsloth extensions (ignored by standard OpenAI clients) ──
top_k: int = Field(20, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling")
@ -339,6 +518,14 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Enable/disable thinking/reasoning mode for supported models",
)
reasoning_effort: Optional[Literal["low", "medium", "high"]] = Field(
None,
description = "[x-unsloth] Reasoning effort level ('low'|'medium'|'high') for Harmony-style reasoning models (e.g. gpt-oss). Overrides enable_thinking when the active model uses reasoning_effort style.",
)
preserve_thinking: Optional[bool] = Field(
None,
description = "[x-unsloth] When true, keep historical <think> blocks from past assistant turns in the prompt (Qwen3.6 templates). Independent of enable_thinking / reasoning_effort.",
)
enable_tools: Optional[bool] = Field(
None,
description = "[x-unsloth] Enable tool calling for supported models",
@ -365,6 +552,10 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
)
cancel_id: Optional[str] = Field(
None,
description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.",
)
# ── External provider routing (x-unsloth extensions) ──────────
provider_id: Optional[str] = Field(
@ -454,3 +645,435 @@ class ChatCompletion(BaseModel):
model: str = "default"
choices: list[CompletionChoice]
usage: CompletionUsage = Field(default_factory = CompletionUsage)
# =====================================================================
# OpenAI Responses API Models (/v1/responses)
# =====================================================================
# ── Request models ──────────────────────────────────────────────
class ResponsesInputTextPart(BaseModel):
"""Text content part in a Responses API message (type=input_text)."""
type: Literal["input_text"]
text: str
class ResponsesInputImagePart(BaseModel):
"""Image content part in a Responses API message (type=input_image)."""
type: Literal["input_image"]
image_url: str = Field(..., description = "data:image/png;base64,... or https://...")
detail: Optional[Literal["auto", "low", "high"]] = "auto"
class ResponsesOutputTextPart(BaseModel):
"""Assistant ``output_text`` content part replayed on subsequent turns.
When a client (OpenAI Codex CLI, OpenAI Python SDK agents) loops on a
stateless Responses endpoint, prior assistant messages are round-tripped
as ``{"role":"assistant","content":[{"type":"output_text","text":...,
"annotations":[],"logprobs":[]}]}``. We preserve the text and ignore
the annotations/logprobs metadata when flattening into Chat Completions.
"""
type: Literal["output_text"]
text: str
annotations: Optional[list] = None
logprobs: Optional[list] = None
model_config = {"extra": "allow"}
class ResponsesUnknownContentPart(BaseModel):
"""Catch-all for content-part types we don't model explicitly.
Keeps validation green when a client sends newer part types (e.g.
``input_audio``, ``input_file``) we haven't mapped; these are silently
skipped during normalisation rather than rejected with a 422.
"""
type: str
model_config = {"extra": "allow"}
ResponsesContentPart = Union[
ResponsesInputTextPart,
ResponsesInputImagePart,
ResponsesOutputTextPart,
ResponsesUnknownContentPart,
]
class ResponsesInputMessage(BaseModel):
"""A single message in the Responses API input array."""
type: Optional[Literal["message"]] = None
role: Literal["system", "user", "assistant", "developer"]
content: Union[str, list[ResponsesContentPart]]
# Codex (gpt-5.3-codex+) attaches a `phase` field ("commentary" |
# "final_answer") to assistant messages and requires clients to preserve
# it on subsequent turns. We accept and round-trip it; llama-server does
# not care about it.
model_config = {"extra": "allow"}
class ResponsesFunctionCallInputItem(BaseModel):
"""A prior assistant function_call being replayed in a multi-turn Responses input.
The Responses API represents tool calls as top-level input items (not
nested inside assistant messages), correlated across turns by ``call_id``.
"""
type: Literal["function_call"]
id: Optional[str] = Field(
None, description = "Item id assigned by the server (e.g. fc_...)"
)
call_id: str = Field(
...,
description = "Correlation id matching a function_call_output on the next turn.",
)
name: str
arguments: str = Field(
..., description = "JSON string of the arguments the model produced."
)
status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
class ResponsesFunctionCallOutputInputItem(BaseModel):
"""A tool result supplied by the client for a prior function_call.
Replaces Chat Completions' ``role="tool"`` message. Correlated to the
originating call by ``call_id``.
"""
type: Literal["function_call_output"]
id: Optional[str] = None
call_id: str
output: Union[str, list] = Field(
..., description = "String or content-array result of the tool call."
)
status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
class ResponsesUnknownInputItem(BaseModel):
"""Catch-all for Responses input item types we don't model explicitly.
Covers ``reasoning`` items (replayed from prior o-series / gpt-5 turns)
and any future item types the client may send. These items are dropped
during normalisation — llama-server-backed GGUFs cannot consume them —
but keeping them in the request-model union stops unrelated turns from
failing validation with a 422.
"""
type: str
model_config = {"extra": "allow"}
def _responses_input_item_discriminator(v: Any) -> str:
"""Route a Responses input item to the correct tagged variant.
Pydantic's default smart-union matching fails when one variant in the
union is tagged with a strict ``Literal`` (``function_call`` /
``function_call_output``) and the incoming dict uses a different
``type`` — the other variants' validation errors are hidden and the
outer ``Union[str, list[...]]`` reports a misleading "Input should be a
valid string" error. An explicit discriminator makes the routing
deterministic and lets us fall through to the catch-all.
"""
if isinstance(v, dict):
t = v.get("type")
r = v.get("role")
else:
t = getattr(v, "type", None)
r = getattr(v, "role", None)
if t == "function_call":
return "function_call"
if t == "function_call_output":
return "function_call_output"
if r is not None or t == "message":
return "message"
return "unknown"
ResponsesInputItem = Annotated[
Union[
Annotated[ResponsesInputMessage, Tag("message")],
Annotated[ResponsesFunctionCallInputItem, Tag("function_call")],
Annotated[ResponsesFunctionCallOutputInputItem, Tag("function_call_output")],
Annotated[ResponsesUnknownInputItem, Tag("unknown")],
],
Discriminator(_responses_input_item_discriminator),
]
class ResponsesFunctionTool(BaseModel):
"""Flat function-tool definition used by the Responses API request.
Unlike Chat Completions (which nests ``{"name": ..., "parameters": ...}``
inside a ``"function"`` key), the Responses API uses a flat shape with
``type``, ``name``, ``description``, ``parameters``, and ``strict`` at the
top level of each tool entry.
"""
type: Literal["function"]
name: str
description: Optional[str] = None
parameters: Optional[dict] = None
strict: Optional[bool] = None
class ResponsesRequest(BaseModel):
"""OpenAI Responses API request."""
model: str = Field("default", description = "Model identifier")
input: Union[str, list[ResponsesInputItem]] = Field(
default = [],
description = "Input text or list of messages / function_call / function_call_output items",
)
instructions: Optional[str] = Field(
None, description = "System / developer instructions"
)
temperature: Optional[float] = Field(None, ge = 0.0, le = 2.0)
top_p: Optional[float] = Field(None, ge = 0.0, le = 1.0)
max_output_tokens: Optional[int] = Field(None, ge = 1)
stream: bool = Field(False, description = "Whether to stream the response via SSE")
# OpenAI function-calling fields — forwarded to llama-server via the
# Chat Completions pass-through (see routes/inference.py). Typed as a
# plain list so built-in tool shapes (``web_search``, ``file_search``,
# ``mcp``, ...) round-trip without validation errors — the translator
# picks out only ``type=="function"`` entries for forwarding.
tools: Optional[list[dict]] = Field(
None,
description = (
"Responses-shape function tool definitions. Entries with "
'`type="function"` are translated to the Chat Completions nested '
"shape before being forwarded to llama-server; other tool types "
"(built-in web_search, file_search, mcp, ...) are accepted for SDK "
"compatibility but ignored on the llama-server passthrough."
),
)
tool_choice: Optional[Any] = Field(
None,
description = (
"'auto' | 'required' | 'none' | {'type': 'function', 'name': ...} — "
"the Responses-shape forcing object is translated to the Chat "
"Completions nested shape internally."
),
)
parallel_tool_calls: Optional[bool] = None
previous_response_id: Optional[str] = None
store: Optional[bool] = None
metadata: Optional[dict] = None
truncation: Optional[Any] = None
user: Optional[str] = None
text: Optional[Any] = None
reasoning: Optional[Any] = None
model_config = {"extra": "allow"}
# ── Response models ─────────────────────────────────────────────
class ResponsesOutputTextContent(BaseModel):
"""A text content block inside an output message."""
type: Literal["output_text"] = "output_text"
text: str
annotations: list = Field(default_factory = list)
class ResponsesOutputMessage(BaseModel):
"""An output message in the Responses API response."""
type: Literal["message"] = "message"
id: str = Field(default_factory = lambda: f"msg_{uuid.uuid4().hex[:12]}")
status: Literal["completed", "in_progress"] = "completed"
role: Literal["assistant"] = "assistant"
content: list[ResponsesOutputTextContent] = Field(default_factory = list)
class ResponsesOutputFunctionCall(BaseModel):
"""A function-call output item in the Responses API response.
Unlike Chat Completions (which nests tool calls inside the assistant
message), the Responses API emits each tool call as its own top-level
``output`` item so clients can correlate results via ``call_id`` on the
next turn.
"""
type: Literal["function_call"] = "function_call"
id: str = Field(default_factory = lambda: f"fc_{uuid.uuid4().hex[:12]}")
call_id: str
name: str
arguments: str = Field(
..., description = "JSON string of the arguments the model produced."
)
status: Literal["completed", "in_progress", "incomplete"] = "completed"
ResponsesOutputItem = Union[ResponsesOutputMessage, ResponsesOutputFunctionCall]
class ResponsesUsage(BaseModel):
"""Token usage for a Responses API response (input_tokens, not prompt_tokens)."""
input_tokens: int = 0
output_tokens: int = 0
total_tokens: int = 0
class ResponsesResponse(BaseModel):
"""Top-level Responses API response object."""
id: str = Field(default_factory = lambda: f"resp_{uuid.uuid4().hex[:12]}")
object: Literal["response"] = "response"
created_at: int = Field(default_factory = lambda: int(time.time()))
status: Literal["completed", "in_progress", "failed"] = "completed"
model: str = "default"
output: list[ResponsesOutputItem] = Field(default_factory = list)
usage: ResponsesUsage = Field(default_factory = ResponsesUsage)
error: Optional[Any] = None
incomplete_details: Optional[Any] = None
instructions: Optional[str] = None
metadata: dict = Field(default_factory = dict)
temperature: Optional[float] = None
top_p: Optional[float] = None
max_output_tokens: Optional[int] = None
previous_response_id: Optional[str] = None
text: Optional[Any] = None
tool_choice: Optional[Any] = None
tools: list = Field(default_factory = list)
truncation: Optional[Any] = None
# =====================================================================
# Anthropic Messages API Models (/v1/messages)
# =====================================================================
# ── Request models ─────────────────────────────────────────────
class AnthropicTextBlock(BaseModel):
type: Literal["text"]
text: str
class AnthropicImageSource(BaseModel):
type: Literal["base64", "url"]
media_type: Optional[str] = None
data: Optional[str] = None
url: Optional[str] = None
class AnthropicImageBlock(BaseModel):
type: Literal["image"]
source: AnthropicImageSource
class AnthropicToolUseBlock(BaseModel):
type: Literal["tool_use"]
id: str
name: str
input: dict
class AnthropicToolResultBlock(BaseModel):
type: Literal["tool_result"]
tool_use_id: str
content: Union[str, list] = ""
AnthropicContentBlock = Union[
AnthropicTextBlock,
AnthropicImageBlock,
AnthropicToolUseBlock,
AnthropicToolResultBlock,
]
class AnthropicMessage(BaseModel):
role: Literal["user", "assistant"]
content: Union[str, list[AnthropicContentBlock]]
class AnthropicTool(BaseModel):
name: str
description: Optional[str] = None
input_schema: dict
class AnthropicMessagesRequest(BaseModel):
model: str = "default"
max_tokens: Optional[int] = None
messages: list[AnthropicMessage]
system: Optional[Union[str, list]] = None
tools: Optional[list[AnthropicTool]] = None
tool_choice: Optional[Any] = None
stream: bool = False
temperature: Optional[float] = None
top_p: Optional[float] = None
top_k: Optional[int] = None
stop_sequences: Optional[list[str]] = None
metadata: Optional[dict] = None
# [x-unsloth] extensions — mirror the OpenAI endpoint convenience fields
min_p: Optional[float] = Field(
None, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold"
)
repetition_penalty: Optional[float] = Field(
None, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty"
)
presence_penalty: Optional[float] = Field(
None, ge = 0.0, le = 2.0, description = "[x-unsloth] Presence penalty"
)
enable_tools: Optional[bool] = None
enabled_tools: Optional[list[str]] = None
session_id: Optional[str] = None
cancel_id: Optional[str] = None
model_config = {"extra": "allow"}
# ── Response models ────────────────────────────────────────────
class AnthropicUsage(BaseModel):
input_tokens: int = 0
output_tokens: int = 0
class AnthropicResponseTextBlock(BaseModel):
type: Literal["text"] = "text"
text: str
class AnthropicResponseToolUseBlock(BaseModel):
type: Literal["tool_use"] = "tool_use"
id: str
name: str
input: dict
AnthropicResponseBlock = Union[
AnthropicResponseTextBlock, AnthropicResponseToolUseBlock
]
class AnthropicMessagesResponse(BaseModel):
id: str = Field(default_factory = lambda: f"msg_{uuid.uuid4().hex[:24]}")
type: Literal["message"] = "message"
role: Literal["assistant"] = "assistant"
content: list[AnthropicResponseBlock] = Field(default_factory = list)
model: str = "default"
stop_reason: Optional[str] = None
stop_sequence: Optional[str] = None
usage: AnthropicUsage = Field(default_factory = AnthropicUsage)

65
studio/backend/models/models.py
View file

@ -213,3 +213,68 @@ class ScanFolderInfo(BaseModel):
id: int = Field(..., description = "Database row ID")
path: str = Field(..., description = "Normalized absolute path")
created_at: str = Field(..., description = "ISO 8601 creation timestamp")
class BrowseEntry(BaseModel):
"""A directory entry surfaced by the folder browser."""
name: str = Field(..., description = "Entry name (basename, not full path)")
has_models: bool = Field(
False,
description = (
"Hint that the directory likely contains models "
"(*.gguf, *.safetensors, config.json, or HF-style "
"`models--*` subfolders). Used by the UI to highlight "
"promising candidates; the scanner itself is authoritative."
),
)
hidden: bool = Field(
False,
description = "Name starts with a dot (e.g. `.cache`)",
)
class BrowseFoldersResponse(BaseModel):
"""Response schema for the folder browser endpoint."""
current: str = Field(..., description = "Absolute path of the directory just listed")
parent: Optional[str] = Field(
None,
description = (
"Parent directory of `current`, or null if `current` is the "
"filesystem root. The frontend uses this to render an `Up` row."
),
)
entries: List[BrowseEntry] = Field(
default_factory = list,
description = (
"Subdirectories of `current`. Sorted with model-bearing "
"directories first, then alphabetically case-insensitive; "
"hidden entries come last within each group."
),
)
suggestions: List[str] = Field(
default_factory = list,
description = (
"Handy starting points (home, HF cache, already-registered "
"scan folders). Rendered as quick-pick chips above the list."
),
)
truncated: bool = Field(
False,
description = (
"True when the listing was capped because the directory had "
"more subfolders than the server is willing to enumerate in "
"one request. The UI should show a hint telling the user to "
"narrow their path."
),
)
model_files_here: int = Field(
0,
description = (
"Count of GGUF/safetensors files immediately inside "
"``current``. Used by the UI to surface a hint on leaf "
"model directories (which otherwise look `empty` because "
"they contain only files, no subdirectories)."
),
)

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

@ -127,6 +127,9 @@ class TrainingStartRequest(BaseModel):
wandb_project: Optional[str] = Field(None, description = "W&B project name")
enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging")
tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory")
resume_from_checkpoint: Optional[str] = Field(
None, description = "Saved training output directory to resume from"
)
# GPU selection
gpu_ids: Optional[List[int]] = Field(
@ -220,6 +223,8 @@ class TrainingRunSummary(BaseModel):
duration_seconds: Optional[float] = None
error_message: Optional[str] = None
loss_sparkline: Optional[List[float]] = None
can_resume: bool = False
resumed_later: bool = False
class TrainingRunListResponse(BaseModel):

73
studio/backend/plugins/data-designer-github-repo-seed/README.md Normal file
View file

@ -0,0 +1,73 @@
# data-designer-github-repo-seed
A Data Designer seed-reader plugin for **Unsloth Studio** that scrapes real
GitHub data (issues, pull requests, commits) from one or more repositories
and hands it to the recipe pipeline as a seed dataset.
Designed to ship with Studio as a default seed source so any user with a
GitHub token can build training datasets straight from live repos.
## What it does
Given a list of `owner/name` repos, a GitHub token, and a per-resource
`limit`, the plugin uses GitHub's GraphQL API to fetch issues, pull
requests, and/or commits, with labels, state, authors, and the first N
comments of each item, and materialises a single JSONL with uniform
columns so the rest of the recipe (LLM text / LLM structured / processors)
can treat it like any other seed table.
| Column | Description |
|---------------|------------------------------------------------|
| `item_type` | `issue` / `pull` / `commit` |
| `repo` | `owner/name` |
| `number` | Issue/PR number, or commit SHA |
| `title` | Title (or commit message headline) |
| `body` | Issue/PR body (or full commit message) |
| `state` | `OPEN` / `CLOSED` / `MERGED` (empty for commit)|
| `author` | GitHub login of the author |
| `created_at` | ISO8601 |
| `closed_at` | ISO8601 (empty for commits) |
| `url` | Permalink |
| `labels` | List of label names |
| `comments` | First N comments concatenated |
## Usage in a recipe
```json
{
"seed_config": {
"source": {
"seed_type": "github_repo",
"repos": ["unslothai/unsloth", "unslothai/unsloth-zoo"],
"token": "",
"item_types": ["issues", "pulls"],
"limit": 100,
"include_comments": true,
"max_comments_per_item": 30
},
"sampling_strategy": "shuffle",
"selection_strategy": null
}
}
```
Leave `token` empty to fall back to the server's `GH_TOKEN` / `GITHUB_TOKEN`
environment variable, useful when the recipe is published and shouldn't
carry a secret.
## Auth
A GitHub personal access token with `public_repo` scope is enough for public
repositories; `repo` scope is required for private ones. GraphQL requests
are rate-limit aware: the client inspects `x-ratelimit-*` headers and
sleeps until reset when the budget drops below a safety threshold.
## Install
Shipped as a default Studio plugin. For development:
```bash
pip install -e .
```
Registered automatically via the `data_designer.plugins` entry point.

25
studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml Normal file
View file

@ -0,0 +1,25 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "data-designer-github-repo-seed"
version = "0.1.0"
description = "Unsloth Studio seed plugin that scrapes GitHub issues, PRs, and commits."
requires-python = ">=3.11"
dependencies = [
"data-designer-engine>=0.5.4,<0.6",
"requests>=2.31",
]
[project.entry-points."data_designer.plugins"]
github_repo_seed = "data_designer_github_repo_seed.plugin:github_repo_seed_plugin"
[tool.setuptools]
package-dir = {"" = "src"}
[tool.setuptools.packages.find]
where = ["src"]

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

@ -0,0 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Intentionally empty. Data-designer loads submodules lazily via qualified names
# (impl_qualified_name / config_qualified_name in plugin.py), so importing this
# package must NOT touch modules that depend on data_designer.engine.* during
# Studio's bootstrap (circular import).

64
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py Normal file
View file

@ -0,0 +1,64 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
from typing import Literal
from pydantic import Field, field_validator, model_validator
from data_designer.config.seed_source import SeedSource
class GitHubRepoSeedSource(SeedSource):
seed_type: Literal["github_repo"] = "github_repo"
repos: list[str] = Field(
default_factory = list,
description = "List of GitHub repositories to scrape, each in `owner/name` form.",
)
token: str = Field(
default = "",
description = "Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.",
)
item_types: list[Literal["issues", "pulls", "commits"]] = Field(
default = ["issues", "pulls"],
description = "Which GitHub item types to fetch per repo.",
)
limit: int = Field(
default = 100,
ge = 1,
le = 5000,
description = "Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).",
)
include_comments: bool = Field(
default = True,
description = "Fetch the first N comments of each issue/PR and include them in the `comments` column.",
)
max_comments_per_item: int = Field(default = 30, ge = 0, le = 200)
@field_validator("repos")
@classmethod
def _validate_repos(cls, v: list[str]) -> list[str]:
out: list[str] = []
for r in v or []:
r = r.strip()
if not r:
continue
if r.count("/") != 1 or not all(r.split("/")):
raise ValueError(f"Each repo must be `owner/name`; got {r!r}")
out.append(r)
return out
@field_validator("item_types")
@classmethod
def _validate_item_types(cls, v: list[str]) -> list[str]:
if not v:
raise ValueError("item_types must not be empty")
return list(dict.fromkeys(v))
@model_validator(mode = "after")
def _ensure_repos(self) -> "GitHubRepoSeedSource":
if not self.repos:
raise ValueError("At least one repo is required")
return self

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

@ -0,0 +1,83 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import hashlib
import tempfile
import threading
from pathlib import Path
from typing import Optional
import data_designer.lazy_heavy_imports as lazy
from data_designer.engine.resources.seed_reader import SeedReader
from .config import GitHubRepoSeedSource
from .scraper import ScrapeConfig, materialize_to_jsonl
# In-process cache mapping a stable config signature to the JSONL materialization
# path. A single recipe job invokes the seed reader multiple times (validation,
# preview, per-column sampling), and the default flow re-scrapes the repo on
# every call: for a 2-repo preview that is ~15s of redundant GitHub GraphQL
# traffic before any generation fires. Memoize the materialization so the second
# and third passes reuse the file the first pass wrote. Cache key excludes the
# raw token and uses a short SHA-256 digest so token values never hit memory
# twice and token rotation invalidates cleanly.
_SCRAPE_CACHE: dict[tuple, str] = {}
_SCRAPE_CACHE_LOCK = threading.Lock()
def _scrape_cache_key(cfg: ScrapeConfig) -> tuple:
token_digest = hashlib.sha256(
(cfg.token or "").encode("utf-8"),
).hexdigest()[:16]
return (
tuple(cfg.repos),
tuple(cfg.item_types),
cfg.limit,
bool(cfg.include_comments),
cfg.max_comments_per_item,
token_digest,
)
def _lookup_cached_scrape(key: tuple) -> Optional[str]:
with _SCRAPE_CACHE_LOCK:
path = _SCRAPE_CACHE.get(key)
if path and Path(path).exists():
return path
# Stale entry (tmp cleanup, user restarted, ...); drop it so the caller
# materializes a fresh file rather than returning a dangling path.
if path:
with _SCRAPE_CACHE_LOCK:
_SCRAPE_CACHE.pop(key, None)
return None
def _store_cached_scrape(key: tuple, path: str) -> None:
with _SCRAPE_CACHE_LOCK:
_SCRAPE_CACHE[key] = path
class GitHubRepoSeedReader(SeedReader[GitHubRepoSeedSource]):
def create_duckdb_connection(self):
return lazy.duckdb.connect()
def get_dataset_uri(self) -> str:
out_dir = Path(tempfile.gettempdir()) / "studio-github-repo-seed"
cfg = ScrapeConfig(
repos = list(self.source.repos),
token = self.source.token,
item_types = list(self.source.item_types),
limit = self.source.limit,
include_comments = self.source.include_comments,
max_comments_per_item = self.source.max_comments_per_item,
)
cache_key = _scrape_cache_key(cfg)
cached_path = _lookup_cached_scrape(cache_key)
if cached_path is not None:
return cached_path
path = materialize_to_jsonl(cfg, out_dir)
_store_cached_scrape(cache_key, str(path))
return str(path)

10
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py Normal file
View file

@ -0,0 +1,10 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from data_designer.plugins.plugin import Plugin, PluginType
github_repo_seed_plugin = Plugin(
impl_qualified_name = "data_designer_github_repo_seed.impl.GitHubRepoSeedReader",
config_qualified_name = "data_designer_github_repo_seed.config.GitHubRepoSeedSource",
plugin_type = PluginType.SEED_READER,
)

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

@ -0,0 +1,236 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Multi-repo GitHub scraper for the Studio seed plugin.
Drives the GraphQL-based scraper in `scraper_impl/` per repo. Each repo is
scraped with a trial_limits cap so we stop at `limit` items per resource.
After scraping, we read the per-resource JSONL shards and flatten them into
a single unified JSONL with stable columns (`item_type`, `repo`, `number`,
`title`, `body`, ...).
"""
from __future__ import annotations
import json
import os
import sys
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
# Defer scraper_impl imports until `scrape()` runs with a resolved token.
_IMPL_DIR = Path(__file__).parent / "scraper_impl"
def _ensure_impl_on_path() -> None:
if str(_IMPL_DIR) not in sys.path:
sys.path.insert(0, str(_IMPL_DIR))
def _load_impl():
_ensure_impl_on_path()
import importlib
gh_client = importlib.import_module("gh_client") # type: ignore
scraper_mod = importlib.import_module("scraper") # type: ignore
return gh_client.GitHubClient, scraper_mod.RepoScraper
@dataclass
class ScrapeConfig:
repos: list[str]
token: str
item_types: list[str]
limit: int
include_comments: bool
max_comments_per_item: int
def _resolve_token(token: str) -> str:
tok = token or os.environ.get("GH_TOKEN", "") or os.environ.get("GITHUB_TOKEN", "")
if not tok:
raise ValueError(
"GitHub token is required. Set it in the recipe config or the GH_TOKEN / GITHUB_TOKEN env var."
)
return tok
def _read_jsonl(path: Path, max_rows: int | None = None):
if not path.exists():
return
with path.open(encoding = "utf-8") as f:
for i, line in enumerate(f):
if not line.strip():
continue
if max_rows is not None and i >= max_rows:
return
try:
yield json.loads(line)
except json.JSONDecodeError:
continue
def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
labels = [
l.get("name")
for l in (r.get("labels", {}) or {}).get("nodes", [])
if l.get("name")
]
comments_nodes = (r.get("comments") or {}).get("nodes") or []
comments_text = ""
if include_comments and comments_nodes:
kept = comments_nodes[:max_c]
comments_text = "\n\n".join(
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}"
for c in kept
)
return {
"item_type": "issue",
"repo": repo,
"number": r.get("number"),
"title": r.get("title") or "",
"body": r.get("body") or "",
"state": r.get("state") or "",
"author": (r.get("author") or {}).get("login", ""),
"created_at": r.get("createdAt") or "",
"closed_at": r.get("closedAt") or "",
"url": r.get("url") or r.get("permalink") or "",
"labels": labels,
"comments": comments_text,
}
def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
labels = [
l.get("name")
for l in (r.get("labels", {}) or {}).get("nodes", [])
if l.get("name")
]
comments_nodes = (r.get("comments") or {}).get("nodes") or []
comments_text = ""
if include_comments and comments_nodes:
kept = comments_nodes[:max_c]
comments_text = "\n\n".join(
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}"
for c in kept
)
return {
"item_type": "pull",
"repo": repo,
"number": r.get("number"),
"title": r.get("title") or "",
"body": r.get("body") or "",
"state": r.get("state") or "",
"author": (r.get("author") or {}).get("login", ""),
"created_at": r.get("createdAt") or "",
"closed_at": r.get("closedAt") or "",
"url": r.get("url") or r.get("permalink") or "",
"labels": labels,
"comments": comments_text,
}
def _flatten_commit_row(r: dict, repo: str) -> dict:
msg = r.get("messageHeadline") or r.get("message") or ""
body = r.get("messageBody") or r.get("message") or msg
author = r.get("author") or {}
return {
"item_type": "commit",
"repo": repo,
"number": r.get("oid") or r.get("sha") or "",
"title": msg,
"body": body,
"state": "",
"author": (author.get("user") or {}).get("login") or author.get("name", ""),
"created_at": (author.get("date") or r.get("committedDate") or ""),
"closed_at": "",
"url": r.get("url") or "",
"labels": [],
"comments": "",
}
def scrape(cfg: ScrapeConfig, base_dir: Path):
token = _resolve_token(cfg.token)
GitHubClient, RepoScraper = _load_impl()
client = GitHubClient(token = token)
base_dir.mkdir(parents = True, exist_ok = True)
# Per-resource trial limits. limit <= 0 means "all": use a very large cap.
effective_limit = cfg.limit if cfg.limit and cfg.limit > 0 else 1_000_000
trial_limits: dict[str, int] = {}
if "issues" in cfg.item_types:
trial_limits["issues"] = effective_limit
if "pulls" in cfg.item_types:
trial_limits["pull_requests"] = effective_limit
if "commits" in cfg.item_types:
trial_limits["commits"] = effective_limit
all_rows: list[dict] = []
for repo in cfg.repos:
owner, name = repo.split("/", 1)
scraper = RepoScraper(
owner = owner,
name = name,
base_dir = base_dir,
client = client,
trial_limits = trial_limits,
light = True,
)
try:
repo_meta = scraper.scrape_repo_meta()
if "issues" in cfg.item_types:
scraper.scrape_issues()
if "pulls" in cfg.item_types:
scraper.scrape_prs()
if "commits" in cfg.item_types:
default_ref = repo_meta.get("defaultBranchRef") or {}
default_branch = (
default_ref.get("name") if isinstance(default_ref, dict) else None
)
branch = (
f"refs/heads/{default_branch}"
if default_branch
else "refs/heads/main"
)
scraper.scrape_commits(branch = branch)
finally:
scraper.close()
read_cap = cfg.limit if cfg.limit and cfg.limit > 0 else None
repo_dir = base_dir / f"{owner}__{name}"
if "issues" in cfg.item_types:
for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap):
all_rows.append(
_flatten_issue_row(
row, repo, cfg.include_comments, cfg.max_comments_per_item
)
)
if "pulls" in cfg.item_types:
for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap):
all_rows.append(
_flatten_pr_row(
row, repo, cfg.include_comments, cfg.max_comments_per_item
)
)
if "commits" in cfg.item_types:
for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap):
all_rows.append(_flatten_commit_row(row, repo))
return all_rows
def materialize_to_jsonl(cfg: ScrapeConfig, out_dir: Path) -> Path:
out_dir.mkdir(parents = True, exist_ok = True)
tag = "-".join(r.replace("/", "__") for r in cfg.repos)[:120]
kinds = "-".join(cfg.item_types)
run_id = f"{int(time.time())}-{uuid.uuid4().hex[:12]}"
fname = f"github_{tag}__{kinds}__{cfg.limit}_{run_id}.jsonl"
out = out_dir / fname
rows = scrape(cfg, out_dir / "raw-runs" / run_id)
with out.open("w", encoding = "utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii = False) + "\n")
return out

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

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

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

@ -0,0 +1,248 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GitHub API client with rate-limit awareness, retry, and dual REST/GraphQL support."""
from __future__ import annotations
import json
import os
import time
import logging
from typing import Any, Dict, Iterable, Iterator, List, Optional
import requests
log = logging.getLogger("gh_client")
GRAPHQL_URL = "https://api.github.com/graphql"
REST_BASE = "https://api.github.com"
BASE_HEADERS = {
"Accept": "application/vnd.github+json",
"X-GitHub-Api-Version": "2022-11-28",
"User-Agent": "github-data-gatherer/1.0",
}
class RateLimitError(Exception):
pass
class GitHubClient:
def __init__(
self,
min_remaining_graphql: int = 100,
min_remaining_rest: int = 100,
token: str | None = None,
):
token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
if not token:
raise RuntimeError("GH_TOKEN not set in environment")
self.session = requests.Session()
self.session.headers.update(
{**BASE_HEADERS, "Authorization": f"Bearer {token}"}
)
self.min_remaining_graphql = min_remaining_graphql
self.min_remaining_rest = min_remaining_rest
self.graphql_remaining: Optional[int] = None
self.graphql_reset: Optional[int] = None
self.rest_remaining: Optional[int] = None
self.rest_reset: Optional[int] = None
self.calls_graphql = 0
self.calls_rest = 0
self.retry_count = 0
def _sleep_until(self, reset_ts: int, buffer_s: int = 10) -> None:
now = int(time.time())
wait = max(0, reset_ts - now) + buffer_s
log.warning("Rate limit hit. Sleeping %ds until reset.", wait)
time.sleep(wait)
def _check_rate_and_wait(self, kind: str) -> None:
if kind == "graphql":
remaining = self.graphql_remaining
reset = self.graphql_reset
min_remaining = self.min_remaining_graphql
else:
remaining = self.rest_remaining
reset = self.rest_reset
min_remaining = self.min_remaining_rest
if remaining is not None and remaining < min_remaining:
if reset:
self._sleep_until(reset)
# Reset remaining so we don't spin
if kind == "graphql":
self.graphql_remaining = None
else:
self.rest_remaining = None
def graphql(
self,
query: str,
variables: Optional[Dict[str, Any]] = None,
max_retries: int = 20,
) -> Dict[str, Any]:
self._check_rate_and_wait("graphql")
backoff = 2
last_err = None
for attempt in range(max_retries):
try:
r = self.session.post(
GRAPHQL_URL,
json = {"query": query, "variables": variables or {}},
timeout = 120,
)
self.calls_graphql += 1
# Update rate info from response headers
rem = r.headers.get("X-RateLimit-Remaining")
rst = r.headers.get("X-RateLimit-Reset")
if rem is not None:
try:
self.graphql_remaining = int(rem)
except ValueError:
pass
if rst is not None:
try:
self.graphql_reset = int(rst)
except ValueError:
pass
if r.status_code in (502, 503, 504):
log.warning("GraphQL %s transient, retrying", r.status_code)
time.sleep(backoff)
backoff = min(backoff * 2, 60)
continue
if r.status_code == 403 or r.status_code == 429:
# Check for secondary/abuse
retry_after = r.headers.get("Retry-After")
if retry_after:
t = int(retry_after)
log.warning("Secondary rate limit. Sleep %ds.", t)
time.sleep(t + 2)
continue
if self.graphql_reset:
self._sleep_until(self.graphql_reset)
continue
time.sleep(60)
continue
r.raise_for_status()
data = r.json()
if "errors" in data and data["errors"]:
# Surface errors but allow partial data
errs = data["errors"]
# Retry on RATE_LIMITED
for e in errs:
if e.get("type") == "RATE_LIMITED":
self._sleep_until(
(self.graphql_reset or int(time.time()) + 60)
)
break
else:
# No rate-limit error, log and return partial
log.warning("GraphQL errors: %s", json.dumps(errs)[:400])
return data
continue
return data
except requests.RequestException as e:
last_err = e
log.warning("GraphQL network error: %s. Retry.", e)
time.sleep(backoff)
backoff = min(backoff * 2, 60)
raise RuntimeError(f"GraphQL failed after {max_retries} retries: {last_err}")
def rest(
self,
method: str,
path: str,
params: Optional[Dict[str, Any]] = None,
json_body: Optional[Dict[str, Any]] = None,
max_retries: int = 6,
) -> requests.Response:
self._check_rate_and_wait("rest")
if path.startswith("http"):
url = path
else:
url = REST_BASE + path
backoff = 2
last_err = None
for attempt in range(max_retries):
try:
r = self.session.request(
method, url, params = params, json = json_body, timeout = 120
)
self.calls_rest += 1
rem = r.headers.get("X-RateLimit-Remaining")
rst = r.headers.get("X-RateLimit-Reset")
if rem is not None:
try:
self.rest_remaining = int(rem)
except ValueError:
pass
if rst is not None:
try:
self.rest_reset = int(rst)
except ValueError:
pass
if r.status_code in (502, 503, 504):
log.warning("REST %s transient, retrying", r.status_code)
time.sleep(backoff)
backoff = min(backoff * 2, 60)
continue
if r.status_code in (403, 429):
retry_after = r.headers.get("Retry-After")
if retry_after:
t = int(retry_after)
log.warning("Secondary rate limit on REST. Sleep %ds.", t)
time.sleep(t + 2)
continue
# Check if primary rate
if self.rest_remaining == 0 and self.rest_reset:
self._sleep_until(self.rest_reset)
continue
log.warning("REST 403/429, sleep 60")
time.sleep(60)
continue
return r
except requests.RequestException as e:
last_err = e
log.warning("REST network error: %s. Retry.", e)
time.sleep(backoff)
backoff = min(backoff * 2, 60)
raise RuntimeError(f"REST failed after {max_retries} retries: {last_err}")
def rest_paginate(
self, path: str, params: Optional[Dict[str, Any]] = None, per_page: int = 100
) -> Iterator[dict]:
params = dict(params or {})
params.setdefault("per_page", per_page)
url = path
while True:
r = self.rest("GET", url, params = params if url == path else None)
if r.status_code != 200:
log.error(
"REST paginate got %s at %s: %s", r.status_code, url, r.text[:200]
)
return
items = r.json()
if isinstance(items, dict):
# Some endpoints return dict with list field
items = items.get("items", [])
for it in items:
yield it
# Follow link header
link = r.headers.get("Link", "")
nxt = None
for part in link.split(","):
if 'rel="next"' in part:
nxt = part.split(";")[0].strip().strip("<>")
break
if not nxt:
return
url = nxt
params = None
def rate_snapshot(self) -> Dict[str, Any]:
r = self.rest("GET", "/rate_limit")
if r.status_code == 200:
return r.json()
return {}

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

@ -0,0 +1,685 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GraphQL queries for GitHub data scraping.
GitHub's GraphQL rejects queries that define unused fragments, so each query
only includes the fragments it actually references.
"""
# ---- Fragments (kept as raw strings, composed per query) ----
F_ACTOR = """
fragment ActorFields on Actor {
__typename
login
url
avatarUrl
... on User { id databaseId name }
... on Bot { id databaseId }
... on Organization { id databaseId name }
}
"""
F_LABEL = """
fragment LabelFields on Label {
id
name
color
description
createdAt
}
"""
F_TIMELINE = """
fragment TimelineItem on IssueTimelineItems {
__typename
... on Node { id }
... on AddedToProjectEvent { createdAt actor { ...ActorFields } }
... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
... on ClosedEvent { createdAt actor { ...ActorFields } stateReason closer { __typename ... on Commit { oid url } ... on PullRequest { number url } } }
... on CommentDeletedEvent { createdAt actor { ...ActorFields } }
... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url repository { nameWithOwner } } ... on PullRequest { number url repository { nameWithOwner } } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } }
... on ConvertedNoteToIssueEvent { createdAt actor { ...ActorFields } }
... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } }
... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } }
... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } }
... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } }
... on LockedEvent { createdAt actor { ...ActorFields } lockReason }
... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } }
... on MentionedEvent { createdAt actor { ...ActorFields } }
... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } }
... on PinnedEvent { createdAt actor { ...ActorFields } }
... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } }
... on RemovedFromProjectEvent { createdAt actor { ...ActorFields } }
... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle }
... on ReopenedEvent { createdAt actor { ...ActorFields } }
... on SubscribedEvent { createdAt actor { ...ActorFields } }
... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } }
... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } }
... on UnlockedEvent { createdAt actor { ...ActorFields } }
... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } }
... on UnpinnedEvent { createdAt actor { ...ActorFields } }
... on UnsubscribedEvent { createdAt actor { ...ActorFields } }
... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration }
}
"""
F_PR_TIMELINE = """
fragment PRTimelineItem on PullRequestTimelineItems {
__typename
... on Node { id }
... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
... on AutoMergeDisabledEvent { createdAt actor { ...ActorFields } reason }
... on AutoMergeEnabledEvent { createdAt actor { ...ActorFields } }
... on AutoRebaseEnabledEvent { createdAt actor { ...ActorFields } }
... on AutoSquashEnabledEvent { createdAt actor { ...ActorFields } }
... on AutomaticBaseChangeFailedEvent { createdAt actor { ...ActorFields } oldBase newBase }
... on AutomaticBaseChangeSucceededEvent { createdAt actor { ...ActorFields } oldBase newBase }
... on BaseRefChangedEvent { createdAt actor { ...ActorFields } previousRefName currentRefName }
... on BaseRefDeletedEvent { createdAt actor { ...ActorFields } baseRefName }
... on BaseRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } }
... on ClosedEvent { createdAt actor { ...ActorFields } stateReason }
... on CommentDeletedEvent { createdAt actor { ...ActorFields } }
... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url } ... on PullRequest { number url } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } }
... on ConvertToDraftEvent { createdAt actor { ...ActorFields } }
... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } }
... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
... on DeployedEvent { createdAt actor { ...ActorFields } }
... on DeploymentEnvironmentChangedEvent { createdAt actor { ...ActorFields } }
... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } }
... on HeadRefDeletedEvent { createdAt actor { ...ActorFields } headRefName }
... on HeadRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } }
... on HeadRefRestoredEvent { createdAt actor { ...ActorFields } }
... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } }
... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } }
... on LockedEvent { createdAt actor { ...ActorFields } lockReason }
... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } }
... on MentionedEvent { createdAt actor { ...ActorFields } }
... on MergedEvent { createdAt actor { ...ActorFields } commit { oid url } mergeRefName }
... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } }
... on PinnedEvent { createdAt actor { ...ActorFields } }
... on PullRequestCommit { commit { oid url message author { user { login } date } committedDate } }
... on PullRequestCommitCommentThread { commit { oid } }
... on PullRequestReview { id databaseId createdAt submittedAt author { ...ActorFields } body state url reactionGroups { content reactors { totalCount } } }
... on PullRequestReviewThread { id isResolved isOutdated path line diffSide }
... on PullRequestRevisionMarker { createdAt lastSeenCommit { oid } }
... on ReadyForReviewEvent { createdAt actor { ...ActorFields } }
... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } }
... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle }
... on ReopenedEvent { createdAt actor { ...ActorFields } }
... on ReviewDismissedEvent { createdAt actor { ...ActorFields } dismissalMessage previousReviewState }
... on ReviewRequestRemovedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } }
... on ReviewRequestedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } }
... on SubscribedEvent { createdAt actor { ...ActorFields } }
... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } }
... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } }
... on UnlockedEvent { createdAt actor { ...ActorFields } }
... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } }
... on UnpinnedEvent { createdAt actor { ...ActorFields } }
... on UnsubscribedEvent { createdAt actor { ...ActorFields } }
... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration }
}
"""
def _q(parts: list[str], body: str) -> str:
return "\n".join(parts + [body])
ISSUES_PAGE_QUERY = _q(
[F_ACTOR, F_LABEL, F_TIMELINE],
"""
query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
pageInfo { hasNextPage endCursor }
totalCount
nodes {
id databaseId number title body state stateReason
createdAt updatedAt closedAt
url
author { ...ActorFields }
editor { ...ActorFields }
labels(first: 50) { nodes { ...LabelFields } }
assignees(first: 20) { nodes { login id } }
milestone { title number state dueOn }
reactionGroups { content reactors { totalCount } }
comments(first: 100) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId createdAt updatedAt url body
author { ...ActorFields }
editor { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
timelineItems(first: 100) {
totalCount
pageInfo { hasNextPage endCursor }
nodes { ...TimelineItem }
}
trackedInIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } }
trackedIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } }
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
PRS_PAGE_QUERY = _q(
[F_ACTOR, F_LABEL, F_PR_TIMELINE],
"""
query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
pageInfo { hasNextPage endCursor }
totalCount
nodes {
id databaseId number title body state isDraft
createdAt updatedAt closedAt mergedAt
url
headRefName headRefOid
baseRefName baseRefOid
additions deletions changedFiles
mergeable merged mergeStateStatus
author { ...ActorFields }
editor { ...ActorFields }
mergedBy { ...ActorFields }
labels(first: 50) { nodes { ...LabelFields } }
assignees(first: 20) { nodes { login id } }
milestone { title number state dueOn }
reactionGroups { content reactors { totalCount } }
closingIssuesReferences(first: 20) { totalCount nodes { number url repository { nameWithOwner } title } }
comments(first: 100) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId createdAt updatedAt url body
author { ...ActorFields }
editor { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
reviewThreads(first: 50) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id isResolved isOutdated path line diffSide
comments(first: 50) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId createdAt updatedAt url body path diffHunk
author { ...ActorFields }
editor { ...ActorFields }
position originalPosition line originalLine
commit { oid }
reactionGroups { content reactors { totalCount } }
}
}
}
}
reviews(first: 50) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId state createdAt submittedAt body url
author { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
commits(first: 100) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
commit {
oid
message
messageHeadline
committedDate
authoredDate
author { name email user { login } date }
committer { name email user { login } date }
additions deletions changedFilesIfAvailable
parents(first: 3) { nodes { oid } }
}
}
}
files(first: 100) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
path additions deletions changeType
}
}
timelineItems(first: 100) {
totalCount
pageInfo { hasNextPage endCursor }
nodes { ...PRTimelineItem }
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
PRS_PAGE_QUERY_LIGHT = _q(
[F_ACTOR, F_LABEL],
"""
query PRsPageLight($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
pageInfo { hasNextPage endCursor }
totalCount
nodes {
id databaseId number title body state isDraft
createdAt updatedAt closedAt mergedAt
url
author { ...ActorFields }
labels(first: 50) { nodes { ...LabelFields } }
comments(first: 30) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId createdAt updatedAt url body
author { ...ActorFields }
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
ISSUES_PAGE_QUERY_LIGHT = _q(
[F_ACTOR, F_LABEL],
"""
query IssuesPageLight($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
pageInfo { hasNextPage endCursor }
totalCount
nodes {
id databaseId number title body state
createdAt updatedAt closedAt
url
author { ...ActorFields }
labels(first: 50) { nodes { ...LabelFields } }
comments(first: 30) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId createdAt updatedAt url body
author { ...ActorFields }
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
ISSUE_COMMENTS_QUERY = _q(
[F_ACTOR],
"""
query IssueComments($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
issueOrPullRequest(number: $number) {
__typename
... on Issue {
comments(first: 100, after: $after) {
pageInfo { hasNextPage endCursor }
nodes {
id databaseId createdAt updatedAt url body
author { ...ActorFields }
editor { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
}
... on PullRequest {
comments(first: 100, after: $after) {
pageInfo { hasNextPage endCursor }
nodes {
id databaseId createdAt updatedAt url body
author { ...ActorFields }
editor { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
ISSUE_TIMELINE_QUERY = _q(
[F_ACTOR, F_TIMELINE],
"""
query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
issue(number: $number) {
timelineItems(first: 100, after: $after) {
pageInfo { hasNextPage endCursor }
nodes { ...TimelineItem }
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
PR_TIMELINE_QUERY = _q(
[F_ACTOR, F_PR_TIMELINE],
"""
query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequest(number: $number) {
timelineItems(first: 100, after: $after) {
pageInfo { hasNextPage endCursor }
nodes { ...PRTimelineItem }
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
PR_COMMITS_QUERY = """
query PRCommits($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequest(number: $number) {
commits(first: 100, after: $after) {
pageInfo { hasNextPage endCursor }
nodes {
commit {
oid message messageHeadline committedDate authoredDate
author { name email user { login } date }
committer { name email user { login } date }
additions deletions changedFilesIfAvailable
parents(first: 3) { nodes { oid } }
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
"""
PR_FILES_QUERY = """
query PRFiles($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequest(number: $number) {
files(first: 100, after: $after) {
pageInfo { hasNextPage endCursor }
nodes { path additions deletions changeType }
}
}
}
rateLimit { cost remaining resetAt }
}
"""
PR_REVIEW_THREADS_QUERY = _q(
[F_ACTOR],
"""
query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
pullRequest(number: $number) {
reviewThreads(first: 50, after: $after) {
pageInfo { hasNextPage endCursor }
nodes {
id isResolved isOutdated path line diffSide
comments(first: 50) {
totalCount
nodes {
id databaseId createdAt updatedAt url body path diffHunk
author { ...ActorFields }
editor { ...ActorFields }
position originalPosition line originalLine
commit { oid }
reactionGroups { content reactors { totalCount } }
}
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
DISCUSSIONS_PAGE_QUERY = _q(
[F_ACTOR, F_LABEL],
"""
query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
discussions(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
pageInfo { hasNextPage endCursor }
totalCount
nodes {
id databaseId number title body
createdAt updatedAt url
author { ...ActorFields }
editor { ...ActorFields }
locked
answerChosenAt
closed closedAt
category { id name emoji description isAnswerable }
labels(first: 30) { nodes { ...LabelFields } }
upvoteCount
answer { id databaseId body author { ...ActorFields } createdAt url }
reactionGroups { content reactors { totalCount } }
comments(first: 50) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId body createdAt updatedAt url
author { ...ActorFields }
editor { ...ActorFields }
upvoteCount
isAnswer
reactionGroups { content reactors { totalCount } }
replies(first: 50) {
totalCount
pageInfo { hasNextPage endCursor }
nodes {
id databaseId body createdAt updatedAt url
author { ...ActorFields }
editor { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
DISCUSSION_COMMENTS_QUERY = _q(
[F_ACTOR],
"""
query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after: String) {
repository(owner: $owner, name: $name) {
discussion(number: $number) {
comments(first: 50, after: $after) {
pageInfo { hasNextPage endCursor }
nodes {
id databaseId body createdAt updatedAt url
author { ...ActorFields }
editor { ...ActorFields }
upvoteCount
isAnswer
reactionGroups { content reactors { totalCount } }
replies(first: 50) {
totalCount
nodes {
id databaseId body createdAt updatedAt url
author { ...ActorFields }
editor { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
DISCUSSION_REPLIES_QUERY = _q(
[F_ACTOR],
"""
query DiscussionReplies($commentId: ID!, $after: String) {
node(id: $commentId) {
... on DiscussionComment {
replies(first: 50, after: $after) {
pageInfo { hasNextPage endCursor }
nodes {
id databaseId body createdAt updatedAt url
author { ...ActorFields }
editor { ...ActorFields }
reactionGroups { content reactors { totalCount } }
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
COMMITS_PAGE_QUERY = """
query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String, $branch: String!) {
repository(owner: $owner, name: $name) {
ref(qualifiedName: $branch) {
target {
... on Commit {
history(first: $first, after: $after) {
pageInfo { hasNextPage endCursor }
totalCount
nodes {
oid
message
messageHeadline
committedDate
authoredDate
url
additions deletions changedFilesIfAvailable
author { name email date user { login id } }
committer { name email date user { login id } }
parents(first: 3) { nodes { oid } }
associatedPullRequests(first: 5) { nodes { number url state } }
}
}
}
}
}
}
rateLimit { cost remaining resetAt }
}
"""
RELEASES_QUERY = _q(
[F_ACTOR],
"""
query Releases($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
releases(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
pageInfo { hasNextPage endCursor }
nodes {
id databaseId name tagName description
createdAt publishedAt updatedAt
isDraft isPrerelease isLatest
url
author { ...ActorFields }
tagCommit { oid url }
reactionGroups { content reactors { totalCount } }
releaseAssets(first: 50) {
nodes { name contentType size downloadUrl createdAt updatedAt }
}
}
}
}
rateLimit { cost remaining resetAt }
}
""",
)
LABELS_QUERY = _q(
[F_LABEL],
"""
query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
labels(first: $first, after: $after) {
pageInfo { hasNextPage endCursor }
nodes { ...LabelFields }
}
}
rateLimit { cost remaining resetAt }
}
""",
)
MILESTONES_QUERY = """
query Milestones($owner: String!, $name: String!, $first: Int!, $after: String) {
repository(owner: $owner, name: $name) {
milestones(first: $first, after: $after) {
pageInfo { hasNextPage endCursor }
nodes {
id number title description state
createdAt updatedAt closedAt dueOn
creator { login }
}
}
}
rateLimit { cost remaining resetAt }
}
"""
REPO_META_QUERY = """
query RepoMeta($owner: String!, $name: String!) {
repository(owner: $owner, name: $name) {
id databaseId name nameWithOwner description url
createdAt updatedAt pushedAt
isArchived isDisabled isFork isPrivate
primaryLanguage { name }
languages(first: 20, orderBy: {field: SIZE, direction: DESC}) {
edges { size node { name } }
totalSize
}
stargazerCount forkCount watchers { totalCount }
diskUsage
licenseInfo { key name }
homepageUrl
defaultBranchRef { name }
}
rateLimit { cost remaining resetAt }
}
"""

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

@ -0,0 +1,756 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Main scraper orchestration. Collects issues, PRs, discussions, commits, releases, etc.
Resumable via state file. Writes JSONL shards under data/{repo}/{resource}.jsonl.
"""
from __future__ import annotations
import argparse
import json
import logging
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple
# Allow running as a module or script
THIS_DIR = Path(__file__).resolve().parent
if str(THIS_DIR) not in sys.path:
sys.path.insert(0, str(THIS_DIR))
from gh_client import GitHubClient
from state_store import JsonlWriter, StateStore
import queries as Q
log = logging.getLogger("scraper")
def ts() -> str:
return time.strftime("%Y-%m-%d %H:%M:%S")
class RepoScraper:
def __init__(
self,
owner: str,
name: str,
base_dir: Path,
client: GitHubClient,
trial_limits: Optional[Dict[str, int]] = None,
light: bool = False,
):
self.owner = owner
self.name = name
self.base_dir = base_dir
self.client = client
self.trial_limits = trial_limits or {}
# When light=True, use trimmed GraphQL queries (no reviewThreads,
# reviews, commits, timelineItems, files) so PR pages can be much
# larger without blowing GitHub's node-count ceiling.
self.light = light
self.repo_dir = base_dir / f"{owner}__{name}"
self.repo_dir.mkdir(parents = True, exist_ok = True)
self.state = StateStore(base_dir / "state" / f"{owner}__{name}.json")
# Writers
self.writers: Dict[str, JsonlWriter] = {}
for key in (
"issues",
"pull_requests",
"discussions",
"commits",
"releases",
"labels",
"milestones",
"pr_extra_comments",
"pr_extra_timeline",
"pr_extra_reviews",
"issue_extra_comments",
"issue_extra_timeline",
"discussion_extra_comments",
"discussion_extra_replies",
"repo_meta",
):
self.writers[key] = JsonlWriter(self.repo_dir / f"{key}.jsonl")
# ----- helpers -----
def _trial_stop(self, key: str, counter: int) -> bool:
lim = self.trial_limits.get(key)
if lim is None:
return False
return counter >= lim
def _log_rate(self, where: str, data: Dict[str, Any]) -> None:
rl = (
data.get("data", {}).get("rateLimit")
if isinstance(data.get("data"), dict)
else None
)
if rl:
log.debug(
"[%s] rate cost=%s remaining=%s resetAt=%s",
where,
rl.get("cost"),
rl.get("remaining"),
rl.get("resetAt"),
)
# ----- repo meta -----
def scrape_repo_meta(self) -> Dict[str, Any]:
data = self.client.graphql(
Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name}
)
self._log_rate("repo_meta", data)
repo = data.get("data", {}).get("repository") or {}
repo["_fetchedAt"] = ts()
self.writers["repo_meta"].write(repo)
return repo
# ----- issues -----
def scrape_issues(self) -> int:
key = "issues"
cursor = self.state.get(f"{key}_cursor")
done = self.state.get(f"{key}_done", False)
if done:
log.info("%s/%s issues already complete", self.owner, self.name)
return 0
total_new = 0
page = 0
# Light query skips heavy nested fields; safe at 50 per page.
# Clamp by trial_limit so e.g. limit=1 asks GitHub for first:1
# instead of fetching a full 50-item page and discarding 49.
page_cap = 50 if self.light else 15
trial_cap = self.trial_limits.get(key)
per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap
while True:
page += 1
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
}
query = Q.ISSUES_PAGE_QUERY_LIGHT if self.light else Q.ISSUES_PAGE_QUERY
data = self.client.graphql(query, vars_)
self._log_rate("issues", data)
repo = (data.get("data") or {}).get("repository") or {}
issues = repo.get("issues") or {}
nodes = issues.get("nodes") or []
for it in nodes:
it["_owner"] = self.owner
it["_repo"] = self.name
it["_fetchedAt"] = ts()
if not self.light:
if it.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_issue_comments(
it["number"], it["comments"]["pageInfo"]["endCursor"]
)
if (
it.get("timelineItems", {})
.get("pageInfo", {})
.get("hasNextPage")
):
self._paginate_issue_timeline(
it["number"],
it["timelineItems"]["pageInfo"]["endCursor"],
)
if self.writers[key].write(it):
total_new += 1
info = issues.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info(
"[%s/%s] issues page %d (+%d) cursor=%s remaining=%s",
self.owner,
self.name,
page,
len(nodes),
str(cursor)[:20],
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
log.info("Trial limit reached for issues (%d)", total_new)
return total_new
if not info.get("hasNextPage"):
self.state.set(f"{key}_done", True)
break
return total_new
def _paginate_issue_comments(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"issueOrPullRequest"
) or {}
comments = item.get("comments") or {}
for c in comments.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
c["_issueNumber"] = number
self.writers["issue_extra_comments"].write(c)
info = comments.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_issue_timeline(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.ISSUE_TIMELINE_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get("issue") or {}
tl = item.get("timelineItems") or {}
for ev in tl.get("nodes") or []:
ev["_owner"] = self.owner
ev["_repo"] = self.name
ev["_issueNumber"] = number
self.writers["issue_extra_timeline"].write(ev)
info = tl.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
# ----- PRs -----
def scrape_prs(self) -> int:
key = "pull_requests"
cursor = self.state.get(f"{key}_cursor")
done = self.state.get(f"{key}_done", False)
if done:
log.info("%s/%s PRs already complete", self.owner, self.name)
return 0
total_new = 0
page = 0
# Heavy nested PR query is capped at 3 per page (GitHub node-count
# ceiling); light query skips reviewThreads/reviews/commits/etc and
# can safely go to 25 per page. Clamp by trial_limit for small
# previews so limit=1 does not fetch a whole 25-item page.
page_cap = 25 if self.light else 3
trial_cap = self.trial_limits.get(key)
per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap
while True:
page += 1
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
}
query = Q.PRS_PAGE_QUERY_LIGHT if self.light else Q.PRS_PAGE_QUERY
data = self.client.graphql(query, vars_)
self._log_rate("prs", data)
repo = (data.get("data") or {}).get("repository") or {}
prs = repo.get("pullRequests") or {}
nodes = prs.get("nodes") or []
for pr in nodes:
pr["_owner"] = self.owner
pr["_repo"] = self.name
pr["_fetchedAt"] = ts()
num = pr["number"]
if not self.light:
if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_comments(
num, pr["comments"]["pageInfo"]["endCursor"]
)
if (
pr.get("timelineItems", {})
.get("pageInfo", {})
.get("hasNextPage")
):
self._paginate_pr_timeline(
num, pr["timelineItems"]["pageInfo"]["endCursor"]
)
if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_commits(
num, pr["commits"]["pageInfo"]["endCursor"]
)
if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_files(
num, pr["files"]["pageInfo"]["endCursor"]
)
if (
pr.get("reviewThreads", {})
.get("pageInfo", {})
.get("hasNextPage")
):
self._paginate_pr_review_threads(
num, pr["reviewThreads"]["pageInfo"]["endCursor"]
)
if self.writers[key].write(pr):
total_new += 1
info = prs.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info(
"[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s",
self.owner,
self.name,
page,
len(nodes),
str(cursor)[:20],
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
log.info("Trial limit reached for PRs (%d)", total_new)
return total_new
if not info.get("hasNextPage"):
self.state.set(f"{key}_done", True)
break
return total_new
def _paginate_pr_comments(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"issueOrPullRequest"
) or {}
comments = item.get("comments") or {}
for c in comments.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
c["_prNumber"] = number
self.writers["pr_extra_comments"].write(c)
info = comments.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_pr_timeline(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
tl = item.get("timelineItems") or {}
for ev in tl.get("nodes") or []:
ev["_owner"] = self.owner
ev["_repo"] = self.name
ev["_prNumber"] = number
self.writers["pr_extra_timeline"].write(ev)
info = tl.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_pr_commits(self, number: int, after: str) -> None:
cur = after
out_key = "pr_extra_commits"
if out_key not in self.writers:
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
cc = item.get("commits") or {}
for c in cc.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
c["_prNumber"] = number
self.writers[out_key].write(c)
info = cc.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_pr_files(self, number: int, after: str) -> None:
cur = after
out_key = "pr_extra_files"
if out_key not in self.writers:
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_FILES_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
ff = item.get("files") or {}
for f in ff.get("nodes") or []:
f["_owner"] = self.owner
f["_repo"] = self.name
f["_prNumber"] = number
# files don't have id, synthesize one
f["_syntheticId"] = f"{self.owner}/{self.name}#{number}:{f.get('path')}"
self.writers[out_key].write(f)
info = ff.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_pr_review_threads(self, number: int, after: str) -> None:
cur = after
out_key = "pr_extra_review_threads"
if out_key not in self.writers:
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
rt = item.get("reviewThreads") or {}
for th in rt.get("nodes") or []:
th["_owner"] = self.owner
th["_repo"] = self.name
th["_prNumber"] = number
self.writers[out_key].write(th)
info = rt.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
# ----- Discussions -----
def scrape_discussions(self) -> int:
key = "discussions"
cursor = self.state.get(f"{key}_cursor")
done = self.state.get(f"{key}_done", False)
if done:
log.info("%s/%s discussions already complete", self.owner, self.name)
return 0
total_new = 0
page = 0
per_page = 15
while True:
page += 1
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
}
data = self.client.graphql(Q.DISCUSSIONS_PAGE_QUERY, vars_)
self._log_rate("discussions", data)
repo = (data.get("data") or {}).get("repository") or {}
dd = repo.get("discussions") or {}
nodes = dd.get("nodes") or []
for d in nodes:
d["_owner"] = self.owner
d["_repo"] = self.name
d["_fetchedAt"] = ts()
num = d["number"]
if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_discussion_comments(
num, d["comments"]["pageInfo"]["endCursor"]
)
# paginate replies per comment if needed
for c in d.get("comments", {}).get("nodes", []) or []:
if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_discussion_replies(
c["id"], c["replies"]["pageInfo"]["endCursor"], num
)
if self.writers[key].write(d):
total_new += 1
info = dd.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info(
"[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s",
self.owner,
self.name,
page,
len(nodes),
str(cursor)[:20],
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
return total_new
if not info.get("hasNextPage"):
self.state.set(f"{key}_done", True)
break
return total_new
def _paginate_discussion_comments(self, number: int, after: str) -> None:
cur = after
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"number": number,
"after": cur,
}
data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_)
disc = ((data.get("data") or {}).get("repository") or {}).get(
"discussion"
) or {}
cc = disc.get("comments") or {}
for c in cc.get("nodes") or []:
c["_owner"] = self.owner
c["_repo"] = self.name
c["_discussionNumber"] = number
self.writers["discussion_extra_comments"].write(c)
info = cc.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_discussion_replies(
self, comment_id: str, after: str, disc_number: int
) -> None:
cur = after
while cur:
vars_ = {
"owner": self.owner,
"name": self.name,
"commentId": comment_id,
"after": cur,
}
data = self.client.graphql(Q.DISCUSSION_REPLIES_QUERY, vars_)
node = (data.get("data") or {}).get("node") or {}
replies = node.get("replies") or {}
for r in replies.get("nodes") or []:
r["_owner"] = self.owner
r["_repo"] = self.name
r["_discussionNumber"] = disc_number
r["_commentId"] = comment_id
self.writers["discussion_extra_replies"].write(r)
info = replies.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
# ----- Commits -----
def scrape_commits(self, branch: str = "refs/heads/main") -> int:
key = "commits"
cursor = self.state.get(f"{key}_cursor")
done = self.state.get(f"{key}_done", False)
if done:
return 0
total_new = 0
page = 0
page_cap = 100
trial_cap = self.trial_limits.get(key)
per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap
while True:
page += 1
vars_ = {
"owner": self.owner,
"name": self.name,
"first": per_page,
"after": cursor,
"branch": branch,
}
data = self.client.graphql(Q.COMMITS_PAGE_QUERY, vars_)
self._log_rate("commits", data)
ref = ((data.get("data") or {}).get("repository") or {}).get("ref") or {}
tgt = ref.get("target") or {}
hist = tgt.get("history") or {}
nodes = hist.get("nodes") or []
for c in nodes:
c["_owner"] = self.owner
c["_repo"] = self.name
c["_fetchedAt"] = ts()
if self.writers[key].write(c):
total_new += 1
info = hist.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
log.info(
"[%s/%s] commits page %d (+%d) remaining=%s",
self.owner,
self.name,
page,
len(nodes),
self.client.graphql_remaining,
)
if self._trial_stop(key, total_new):
return total_new
if not info.get("hasNextPage"):
self.state.set(f"{key}_done", True)
break
return total_new
# ----- Releases/Labels/Milestones -----
def scrape_releases(self) -> int:
return self._scrape_simple("releases", Q.RELEASES_QUERY, "releases")
def scrape_labels(self) -> int:
return self._scrape_simple("labels", Q.LABELS_QUERY, "labels")
def scrape_milestones(self) -> int:
return self._scrape_simple("milestones", Q.MILESTONES_QUERY, "milestones")
def _scrape_simple(self, key: str, query: str, field: str) -> int:
cursor = self.state.get(f"{key}_cursor")
done = self.state.get(f"{key}_done", False)
if done:
return 0
total_new = 0
while True:
vars_ = {
"owner": self.owner,
"name": self.name,
"first": 50,
"after": cursor,
}
data = self.client.graphql(query, vars_)
repo = (data.get("data") or {}).get("repository") or {}
col = repo.get(field) or {}
for it in col.get("nodes") or []:
it["_owner"] = self.owner
it["_repo"] = self.name
it["_fetchedAt"] = ts()
if self.writers[key].write(it):
total_new += 1
info = col.get("pageInfo") or {}
cursor = info.get("endCursor")
self.state.set(f"{key}_cursor", cursor)
if self._trial_stop(key, total_new):
return total_new
if not info.get("hasNextPage"):
self.state.set(f"{key}_done", True)
break
log.info("[%s/%s] %s done +%d", self.owner, self.name, key, total_new)
return total_new
def close(self) -> None:
for w in self.writers.values():
try:
w.close()
except Exception:
pass
def setup_logging(log_file: Path) -> None:
log_file.parent.mkdir(parents = True, exist_ok = True)
fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
handlers = [
logging.StreamHandler(sys.stdout),
logging.FileHandler(log_file, mode = "a", encoding = "utf-8"),
]
logging.basicConfig(level = logging.INFO, format = fmt, handlers = handlers, force = True)
def main():
ap = argparse.ArgumentParser()
ap.add_argument(
"--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper"
)
ap.add_argument(
"--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"]
)
ap.add_argument("--trial", action = "store_true", help = "Small trial run")
ap.add_argument(
"--only",
nargs = "+",
default = None,
help = "Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta",
)
ap.add_argument(
"--hf-upload-interval",
type = int,
default = 900,
help = "Seconds between HF uploads (0 to disable)",
)
args = ap.parse_args()
base = Path(args.base_dir)
data_dir = base / "data"
data_dir.mkdir(parents = True, exist_ok = True)
setup_logging(base / "logs" / f"scraper_{time.strftime('%Y%m%d_%H%M%S')}.log")
log.info("Scraper starting: repos=%s trial=%s", args.repos, args.trial)
client = GitHubClient(min_remaining_graphql = 80, min_remaining_rest = 80)
rl = client.rate_snapshot()
log.info(
"Rate limit snapshot: %s",
json.dumps(rl.get("resources", {}), default = str)[:400],
)
# Start HF uploader in background if requested
uploader = None
if args.hf_upload_interval > 0:
from hf_uploader import HFUploader
uploader = HFUploader(data_dir, interval_s = args.hf_upload_interval)
uploader.start()
trial_limits = None
if args.trial:
trial_limits = {
"issues": 5,
"pull_requests": 5,
"discussions": 3,
"commits": 20,
"releases": 3,
"labels": 20,
"milestones": 20,
}
only = set(args.only or [])
try:
for repo_spec in args.repos:
owner, name = repo_spec.split("/")
scraper = RepoScraper(owner, name, data_dir, client, trial_limits)
try:
repo_meta: Dict[str, Any] = {}
if not only or "meta" in only or "commits" in only:
repo_meta = scraper.scrape_repo_meta()
if not only or "labels" in only:
scraper.scrape_labels()
if not only or "milestones" in only:
scraper.scrape_milestones()
if not only or "releases" in only:
scraper.scrape_releases()
if not only or "discussions" in only:
scraper.scrape_discussions()
if not only or "issues" in only:
scraper.scrape_issues()
if not only or "pulls" in only:
scraper.scrape_prs()
if not only or "commits" in only:
default_ref = repo_meta.get("defaultBranchRef") or {}
default_branch = (
default_ref.get("name")
if isinstance(default_ref, dict)
else None
)
branch = (
f"refs/heads/{default_branch}"
if default_branch
else "refs/heads/main"
)
scraper.scrape_commits(branch = branch)
finally:
scraper.close()
finally:
if uploader:
log.info("Stopping uploader and final sync...")
uploader.stop(final_upload = True)
log.info(
"Scraper complete. GraphQL calls=%d REST calls=%d",
client.calls_graphql,
client.calls_rest,
)
if __name__ == "__main__":
main()

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

@ -0,0 +1,105 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Checkpoint state management for resumable scraping."""
from __future__ import annotations
import json
import os
import threading
from pathlib import Path
from typing import Any, Dict
class StateStore:
def __init__(self, path: str | Path):
self.path = Path(path)
self.path.parent.mkdir(parents = True, exist_ok = True)
self._lock = threading.Lock()
self._data: Dict[str, Any] = {}
if self.path.exists():
try:
with self.path.open() as f:
self._data = json.load(f)
except Exception:
self._data = {}
def get(self, key: str, default: Any = None) -> Any:
with self._lock:
return self._data.get(key, default)
def set(self, key: str, value: Any) -> None:
with self._lock:
self._data[key] = value
self._flush()
def update(self, key: str, **kwargs) -> None:
with self._lock:
sub = dict(self._data.get(key, {}))
sub.update(kwargs)
self._data[key] = sub
self._flush()
def all(self) -> Dict[str, Any]:
with self._lock:
return dict(self._data)
def _flush(self) -> None:
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
with tmp.open("w") as f:
json.dump(self._data, f, indent = 2, default = str)
os.replace(tmp, self.path)
class JsonlWriter:
"""Append-only JSONL writer, thread-safe, with line buffering."""
def __init__(self, path: str | Path):
self.path = Path(path)
self.path.parent.mkdir(parents = True, exist_ok = True)
self._lock = threading.Lock()
self._fh = self.path.open("a", buffering = 1)
self._count_seen_keys: set[str] = set()
# Preload seen keys if file exists (for dedup across resumes)
if self.path.exists() and self.path.stat().st_size > 0:
try:
with self.path.open() as f:
for line in f:
try:
obj = json.loads(line)
k = self._key(obj)
if k is not None:
self._count_seen_keys.add(k)
except Exception:
pass
except Exception:
pass
def _key(self, obj: dict) -> str | None:
for k in ("id", "node_id", "number", "sha", "url"):
if k in obj:
return f"{k}:{obj[k]}"
return None
def has(self, key: str) -> bool:
return key in self._count_seen_keys
def write(self, obj: dict) -> bool:
"""Return True if newly written, False if already present."""
k = self._key(obj)
with self._lock:
if k is not None and k in self._count_seen_keys:
return False
if k is not None:
self._count_seen_keys.add(k)
self._fh.write(json.dumps(obj, default = str, ensure_ascii = False))
self._fh.write("\n")
self._fh.flush()
return True
def close(self) -> None:
try:
self._fh.close()
except Exception:
pass

8
studio/backend/requirements/extras-no-deps.txt
View file

@ -2,9 +2,13 @@
descript-audio-codec
descript-audiotools
julius
torchcodec
torchcodec==0.10.0
snac
# peft 0.19.0 causes export subprocess shutdown issues in Studio;
# installing with --no-deps to avoid pulling in torch>=0.11.0
peft==0.18.1
# TRL and related packages
trl==0.23.1
git+https://github.com/meta-pytorch/OpenEnv.git
@ -13,4 +17,4 @@ torch-c-dlpack-ext
sentence_transformers==5.2.0
transformers==4.57.6
pytorch_tokenizers
kernels
kernels==0.12.1

3
studio/backend/requirements/single-env/data-designer-deps.txt
View file

@ -19,7 +19,8 @@ ruff<1,>=0.14.10
scipy<2,>=1.11.0
sqlfluff<4,>=3.2.0
tiktoken<1,>=0.8.0
# Unstructured-seed plugin deps (plugin installed with --no-deps)
# Local seed plugin deps (plugins installed with --no-deps)
requests>=2.31
pymupdf>=1.24.0
pymupdf4llm>=0.0.17
mammoth>=1.8.0

2
studio/backend/routes/__init__.py
View file

@ -8,6 +8,7 @@ API Routes
from routes.training import router as training_router
from routes.models import router as models_router
from routes.inference import router as inference_router
from routes.inference import studio_router as inference_studio_router
from routes.datasets import router as datasets_router
from routes.auth import router as auth_router
from routes.data_recipe import router as data_recipe_router
@ -19,6 +20,7 @@ __all__ = [
"training_router",
"models_router",
"inference_router",
"inference_studio_router",
"datasets_router",
"auth_router",
"data_recipe_router",

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

@ -7,11 +7,18 @@ Authentication API routes
from fastapi import APIRouter, Depends, HTTPException, status
from datetime import datetime, timedelta, timezone
from models.auth import (
ApiKeyListResponse,
ApiKeyResponse,
AuthLoginRequest,
RefreshTokenRequest,
AuthStatusResponse,
ChangePasswordRequest,
CreateApiKeyRequest,
CreateApiKeyResponse,
DesktopLoginRequest,
RefreshTokenRequest,
)
from models.users import Token
from auth import storage, hashing
@ -74,6 +81,24 @@ async def login(payload: AuthLoginRequest) -> Token:
)
@router.post("/desktop-login", response_model = Token)
async def desktop_login(payload: DesktopLoginRequest) -> Token:
"""Exchange a local desktop secret for normal admin-subject tokens."""
username = storage.validate_desktop_secret(payload.secret)
if username is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Desktop authentication failed",
)
return Token(
access_token = create_access_token(subject = username, desktop = True),
refresh_token = create_refresh_token(subject = username, desktop = True),
token_type = "bearer",
must_change_password = False,
)
@router.post("/refresh", response_model = Token)
async def refresh(payload: RefreshTokenRequest) -> Token:
"""
@ -81,7 +106,7 @@ async def refresh(payload: RefreshTokenRequest) -> Token:
The refresh token itself is reusable until it expires (7 days).
"""
new_access_token, username = refresh_access_token(payload.refresh_token)
new_access_token, username, is_desktop = refresh_access_token(payload.refresh_token)
if new_access_token is None or username is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
@ -92,7 +117,9 @@ async def refresh(payload: RefreshTokenRequest) -> Token:
access_token = new_access_token,
refresh_token = payload.refresh_token,
token_type = "bearer",
must_change_password = storage.requires_password_change(username),
must_change_password = False
if is_desktop
else storage.requires_password_change(username),
)
@ -131,3 +158,68 @@ async def change_password(
token_type = "bearer",
must_change_password = False,
)
# ---------------------------------------------------------------------------
# API key management
# ---------------------------------------------------------------------------
def _row_to_api_key_response(row: dict) -> ApiKeyResponse:
return ApiKeyResponse(
id = row["id"],
name = row["name"],
key_prefix = row["key_prefix"],
created_at = row["created_at"],
last_used_at = row.get("last_used_at"),
expires_at = row.get("expires_at"),
is_active = bool(row["is_active"]),
)
@router.post("/api-keys", response_model = CreateApiKeyResponse)
async def create_api_key(
payload: CreateApiKeyRequest,
current_subject: str = Depends(get_current_subject),
) -> CreateApiKeyResponse:
"""Create a new API key. The raw key is returned once and cannot be retrieved later."""
expires_at = None
if payload.expires_in_days is not None:
expires_at = (
datetime.now(timezone.utc) + timedelta(days = payload.expires_in_days)
).isoformat()
raw_key, row = storage.create_api_key(
username = current_subject,
name = payload.name,
expires_at = expires_at,
)
return CreateApiKeyResponse(
key = raw_key,
api_key = _row_to_api_key_response(row),
)
@router.get("/api-keys", response_model = ApiKeyListResponse)
async def list_api_keys(
current_subject: str = Depends(get_current_subject),
) -> ApiKeyListResponse:
"""List all API keys for the authenticated user (raw keys are never exposed)."""
rows = storage.list_api_keys(current_subject)
return ApiKeyListResponse(
api_keys = [_row_to_api_key_response(r) for r in rows],
)
@router.delete("/api-keys/{key_id}")
async def revoke_api_key(
key_id: int,
current_subject: str = Depends(get_current_subject),
) -> dict:
"""Revoke (soft-delete) an API key."""
if not storage.revoke_api_key(current_subject, key_id):
raise HTTPException(
status_code = status.HTTP_404_NOT_FOUND,
detail = "API key not found",
)
return {"detail": "API key revoked"}

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

@ -5,7 +5,10 @@
from __future__ import annotations
from typing import Any
import copy
from datetime import datetime, timedelta, timezone
from typing import Any, Optional
from urllib.parse import urlparse
from fastapi import APIRouter, HTTPException, Query, Request
from fastapi.responses import JSONResponse, StreamingResponse
@ -26,6 +29,308 @@ from models.data_recipe import (
router = APIRouter()
def _resolve_local_v1_endpoint(request: Request) -> str:
"""Return the loopback /v1 URL for the actual backend listen port.
Resolution order:
1. ``app.state.server_port`` - explicitly published by run.py after
the uvicorn server has bound. This is the most reliable source
because it survives reverse proxies, TLS terminators and tunnels.
2. ``request.scope["server"]`` - the real (host, port) tuple uvicorn
sets when the request is dispatched. Used when Studio is started
outside ``run_server`` (e.g. ``uvicorn studio.backend.main:app``).
3. ``request.base_url`` parsed - last resort for test fixtures that
do not route through a live uvicorn server.
"""
port: Any = getattr(request.app.state, "server_port", None)
if not isinstance(port, int) or port <= 0:
server = request.scope.get("server")
if (
isinstance(server, tuple)
and len(server) >= 2
and isinstance(server[1], int)
and server[1] > 0
):
port = server[1]
else:
parsed = urlparse(str(request.base_url))
port = parsed.port if parsed.port is not None else 8888
return f"http://127.0.0.1:{int(port)}/v1"
def _request_has_desktop_access_token(request: Request) -> bool:
auth_header = request.headers.get("authorization")
if not auth_header:
return False
parts = auth_header.split(None, 1)
if len(parts) != 2 or parts[0].lower() != "bearer":
return False
from auth.authentication import is_desktop_access_token
return is_desktop_access_token(parts[1])
def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]:
"""Return the set of model_aliases that are actually referenced by an
LLM column. Used to narrow the "Chat model loaded" gate so that orphan
model_config nodes on the canvas do not block unrelated recipe runs.
The ``llm-`` prefix matches the existing convention in
``core/data_recipe/service.py::_recipe_has_llm_columns`` and covers all
LLM column types emitted by the frontend (llm-text, llm-code,
llm-structured, llm-judge).
"""
aliases: set[str] = set()
for column in recipe.get("columns", []):
if not isinstance(column, dict):
continue
column_type = column.get("column_type")
if not isinstance(column_type, str) or not column_type.startswith("llm-"):
continue
alias = column.get("model_alias")
if isinstance(alias, str) and alias:
aliases.add(alias)
return aliases
def _inject_local_structured_response_format(
recipe: dict[str, Any], local_provider_names: set[str]
) -> None:
"""For each llm-structured column that targets a local-provider model_config,
clone the model_config and inject an OpenAI ``response_format`` with the
column's ``output_format`` JSON schema. The column is rewritten to point at
the clone so llm-text / llm-judge columns that share the same alias keep
free-form sampling.
Without this, data_designer only injects a prompt-level "return JSON in a
```json fence" instruction. Small GGUF models frequently break format,
wasting the full ``max_tokens`` budget per row and then failing to parse.
Forwarding ``response_format`` lets llama-server apply grammar-constrained
sampling from the JSON schema, which guarantees a parseable response and
terminates early.
"""
columns = recipe.get("columns")
model_configs = recipe.get("model_configs")
if not isinstance(columns, list) or not isinstance(model_configs, list):
return
# alias -> model_config (only configs referencing a local provider qualify).
alias_to_local_mc: dict[str, dict[str, Any]] = {}
for mc in model_configs:
if not isinstance(mc, dict):
continue
if mc.get("provider") in local_provider_names and isinstance(
mc.get("alias"), str
):
alias_to_local_mc[mc["alias"]] = mc
if not alias_to_local_mc:
return
# Clone per (alias, column) so each llm-structured column gets its own
# schema without leaking response_format onto other columns that share the
# same base alias.
seen_clone_aliases: set[str] = {
mc.get("alias") for mc in model_configs if isinstance(mc.get("alias"), str)
}
new_configs: list[dict[str, Any]] = []
for column in columns:
if not isinstance(column, dict):
continue
if column.get("column_type") != "llm-structured":
continue
alias = column.get("model_alias")
if not isinstance(alias, str) or alias not in alias_to_local_mc:
continue
output_format = column.get("output_format")
if not isinstance(output_format, dict) or not output_format:
continue
base_mc = alias_to_local_mc[alias]
column_name = column.get("name") or "structured"
clone_alias_base = f"{alias}__{column_name}_structured"
clone_alias = clone_alias_base
counter = 1
while clone_alias in seen_clone_aliases:
counter += 1
clone_alias = f"{clone_alias_base}_{counter}"
seen_clone_aliases.add(clone_alias)
clone = copy.deepcopy(base_mc)
clone["alias"] = clone_alias
params = clone.get("inference_parameters")
if not isinstance(params, dict):
params = {}
clone["inference_parameters"] = params
# data_designer's BaseInferenceParams is a pydantic model with
# extra="forbid", so response_format cannot sit at the top level of
# inference_parameters. It does expose an `extra_body: dict` pass-
# through that the OpenAI client spreads into the request body at the
# top level, which is where llama-server reads response_format from.
# llama.cpp server shape (tools/server/README.md): the schema sits
# directly under response_format, not nested in a json_schema object
# the way OpenAI's Chat Completions API expects. llama-server converts
# the schema to a GBNF grammar and applies it during sampling.
extra_body = params.get("extra_body")
if not isinstance(extra_body, dict):
extra_body = {}
extra_body["response_format"] = {
"type": "json_schema",
"schema": output_format,
}
params["extra_body"] = extra_body
new_configs.append(clone)
column["model_alias"] = clone_alias
if new_configs:
model_configs.extend(new_configs)
def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optional[int]:
"""
Mutate recipe dict in-place: for any provider with is_local=True,
fill in the endpoint pointing at this server and inject a short-lived
internal sk-unsloth-* API key for workflow auth.
Returns the row id of the minted internal key (so the caller can
revoke it on job completion) or ``None`` when no local provider is
actually reachable from an LLM column.
"""
providers = recipe.get("model_providers")
if not providers:
return None
# Collect local providers and pop is_local from ALL dicts unconditionally.
# Strict `is True` guard so malformed payloads (is_local: 1,
# is_local: "true") do not accidentally trigger the loopback rewrite.
local_indices: list[int] = []
for i, provider in enumerate(providers):
if not isinstance(provider, dict):
continue
is_local = provider.pop("is_local", None)
if is_local is True:
local_indices.append(i)
if not local_indices:
return None
endpoint = _resolve_local_v1_endpoint(request)
# Only gate on model-loaded if a local provider is actually reachable
# from an LLM column through a model_config. Orphan model_config nodes
# that reference a local provider but that no LLM column uses should
# not block runs; the recipe would never call /v1 for them.
local_names = {
providers[i].get("name") for i in local_indices if providers[i].get("name")
}
used_aliases = _used_llm_model_aliases(recipe)
referenced_providers = {
mc.get("provider")
for mc in recipe.get("model_configs", [])
if (
isinstance(mc, dict)
and mc.get("provider")
and mc.get("alias") in used_aliases
)
}
token = ""
internal_key_id: Optional[int] = None
if local_names & referenced_providers:
# Verify a model is loaded.
# NOTE: This is a point-in-time check (TOCTOU). The model could be unloaded
# or swapped after this check but before the recipe subprocess calls /v1.
# The inference endpoint returns a clear 400 in that case.
#
# Imports are deferred to avoid circular dependencies with inference modules.
from routes.inference import get_llama_cpp_backend
from core.inference import get_inference_backend
llama = get_llama_cpp_backend()
model_loaded = llama.is_loaded
if not model_loaded:
backend = get_inference_backend()
model_loaded = bool(backend.active_model_name)
if not model_loaded:
raise ValueError(
"No model loaded in Chat. Load a model first, then run the recipe."
)
from auth import storage # deferred: avoids circular import
# Mint an internal sk-unsloth-* key scoped to this workflow run.
# Uses the unified API-key issuance path (one mint/revoke/verify
# surface instead of a second JWT code path). The key is marked
# internal so it is hidden from the user's API-key list, and the
# caller revokes it when the job terminates.
expires_at = (datetime.now(timezone.utc) + timedelta(hours = 24)).isoformat()
token, row = storage.create_api_key(
username = "unsloth",
name = "data-recipe workflow",
expires_at = expires_at,
internal = True,
)
internal_key_id = int(row["id"])
# Defensively strip any stale "external"-only fields the frontend may
# have left on the dict (extra_headers/extra_body/api_key_env). The UI
# hides these inputs in local mode but the payload builder still serializes
# them, so a previously external provider that flipped to local can carry
# invalid JSON or rogue auth headers into the local /v1 call.
for i in local_indices:
providers[i]["endpoint"] = endpoint
providers[i]["api_key"] = token
providers[i]["provider_type"] = "openai"
providers[i].pop("api_key_env", None)
providers[i].pop("extra_headers", None)
providers[i].pop("extra_body", None)
# Force skip_health_check on any model_config that references a local
# provider. The local /v1/models endpoint only lists the real loaded
# model (e.g. "unsloth/llama-3.2-1b") and not the placeholder "local"
# that the recipe sends as the model id, so data_designer's pre-flight
# health check would otherwise fail before the first completion call.
# The backend route ignores the model id field in chat completions, so
# skipping the check is safe.
for mc in recipe.get("model_configs", []):
if not isinstance(mc, dict):
continue
if mc.get("provider") in local_names:
mc["skip_health_check"] = True
# Disable thinking for data-recipe inference on local providers.
# Reasoning models emit a <think>...</think> preamble before the
# answer, which roughly doubles generated token count per row and
# pushes the visible answer past data_designer's json-fence
# regex. Forward chat_template_kwargs={enable_thinking: False}
# through the OpenAI SDK's extra_body passthrough so llama-server
# renders the template without the reasoning preamble. Free-form
# llm-text columns benefit from the latency cut, and structured
# columns also stop leaking think tags into the grammar-
# constrained JSON (llama-server's GBNF path still enforces the
# schema either way).
params = mc.get("inference_parameters")
if not isinstance(params, dict):
params = {}
mc["inference_parameters"] = params
extra_body = params.get("extra_body")
if not isinstance(extra_body, dict):
extra_body = {}
tpl_kwargs = extra_body.get("chat_template_kwargs")
if not isinstance(tpl_kwargs, dict):
tpl_kwargs = {}
tpl_kwargs.setdefault("enable_thinking", False)
extra_body["chat_template_kwargs"] = tpl_kwargs
params["extra_body"] = extra_body
# Forward each llm-structured column's output_format as an OpenAI
# response_format so llama-server uses grammar-constrained sampling and
# small GGUFs stop wasting the full max_tokens budget on broken JSON.
_inject_local_structured_response_format(recipe, local_names)
return internal_key_id
def _normalize_run_name(value: Any) -> str | None:
if value is None:
return None
@ -40,7 +345,7 @@ def _normalize_run_name(value: Any) -> str | None:
@router.post("/jobs", response_class = JSONResponse, response_model = JobCreateResponse)
def create_job(payload: RecipePayload):
def create_job(payload: RecipePayload, request: Request):
recipe = payload.recipe
if not recipe.get("columns"):
raise HTTPException(status_code = 400, detail = "Recipe must include columns.")
@ -67,17 +372,50 @@ def create_job(payload: RecipePayload):
status_code = 400, detail = f"invalid run_config: {exc}"
) from exc
mgr = get_job_manager()
try:
job_id = mgr.start(recipe = recipe, run = run)
except RuntimeError as exc:
raise HTTPException(status_code = 409, detail = str(exc)) from exc
internal_api_key_id = _inject_local_providers(recipe, request)
except ValueError as exc:
raise HTTPException(status_code = 400, detail = str(exc)) from exc
# Single try block covers get_job_manager() AND mgr.start() so a workflow
# key minted above never outlives the request even when an unexpected
# exception type (TypeError from a stale kwarg, OSError from a queue
# write, etc.) bubbles up. Without the bare except, such exceptions let
# the sk-unsloth-* key live until its 24h TTL.
try:
mgr = get_job_manager()
job_id = mgr.start(
recipe = recipe,
run = run,
internal_api_key_id = internal_api_key_id,
)
except RuntimeError as exc:
if internal_api_key_id is not None:
_revoke_internal_api_key_safe(internal_api_key_id)
raise HTTPException(status_code = 409, detail = str(exc)) from exc
except ValueError as exc:
if internal_api_key_id is not None:
_revoke_internal_api_key_safe(internal_api_key_id)
raise HTTPException(status_code = 400, detail = str(exc)) from exc
except Exception:
if internal_api_key_id is not None:
_revoke_internal_api_key_safe(internal_api_key_id)
raise
return {"job_id": job_id}
def _revoke_internal_api_key_safe(key_id: int) -> None:
"""Best-effort revoke of a workflow-minted key; swallow any error so
that revocation failures never mask the caller's own error path."""
try:
from auth import storage # deferred: avoids circular import
storage.revoke_internal_api_key(key_id)
except Exception:
pass
@router.get("/jobs/{job_id}/status")
def job_status(job_id: str):
mgr = get_job_manager()

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

@ -8,6 +8,7 @@ from __future__ import annotations
import base64
import binascii
import json
import os
import re
from itertools import islice
from pathlib import Path
@ -627,3 +628,14 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
split = None,
subset = None,
)
@router.get("/seed/github/env-token")
def get_github_env_token_status() -> dict:
"""Report whether the server has a GH_TOKEN / GITHUB_TOKEN env var.
The value is never returned; the UI uses this to tell the user they
can leave the token field blank.
"""
has_token = bool(os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN"))
return {"has_token": has_token}

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

@ -14,10 +14,63 @@ from core.data_recipe.service import (
create_data_designer,
validate_recipe,
)
from loggers import get_logger
from models.data_recipe import RecipePayload, ValidateError, ValidateResponse
logger = get_logger(__name__)
router = APIRouter()
_GITHUB_VALIDATE_NOTE = "Recipe shape is valid. GitHub access and rate limits are checked when the run starts."
_GITHUB_ITEM_TYPES = {"issues", "pulls", "commits"}
def _github_seed_source(recipe: dict[str, Any]) -> dict[str, Any] | None:
seed_config = recipe.get("seed_config")
if not isinstance(seed_config, dict):
return None
source = seed_config.get("source")
if not isinstance(source, dict) or source.get("seed_type") != "github_repo":
return None
return source
def _validate_github_seed_static(source: dict[str, Any]) -> list[ValidateError]:
errors: list[ValidateError] = []
repos = source.get("repos")
if not isinstance(repos, list) or not repos:
errors.append(ValidateError(message = "GitHub seed requires at least one repo."))
else:
for repo in repos:
if not isinstance(repo, str) or not repo.strip() or "/" not in repo:
errors.append(
ValidateError(message = "GitHub repos must be owner/name strings.")
)
break
item_types = source.get("item_types")
if not isinstance(item_types, list) or not item_types:
errors.append(
ValidateError(message = "GitHub seed requires at least one item type.")
)
else:
invalid_items = [item for item in item_types if item not in _GITHUB_ITEM_TYPES]
if invalid_items:
errors.append(
ValidateError(
message = "GitHub item types must be issues, pulls, or commits."
)
)
try:
limit = int(source.get("limit"))
except (TypeError, ValueError):
limit = 0
if limit < 1 or limit > 5000:
errors.append(ValidateError(message = "GitHub limit must be from 1 to 5000."))
return errors
def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
try:
@ -68,6 +121,20 @@ def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
return errors
def _patch_local_providers(recipe: dict[str, Any]) -> None:
"""Strip is_local and fill a dummy endpoint so validation doesn't choke.
Uses a strict `is True` check to match _inject_local_providers in
jobs.py - malformed payloads with truthy but non-boolean is_local
values should not be treated as local.
"""
for provider in recipe.get("model_providers", []):
if not isinstance(provider, dict):
continue
if provider.pop("is_local", None) is True:
provider["endpoint"] = "http://127.0.0.1"
@router.post("/validate", response_model = ValidateResponse)
def validate(payload: RecipePayload) -> ValidateResponse:
recipe = payload.recipe
@ -77,6 +144,40 @@ def validate(payload: RecipePayload) -> ValidateResponse:
errors = [ValidateError(message = "Recipe must include columns.")],
)
_patch_local_providers(recipe)
github_source = _github_seed_source(recipe)
if github_source is not None:
static_errors = _validate_github_seed_static(github_source)
if static_errors:
return ValidateResponse(valid = False, errors = static_errors)
try:
build_config_builder(recipe)
except ModuleNotFoundError as exc:
# data_designer is an optional runtime dep. Static validation
# already passed; live access + full config validation are
# deferred to run start (per _GITHUB_VALIDATE_NOTE), so a missing
# optional import at validate time should not block the recipe.
# Restrict the bypass to the data_designer module specifically so
# other ImportErrors (e.g. broken internal imports or missing
# transitive deps after a package upgrade) still surface as
# validation failures instead of being silently swallowed.
if not (exc.name or "").startswith("data_designer"):
raise
logger.debug(
"data_designer not installed; deferring full config "
"validation to run start",
missing_module = exc.name,
)
except Exception as exc:
detail = str(exc).strip() or "Validation failed."
return ValidateResponse(
valid = False,
errors = [ValidateError(message = detail)],
raw_detail = detail,
)
return ValidateResponse(valid = True, raw_detail = _GITHUB_VALIDATE_NOTE)
try:
validate_recipe(recipe)
except RuntimeError as exc:

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

@ -11,10 +11,55 @@ import json
import sys
from pathlib import Path
from uuid import uuid4
from fastapi import APIRouter, Depends, HTTPException, UploadFile
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile
import re as _re
import structlog
from loggers import get_logger
_VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")
def _is_valid_repo_id(repo_id: str) -> bool:
return bool(_VALID_REPO_ID.fullmatch(repo_id))
_dataset_size_cache: dict[str, int] = {}
def _get_dataset_size_cached(repo_id: str) -> int:
if repo_id in _dataset_size_cache:
return _dataset_size_cache[repo_id]
try:
from huggingface_hub import dataset_info as hf_dataset_info
info = hf_dataset_info(repo_id, token = None, files_metadata = True)
total = sum(s.size for s in info.siblings if getattr(s, "size", None))
_dataset_size_cache[repo_id] = total
return total
except Exception:
return 0
def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
"""Pick the most useful on-disk path for a HF cache repo dir.
Mirrors the helper in routes/models.py: prefer the most-recent
snapshot dir, fall back to the cache repo root, return resolved
realpath. Duplicated here to keep routes/datasets.py self-contained.
"""
try:
snapshots_dir = repo_dir / "snapshots"
if snapshots_dir.is_dir():
snaps = [s for s in snapshots_dir.iterdir() if s.is_dir()]
if snaps:
latest = max(snaps, key = lambda s: s.stat().st_mtime)
return str(latest.resolve())
return str(repo_dir.resolve())
except Exception:
return None
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
@ -308,6 +353,89 @@ def list_local_datasets(
return LocalDatasetsResponse(datasets = _build_local_dataset_items())
@router.get("/download-progress")
async def get_dataset_download_progress(
repo_id: str = Query(
..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"
),
current_subject: str = Depends(get_current_subject),
):
"""Return download progress for a HuggingFace dataset repo.
Mirrors ``GET /api/models/download-progress`` but scans the
``datasets--owner--name`` cache directory under HF_HUB_CACHE.
Modern ``datasets``/``huggingface_hub`` caches both raw model and
raw dataset blobs in HF_HUB_CACHE; the ``datasets`` library writes
its processed Arrow shards elsewhere, but the in-progress *download*
bytes are observable here. Returns ``cache_path`` so the UI can
show users where the dataset blobs landed on disk.
"""
_empty = {
"downloaded_bytes": 0,
"expected_bytes": 0,
"progress": 0,
"cache_path": None,
}
try:
if not _is_valid_repo_id(repo_id):
return _empty
from huggingface_hub import constants as hf_constants
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"datasets--{repo_id.replace('/', '--')}".lower()
completed_bytes = 0
in_progress_bytes = 0
cache_path: Optional[str] = None
if cache_dir.is_dir():
for entry in cache_dir.iterdir():
if entry.name.lower() != target:
continue
cache_path = _resolve_hf_cache_realpath(entry)
blobs_dir = entry / "blobs"
if not blobs_dir.is_dir():
break
for f in blobs_dir.iterdir():
if not f.is_file():
continue
if f.name.endswith(".incomplete"):
in_progress_bytes += f.stat().st_size
else:
completed_bytes += f.stat().st_size
break
downloaded_bytes = completed_bytes + in_progress_bytes
if downloaded_bytes == 0:
return {**_empty, "cache_path": cache_path}
expected_bytes = _get_dataset_size_cached(repo_id)
if expected_bytes <= 0:
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": 0,
"progress": 0,
"cache_path": cache_path,
}
# Same 95% completion threshold as the model endpoint -- HF blob
# dedup makes completed_bytes drift slightly under expected_bytes,
# and inter-file gaps would otherwise look like "done".
if completed_bytes >= expected_bytes * 0.95:
progress = 1.0
else:
progress = min(downloaded_bytes / expected_bytes, 0.99)
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": expected_bytes,
"progress": round(progress, 3),
"cache_path": cache_path,
}
except Exception as e:
logger.warning(f"Error checking dataset download progress for {repo_id}: {e}")
return _empty
@router.post("/check-format", response_model = CheckFormatResponse)
def check_format(
request: CheckFormatRequest,

215
studio/backend/routes/export.py
View file

@ -5,9 +5,15 @@
Export API routes: checkpoint discovery and model export operations.
"""
import asyncio
import json
import sys
import time
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Query
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from fastapi.responses import StreamingResponse
import structlog
from loggers import get_logger
@ -97,7 +103,11 @@ async def load_checkpoint(
logger.warning("Could not stop training: %s", e)
backend = get_export_backend()
success, message = backend.load_checkpoint(
# load_checkpoint spawns and waits on a subprocess and can take
# minutes. Run it in a worker thread so the event loop stays
# free to serve the live log SSE stream concurrently.
success, message = await asyncio.to_thread(
backend.load_checkpoint,
checkpoint_path = request.checkpoint_path,
max_seq_length = request.max_seq_length,
load_in_4bit = request.load_in_4bit,
@ -129,7 +139,7 @@ async def cleanup_export_memory(
"""
try:
backend = get_export_backend()
success = backend.cleanup_memory()
success = await asyncio.to_thread(backend.cleanup_memory)
if not success:
raise HTTPException(
@ -173,6 +183,17 @@ async def get_export_status(
)
def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]:
"""Wrap the resolved on-disk export path into the details dict the
frontend reads to populate the Export Complete screen. Returns None
when the export had no local component (Hub-only push) so the
Pydantic field stays absent rather than ``{"output_path": null}``.
"""
if not output_path:
return None
return {"output_path": output_path}
@router.post("/export/merged", response_model = ExportOperationResponse)
async def export_merged_model(
request: ExportMergedModelRequest,
@ -185,7 +206,8 @@ async def export_merged_model(
"""
try:
backend = get_export_backend()
success, message = backend.export_merged_model(
success, message, output_path = await asyncio.to_thread(
backend.export_merged_model,
save_directory = request.save_directory,
format_type = request.format_type,
push_to_hub = request.push_to_hub,
@ -197,7 +219,11 @@ async def export_merged_model(
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success = True, message = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
@ -220,7 +246,8 @@ async def export_base_model(
"""
try:
backend = get_export_backend()
success, message = backend.export_base_model(
success, message, output_path = await asyncio.to_thread(
backend.export_base_model,
save_directory = request.save_directory,
push_to_hub = request.push_to_hub,
repo_id = request.repo_id,
@ -232,7 +259,11 @@ async def export_base_model(
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success = True, message = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
@ -255,7 +286,8 @@ async def export_gguf(
"""
try:
backend = get_export_backend()
success, message = backend.export_gguf(
success, message, output_path = await asyncio.to_thread(
backend.export_gguf,
save_directory = request.save_directory,
quantization_method = request.quantization_method,
push_to_hub = request.push_to_hub,
@ -266,7 +298,11 @@ async def export_gguf(
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success = True, message = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
@ -289,7 +325,8 @@ async def export_lora_adapter(
"""
try:
backend = get_export_backend()
success, message = backend.export_lora_adapter(
success, message, output_path = await asyncio.to_thread(
backend.export_lora_adapter,
save_directory = request.save_directory,
push_to_hub = request.push_to_hub,
repo_id = request.repo_id,
@ -300,7 +337,11 @@ async def export_lora_adapter(
if not success:
raise HTTPException(status_code = 400, detail = message)
return ExportOperationResponse(success = True, message = message)
return ExportOperationResponse(
success = True,
message = message,
details = _export_details(output_path),
)
except HTTPException:
raise
except Exception as e:
@ -309,3 +350,155 @@ async def export_lora_adapter(
status_code = 500,
detail = f"Failed to export LoRA adapter: {str(e)}",
)
# ─────────────────────────────────────────────────────────────────────
# Live export log stream (Server-Sent Events)
# ─────────────────────────────────────────────────────────────────────
#
# The export worker subprocess redirects its stdout/stderr into a pipe
# that a reader thread forwards to the orchestrator as log entries (see
# core/export/worker.py::_setup_log_capture and
# core/export/orchestrator.py::_append_log). This endpoint streams
# those entries to the browser so the export dialog can show a live
# terminal-style output panel while load_checkpoint / export_merged /
# export_gguf / export_lora / export_base run.
#
# Shape follows the training progress SSE endpoint
# (routes/training.py::stream_training_progress): each event carries
# `id`, `event`, and `data` fields, the stream starts with a `retry:`
# directive, and `Last-Event-ID` is honored on reconnect.
def _format_sse(data: str, event: str, event_id: Optional[int] = None) -> str:
"""Format a single SSE message with id/event/data fields."""
lines = []
if event_id is not None:
lines.append(f"id: {event_id}")
lines.append(f"event: {event}")
lines.append(f"data: {data}")
lines.append("")
lines.append("")
return "\n".join(lines)
@router.get("/logs/stream")
async def stream_export_logs(
request: Request,
since: Optional[int] = Query(
None,
description = "Return log entries with seq strictly greater than this cursor.",
),
current_subject: str = Depends(get_current_subject),
):
"""
Stream live stdout/stderr output from the export worker subprocess
as Server-Sent Events.
Events:
- `log` : a single log line (data: {"stream","line","ts"})
- `heartbeat`: periodic keepalive when no new lines are available
- `complete` : emitted once the export worker is idle and no new
lines arrived for ~1 second. Clients should close.
- `error` : unrecoverable server-side error
The `id:` field on each event is the log entry's monotonic seq
number so the browser can resume via `Last-Event-ID` on reconnect.
"""
backend = get_export_backend()
# Determine starting cursor. Explicit `since` wins, then
# Last-Event-ID header on reconnect, otherwise start from the
# run-start snapshot captured by clear_logs() so the client sees
# every line emitted since the current run began -- even if the
# SSE connection opened after the POST that kicked off the export.
# Using get_current_log_seq() here would lose the early bootstrap
# lines that arrive in the gap between POST and SSE connect.
last_event_id = request.headers.get("last-event-id")
if since is None and last_event_id is not None:
try:
since = int(last_event_id)
except ValueError:
pass
if since is None:
cursor = backend.get_run_start_seq()
else:
cursor = max(0, int(since))
async def event_generator() -> AsyncGenerator[str, None]:
nonlocal cursor
# Tell the browser to reconnect after 3 seconds if the
# connection drops mid-export.
yield "retry: 3000\n\n"
last_yield = time.monotonic()
idle_since: Optional[float] = None
try:
while True:
if await request.is_disconnected():
return
entries, new_cursor = backend.get_logs_since(cursor)
if entries:
for entry in entries:
payload = json.dumps(
{
"stream": entry.get("stream", "stdout"),
"line": entry.get("line", ""),
"ts": entry.get("ts"),
}
)
yield _format_sse(
payload,
event = "log",
event_id = int(entry.get("seq", 0)),
)
cursor = new_cursor
last_yield = time.monotonic()
idle_since = None
else:
now = time.monotonic()
if now - last_yield > 10.0:
yield _format_sse("{}", event = "heartbeat")
last_yield = now
if not backend.is_export_active():
# Give the reader thread a moment to drain any
# trailing lines the worker process printed
# just before signalling done.
if idle_since is None:
idle_since = now
elif now - idle_since > 1.0:
yield _format_sse(
"{}",
event = "complete",
event_id = cursor,
)
return
else:
idle_since = None
await asyncio.sleep(0.1)
except asyncio.CancelledError:
# Client disconnected mid-yield. Don't re-raise, just end
# the generator cleanly so StreamingResponse finalizes.
return
except Exception as exc:
logger.error("Export log stream failed: %s", exc, exc_info = True)
try:
yield _format_sse(
json.dumps({"error": str(exc)}),
event = "error",
)
except Exception:
pass
return StreamingResponse(
event_generator(),
media_type = "text/event-stream",
headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)

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

File diff suppressed because it is too large Load diff

1516
studio/backend/routes/models.py
View file

File diff suppressed because it is too large Load diff

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

@ -25,6 +25,12 @@ if str(backend_path) not in sys.path:
# Import backend functions
try:
from core.training import get_training_backend
from core.training.resume import (
can_resume_run,
get_resume_checkpoint_path,
normalize_resume_output_dir,
)
from storage.studio_db import get_resumable_run_by_output_dir
from utils.models.model_config import load_model_defaults
from utils.paths import resolve_dataset_path
except ImportError:
@ -33,6 +39,12 @@ except ImportError:
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
from core.training import get_training_backend
from core.training.resume import (
can_resume_run,
get_resume_checkpoint_path,
normalize_resume_output_dir,
)
from storage.studio_db import get_resumable_run_by_output_dir
from utils.models.model_config import load_model_defaults
from utils.paths import resolve_dataset_path
@ -152,6 +164,28 @@ async def start_training(
request.local_eval_datasets = _validate_local_dataset_paths(
request.local_eval_datasets, "Local eval dataset"
)
resume_output_dir: Optional[str] = None
if request.resume_from_checkpoint:
try:
resume_output_dir = normalize_resume_output_dir(
request.resume_from_checkpoint
)
except ValueError as e:
raise HTTPException(status_code = 400, detail = str(e))
resume_run = get_resumable_run_by_output_dir(resume_output_dir)
if not resume_run or not can_resume_run(resume_run):
raise HTTPException(
status_code = 400,
detail = "Resume checkpoint must belong to a stopped run with saved trainer state.",
)
resume_checkpoint = get_resume_checkpoint_path(resume_output_dir)
if not resume_checkpoint:
raise HTTPException(
status_code = 400,
detail = "Resume checkpoint must include saved trainer state.",
)
request.resume_from_checkpoint = resume_checkpoint
# Convert request to kwargs for backend
training_kwargs = {
@ -209,6 +243,8 @@ async def start_training(
"wandb_project": request.wandb_project or "",
"enable_tensorboard": request.enable_tensorboard,
"tensorboard_dir": request.tensorboard_dir or "",
"output_dir": resume_output_dir,
"resume_from_checkpoint": request.resume_from_checkpoint,
"trust_remote_code": request.trust_remote_code,
"gpu_ids": request.gpu_ids,
}
@ -437,6 +473,9 @@ async def get_training_status(
"loss": getattr(progress, "loss", None),
"learning_rate": getattr(progress, "learning_rate", None),
}
output_dir = getattr(backend, "_output_dir", None)
if output_dir:
details["output_dir"] = output_dir
# Build metric history for chart recovery after SSE reconnection
metric_history = None

13
studio/backend/routes/training_history.py
View file

@ -11,6 +11,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query
from loggers import get_logger
from auth.authentication import get_current_subject
from core.training.resume import can_resume_run
from models import (
TrainingRunDeleteResponse,
TrainingRunDetailResponse,
@ -34,7 +35,10 @@ async def list_training_runs(
"""List training runs, newest first."""
result = list_runs(limit = limit, offset = offset)
return TrainingRunListResponse(
runs = [TrainingRunSummary(**r) for r in result["runs"]],
runs = [
TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)})
for r in result["runs"]
],
total = result["total"],
)
@ -58,7 +62,12 @@ async def get_training_run_detail(
metrics_data = get_run_metrics(run_id)
return TrainingRunDetailResponse(
run = TrainingRunSummary(**{k: v for k, v in run.items() if k != "config_json"}),
run = TrainingRunSummary(
**{
**{k: v for k, v in run.items() if k != "config_json"},
"can_resume": can_resume_run(run),
}
),
config = config,
metrics = TrainingRunMetrics(**metrics_data),
)

42
studio/backend/run.py
View file

@ -244,10 +244,12 @@ _shutdown_event = None
def run_server(
host: str = "0.0.0.0",
host: str = "127.0.0.1",
port: int = 8888,
frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
silent: bool = False,
api_only: bool = False,
llama_parallel_slots: int = 1,
):
"""
Start the FastAPI server.
@ -257,6 +259,8 @@ def run_server(
port: Port to bind to (auto-increments if in use)
frontend_path: Path to frontend build directory (optional)
silent: Suppress startup messages
api_only: Run API server only, no frontend serving (for Tauri desktop app)
llama_parallel_slots: Number of parallel slots for llama-server
Note:
Signal handlers are NOT registered here so that embedders
@ -273,6 +277,10 @@ def run_server(
except Exception:
pass
# Set env var BEFORE importing main so CORS middleware picks it up
if api_only:
os.environ["UNSLOTH_API_ONLY"] = "1"
import nest_asyncio
nest_asyncio.apply()
@ -308,8 +316,12 @@ def run_server(
print("=" * 50)
print("")
# Setup frontend if path provided
if frontend_path:
# Output port for Tauri to parse when in api-only mode
if api_only:
print(f"TAURI_PORT={port}", flush = True)
# Setup frontend if path provided (skip in api-only mode)
if frontend_path and not api_only:
if setup_frontend(app, frontend_path):
if not silent:
print(f"[OK] Frontend loaded from {frontend_path}")
@ -324,6 +336,15 @@ def run_server(
_server = uvicorn.Server(config)
_shutdown_event = Event()
# Expose the actual bound port so request-handling code can build
# loopback URLs that point at the real backend, not whatever port a
# reverse proxy or tunnel exposed in the request URL. Only publish
# an explicit value when we know the concrete port; for ephemeral
# binds (port==0) leave it unset and let request handlers fall back
# to the ASGI request scope or request.base_url.
app.state.server_port = port if port and port > 0 else None
app.state.llama_parallel_slots = llama_parallel_slots
# Run server in a daemon thread
def _run():
asyncio.run(_server.serve())
@ -371,7 +392,11 @@ if __name__ == "__main__":
pass
parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server")
parser.add_argument("--host", default = "0.0.0.0", help = "Host to bind to")
parser.add_argument(
"--host",
default = "127.0.0.1",
help = "Host to bind to (default: 127.0.0.1; use 0.0.0.0 for network/cloud access)",
)
parser.add_argument("--port", type = int, default = 8888, help = "Port to bind to")
parser.add_argument(
"--frontend",
@ -380,10 +405,17 @@ if __name__ == "__main__":
help = "Path to frontend build",
)
parser.add_argument("--silent", action = "store_true", help = "Suppress output")
parser.add_argument(
"--api-only",
action = "store_true",
help = "API server only, no frontend (for Tauri)",
)
args = parser.parse_args()
kwargs = dict(host = args.host, port = args.port, silent = args.silent)
kwargs = dict(
host = args.host, port = args.port, silent = args.silent, api_only = args.api_only
)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)

33
studio/backend/state/tool_policy.py Normal file
View file

@ -0,0 +1,33 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Process-level server-side tool policy.
Set by `unsloth run` at startup; consulted by the inference route gates.
None -> no CLI override (default). Per-request `enable_tools` is honored.
True -> CLI forced tools on for every request.
False -> CLI forced tools off for every request.
"""
from typing import Optional
_tool_policy: Optional[bool] = None
def get_tool_policy() -> Optional[bool]:
return _tool_policy
def set_tool_policy(value: Optional[bool]) -> None:
if value is not None and not isinstance(value, bool):
raise TypeError(
f"tool_policy must be Optional[bool], got {type(value).__name__}"
)
global _tool_policy
_tool_policy = value
def reset_tool_policy() -> None:
global _tool_policy
_tool_policy = None

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

@ -267,10 +267,23 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
total = conn.execute("SELECT COUNT(*) FROM training_runs").fetchone()[0]
rows = conn.execute(
"""
SELECT id, status, model_name, dataset_name, started_at, ended_at,
total_steps, final_step, final_loss, output_dir,
duration_seconds, error_message, loss_sparkline
FROM training_runs
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
r.ended_at, r.total_steps, r.final_step, r.final_loss,
r.output_dir, r.duration_seconds, r.error_message,
r.loss_sparkline,
CASE
WHEN r.status = 'stopped'
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
END AS resumed_later
FROM training_runs r
ORDER BY started_at DESC
LIMIT ? OFFSET ?
""",
@ -297,7 +310,26 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
def get_run(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute("SELECT * FROM training_runs WHERE id = ?", (id,)).fetchone()
row = conn.execute(
"""
SELECT r.*,
CASE
WHEN r.status = 'stopped'
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
END AS resumed_later
FROM training_runs r
WHERE r.id = ?
""",
(id,),
).fetchone()
if row is None:
return None
run = dict(row)
@ -313,6 +345,45 @@ def get_run(id: str) -> Optional[dict]:
conn.close()
def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT r.*,
0 AS resumed_later
FROM training_runs r
WHERE r.output_dir = ?
AND r.status = 'stopped'
AND NOT EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.started_at > r.started_at
)
ORDER BY r.started_at DESC
LIMIT 1
""",
(output_dir,),
).fetchone()
if row is None:
return None
run = dict(row)
sparkline = run.get("loss_sparkline")
if sparkline:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug(
"Failed to parse loss_sparkline for output_dir %s", output_dir
)
run["loss_sparkline"] = None
return run
finally:
conn.close()
def get_run_metrics(id: str) -> dict:
"""Return metric arrays for a run, using paired step arrays per metric."""
conn = get_connection()

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

@ -3,14 +3,136 @@
"""
Shared pytest configuration for the backend test suite.
Ensures that the backend root is on sys.path so that
`import utils.utils` (and similar flat imports) resolve correctly.
Responsibilities:
1. Put the backend root on sys.path so `from models.inference import ...`
(and similar flat imports) resolve in test modules — mirrors how the
app itself is launched.
2. Provide a hybrid ``studio_server`` session fixture for end-to-end tests
(see ``test_studio_api.py``). The fixture supports two invocation modes:
a. **External server.** If ``UNSLOTH_E2E_BASE_URL`` is set, tests point
at an already-running Studio instance. ``UNSLOTH_E2E_API_KEY`` must
also be set. This is the fast-iteration mode: start the server once
with ``unsloth studio run ...``, then run pytest against it many
times with no per-run GGUF load cost.
b. **Fixture-managed server.** Otherwise, the fixture launches a fresh
server via ``_start_server`` and tears it down at session end. This
is the one-shot mode for CI or a clean-slate verification run.
The model / variant for mode (b) come from ``--unsloth-model`` /
``--unsloth-gguf-variant`` pytest options, then ``UNSLOTH_E2E_MODEL`` /
``UNSLOTH_E2E_VARIANT`` env vars, then the defaults in
``test_studio_api.py``.
"""
import os
import sys
from pathlib import Path
import pytest
# Add backend root to sys.path (mirrors how the app itself is launched)
_backend_root = Path(__file__).resolve().parent.parent
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
# ── Pytest CLI options ───────────────────────────────────────────────
def pytest_addoption(parser):
group = parser.getgroup(
"unsloth-e2e",
"Unsloth Studio end-to-end test options",
)
group.addoption(
"--unsloth-model",
action = "store",
default = None,
help = (
"GGUF model id used when starting a server for e2e tests. "
"Ignored if UNSLOTH_E2E_BASE_URL is set. Overrides "
"UNSLOTH_E2E_MODEL env var. Defaults to test_studio_api.py's "
"DEFAULT_MODEL."
),
)
group.addoption(
"--unsloth-gguf-variant",
action = "store",
default = None,
help = (
"GGUF variant used when starting a server for e2e tests. "
"Ignored if UNSLOTH_E2E_BASE_URL is set. Overrides "
"UNSLOTH_E2E_VARIANT env var. Defaults to test_studio_api.py's "
"DEFAULT_VARIANT."
),
)
# ── E2E server fixtures ──────────────────────────────────────────────
@pytest.fixture(scope = "session")
def studio_server(request):
"""Yield ``(base_url, api_key)`` for e2e tests.
Resolution order:
1. If ``UNSLOTH_E2E_BASE_URL`` is set → point at that server,
require ``UNSLOTH_E2E_API_KEY`` alongside (skip if missing).
2. Otherwise → start a fresh ``unsloth studio run`` subprocess via
the existing ``_start_server`` helper in ``test_studio_api.py``
and tear it down on session teardown.
Session-scoped so the expensive GGUF load happens at most once per
pytest invocation. Lazily instantiated — tests that don't request
the fixture (e.g. the unit tests in ``test_anthropic_messages.py``
or ``test_help_output``) do not trigger server startup.
"""
external_url = os.environ.get("UNSLOTH_E2E_BASE_URL")
if external_url:
api_key = os.environ.get("UNSLOTH_E2E_API_KEY")
if not api_key:
pytest.skip(
"UNSLOTH_E2E_BASE_URL is set but UNSLOTH_E2E_API_KEY is "
"missing — tests that require auth cannot run against an "
"external server without it.",
)
yield external_url, api_key
return
# Lazy import: pytest has already loaded test_studio_api into
# sys.modules by the time any test requests this fixture, so this
# is a cache hit, not a re-execution.
import test_studio_api as _e2e
model = (
request.config.getoption("--unsloth-model")
or os.environ.get("UNSLOTH_E2E_MODEL")
or _e2e.DEFAULT_MODEL
)
variant = (
request.config.getoption("--unsloth-gguf-variant")
or os.environ.get("UNSLOTH_E2E_VARIANT")
or _e2e.DEFAULT_VARIANT
)
proc, api_key = _e2e._start_server(model, variant)
try:
yield f"http://{_e2e.HOST}:{_e2e.PORT}", api_key
finally:
_e2e._kill_server(proc)
@pytest.fixture
def base_url(studio_server):
"""Base URL for the e2e Studio server (from ``studio_server``)."""
return studio_server[0]
@pytest.fixture
def api_key(studio_server):
"""API key for the e2e Studio server (from ``studio_server``)."""
return studio_server[1]

1013
studio/backend/tests/test_anthropic_messages.py Normal file
View file

File diff suppressed because it is too large Load diff

86
studio/backend/tests/test_browse_folders_route.py Normal file
View file

@ -0,0 +1,86 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import os
import sys
import types
from pathlib import Path
import pytest
from fastapi import HTTPException
# Keep this test runnable in lightweight environments where optional logging
# deps are not installed.
if "structlog" not in sys.modules:
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
sys.modules["structlog"] = types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
)
import routes.models as models_route
def test_resolve_browse_target_returns_allowed_directory(tmp_path):
allowed = tmp_path / "allowed"
target = allowed / "models" / "nested"
target.mkdir(parents = True)
resolved = models_route._resolve_browse_target(str(target), [allowed])
assert resolved == target.resolve()
def test_resolve_browse_target_rejects_outside_allowlist(tmp_path):
allowed = tmp_path / "allowed"
disallowed = tmp_path / "disallowed"
allowed.mkdir()
disallowed.mkdir()
with pytest.raises(HTTPException) as exc_info:
models_route._resolve_browse_target(str(disallowed), [allowed])
assert exc_info.value.status_code == 403
def test_resolve_browse_target_rejects_file_path(tmp_path):
allowed = tmp_path / "allowed"
allowed.mkdir()
model_file = allowed / "model.gguf"
model_file.write_text("gguf")
with pytest.raises(HTTPException) as exc_info:
models_route._resolve_browse_target(str(model_file), [allowed])
assert exc_info.value.status_code == 400
def test_resolve_browse_target_allows_symlink_into_other_allowed_root(tmp_path):
home_root = tmp_path / "home"
scan_root = tmp_path / "scan"
target = scan_root / "nested"
home_root.mkdir()
target.mkdir(parents = True)
(home_root / "scan-link").symlink_to(scan_root, target_is_directory = True)
resolved = models_route._resolve_browse_target(
str(home_root / "scan-link" / "nested"),
[home_root, scan_root],
)
assert resolved == target.resolve()
@pytest.mark.skipif(os.altsep is not None, reason = "POSIX-only path semantics")
def test_resolve_browse_target_allows_backslash_in_posix_segment(tmp_path):
allowed = tmp_path / "allowed"
target = allowed / r"dir\name"
target.mkdir(parents = True)
resolved = models_route._resolve_browse_target(str(target), [allowed])
assert resolved == target.resolve()

120
studio/backend/tests/test_cache_case_resolution.py Normal file
View file

@ -0,0 +1,120 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from pathlib import Path
import sys
import types
# Keep this test runnable in lightweight environments where optional logging
# deps are not installed.
if "structlog" not in sys.modules:
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
sys.modules["structlog"] = types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
)
from utils.paths.path_utils import (
resolve_cached_repo_id_case,
get_cache_case_resolution_stats,
reset_cache_case_resolution_state,
)
import utils.paths.path_utils as path_utils
def _mk_cache_repo(cache_root: Path, repo_id: str) -> Path:
repo_dir = cache_root / f"models--{repo_id.replace('/', '--')}"
repo_dir.mkdir(parents = True, exist_ok = True)
return repo_dir
def test_resolve_cached_repo_id_case_exact_hit(tmp_path, monkeypatch):
reset_cache_case_resolution_state()
_mk_cache_repo(tmp_path, "Org/Model")
monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
resolved = resolve_cached_repo_id_case("Org/Model")
assert resolved == "Org/Model"
stats = get_cache_case_resolution_stats()
assert stats["calls"] == 1
assert stats["exact_hits"] == 1
assert stats["variant_hits"] == 0
def test_resolve_cached_repo_id_case_variant_hit(tmp_path, monkeypatch):
reset_cache_case_resolution_state()
_mk_cache_repo(tmp_path, "Org/Model")
monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
resolved = resolve_cached_repo_id_case("org/model")
assert resolved == "Org/Model"
stats = get_cache_case_resolution_stats()
assert stats["variant_hits"] == 1
assert stats["tie_breaks"] == 0
def test_resolve_cached_repo_id_case_tie_break_deterministic(tmp_path, monkeypatch):
reset_cache_case_resolution_state()
_mk_cache_repo(tmp_path, "Org/Model")
_mk_cache_repo(tmp_path, "org/model")
monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
resolved = resolve_cached_repo_id_case("oRg/mOdEl")
# Deterministic rule: lexical sort of candidate repo ids.
assert resolved == "Org/Model"
stats = get_cache_case_resolution_stats()
assert stats["variant_hits"] == 1
assert stats["tie_breaks"] == 1
def test_resolve_cached_repo_id_case_no_cache_fallback(tmp_path, monkeypatch):
reset_cache_case_resolution_state()
monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
resolved = resolve_cached_repo_id_case("Org/Missing")
assert resolved == "Org/Missing"
stats = get_cache_case_resolution_stats()
assert stats["fallbacks"] == 1
assert stats["variant_hits"] == 0
assert stats["exact_hits"] == 0
def test_resolve_cached_repo_id_case_memoization(tmp_path, monkeypatch):
reset_cache_case_resolution_state()
_mk_cache_repo(tmp_path, "Org/Model")
monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
first = resolve_cached_repo_id_case("org/model")
second = resolve_cached_repo_id_case("org/model")
assert first == "Org/Model"
assert second == "Org/Model"
stats = get_cache_case_resolution_stats()
assert stats["calls"] == 2
assert stats["variant_hits"] == 1
assert stats["memo_hits"] == 1
def test_resolve_cached_repo_id_case_late_cache_population(tmp_path, monkeypatch):
"""Regression guard: memoized fallback should not hide a later cache variant."""
reset_cache_case_resolution_state()
monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
first = resolve_cached_repo_id_case("org/model")
assert first == "org/model"
# Simulate cache being populated after first miss (e.g. another code path/download).
_mk_cache_repo(tmp_path, "Org/Model")
second = resolve_cached_repo_id_case("org/model")
# Desired behavior: second lookup should pick up the now-existing variant.
assert second == "Org/Model"

398
studio/backend/tests/test_cached_gguf_routes.py Normal file
View file

@ -0,0 +1,398 @@
# 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
from pathlib import Path
from types import SimpleNamespace
# Keep this test runnable in lightweight environments where optional logging
# deps are not installed.
if "structlog" not in sys.modules:
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
sys.modules["structlog"] = types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
)
import routes.models as models_route
def _repo(
repo_id: str,
files: list[SimpleNamespace],
repo_path: Path,
*,
revisions: list[SimpleNamespace] | None = None,
) -> SimpleNamespace:
return SimpleNamespace(
repo_id = repo_id,
repo_type = "model",
repo_path = repo_path,
revisions = revisions or [SimpleNamespace(files = files)],
)
def _file(
name: str,
size_on_disk: int,
*,
blob_path: str | None = None,
) -> SimpleNamespace:
return SimpleNamespace(
file_name = name,
size_on_disk = size_on_disk,
blob_path = blob_path,
)
def test_iter_gguf_paths_matches_extension_case_insensitively(tmp_path):
nested = tmp_path / "snapshots" / "rev"
nested.mkdir(parents = True)
lower = nested / "Q4_K_M.gguf"
upper = nested / "Q8_0.GGUF"
other = nested / "README.md"
lower.write_text("a")
upper.write_text("b")
other.write_text("c")
result = sorted(path.name for path in models_route._iter_gguf_paths(tmp_path))
assert result == ["Q4_K_M.gguf", "Q8_0.GGUF"]
def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(
monkeypatch, tmp_path
):
repo = _repo(
"HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive",
[_file("Q4_K_M.gguf", 5_000), _file("README.md", 10)],
tmp_path / "models--HauhauCS--Gemma",
)
scan = SimpleNamespace(repos = [repo])
monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [scan])
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive",
"size_bytes": 5_000,
"cache_path": str(repo.repo_path),
}
]
def test_list_cached_gguf_matches_extension_case_insensitively(monkeypatch, tmp_path):
repo = _repo(
"Org/Model-Without-Suffix",
[_file("Q8_0.GGUF", 7_000)],
tmp_path / "models--Org--Model-Without-Suffix",
)
scan = SimpleNamespace(repos = [repo])
monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [scan])
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/Model-Without-Suffix",
"size_bytes": 7_000,
"cache_path": str(repo.repo_path),
}
]
def test_list_cached_gguf_skips_repos_without_positive_gguf_size(monkeypatch, tmp_path):
missing = _repo(
"Org/ReadmeOnly",
[_file("README.md", 10)],
tmp_path / "models--Org--ReadmeOnly",
)
zero = _repo(
"Org/ZeroSize",
[_file("Q4_K_M.gguf", 0)],
tmp_path / "models--Org--ZeroSize",
)
scan = SimpleNamespace(repos = [missing, zero])
monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [scan])
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == []
def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(
monkeypatch, tmp_path
):
smaller = _repo(
"Org/Dupe",
[_file("Q4_K_M.gguf", 2_000)],
tmp_path / "models--Org--Dupe-a",
)
larger = _repo(
"org/dupe",
[_file("Q4_K_M.gguf", 5_000), _file("Q6_K.gguf", 1_000)],
tmp_path / "models--Org--Dupe-b",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [
SimpleNamespace(repos = [smaller]),
SimpleNamespace(repos = [larger]),
],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "org/dupe",
"size_bytes": 6_000,
"cache_path": str(larger.repo_path),
}
]
def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp_path):
shared = "blobs/shared-q4"
repo = _repo(
"Org/SharedBlobRepo",
[],
tmp_path / "models--Org--SharedBlobRepo",
revisions = [
SimpleNamespace(files = [_file("Q4_K_M.gguf", 5_000, blob_path = shared)]),
SimpleNamespace(files = [_file("Q4_K_M.gguf", 5_000, blob_path = shared)]),
],
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/SharedBlobRepo",
"size_bytes": 5_000,
"cache_path": str(repo.repo_path),
}
]
def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(
monkeypatch, tmp_path
):
mixed = _repo(
"Org/MixedRepo",
[
_file("Q4_K_M.gguf", 5_000),
_file("model.safetensors", 10_000),
],
tmp_path / "models--Org--MixedRepo",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [mixed])],
)
result = asyncio.run(models_route.list_cached_models(current_subject = "test-user"))
assert result["cached"] == []
def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(
monkeypatch, tmp_path
):
"""Mirror of the _skips_ test: the mixed repo should still surface in
cached-gguf so the picker can show it as a GGUF download."""
mixed = _repo(
"Org/MixedRepo",
[
_file("Q4_K_M.gguf", 5_000),
_file("model.safetensors", 10_000),
],
tmp_path / "models--Org--MixedRepo",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [mixed])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/MixedRepo",
"size_bytes": 5_000,
"cache_path": str(mixed.repo_path),
}
]
def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path):
"""A partial/interrupted GGUF download has ``size_on_disk = None``. The
route must treat the unknown bytes as zero instead of raising TypeError
out of ``sum()`` and wiping the entire response."""
partial = _repo(
"Org/PartialDownload",
[_file("Q4_K_M.gguf", None), _file("Q6_K.gguf", 5_000)],
tmp_path / "models--Org--PartialDownload",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [partial])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/PartialDownload",
"size_bytes": 5_000,
"cache_path": str(partial.repo_path),
}
]
def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(
monkeypatch, tmp_path
):
"""One repo raising during classification must not poison the response
for every other repo in the scan."""
class _ExplodingRepo:
repo_id = "Org/Broken"
repo_type = "model"
repo_path = tmp_path / "models--Org--Broken"
@property
def revisions(self):
raise RuntimeError("boom")
healthy = _repo(
"Org/Healthy",
[_file("Q4_K_M.gguf", 5_000)],
tmp_path / "models--Org--Healthy",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [_ExplodingRepo(), healthy])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/Healthy",
"size_bytes": 5_000,
"cache_path": str(healthy.repo_path),
}
]
def test_list_cached_gguf_skips_repo_with_only_mmproj_gguf(monkeypatch, tmp_path):
"""A repo whose only ``.gguf`` artifact is an mmproj vision adapter
must not be classified as a GGUF repo: the variant selector filters
mmproj out and the picker would otherwise show zero variants."""
mmproj_only = _repo(
"Org/MmprojOnly",
[
_file("mmproj-Q8_0.gguf", 5_000),
_file("model.safetensors", 10_000),
],
tmp_path / "models--Org--MmprojOnly",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [mmproj_only])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == []
def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp_path):
"""Mirror of the cached-gguf skip: a safetensors repo with an
auxiliary mmproj vision adapter must still surface in cached-models
so the user can load it as a normal model."""
mmproj_aux = _repo(
"Org/MmprojAux",
[
_file("mmproj-Q8_0.gguf", 5_000),
_file("model.safetensors", 10_000),
],
tmp_path / "models--Org--MmprojAux",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [mmproj_aux])],
)
result = asyncio.run(models_route.list_cached_models(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/MmprojAux",
"size_bytes": 15_000,
}
]
def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(
monkeypatch, tmp_path
):
"""A vision-capable GGUF repo (main weight + mmproj adapter) is still
a GGUF repo. The reported size is the main weight size; mmproj is
excluded from the GGUF-size accounting because it is filtered out at
classification time."""
vision_repo = _repo(
"Org/VisionGguf",
[
_file("Q4_K_M.gguf", 5_000),
_file("mmproj-Q8_0.gguf", 1_000),
],
tmp_path / "models--Org--VisionGguf",
)
monkeypatch.setattr(
models_route,
"_all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [vision_repo])],
)
result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/VisionGguf",
"size_bytes": 5_000,
"cache_path": str(vision_repo.repo_path),
}
]

91
studio/backend/tests/test_data_recipe_github_progress.py Normal file
View file

@ -0,0 +1,91 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from core.data_recipe.jobs.parse import apply_update, parse_log_message
from core.data_recipe.jobs.types import Job
from routes.data_recipe.validate import _GITHUB_VALIDATE_NOTE, validate
from models.data_recipe import RecipePayload
def test_github_page_log_updates_source_progress_without_cursor():
job = Job(job_id = "job-1")
job.source_progress_estimated_total = 200
update = parse_log_message(
"[unslothai/unsloth] issues page 2 (+15) cursor=abc123 remaining=2960"
)
assert update is not None
apply_update(job, update)
progress = job.source_progress
assert progress is not None
assert progress.source == "github"
assert progress.status == "fetching"
assert progress.repo == "unslothai/unsloth"
assert progress.resource == "issues"
assert progress.page == 2
assert progress.page_items == 15
assert progress.fetched_items == 15
assert progress.estimated_total == 200
assert progress.rate_remaining == 2960
assert progress.message is not None
assert "cursor" not in progress.message
assert "abc123" not in progress.message
def test_github_rate_limit_log_updates_source_progress():
job = Job(job_id = "job-1")
update = parse_log_message("Rate limit hit. Sleeping 123s until reset.")
assert update is not None
apply_update(job, update)
progress = job.source_progress
assert progress is not None
assert progress.status == "rate_limited"
assert progress.retry_after_sec == 123
assert "resume automatically" in (progress.message or "")
def test_github_real_sample_prs_and_trial_limit_are_parsed():
job = Job(job_id = "job-1")
for message in (
"[unslothai/unsloth] PRs page 4 (+25) cursor=abc123 remaining=4983",
"Trial limit reached for PRs (100)",
):
update = parse_log_message(message)
assert update is not None
apply_update(job, update)
progress = job.source_progress
assert progress is not None
assert progress.repo == "unslothai/unsloth"
assert progress.resource == "pulls"
assert progress.page == 4
assert progress.fetched_items == 25
assert progress.rate_remaining == 4983
assert progress.message == "GitHub pulls trial limit reached (100)."
def test_github_validate_skips_live_access_with_honest_note():
response = validate(
RecipePayload(
recipe = {
"seed_config": {
"source": {
"seed_type": "github_repo",
"repos": ["unslothai/unsloth"],
"item_types": ["issues"],
"limit": 1,
}
},
"columns": [{"column_type": "expression", "name": "x", "expr": "1"}],
}
)
)
assert response.valid is True
assert response.raw_detail == _GITHUB_VALIDATE_NOTE

598
studio/backend/tests/test_desktop_auth.py Normal file
View file

@ -0,0 +1,598 @@
import importlib.util
import asyncio
import hashlib
import json
import os
import platform
import secrets
import sqlite3
import subprocess
import sys
from pathlib import Path
from types import SimpleNamespace
import jwt
import pytest
from fastapi import APIRouter, FastAPI
from fastapi.security import HTTPAuthorizationCredentials
from fastapi.testclient import TestClient
from auth import storage
@pytest.fixture(autouse = True)
def isolated_auth_db(tmp_path, monkeypatch):
monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password")
monkeypatch.setattr(storage, "_bootstrap_password", None)
monkeypatch.setattr(storage, "_api_key_pbkdf2_salt_cache", None)
yield
def seed_user(*, must_change_password = False):
storage.create_initial_user(
username = storage.DEFAULT_ADMIN_USERNAME,
password = "human-password-123",
jwt_secret = secrets.token_urlsafe(64),
must_change_password = must_change_password,
)
def auth_client():
route_path = Path(__file__).resolve().parents[1] / "routes" / "auth.py"
spec = importlib.util.spec_from_file_location("_desktop_auth_route", route_path)
auth_route = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(auth_route)
app = FastAPI()
app.include_router(auth_route.router, prefix = "/api/auth")
return TestClient(app)
def data_recipe_jobs_module():
route_path = (
Path(__file__).resolve().parents[1] / "routes" / "data_recipe" / "jobs.py"
)
spec = importlib.util.spec_from_file_location(
"_desktop_data_recipe_jobs", route_path
)
jobs_route = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(jobs_route)
return jobs_route
def local_recipe():
return {
"model_providers": [{"name": "local", "is_local": True}],
"model_configs": [{"alias": "local-model", "provider": "local"}],
"columns": [{"column_type": "llm-text", "model_alias": "local-model"}],
}
def local_recipe_request(token):
return SimpleNamespace(
headers = {"authorization": f"Bearer {token}"},
app = SimpleNamespace(state = SimpleNamespace(server_port = 8888)),
scope = {},
base_url = "http://testserver/",
)
@pytest.fixture
def loaded_local_model(monkeypatch):
inference_module = SimpleNamespace(
get_llama_cpp_backend = lambda: SimpleNamespace(is_loaded = True),
)
monkeypatch.setitem(sys.modules, "routes.inference", inference_module)
def test_desktop_secret_round_trip_uses_real_admin_subject():
seed_user()
raw = storage.create_desktop_secret()
assert raw.startswith("desktop-")
assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME
assert storage.validate_desktop_secret(raw + "x") is None
def test_create_desktop_secret_rotates_old_secret():
seed_user()
old = storage.create_desktop_secret()
new = storage.create_desktop_secret()
assert old != new
assert storage.validate_desktop_secret(old) is None
assert storage.validate_desktop_secret(new) == storage.DEFAULT_ADMIN_USERNAME
def test_clear_desktop_secret_invalidates_secret():
seed_user()
raw = storage.create_desktop_secret()
storage.clear_desktop_secret()
assert storage.validate_desktop_secret(raw) is None
def test_ensure_default_admin_does_not_recreate_bootstrap_for_existing_admin():
seed_user()
created = storage.ensure_default_admin()
assert created is False
assert not storage._BOOTSTRAP_PW_PATH.exists()
def test_ensure_default_admin_loads_existing_bootstrap_after_restart(monkeypatch):
created = storage.ensure_default_admin()
bootstrap_pw = storage._BOOTSTRAP_PW_PATH.read_text().strip()
monkeypatch.setattr(storage, "_bootstrap_password", None)
created_again = storage.ensure_default_admin()
assert created is True
assert storage._BOOTSTRAP_PW_PATH.exists()
assert created_again is False
assert storage.get_bootstrap_password() == bootstrap_pw
def test_ensure_default_admin_does_not_generate_for_empty_existing_bootstrap():
seed_user()
storage._BOOTSTRAP_PW_PATH.write_text(" \n")
created = storage.ensure_default_admin()
assert created is False
assert storage._BOOTSTRAP_PW_PATH.read_text() == " \n"
assert storage.get_bootstrap_password() is None
def test_web_login_token_has_no_desktop_marker_and_keeps_password_gate():
seed_user(must_change_password = True)
client = auth_client()
response = client.post(
"/api/auth/login",
json = {
"username": storage.DEFAULT_ADMIN_USERNAME,
"password": "human-password-123",
},
)
assert response.status_code == 200
body = response.json()
assert body["must_change_password"] is True
payload = jwt.decode(
body["access_token"],
storage.get_jwt_secret(storage.DEFAULT_ADMIN_USERNAME),
algorithms = ["HS256"],
)
assert payload["sub"] == storage.DEFAULT_ADMIN_USERNAME
assert "desktop" not in payload
gated = client.post(
"/api/auth/api-keys",
headers = {"Authorization": f"Bearer {body['access_token']}"},
json = {"name": "web"},
)
assert gated.status_code == 403
def test_desktop_login_mints_admin_token_without_clearing_web_password_change():
seed_user(must_change_password = True)
raw = storage.create_desktop_secret()
client = auth_client()
response = client.post("/api/auth/desktop-login", json = {"secret": raw})
assert response.status_code == 200
body = response.json()
assert body["access_token"]
assert body["refresh_token"]
assert body["token_type"] == "bearer"
assert body["must_change_password"] is False
assert storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME) is True
payload = jwt.decode(
body["access_token"],
storage.get_jwt_secret(storage.DEFAULT_ADMIN_USERNAME),
algorithms = ["HS256"],
)
assert payload["sub"] == storage.DEFAULT_ADMIN_USERNAME
assert payload["desktop"] is True
def test_desktop_refresh_preserves_desktop_marker():
seed_user(must_change_password = True)
raw = storage.create_desktop_secret()
client = auth_client()
login_body = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()
response = client.post(
"/api/auth/refresh",
json = {"refresh_token": login_body["refresh_token"]},
)
assert response.status_code == 200
body = response.json()
assert body["must_change_password"] is False
payload = jwt.decode(
body["access_token"],
storage.get_jwt_secret(storage.DEFAULT_ADMIN_USERNAME),
algorithms = ["HS256"],
)
assert payload["sub"] == storage.DEFAULT_ADMIN_USERNAME
assert payload["desktop"] is True
def test_desktop_session_uses_real_admin_identity_for_api_keys():
seed_user(must_change_password = True)
raw = storage.create_desktop_secret()
client = auth_client()
token = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()[
"access_token"
]
response = client.post(
"/api/auth/api-keys",
headers = {"Authorization": f"Bearer {token}"},
json = {"name": "desktop"},
)
assert response.status_code == 200
rows = storage.list_api_keys(storage.DEFAULT_ADMIN_USERNAME)
assert [row["name"] for row in rows] == ["desktop"]
def test_local_recipe_token_authenticates_as_admin_for_desktop_user(loaded_local_model):
# _inject_local_providers mints an internal sk-unsloth-* API key (not a
# forwarded JWT). The unified API-key path validates as the real admin
# user regardless of whether the incoming session was desktop or web.
from auth.authentication import create_access_token, get_current_subject
seed_user(must_change_password = True)
jobs_route = data_recipe_jobs_module()
incoming_token = create_access_token(
subject = storage.DEFAULT_ADMIN_USERNAME,
desktop = True,
)
recipe = local_recipe()
jobs_route._inject_local_providers(recipe, local_recipe_request(incoming_token))
local_token = recipe["model_providers"][0]["api_key"]
assert local_token.startswith(storage.API_KEY_PREFIX)
credentials = HTTPAuthorizationCredentials(
scheme = "Bearer",
credentials = local_token,
)
assert (
asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
)
def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_model):
# Mirror of the desktop variant: API-key issuance is identical for web
# and desktop incoming tokens; auth via get_current_subject works the same.
from auth.authentication import create_access_token, get_current_subject
seed_user(must_change_password = False)
jobs_route = data_recipe_jobs_module()
incoming_token = create_access_token(subject = storage.DEFAULT_ADMIN_USERNAME)
recipe = local_recipe()
jobs_route._inject_local_providers(recipe, local_recipe_request(incoming_token))
local_token = recipe["model_providers"][0]["api_key"]
assert local_token.startswith(storage.API_KEY_PREFIX)
credentials = HTTPAuthorizationCredentials(
scheme = "Bearer",
credentials = local_token,
)
assert (
asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
)
def test_desktop_login_rejects_invalid_secret():
seed_user(must_change_password = False)
client = auth_client()
response = client.post(
"/api/auth/desktop-login",
json = {"secret": "desktop-invalid"},
)
assert response.status_code == 401
def test_write_desktop_secret_file_is_0600_on_unix(tmp_path):
from unsloth_cli.commands import studio as studio_cli
path = tmp_path / ".desktop_secret"
if platform.system() != "Windows":
path.write_text("old-secret")
os.chmod(path, 0o644)
studio_cli._write_auth_secret(path, "desktop-secret")
assert path.read_text() == "desktop-secret"
if platform.system() != "Windows":
assert oct(path.stat().st_mode & 0o777) == "0o600"
def test_reset_password_removes_desktop_secret_files(tmp_path, monkeypatch):
from typer.testing import CliRunner
from unsloth_cli.commands import studio as studio_cli
auth_dir = tmp_path / "auth"
auth_dir.mkdir()
(auth_dir / "auth.db").write_text("db")
(auth_dir / ".bootstrap_password").write_text("boot")
(auth_dir / ".desktop_secret").write_text("new")
monkeypatch.setattr(studio_cli, "STUDIO_HOME", tmp_path)
result = CliRunner().invoke(studio_cli.studio_app, ["reset-password"])
assert result.exit_code == 0
assert not (auth_dir / "auth.db").exists()
assert not (auth_dir / ".bootstrap_password").exists()
assert not (auth_dir / ".desktop_secret").exists()
def test_reset_password_removes_desktop_secret_files_without_db(tmp_path, monkeypatch):
from typer.testing import CliRunner
from unsloth_cli.commands import studio as studio_cli
auth_dir = tmp_path / "auth"
auth_dir.mkdir()
(auth_dir / ".desktop_secret").write_text("new")
monkeypatch.setattr(studio_cli, "STUDIO_HOME", tmp_path)
result = CliRunner().invoke(studio_cli.studio_app, ["reset-password"])
assert result.exit_code == 0
assert not (auth_dir / ".desktop_secret").exists()
def test_desktop_capabilities_json_reports_rollout_safe_flags():
from typer.testing import CliRunner
import unsloth_cli.commands.studio as studio_cli
result = CliRunner().invoke(
studio_cli.studio_app,
["desktop-capabilities", "--json"],
)
assert result.exit_code == 0
body = json.loads(result.output)
assert body["desktop_protocol_version"] == 1
assert body["supports_provision_desktop_auth"] is True
assert body["supports_api_only"] is True
assert isinstance(body["version"], str)
def test_health_response_reports_desktop_capability_fields(monkeypatch):
router_stub = SimpleNamespace(
auth_router = APIRouter(),
data_recipe_router = APIRouter(),
datasets_router = APIRouter(),
export_router = APIRouter(),
inference_router = APIRouter(),
inference_studio_router = APIRouter(),
models_router = APIRouter(),
training_history_router = APIRouter(),
training_router = APIRouter(),
)
monkeypatch.setitem(sys.modules, "routes", router_stub)
import studio.backend.main as backend_main
monkeypatch.setattr(backend_main._hw_module, "CHAT_ONLY", False)
body = asyncio.run(backend_main.health_check())
assert body["desktop_protocol_version"] == 1
assert body["supports_desktop_auth"] is True
def test_provision_desktop_auth_writes_secret_and_creates_db_without_backend_deps(
tmp_path,
monkeypatch,
):
auth_dir = tmp_path / "auth"
auth_dir.mkdir()
code = """
import builtins
import sys
from pathlib import Path
from typer.testing import CliRunner
studio_home = Path(sys.argv[1])
real_import = builtins.__import__
def guarded_import(name, *args, **kwargs):
blocked = ("auth", "fastapi", "structlog", "utils")
if name in blocked or name.startswith(("auth.", "utils.")):
raise ModuleNotFoundError(name)
return real_import(name, *args, **kwargs)
builtins.__import__ = guarded_import
from unsloth_cli.commands import studio as studio_cli
studio_cli.STUDIO_HOME = studio_home
result = CliRunner().invoke(studio_cli.studio_app, ["provision-desktop-auth"])
if result.exit_code != 0:
print(result.output)
if result.exception is not None:
raise result.exception
raise SystemExit(result.exit_code)
"""
result = subprocess.run(
[sys.executable, "-c", code, str(tmp_path)],
cwd = Path(__file__).resolve().parents[3],
env = {**os.environ, "PYTHONPATH": "."},
text = True,
capture_output = True,
)
assert result.returncode == 0, result.stderr + result.stdout
secret = (auth_dir / ".desktop_secret").read_text()
assert secret.startswith("desktop-")
conn = sqlite3.connect(auth_dir / "auth.db")
conn.row_factory = sqlite3.Row
try:
user = conn.execute(
"""
SELECT username, password_salt, password_hash, must_change_password
FROM auth_user
"""
).fetchone()
app_secrets = {
row["key"]: row["value"]
for row in conn.execute("SELECT key, value FROM app_secrets")
}
refresh_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")
}
finally:
conn.close()
bootstrap_password = (auth_dir / ".bootstrap_password").read_text().strip()
bootstrap_hash = hashlib.pbkdf2_hmac(
"sha256",
bootstrap_password.encode("utf-8"),
user["password_salt"].encode("utf-8"),
100_000,
).hex()
assert bootstrap_password
assert user["username"] == "unsloth"
assert user["must_change_password"] == 1
assert bootstrap_hash == user["password_hash"]
assert len(app_secrets["api_key_pbkdf2_salt"]) == 64
assert len(app_secrets["desktop_secret_hash"]) == 64
assert app_secrets["desktop_secret_created_at"]
assert "is_desktop" in refresh_columns
monkeypatch.setattr(storage, "DB_PATH", auth_dir / "auth.db")
monkeypatch.setattr(storage, "_api_key_pbkdf2_salt_cache", None)
assert storage.validate_desktop_secret(secret) == storage.DEFAULT_ADMIN_USERNAME
assert storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME) is True
def test_provision_desktop_auth_keeps_existing_admin_password(tmp_path, monkeypatch):
from typer.testing import CliRunner
from unsloth_cli.commands import studio as studio_cli
auth_dir = tmp_path / "auth"
auth_dir.mkdir()
monkeypatch.setattr(studio_cli, "STUDIO_HOME", tmp_path)
conn = sqlite3.connect(auth_dir / "auth.db")
try:
conn.execute(
"""
CREATE TABLE auth_user (
id INTEGER PRIMARY KEY,
username TEXT UNIQUE NOT NULL,
password_salt TEXT NOT NULL,
password_hash TEXT NOT NULL,
jwt_secret TEXT NOT NULL,
must_change_password INTEGER NOT NULL DEFAULT 0
)
"""
)
conn.execute(
"""
INSERT INTO auth_user (
username, password_salt, password_hash, jwt_secret, must_change_password
)
VALUES (?, ?, ?, ?, ?)
""",
("unsloth", "existing-salt", "existing-hash", "existing-jwt", 0),
)
conn.commit()
finally:
conn.close()
result = CliRunner().invoke(studio_cli.studio_app, ["provision-desktop-auth"])
assert result.exit_code == 0
assert not (auth_dir / ".bootstrap_password").exists()
conn = sqlite3.connect(auth_dir / "auth.db")
conn.row_factory = sqlite3.Row
try:
user = conn.execute(
"""
SELECT password_salt, password_hash, jwt_secret, must_change_password
FROM auth_user WHERE username = ?
""",
("unsloth",),
).fetchone()
finally:
conn.close()
assert dict(user) == {
"password_salt": "existing-salt",
"password_hash": "existing-hash",
"jwt_secret": "existing-jwt",
"must_change_password": 0,
}
def test_update_password_clears_desktop_secret():
seed_user()
raw = storage.create_desktop_secret()
assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME
changed = storage.update_password(
storage.DEFAULT_ADMIN_USERNAME, "new-admin-password"
)
assert changed is True
assert storage.validate_desktop_secret(raw) is None
def test_update_password_on_unknown_user_leaves_desktop_secret_intact():
seed_user()
raw = storage.create_desktop_secret()
changed = storage.update_password("not-a-user", "irrelevant")
assert changed is False
assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME
def test_desktop_auth_provision_has_bounded_timeout():
rs_path = (
Path(__file__).resolve().parents[3]
/ "studio"
/ "src-tauri"
/ "src"
/ "desktop_auth.rs"
)
src = rs_path.read_text()
start = src.index("async fn provision_desktop_auth(")
depth = 0
body_start = src.index("{", start)
body_end = None
for i in range(body_start, len(src)):
c = src[i]
if c == "{":
depth += 1
elif c == "}":
depth -= 1
if depth == 0:
body_end = i + 1
break
assert body_end is not None
body = src[start:body_end]
assert "tokio::time::timeout" in body
import re
m = re.search(r"Duration::from_secs\(\s*(\d+)\s*\)", body)
assert m is not None
seconds = int(m.group(1))
assert 5 <= seconds <= 120

179
studio/backend/tests/test_export_log_cursor.py Normal file
View file

@ -0,0 +1,179 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Regression tests for the export log ring-buffer cursor semantics.
Context: the live export log SSE stream has a race where the frontend
opens the SSE connection AFTER the POST that starts the export. Any
lines the worker subprocess emits during the gap between POST and SSE
connect get buffered with seqs 1..k, and then the SSE default cursor
`get_current_log_seq()` returns k -- so lines 1..k are forever
unreachable to that client.
Fix: `clear_logs()` snapshots the pre-run seq into `_run_start_seq`
(exposed via `get_run_start_seq()`), and `routes/export.py` defaults
the SSE cursor to that snapshot instead of the current seq. Every line
appended during the current run has seq strictly greater than the
snapshot, so the client sees the full run regardless of when it
connects.
These tests exercise the orchestrator-side contract only (no
subprocess, no FastAPI, no frontend). The routes-level integration
with get_run_start_seq() is a one-line edit covered by manual testing
and the frontend build.
"""
from __future__ import annotations
import sys
import types
from pathlib import Path
import pytest
# Backend root on sys.path so `from core.export.orchestrator import ...`
# and friends resolve without the studio app bootstrap.
_BACKEND_DIR = Path(__file__).resolve().parent.parent
if str(_BACKEND_DIR) not in sys.path:
sys.path.insert(0, str(_BACKEND_DIR))
# ExportOrchestrator imports structlog and a few heavy modules at the
# top of orchestrator.py. Stub the ones we don't need in these unit
# tests so the import succeeds on machines without the full studio
# venv.
_loggers_stub = types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
# structlog is only used for a module-level import; a bare stub is
# enough because we never call into it in these tests.
sys.modules.setdefault("structlog", types.ModuleType("structlog"))
# utils.paths.outputs_root is only called inside scan_checkpoints which
# we don't hit in these tests. Provide a stub module so the top-level
# import in orchestrator.py resolves.
_utils_pkg = types.ModuleType("utils")
_utils_pkg.__path__ = [] # mark as package
_utils_paths_stub = types.ModuleType("utils.paths")
_utils_paths_stub.outputs_root = lambda: Path("/tmp")
sys.modules.setdefault("utils", _utils_pkg)
sys.modules.setdefault("utils.paths", _utils_paths_stub)
@pytest.fixture
def orchestrator():
"""Fresh ExportOrchestrator with only the log-buffer state exercised."""
from core.export.orchestrator import ExportOrchestrator
return ExportOrchestrator()
def _append(orch, line: str, stream: str = "stdout") -> None:
"""Shortcut for simulating a worker log message."""
orch._append_log({"type": "log", "stream": stream, "line": line, "ts": 0.0})
# ---------------------------------------------------------------------------
# clear_logs() semantics
# ---------------------------------------------------------------------------
def test_run_start_seq_is_zero_before_any_logs(orchestrator) -> None:
"""A brand-new orchestrator must report run_start_seq == 0 so a
first SSE connection picks up every line from seq 1 onward."""
assert orchestrator.get_run_start_seq() == 0
def test_clear_logs_snapshots_current_seq(orchestrator) -> None:
"""clear_logs() must capture _log_seq BEFORE clearing the buffer,
so subsequent runs can anchor their SSE cursor at the snapshot."""
_append(orchestrator, "old run line 1")
_append(orchestrator, "old run line 2")
_append(orchestrator, "old run line 3")
assert orchestrator.get_current_log_seq() == 3
orchestrator.clear_logs()
assert orchestrator.get_run_start_seq() == 3
assert orchestrator.get_current_log_seq() == 3 # seq counter preserved
# ---------------------------------------------------------------------------
# Race regression: SSE connects AFTER lines have been emitted
# ---------------------------------------------------------------------------
def test_sse_default_cursor_catches_all_current_run_lines(orchestrator) -> None:
"""Simulate the POST-then-SSE race: worker starts emitting lines
immediately after clear_logs(), SSE connects several lines later.
Using get_run_start_seq() as the default cursor MUST return every
line emitted since clear_logs() ran.
Pre-fix, the SSE defaulted to get_current_log_seq() at connect
time, which would return the last-seen seq and miss lines N+1..M.
"""
# Previous run leaves some buffered lines.
_append(orchestrator, "previous run line A")
_append(orchestrator, "previous run line B")
# New run starts: orchestrator clears the buffer and snapshots seq.
orchestrator.clear_logs()
run_start = orchestrator.get_run_start_seq()
# Worker emits early lines BEFORE the SSE connects.
_append(orchestrator, "Importing Unsloth...")
_append(orchestrator, "Loading checkpoint: /foo/bar")
_append(orchestrator, "Starting export...")
# SSE connects now and asks "give me everything after the run
# start cursor".
entries, new_cursor = orchestrator.get_logs_since(run_start)
# All three early lines must be present. Pre-fix this was [].
lines = [e["line"] for e in entries]
assert lines == [
"Importing Unsloth...",
"Loading checkpoint: /foo/bar",
"Starting export...",
]
assert new_cursor == entries[-1]["seq"]
def test_sse_default_cursor_excludes_previous_run(orchestrator) -> None:
"""After clear_logs(), lines from the PREVIOUS run must not leak
into the new run's SSE stream. Pre-fix this worked correctly
(clear_logs cleared the deque); the fix must preserve it.
"""
_append(orchestrator, "previous run line 1")
_append(orchestrator, "previous run line 2")
_append(orchestrator, "previous run line 3")
assert orchestrator.get_current_log_seq() == 3
orchestrator.clear_logs()
run_start = orchestrator.get_run_start_seq()
_append(orchestrator, "new run line")
entries, _ = orchestrator.get_logs_since(run_start)
assert [e["line"] for e in entries] == ["new run line"]
def test_clear_logs_twice_advances_run_start(orchestrator) -> None:
"""Back-to-back clear_logs() calls (e.g. cleanup -> load ->
export in the same dialog session) must each re-anchor run_start
at the current seq, so successive runs each start with a fresh
low-water mark."""
_append(orchestrator, "run 1 line a")
_append(orchestrator, "run 1 line b")
orchestrator.clear_logs()
assert orchestrator.get_run_start_seq() == 2
_append(orchestrator, "run 2 line a")
_append(orchestrator, "run 2 line b")
_append(orchestrator, "run 2 line c")
orchestrator.clear_logs()
assert orchestrator.get_run_start_seq() == 5

326
studio/backend/tests/test_gpu_selection.py
View file

@ -746,7 +746,15 @@ class TestRouteErrors(unittest.TestCase):
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route.load_model(request, current_subject = "test-user")
inference_route.load_model(
request,
SimpleNamespace(
app = SimpleNamespace(
state = SimpleNamespace(llama_parallel_slots = 1),
),
),
current_subject = "test-user",
)
)
self.assertEqual(exc_info.exception.status_code, 400)
@ -886,7 +894,15 @@ class TestRouteErrors(unittest.TestCase):
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route.load_model(request, current_subject = "test-user")
inference_route.load_model(
request,
SimpleNamespace(
app = SimpleNamespace(
state = SimpleNamespace(llama_parallel_slots = 1),
),
),
current_subject = "test-user",
)
)
self.assertEqual(exc_info.exception.status_code, 400)
@ -942,7 +958,15 @@ class TestRouteErrors(unittest.TestCase):
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route.load_model(request, current_subject = "test-user")
inference_route.load_model(
request,
SimpleNamespace(
app = SimpleNamespace(
state = SimpleNamespace(llama_parallel_slots = 1),
),
),
current_subject = "test-user",
)
)
self.assertEqual(exc_info.exception.status_code, 400)
@ -1025,6 +1049,182 @@ class TestMinGpuVram(unittest.TestCase):
class TestPerGpuFitGuardAllCounts(unittest.TestCase):
def test_training_estimate_resolves_attention_without_raising(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (8 * (1024**3), "config"),
),
patch(
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
return_value = "unsloth/test",
),
patch(
"utils.hardware.hardware._load_config_for_gpu_estimate",
return_value = SimpleNamespace(
hidden_size = 4096,
num_hidden_layers = 32,
num_attention_heads = 32,
num_key_value_heads = 8,
intermediate_size = 14336,
vocab_size = 128256,
tie_word_embeddings = False,
),
),
patch(
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
return_value = "eager",
),
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
):
_, metadata = estimate_required_model_memory_gb(
"unsloth/test",
training_type = "LoRA/QLoRA",
load_in_4bit = True,
)
self.assertEqual(metadata.get("estimation_mode"), "detailed")
self.assertEqual(metadata.get("attention_implementation"), "eager")
def test_training_estimate_falls_back_when_attention_resolution_fails(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (8 * (1024**3), "config"),
),
patch(
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
return_value = "unsloth/test",
),
patch(
"utils.hardware.hardware._load_config_for_gpu_estimate",
return_value = SimpleNamespace(
hidden_size = 4096,
num_hidden_layers = 32,
num_attention_heads = 32,
num_key_value_heads = 8,
intermediate_size = 14336,
vocab_size = 128256,
tie_word_embeddings = False,
),
),
patch(
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
side_effect = RuntimeError("attention unavailable"),
),
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
):
_, metadata = estimate_required_model_memory_gb(
"unsloth/test",
training_type = "LoRA/QLoRA",
load_in_4bit = True,
)
self.assertEqual(metadata.get("estimation_mode"), "detailed")
self.assertEqual(
metadata.get("attention_implementation"),
"eager",
)
def test_attention_resolver_does_not_mutate_loaded_config(self):
from utils.hardware import hardware as hardware_module
config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
)
def _stub_resolver(model_class, cfg):
cfg._attn_implementation = "eager"
return "eager"
with patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
):
hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertFalse(hasattr(config, "_attn_implementation"))
def test_attention_resolver_handles_missing_model_mapping(self):
from utils.hardware import hardware as hardware_module
config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
)
captured = {}
def _stub_resolver(model_class, cfg):
captured["model_class"] = model_class
return "eager"
from transformers import AutoModel, AutoModelForCausalLM
with (
patch.object(AutoModelForCausalLM, "_model_mapping", new = None),
patch.object(AutoModel, "_model_mapping", new = None),
patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
),
):
result = hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertEqual(result, "eager")
self.assertIsNone(captured["model_class"])
def test_attention_resolver_does_not_mutate_nested_text_config(self):
from utils.hardware import hardware as hardware_module
text_config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
)
config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
text_config = text_config,
)
def _stub_resolver(model_class, cfg):
cfg._attn_implementation = "eager"
inner = getattr(cfg, "text_config", None)
if inner is not None:
inner._attn_implementation = "eager"
return "eager"
with patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
):
hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertFalse(hasattr(config, "_attn_implementation"))
self.assertFalse(hasattr(text_config, "_attn_implementation"))
def test_min_per_gpu_generated_for_all_visible_counts(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
@ -1101,3 +1301,123 @@ class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase):
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU):
with self.assertRaisesRegex(ValueError, "only supported on CUDA"):
prepare_gpu_selection([0], model_name = "unsloth/test")
class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase):
def _run(
self,
model_path,
*,
config_bytes,
local_bytes,
safetensors_params = None,
config = object(),
):
from utils.hardware import hardware as hardware_module
with (
patch.object(
hardware_module,
"_resolve_model_identifier_for_gpu_estimate",
return_value = model_path,
),
patch.object(
hardware_module,
"_get_hf_safetensors_total_params",
return_value = safetensors_params,
),
patch.object(
hardware_module,
"_load_config_for_gpu_estimate",
return_value = config,
),
patch.object(
hardware_module,
"_estimate_fp16_model_size_bytes_from_config",
return_value = config_bytes,
),
patch.object(
hardware_module,
"_get_local_weight_size_bytes",
return_value = local_bytes,
),
):
return hardware_module.estimate_fp16_model_size_bytes(model_path)
def test_local_weight_bytes_preferred_when_larger_than_config(self):
bytes_, src = self._run(
"/local/vlm",
config_bytes = 2 * (1 << 30),
local_bytes = 20 * (1 << 30),
)
self.assertEqual(bytes_, 20 * (1 << 30))
self.assertEqual(src, "weight_bytes")
def test_config_bytes_preferred_when_larger_than_local(self):
bytes_, src = self._run(
"/local/text-only",
config_bytes = 20 * (1 << 30),
local_bytes = 2 * (1 << 30),
)
self.assertEqual(bytes_, 20 * (1 << 30))
self.assertEqual(src, "config")
def test_config_bytes_returned_when_no_local_weights(self):
bytes_, src = self._run(
"/local/no-weights",
config_bytes = 5 * (1 << 30),
local_bytes = None,
)
self.assertEqual(bytes_, 5 * (1 << 30))
self.assertEqual(src, "config")
def test_local_bytes_returned_when_config_resolution_fails(self):
bytes_, src = self._run(
"/local/no-config",
config_bytes = None,
local_bytes = 7 * (1 << 30),
config = None,
)
self.assertEqual(bytes_, 7 * (1 << 30))
self.assertEqual(src, "weight_bytes")
def test_equal_local_and_config_keeps_config_label(self):
# why: tie-breaker is "local must be strictly larger" so an exact
# match keeps the config-derived path.
same = 8 * (1 << 30)
bytes_, src = self._run(
"/local/equal",
config_bytes = same,
local_bytes = same,
)
self.assertEqual(bytes_, same)
self.assertEqual(src, "config")
def test_remote_safetensors_path_unaffected_by_local_weights(self):
from utils.hardware import hardware as hardware_module
with (
patch.object(
hardware_module,
"_resolve_model_identifier_for_gpu_estimate",
return_value = "owner/repo",
),
patch.object(
hardware_module,
"_get_hf_safetensors_total_params",
return_value = 1_000_000_000,
),
patch.object(
hardware_module,
"_load_config_for_gpu_estimate",
) as mock_load,
patch.object(
hardware_module,
"_get_local_weight_size_bytes",
) as mock_local,
):
bytes_, src = hardware_module.estimate_fp16_model_size_bytes("owner/repo")
self.assertEqual(bytes_, 2 * 1_000_000_000)
self.assertEqual(src, "safetensors")
mock_load.assert_not_called()
mock_local.assert_not_called()

98
studio/backend/tests/test_host_defaults.py Normal file
View file

@ -0,0 +1,98 @@
# 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 that Unsloth Studio defaults to 127.0.0.1 (loopback) not 0.0.0.0.
Uses AST parsing to inspect source-level defaults without requiring the
full studio venv (run.py has heavy dependencies like structlog/uvicorn).
"""
import ast
from pathlib import Path
_RUN_PY = Path(__file__).resolve().parent.parent / "run.py"
def _parse_function_param_defaults(source: str, func_name: str) -> dict:
"""Return {param_name: default_value} for a named function in *source*.
Only handles ast.Constant defaults (strings, ints, bools).
"""
tree = ast.parse(source)
for node in ast.walk(tree):
if (
isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
and node.name == func_name
):
result = {}
all_args = node.args.args
defaults = node.args.defaults
# Defaults are right-aligned against the args list
offset = len(all_args) - len(defaults)
for i, default in enumerate(defaults):
arg_name = all_args[offset + i].arg
if isinstance(default, ast.Constant):
result[arg_name] = default.value
return result
return {}
def _parse_argparse_add_argument_default(source: str, option_name: str):
"""Return the 'default' kwarg value for add_argument(option_name, ...) in *source*.
Walks the entire module so the call can live in __main__ or in a helper
function — only handles ast.Constant defaults.
"""
tree = ast.parse(source)
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
if not (isinstance(func, ast.Attribute) and func.attr == "add_argument"):
continue
if not node.args:
continue
first_arg = node.args[0]
if not (isinstance(first_arg, ast.Constant) and first_arg.value == option_name):
continue
for kw in node.keywords:
if kw.arg == "default" and isinstance(kw.value, ast.Constant):
return kw.value.value
return None
def test_run_server_default_host_is_loopback():
"""run_server() parameter default for 'host' must be 127.0.0.1, not 0.0.0.0.
Binding to 0.0.0.0 by default exposes the service on all network
interfaces, contradicting the documented "privacy first / 100% local"
guarantee. Loopback (127.0.0.1) is the least-permissive default;
users who need network access can pass -H 0.0.0.0 explicitly.
"""
source = _RUN_PY.read_text()
defaults = _parse_function_param_defaults(source, "run_server")
assert (
"host" in defaults
), "run_server() must have a 'host' parameter with a default"
host_default = defaults["host"]
assert host_default == "127.0.0.1", (
f"run_server() host default must be '127.0.0.1' (loopback) "
f"but got '{host_default}'. Binding to '{host_default}' by default "
f"exposes the service beyond localhost."
)
def test_argparse_default_host_is_loopback():
"""argparse --host add_argument default must be 127.0.0.1.
When run.py is invoked directly (python run.py), the argparse default
should match the function default so direct execution is equally safe.
"""
source = _RUN_PY.read_text()
host_default = _parse_argparse_add_argument_default(source, "--host")
assert (
host_default is not None
), "Could not find add_argument('--host', ...) in run.py"
assert (
host_default == "127.0.0.1"
), f"run.py argparse --host default must be '127.0.0.1', got '{host_default}'"

1340
studio/backend/tests/test_kv_cache_estimation.py
View file

File diff suppressed because it is too large Load diff

243
studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py Normal file
View file

@ -0,0 +1,243 @@
# 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 the cache-aware disk-space preflight in
``LlamaCppBackend.load_model``.
The preflight used to compare the repo's total GGUF download size against
free disk without accounting for bytes already present in the Hugging
Face cache. That made re-loading a cached large model (e.g.
``unsloth/MiniMax-M2.7-GGUF`` at 131 GB) fail cold whenever free disk was
below the full weight footprint, even though nothing needed
downloading.
These tests exercise the preflight arithmetic in isolation by driving
``get_paths_info`` and ``try_to_load_from_cache`` through ``mock.patch``.
No network, GPU, or subprocess use.
Cross-platform: Linux, macOS, Windows, WSL.
"""
from __future__ import annotations
import sys
import tempfile
import types as _types
from pathlib import Path
from unittest.mock import patch
import pytest
# ---------------------------------------------------------------------------
# Stub heavy / unavailable external dependencies before importing the
# module under test. Same pattern as test_kv_cache_estimation.py.
# ---------------------------------------------------------------------------
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# loggers
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
# structlog
_structlog_stub = _types.ModuleType("structlog")
sys.modules.setdefault("structlog", _structlog_stub)
# httpx
_httpx_stub = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
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)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
GIB = 1024**3
class _FakePathInfo:
"""Mimics huggingface_hub's RepoFile-ish return type from get_paths_info."""
def __init__(self, path: str, size: int):
self.path = path
self.size = size
def _preflight(
repo_files,
cached_files,
free_bytes,
hf_repo = "unsloth/Example-GGUF",
hf_token = None,
):
"""Run the preflight arithmetic as written in llama_cpp.py and return
the decision outcome as a dict.
``repo_files``: list of (filename, remote_bytes).
``cached_files``: dict {filename: on_disk_bytes} for files already in cache.
``free_bytes``: value returned by shutil.disk_usage(cache_dir).free.
"""
import os
import shutil
path_infos = [_FakePathInfo(name, size) for name, size in repo_files]
with tempfile.TemporaryDirectory() as tmp:
# Create SPARSE files for the cached ones so os.path.exists /
# os.path.getsize pass without actually allocating bytes on disk.
# This is critical when simulating multi-GB models.
cache_paths = {}
for name, sz in cached_files.items():
p = Path(tmp) / name.replace("/", "_")
with open(p, "wb") as fh:
if sz > 0:
fh.truncate(sz) # sparse allocation: no data blocks written
cache_paths[name] = str(p)
def fake_try_to_load_from_cache(repo_id, filename):
return cache_paths.get(filename)
# Mirror the same variable names and control flow as the real code
# so behavioral drift is caught immediately.
total_bytes = sum((p.size or 0) for p in path_infos)
already_cached_bytes = 0
for p in path_infos:
if not p.size:
continue
cached_path = fake_try_to_load_from_cache(hf_repo, p.path)
if isinstance(cached_path, str) and os.path.exists(cached_path):
try:
on_disk = os.path.getsize(cached_path)
except OSError:
on_disk = 0
if on_disk >= p.size:
already_cached_bytes += p.size
total_download_bytes = max(0, total_bytes - already_cached_bytes)
needed_download = total_download_bytes > free_bytes
return {
"total_bytes": total_bytes,
"already_cached_bytes": already_cached_bytes,
"total_download_bytes": total_download_bytes,
"would_raise_disk_error": (needed_download and total_download_bytes > 0),
}
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestCacheAwarePreflight:
def test_fully_cached_model_does_not_require_disk(self):
"""The MiniMax case: 131 GB weights cached, only 36 GB free.
Preflight must not raise."""
shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)]
cached = {name: size for name, size in shards}
out = _preflight(
repo_files = shards,
cached_files = cached,
free_bytes = 36 * GIB,
)
assert out["total_download_bytes"] == 0
assert out["already_cached_bytes"] == 140 * GIB
assert out["would_raise_disk_error"] is False
def test_partial_cache_only_counts_remaining_bytes(self):
"""Two of four shards cached: preflight against remaining 70 GB."""
shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)]
cached = {
shards[0][0]: shards[0][1],
shards[1][0]: shards[1][1],
}
out = _preflight(
repo_files = shards,
cached_files = cached,
free_bytes = 80 * GIB,
)
assert out["already_cached_bytes"] == 70 * GIB
assert out["total_download_bytes"] == 70 * GIB
assert out["would_raise_disk_error"] is False
def test_partial_cache_insufficient_disk_for_rest_still_raises(self):
"""Two of four shards cached; remaining 70 GB still bigger than
free disk -> preflight correctly wants to raise."""
shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)]
cached = {
shards[0][0]: shards[0][1],
shards[1][0]: shards[1][1],
}
out = _preflight(
repo_files = shards,
cached_files = cached,
free_bytes = 50 * GIB,
)
assert out["total_download_bytes"] == 70 * GIB
assert out["would_raise_disk_error"] is True
def test_nothing_cached_preserves_existing_behavior(self):
"""Cold-cache path still compares full download vs free disk."""
shards = [("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)]
out = _preflight(
repo_files = shards,
cached_files = {},
free_bytes = 50 * GIB,
)
assert out["already_cached_bytes"] == 0
assert out["total_download_bytes"] == 40 * GIB
assert out["would_raise_disk_error"] is False
def test_incomplete_cached_blob_is_not_credited(self):
"""A partial file on disk (e.g. interrupted download) is not
counted as cached -- we still require bytes for it."""
shards = [("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)]
partial = {"UD-Q4_K_XL/shard-0.gguf": 10 * GIB}
out = _preflight(
repo_files = shards,
cached_files = partial,
free_bytes = 50 * GIB,
)
assert out["already_cached_bytes"] == 0
assert out["total_download_bytes"] == 40 * GIB
assert out["would_raise_disk_error"] is False
def test_zero_size_path_infos_do_not_crash(self):
"""A path_info with size=0 should not be credited or break the
arithmetic."""
shards = [("mmproj.gguf", 0), ("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)]
out = _preflight(
repo_files = shards,
cached_files = {},
free_bytes = 50 * GIB,
)
assert out["already_cached_bytes"] == 0
assert out["total_bytes"] == 40 * GIB

393
studio/backend/tests/test_llama_cpp_context_fit.py Normal file
View file

@ -0,0 +1,393 @@
# 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 the GGUF load-time context auto-fit decision.
Guards two regressions in ``LlamaCppBackend.load_model``:
1. **Auto mode on weights-exceed-VRAM** (``n_ctx == 0``): when the model
weights alone exceed 90% of every GPU subset's free memory, the
auto-pick loop used to exit without matching, leaving
``effective_ctx`` at the model's native context (e.g. 196608 for
MiniMax-M2.7). The intended default per Studio's UI spec is 4096 so
the slider lands on a usable value; the user can still drag higher
and trigger ``--fit on`` with a warning.
2. **Explicit ctx silently shrunk when KV overflows**: with fittable
weights but a requested ctx whose KV cache pushes total memory over
90% of VRAM, the old code binary-searched a smaller ctx and emitted
``-c <capped> -ngl -1`` without informing the caller. The UI had
already surfaced its "might be slower" warning and expects the user's
explicit ctx to be honored with ``--fit on`` flexing ``-ngl`` instead.
Tests avoid GPU probing, subprocess spawning, and GGUF I/O by driving the
post-metadata decision block directly against a stubbed instance.
Requires no GPU, network, or external libraries beyond pytest.
Cross-platform: Linux, macOS, Windows, WSL.
"""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# Stub heavy / unavailable external dependencies before importing the
# module under test. Same pattern as test_kv_cache_estimation.py.
# ---------------------------------------------------------------------------
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# loggers
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
# structlog
_structlog_stub = _types.ModuleType("structlog")
sys.modules.setdefault("structlog", _structlog_stub)
# httpx
_httpx_stub = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
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 core.inference.llama_cpp import LlamaCppBackend
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
GIB = 1024**3
FALLBACK_CTX = 4096
def _make_backend(
native_ctx = 131072,
n_layers = 80,
n_kv_heads = 8,
n_heads = 64,
kv_key_length = 128,
kv_value_length = 128,
):
"""Create a LlamaCppBackend instance with GGUF metadata fields set and
the helpers used by the decision block stubbed out."""
inst = LlamaCppBackend.__new__(LlamaCppBackend)
inst._context_length = native_ctx
inst._n_layers = n_layers
inst._n_kv_heads = n_kv_heads
inst._n_heads = n_heads
inst._embedding_length = 8192
inst._kv_key_length = kv_key_length
inst._kv_value_length = kv_value_length
inst._kv_lora_rank = None
inst._sliding_window = None
inst._sliding_window_pattern = None
inst._ssm_inner_size = None
inst._full_attention_interval = None
inst._key_length_mla = None
inst._n_kv_heads_by_layer = None
inst._kv_key_length_swa = None
inst._kv_value_length_swa = None
return inst
def _drive(
n_ctx,
model_gib,
gpus,
native_ctx = 131072,
kv_per_token_bytes = 325_000,
can_estimate_kv = True,
):
"""Drive the post-metadata portion of load_model with stubbed inputs.
Mirrors the decision block at llama_cpp.py:1137-1296 so we can assert
the command that would be built, without subprocesses or GPU probes.
"""
inst = _make_backend(native_ctx = native_ctx)
model_size = int(model_gib * GIB)
cache_type_kv = None
def fake_estimate(n_ctx_, _type = None, **_kwargs):
return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes
inst._estimate_kv_cache_bytes = fake_estimate
inst._can_estimate_kv = lambda: can_estimate_kv
context_length = inst._context_length
effective_ctx = n_ctx if n_ctx > 0 else (context_length or 0)
max_available_ctx = context_length or effective_ctx
if n_ctx > 0:
effective_ctx = n_ctx
elif context_length is not None:
effective_ctx = context_length
else:
effective_ctx = 0
original_ctx = effective_ctx
max_available_ctx = context_length or effective_ctx
gpu_indices, use_fit = None, True
explicit_ctx = n_ctx > 0
if gpus and inst._can_estimate_kv() and effective_ctx > 0:
native_ctx_for_cap = context_length or effective_ctx
if native_ctx_for_cap > 0:
ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True)
best_cap = 0
for n_gpus in range(1, len(ranked_for_cap) + 1):
subset = ranked_for_cap[:n_gpus]
pool_mib = sum(free for _, free in subset)
capped = inst._fit_context_to_vram(
native_ctx_for_cap,
pool_mib,
model_size,
cache_type_kv,
)
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * 0.90:
best_cap = max(best_cap, capped)
if best_cap > 0:
max_available_ctx = best_cap
if explicit_ctx:
requested_total = model_size + inst._estimate_kv_cache_bytes(
effective_ctx, cache_type_kv
)
gpu_indices, use_fit = inst._select_gpus(requested_total, gpus)
else:
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
matched = False
for n_gpus in range(1, len(ranked) + 1):
subset = ranked[:n_gpus]
pool_mib = sum(free for _, free in subset)
capped = inst._fit_context_to_vram(
effective_ctx,
pool_mib,
model_size,
cache_type_kv,
)
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * 0.90:
effective_ctx = capped
gpu_indices = sorted(idx for idx, _ in subset)
use_fit = False
matched = True
break
if not matched:
effective_ctx = min(FALLBACK_CTX, effective_ctx)
elif gpus:
gpu_indices, use_fit = inst._select_gpus(model_size, gpus)
if use_fit and not explicit_ctx:
effective_ctx = (
min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX
)
return {
"c_arg": effective_ctx if effective_ctx > 0 else 0,
"use_fit": use_fit,
"gpu_indices": gpu_indices,
"max_available_ctx": max_available_ctx,
"original_ctx": original_ctx,
}
# ---------------------------------------------------------------------------
# Auto mode, model weights exceed VRAM (Bug A guard)
# ---------------------------------------------------------------------------
class TestAutoModeWeightsExceedVRAM:
"""``n_ctx == 0`` on a model whose weights don't fit anywhere."""
def test_minimax_like_single_gpu(self):
plan = _drive(
n_ctx = 0,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
assert plan["gpu_indices"] is None
# UI slider ceiling stays at native: user can still drag higher
# and get the "might be slower" path.
assert plan["max_available_ctx"] == 196608
def test_multi_gpu_all_subsets_fail(self):
plan = _drive(
n_ctx = 0,
model_gib = 400,
gpus = [(0, 80_000), (1, 80_000), (2, 80_000), (3, 80_000)],
native_ctx = 131072,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
assert plan["gpu_indices"] is None
def test_no_kv_metadata_auto(self):
"""File-size-only fallback path also defaults to 4096."""
plan = _drive(
n_ctx = 0,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
can_estimate_kv = False,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
# ---------------------------------------------------------------------------
# Explicit ctx, KV overflows fittable weights (Bug B guard)
# ---------------------------------------------------------------------------
class TestExplicitCtxRespectsUser:
"""``n_ctx > 0`` must never be silently shrunk."""
def test_fittable_weights_oversized_kv(self):
# 8 GB weights + 131k ctx KV on 24 GB VRAM.
# Budget = 21.6 GB, KV at 131k >> 13.6 GB remaining, so
# _select_gpus flips use_fit=True.
plan = _drive(
n_ctx = 131072,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
)
assert plan["c_arg"] == 131072
assert plan["use_fit"] is True
assert plan["gpu_indices"] is None
def test_explicit_that_fits_uses_ngl(self):
plan = _drive(
n_ctx = 8192,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
)
assert plan["c_arg"] == 8192
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
def test_explicit_on_weights_exceed_vram(self):
# User drags the slider to 32k on a too-big model: honored.
plan = _drive(
n_ctx = 32768,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
)
assert plan["c_arg"] == 32768
assert plan["use_fit"] is True
def test_explicit_at_fallback_on_too_big(self):
plan = _drive(
n_ctx = FALLBACK_CTX,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
def test_explicit_below_floor_honored(self):
# 2048 is below --fit-ctx default; still honored since user set it.
plan = _drive(
n_ctx = 2048,
model_gib = 8,
gpus = [(0, 24_000)],
)
assert plan["c_arg"] == 2048
# ---------------------------------------------------------------------------
# Non-regression: fittable + auto still auto-picks largest fitting ctx
# ---------------------------------------------------------------------------
class TestFittableAutoPickRegressions:
def test_small_model_one_gpu(self):
plan = _drive(
n_ctx = 0,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
kv_per_token_bytes = 8192,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
assert plan["c_arg"] > FALLBACK_CTX
def test_medium_model_needs_multi_gpu(self):
plan = _drive(
n_ctx = 0,
model_gib = 60,
gpus = [(0, 40_000), (1, 40_000)],
native_ctx = 131072,
kv_per_token_bytes = 8192,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0, 1]
def test_no_kv_metadata_fittable_auto(self):
plan = _drive(
n_ctx = 0,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
can_estimate_kv = False,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
# ---------------------------------------------------------------------------
# Platform-agnostic input shape
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("platform_tag", ["linux", "windows", "mac", "rocm"])
def test_identical_decision_across_platforms(platform_tag):
"""The decision function takes ``[(gpu_idx, free_mib), ...]`` regardless
of how upstream (nvidia-smi / nvidia-smi.exe / Metal / rocm-smi) produced
it. Identical inputs must yield identical plans."""
plan_a = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
plan_b = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
assert plan_a == plan_b, platform_tag

258
studio/backend/tests/test_llama_cpp_load_progress.py Normal file
View file

@ -0,0 +1,258 @@
# 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 ``LlamaCppBackend.load_progress()``.
The chat settings flow and the training overlay both show a generic
"Starting model..." spinner during the window after a GGUF download
finishes and before llama-server reports healthy. For small models
that window is a second or two and nobody notices. For large MoE GGUFs
(MiniMax-M2.7, Qwen3.5-397B-A17B, etc.) the llama-server process spends
minutes in kernel state D, paging tens or hundreds of GB of shards
into the page cache. The UI has no way to show a real progress bar,
rate, or ETA during that window.
``load_progress()`` samples ``/proc/<pid>/status VmRSS`` (what the
kernel has actually paged in) against the total shard file size on
disk, so the frontend can render a real bar plus rate/ETA. This
module pins that contract:
* returns ``None`` when no load is in flight
* returns ``{"phase": "mmap", ...}`` while the subprocess is alive
but ``_healthy`` is False
* returns ``{"phase": "ready", ...}`` once ``_healthy`` flips
* ``bytes_total`` is derived from the resolved on-disk path
(which the paired fix assigns to ``self._gguf_path`` on both the
local-GGUF and HF-download code paths)
* ``bytes_loaded`` is VmRSS in bytes, capped by total, rounded
* ``fraction`` is clamped to 0..1 and rounded to 4 decimal places
Linux-only via ``/proc``. On platforms without ``/proc`` the method
returns ``None`` instead of raising.
Cross-platform test: skips cleanly on macOS / Windows if ``/proc`` is
not available.
"""
from __future__ import annotations
import os
import sys
import tempfile
import types as _types
from pathlib import Path
from unittest.mock import patch
import pytest
# ---------------------------------------------------------------------------
# Stub heavy / unavailable external dependencies before importing the
# module under test. Same pattern as test_kv_cache_estimation.py.
# ---------------------------------------------------------------------------
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_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)
_httpx_stub = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
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 core.inference.llama_cpp import LlamaCppBackend
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_instance():
inst = LlamaCppBackend.__new__(LlamaCppBackend)
inst._process = None
inst._gguf_path = None
inst._healthy = False
return inst
class _FakeProc:
"""Minimal stand-in for subprocess.Popen that just carries a pid."""
def __init__(self, pid: int):
self.pid = pid
def _write_sparse_file(path: Path, size_bytes: int) -> None:
"""Create a sparse file of the given size without allocating blocks."""
with open(path, "wb") as fh:
if size_bytes > 0:
fh.truncate(size_bytes)
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestLoadProgressEmptyStates:
def test_returns_none_when_no_process(self):
inst = _make_instance()
assert inst.load_progress() is None
def test_returns_none_when_process_has_no_pid(self):
inst = _make_instance()
inst._process = _FakeProc(pid = None) # type: ignore[arg-type]
assert inst.load_progress() is None
class TestLoadProgressSingleShard:
def test_mmap_phase_for_alive_but_unhealthy(self, tmp_path):
"""VmRSS below total -> phase='mmap', fraction reflects progress."""
gguf = tmp_path / "model.gguf"
_write_sparse_file(gguf, 40 * 1024**3) # 40 GB
inst = _make_instance()
inst._process = _FakeProc(pid = os.getpid()) # use our own pid
inst._gguf_path = str(gguf)
inst._healthy = False
# Patch /proc read to claim 10 GB RSS.
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
return io.StringIO(f"Name:\ttest\nVmRSS:\t{10 * 1024 ** 2}\tkB\n")
return open(path, *args, **kwargs) # fall through
with patch("builtins.open", side_effect = fake_open):
out = inst.load_progress()
assert out is not None
assert out["phase"] == "mmap"
assert out["bytes_total"] == 40 * 1024**3
assert out["bytes_loaded"] == 10 * 1024**3
assert 0.24 < out["fraction"] < 0.26 # ~25%
def test_ready_phase_when_healthy(self, tmp_path):
gguf = tmp_path / "model.gguf"
_write_sparse_file(gguf, 8 * 1024**3)
inst = _make_instance()
inst._process = _FakeProc(pid = os.getpid())
inst._gguf_path = str(gguf)
inst._healthy = True
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
return io.StringIO(f"VmRSS:\t{8 * 1024 ** 2}\tkB\n")
return open(path, *args, **kwargs)
with patch("builtins.open", side_effect = fake_open):
out = inst.load_progress()
assert out is not None
assert out["phase"] == "ready"
assert out["bytes_total"] == 8 * 1024**3
assert out["bytes_loaded"] == 8 * 1024**3
assert out["fraction"] == 1.0
class TestLoadProgressMultiShard:
"""Shard-aware total: for ``*-00001-of-00004.gguf`` primaries the
method sums sibling files with the same prefix."""
def test_sharded_total_aggregates_siblings(self, tmp_path):
for i in range(1, 5):
_write_sparse_file(
tmp_path / f"model-{i:05d}-of-00004.gguf",
size_bytes = 20 * 1024**3,
)
# Drop an unrelated .gguf in the same folder -- must not be counted.
_write_sparse_file(tmp_path / "mmproj-BF16.gguf", 2 * 1024**3)
inst = _make_instance()
inst._process = _FakeProc(pid = os.getpid())
inst._gguf_path = str(tmp_path / "model-00001-of-00004.gguf")
inst._healthy = False
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
return io.StringIO("VmRSS:\t0\tkB\n")
return open(path, *args, **kwargs)
with patch("builtins.open", side_effect = fake_open):
out = inst.load_progress()
assert out is not None
assert out["bytes_total"] == 80 * 1024**3 # 4 x 20 GB, no mmproj
class TestLoadProgressDegradation:
"""Broken / unusual inputs never raise; they produce best-effort output."""
def test_missing_gguf_path_still_reports_rss(self, tmp_path):
inst = _make_instance()
inst._process = _FakeProc(pid = os.getpid())
inst._gguf_path = None
inst._healthy = False
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
return io.StringIO("VmRSS:\t1024\tkB\n")
return open(path, *args, **kwargs)
with patch("builtins.open", side_effect = fake_open):
out = inst.load_progress()
assert out is not None
assert out["phase"] == "mmap"
assert out["bytes_total"] == 0
assert out["bytes_loaded"] == 1024 * 1024
assert out["fraction"] == 0.0
def test_unreadable_proc_returns_none(self, tmp_path):
inst = _make_instance()
# Pid that doesn't exist -> /proc read fails.
inst._process = _FakeProc(pid = 999_999_999)
inst._gguf_path = str(tmp_path / "model.gguf") # doesn't need to exist
inst._healthy = False
out = inst.load_progress()
# FileNotFoundError on /proc path -> load_progress returns None.
assert out is None

202
studio/backend/tests/test_llama_cpp_load_progress_live.py Normal file
View file

@ -0,0 +1,202 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Live, no-mock integration test for ``LlamaCppBackend.load_progress()``.
The companion files (``test_llama_cpp_load_progress.py`` and
``test_llama_cpp_load_progress_matrix.py``) patch ``builtins.open`` to
feed synthetic VmRSS values. This file is the opposite: it uses **real**
subprocesses, **real** file sizes, and the **real** ``/proc``
interface. It is the sanity check that the contract we keep in the
mocked tests still maps to what the kernel actually returns on a live
Linux system.
Why both: the mocked tests can be fooled by a buggy implementation that
parses ``/proc`` output in a format the kernel no longer uses, or that
makes assumptions about ``Path.stat()`` vs ``os.path.getsize``. This
file hits the real APIs so any format drift gets caught.
Skipped cleanly on non-Linux (no ``/proc``).
"""
from __future__ import annotations
import os
import subprocess
import sys
import time
import types as _types
from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# Same stubs as the matrix file (keep self-contained so the file can be
# run standalone as well as via the full suite).
# ---------------------------------------------------------------------------
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_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)
_httpx_stub = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **k: None})
_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 core.inference.llama_cpp import LlamaCppBackend
pytestmark = pytest.mark.skipif(
not Path("/proc").exists(),
reason = "live /proc test is Linux-only",
)
def _make_backend(pid: int, gguf_path: str, healthy: bool = False):
inst = LlamaCppBackend.__new__(LlamaCppBackend)
inst._process = type("P", (), {"pid": pid})()
inst._gguf_path = gguf_path
inst._healthy = healthy
return inst
def test_live_rss_matches_kernel_vmrss(tmp_path):
"""Spawn a real child, let it allocate real bytes, confirm
``bytes_loaded`` tracks the kernel's VmRSS within a sane tolerance."""
# Child that allocates ~100 MB of zero'd bytes and then idles.
script = tmp_path / "burn.py"
script.write_text(
"import time, sys\n"
"buf = bytearray(100 * 1024 * 1024)\n" # 100 MB
"# touch every page so RSS actually grows\n"
"for i in range(0, len(buf), 4096):\n"
" buf[i] = 1\n"
"sys.stdout.write('ready\\n')\n"
"sys.stdout.flush()\n"
"time.sleep(10)\n"
)
proc = subprocess.Popen(
[sys.executable, str(script)],
stdout = subprocess.PIPE,
stderr = subprocess.PIPE,
)
try:
# Wait for the child to finish touching pages.
ready = proc.stdout.readline()
assert ready.strip() == b"ready"
# Create a fake 200 MB sparse gguf so bytes_total is concrete.
gguf = tmp_path / "model.gguf"
with open(gguf, "wb") as f:
f.truncate(200 * 1024 * 1024)
inst = _make_backend(proc.pid, str(gguf), healthy = False)
out = inst.load_progress()
assert out is not None, "load_progress returned None for live pid"
assert out["phase"] == "mmap"
assert out["bytes_total"] == 200 * 1024 * 1024
# VmRSS for the Python child includes the interpreter + the 100MB
# buffer, so a realistic floor is 50 MB and ceiling is 200 MB.
assert (
out["bytes_loaded"] >= 50 * 1024 * 1024
), f"bytes_loaded unexpectedly low: {out['bytes_loaded']}"
assert out["bytes_loaded"] <= 200 * 1024 * 1024
assert 0.0 < out["fraction"] <= 1.0
finally:
proc.terminate()
try:
proc.wait(timeout = 5)
except subprocess.TimeoutExpired:
proc.kill()
def test_live_ready_phase_when_healthy(tmp_path):
gguf = tmp_path / "m.gguf"
with open(gguf, "wb") as f:
f.truncate(1 * 1024 * 1024)
inst = _make_backend(os.getpid(), str(gguf), healthy = True)
out = inst.load_progress()
assert out is not None
assert out["phase"] == "ready"
assert out["bytes_total"] == 1 * 1024 * 1024
# Self-pid RSS is well above 1 MiB for CPython; fraction caps at 1.
assert out["fraction"] == 1.0
def test_live_dead_pid_returns_none(tmp_path):
"""A recently-dead pid may linger in /proc for ms; use a clearly
invalid id so the read reliably fails."""
gguf = tmp_path / "m.gguf"
gguf.touch()
inst = _make_backend(9_999_999_999, str(gguf), healthy = False)
out = inst.load_progress()
assert out is None
def test_live_shard_aggregation_counts_real_files(tmp_path):
"""With 4 real sibling shards on disk, ``bytes_total`` equals their
summed size to the byte."""
shard_size = 7 * 1024 * 1024 # 7 MB each
for i in range(1, 5):
f = tmp_path / f"model-{i:05d}-of-00004.gguf"
with open(f, "wb") as fh:
fh.truncate(shard_size)
# Unrelated file in same dir -- must not be counted.
with open(tmp_path / "config.json", "wb") as fh:
fh.truncate(123)
inst = _make_backend(
os.getpid(),
str(tmp_path / "model-00001-of-00004.gguf"),
healthy = False,
)
out = inst.load_progress()
assert out is not None
assert out["bytes_total"] == 4 * shard_size
def test_live_repeated_polling_stays_sane(tmp_path):
"""Sampling the same backend 20 times should not raise or produce
non-numeric output, even under normal kernel RSS jitter."""
gguf = tmp_path / "m.gguf"
with open(gguf, "wb") as f:
f.truncate(500 * 1024 * 1024)
inst = _make_backend(os.getpid(), str(gguf), healthy = False)
seen = []
for _ in range(20):
out = inst.load_progress()
assert out is not None
assert isinstance(out["bytes_loaded"], int)
assert isinstance(out["bytes_total"], int)
assert 0.0 <= out["fraction"] <= 1.0
seen.append(out["bytes_loaded"])
time.sleep(0.01)
# RSS of a healthy Python process doesn't go below ~5 MB.
assert min(seen) > 1 * 1024 * 1024

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