commit e93507a09cf3d65e78c3d0a7581652d0cbbb4b43 Author: wehub-resource-sync Date: Mon Jul 13 12:59:56 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..5f04b5e --- /dev/null +++ b/.gitattributes @@ -0,0 +1,13 @@ +# Normalize Python files to LF line endings +*.py text eol=lf + +# Always check out shell scripts with LF endings. Without this rule a Windows +# clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks +# them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -"). +*.sh text eol=lf + +# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather +# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files +# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts) +# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF. +studio/frontend/** text=auto eol=lf diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..a969403 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,62 @@ +# Inspired from https://github.com/vllm-project/vllm/blob/main/.github/CODEOWNERS + +/unsloth/models/loader.py @danielhanchen @mmathew23 +/unsloth/models/llama.py @Datta0 @danielhanchen @mmathew23 +/unsloth/models/rl.py @Datta0 @pluesclues @danielhanchen +/unsloth/models/rl_replacements.py @Datta0 @pluesclues @danielhanchen +/unsloth/trainer.py @danielhanchen +/unsloth/models/sentence_transformer.py @Etherll @danielhanchen +/unsloth/save.py @danielhanchen +/unsloth/tokenizer_utils.py @mmathew23 @danielhanchen +/unsloth/chat_templates.py @danielhanchen +/unsloth/ollama_template_mappers.py @danielhanchen +/unsloth/kernels/moe/*.py @Datta0 +/unsloth/import_fixes.py @danielhanchen +/unsloth/device_type.py @danielhanchen +/unsloth/_auto_install.py @danielhanchen +/unsloth/dataprep/*.py @danielhanchen +/unsloth/kernels/cross_entropy_loss.py @danielhanchen +/unsloth/kernels/fast_lora.py @danielhanchen +/unsloth/kernels/flex_attention.py @danielhanchen +/unsloth/kernels/fp8.py @Datta0 +/unsloth/kernels/geglu.py @danielhanchen +/unsloth/kernels/layernorm.py @danielhanchen +/unsloth/kernels/rms_layernorm.py @danielhanchen +/unsloth/kernels/rope_embedding.py @danielhanchen +/unsloth/kernels/swiglu.py @danielhanchen +/unsloth/kernels/utils.py @danielhanchen @Datta0 +/unsloth/models/_utils.py @danielhanchen @mmathew23 +/unsloth/models/cohere.py @danielhanchen +/unsloth/models/dpo.py @danielhanchen +/unsloth/models/falcon_h1.py @danielhanchen +/unsloth/models/gemma.py @danielhanchen +/unsloth/models/gemma2.py @danielhanchen +/unsloth/models/glm4_moe.py @Datta0 +/unsloth/models/granite.py @danielhanchen +/unsloth/models/llama4.py @danielhanchen +/unsloth/models/loader_utils.py @Datta0 @danielhanchen +/unsloth/models/mapper.py @danielhanchen +/unsloth/models/mistral.py @danielhanchen +/unsloth/models/qwen2.py @danielhanchen +/unsloth/models/qwen3.py @Datta0 +/unsloth/models/qwen3_moe.py @Datta0 +/unsloth/models/vision.py @mmathew23 @danielhanchen +/unsloth/utils/attention_dispatch.py @mmathew23 +/unsloth/utils/hf_hub.py @mmathew23 +/unsloth/utils/packing.py @mmathew23 + +/cli/ @Manan17 +/studio/frontend/ @Shine1i @Manan17 +/studio/frontend/public/ @Shine1i +/studio/backend/ +/studio/backend/core/data_recipe/ +/studio/backend/tests/ @danielhanchen +/tests/ @danielhanchen +/scripts/ @danielhanchen + +# Snapshot data for the notebook linter / Colab oracle. Drift in these +# files changes the pin floor for every Unsloth notebook, so refreshes +# must be reviewed by the notebook owners directly. CODEOWNERS later +# wins, so this overrides the broader /scripts/ rule above. +/scripts/data/colab_*.txt @danielhanchen @shimmyshimmer +/scripts/data/colab_*.json @danielhanchen @shimmyshimmer diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 0000000..ae5dade --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1,13 @@ +# These are supported funding model platforms + +github: unslothai +patreon: # Replace with a single Patreon username +open_collective: # Replace with a single Open Collective username +ko_fi: # unsloth +tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel +community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry +liberapay: # Replace with a single Liberapay username +issuehunt: # Replace with a single IssueHunt username +otechie: # Replace with a single Otechie username +lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry +custom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2'] diff --git a/.github/ISSUE_TEMPLATE/bug---issue.md b/.github/ISSUE_TEMPLATE/bug---issue.md new file mode 100644 index 0000000..ffa3d3c --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug---issue.md @@ -0,0 +1,22 @@ +--- +name: Bug / Issue +about: Bug / Issue +title: "[Bug] Please fill in your issue title here." +labels: bug +assignees: '' + +--- +Note: Please do not remove the questions. Answer beside them. +1. Did you update? `pip install --upgrade unsloth unsloth_zoo` +2. `Colab` or `Kaggle` or local / cloud +3. Number GPUs used, use `nvidia-smi` +4. Which notebook? Please link! +5. Which Unsloth version, TRL version, transformers version, PyTorch version? +6. Which trainer? `SFTTrainer`, `GRPOTrainer` etc + +```python +Put Minimal code to reproduce error here ###Remove Hugging Face token### +###Please make sure to check formatting properly, edit if needed.### +``` + +🦥 You can also ask via our Reddit page: https://reddit.com/r/unsloth/ diff --git a/.github/ISSUE_TEMPLATE/feature-request.md b/.github/ISSUE_TEMPLATE/feature-request.md new file mode 100644 index 0000000..5ea70a8 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature-request.md @@ -0,0 +1,21 @@ +--- +name: Feature Request +about: New features, model support, ideas +title: "[Feature]" +labels: feature request +assignees: '' + +--- + +For new models, have you tried: +```python +from unsloth import FastModel +model, tokenizer = FastModel.from_pretrained( + "microsoft/Phi-4-multimodal-instruct", + trust_remote_code = True, +) +from transformers import AutoModelForSequenceClassification +model, tokenizer = FastModel.from_pretrained( + auto_model = AutoModelForSequenceClassification, +) +``` diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..4908382 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,100 @@ +--- +version: 2 +updates: + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + cooldown: + # github-actions refs are git tags / SHAs, not semver -- the + # `semver-minor-days` / `semver-patch-days` knobs are rejected + # by Dependabot's validator for this ecosystem. Only the + # `default-days` floor applies. + default-days: 7 + groups: + actions: + patterns: ["*"] + actions-security: + applies-to: security-updates + patterns: ["*"] + + # Removed a stray `package-ecosystem: "bun"` entry for + # /studio/frontend: that path has no bun.lock / bun.lockb, so + # Dependabot's bun ecosystem silently no-ops on it. The actual + # lockfile committed at /studio/frontend is package-lock.json + # (npm), and the npm entry further below already catches + # npm_and_yarn security advisories for that directory. Version + # updates for /studio/frontend stay suppressed (open-pull- + # requests-limit: 0 in that entry) -- security PRs flow through + # regardless. Add a real bun entry IF and WHEN bun.lock lands. + + - package-ecosystem: "npm" + directory: "/studio/backend/core/data_recipe/oxc-validator" + schedule: + interval: "weekly" + cooldown: + default-days: 7 + semver-minor-days: 3 + semver-patch-days: 3 + groups: + npm-oxc-validator: + patterns: ["*"] + npm-oxc-validator-security: + applies-to: security-updates + patterns: ["*"] + + # pip + cargo grouped weekly; the *-security siblings batch + # advisories that would otherwise each open their own PR. + - package-ecosystem: "pip" + directory: "/" + schedule: + interval: "weekly" + open-pull-requests-limit: 5 + cooldown: + default-days: 7 + groups: + python: + patterns: ["*"] + python-security: + applies-to: security-updates + patterns: ["*"] + + - package-ecosystem: "cargo" + directory: "/studio/src-tauri" + schedule: + interval: "weekly" + cooldown: + default-days: 7 + semver-minor-days: 3 + semver-patch-days: 3 + groups: + cargo-tauri: + patterns: ["*"] + cargo-tauri-security: + applies-to: security-updates + patterns: ["*"] + + # /studio/frontend npm dependencies. Version-update PRs are + # deliberately suppressed (open-pull-requests-limit: 0) -- the + # frontend dep tree is large, the lockfile is the authoritative + # pin, and `min-release-age=7` in studio/frontend/.npmrc already + # blocks fresh tarballs at install time. Security advisories + # arrive via GitHub's npm_and_yarn channel and are NOT capped by + # `open-pull-requests-limit` per Dependabot's documented + # behaviour; they flow through this entry, group together, and + # still respect the cooldown below so we never ingest a tarball + # that was hot-published less than 3 days ago. + - package-ecosystem: "npm" + directory: "/studio/frontend" + schedule: + interval: "weekly" + open-pull-requests-limit: 0 + cooldown: + default-days: 7 + semver-minor-days: 3 + semver-patch-days: 3 + groups: + npm-frontend-security: + applies-to: security-updates + patterns: ["*"] +... diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh new file mode 100755 index 0000000..f4189a1 --- /dev/null +++ b/.github/scripts/agent-guides-drive.sh @@ -0,0 +1,682 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Drive one coding agent against the running `unsloth run` server for the +# Local Agent Guides CI. All failures from here are failure class (c) +# "guide drift": the server preflight already passed and the agent CLI +# already installed, so a failure here means the documented recipe in +# unsloth_cli/commands/start.py no longer produces a working flow. +# +# Self-updating: for all six agents (claude, codex, hermes, openclaw, +# opencode, pi) we obtain the exact env + command from +# `unsloth start --no-launch` and run THAT, so a recipe change is +# exercised automatically. +# +# Every agent invocation is wrapped in `timeout` so a headless-TTY prompt +# can never hang the runner -- a timeout is reported as guide drift with a +# distinct message. +# +# Usage: +# agent-guides-drive.sh connection +# agent-guides-drive.sh file-edit +# agent-guides-drive.sh attribution-ab claude +# +# Required env (exported by serve-unsloth-run.sh): +# UNSLOTH_BASE_URL UNSLOTH_API_KEY UNSLOTH_MODEL_ID +# UNSLOTH_LLAMA_LOG_DIR AGENT_INVOKE_TIMEOUT UNSLOTH_SEED +set -uo pipefail + +MODE="${1:?usage: agent-guides-drive.sh }" +AGENT="${2:?usage: agent-guides-drive.sh }" + +: "${UNSLOTH_BASE_URL:?serve step did not export UNSLOTH_BASE_URL}" +: "${UNSLOTH_API_KEY:?serve step did not export UNSLOTH_API_KEY}" +: "${UNSLOTH_MODEL_ID:?serve step did not export UNSLOTH_MODEL_ID}" +# Determinism (seed/temp) is applied at the server level by +# serve-unsloth-run.sh --extra; agents inherit it through the API. +TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}" + +# Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner +# IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless +# to the other agents, which ignore it. +export IS_SANDBOX=1 + +# Absolute paths anchored at the repo root (this script lives in +# .github/scripts/). Everything writes here regardless of the current working +# directory, so the file-edit mode can `cd` into a scratch work dir without +# breaking log/redaction writes. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +LOGS_DIR="$REPO_ROOT/logs" +REDACTED_DIR="$REPO_ROOT/redacted-configs" +WORKDIR_BASE="$REPO_ROOT/agent-workdir" +CACHE_HELPER="$SCRIPT_DIR/assert-prompt-cache.sh" +mkdir -p "$LOGS_DIR" "$REDACTED_DIR" +CONNECT_REF="unsloth_cli/commands/start.py" + +# Prefill-shrinking flags for Claude Code. The heavyweight agents send +# multi-thousand-token system prompts + full tool schemas, which on a CPU-only +# runner is minutes of prefill per model round-trip (~16 tok/s for a 4B model). +# Replacing the ~5.7k default system prompt with a tiny one (--system-prompt-file) +# and restricting tools cuts the prefill to a few hundred tokens so it completes +# quickly on CPU. These only shape the request size; the start.py recipe +# (endpoint, auth, model) is still exercised end to end. +# +# The bulk of Claude Code's prompt is the built-in tool JSON schemas: measured +# via `claude -p /context`, the default prompt is ~28k tokens of which ~18k is +# "System tools" alone. --allowedTools/--disallowedTools only gate PERMISSION to +# call a tool; they do NOT remove its schema from what is sent to the model, so +# the earlier whitelist left the full ~18k in the prompt and CPU prefill +# (~16 tok/s) overran claude's own request timeout into a retry loop. --tools is +# the flag that restricts which schemas are sent. (The ~8k "Memory files" chunk +# is auto-loaded CLAUDE.md; the unsloth repo ships none, so it is 0 in CI.) +# +# Connection probe: --tools "" sends ZERO tool schemas, leaving ~20 tokens total +# (a one-line --system-prompt-file + the user turn), which prefills instantly. +CLAUDE_CONNECT_FLAGS=( + --system-prompt-file "$SCRIPT_DIR/ci-connect-prompt.txt" + --tools "" +) +# File-edit: the task needs the file/shell tools, so send only those schemas +# (~2.3k tokens vs ~18k for the full set). +CLAUDE_EDIT_FLAGS=( + --system-prompt-file "$SCRIPT_DIR/ci-min-system-prompt.txt" + --tools "Bash,Edit,Write,Read" +) + +guide_fail() { + echo "::error::[guide drift] agent=${AGENT}: $* (preflight passed + install OK, so the documented flow in ${CONNECT_REF} drifted)." >&2 + exit 1 +} + +# Redact the API key from any file we are about to keep as an artifact. +# Portable across GNU sed (Linux runners) and BSD sed (macOS), so the +# redaction is never silently skipped. +redact() { + local f + for f in "$@"; do + [ -f "$f" ] || continue + if sed --version >/dev/null 2>&1; then + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + else + sed -i '' "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + fi + done +} + +# Print a file to the log with the key scrubbed, without mutating it (the raw file is +# still needed to parse the real env). Use this instead of `cat` for any transcript that +# carries an `export UNSLOTH_API_KEY=...` line, so a live key never reaches Actions logs. +cat_redacted() { + sed "s#${UNSLOTH_API_KEY}##g" "$1" +} + +# A reply must be non-empty and free of connection/auth errors. +assert_reply() { + local out="$1" + if [ ! -s "$out" ]; then + guide_fail "agent produced an EMPTY reply" + fi + if grep -qiE 'connection refused|connection error|econnrefused|fetch failed|http 4[0-9][0-9]|unauthorized|invalid api key|authentication failed' "$out"; then + guide_fail "agent reply contained a connection/auth error: $(grep -iE 'connection|unauthorized|auth|http 4' "$out" | head -1)" + fi + echo "[$AGENT] reply (first 20 lines):" + head -20 "$out" +} + +# Run a command under a hard timeout; map 124 to a guide-drift hang message. +run_timed() { # $1=outfile, rest=command + local out="$1"; shift + timeout "$TIMEOUT" "$@" > "$out" 2>&1 + local rc=$? + if [ "$rc" -eq 124 ]; then + redact "$out" # guide_fail exits below, so scrub the transcript here too + echo "[$AGENT] last 40 lines before timeout:"; tail -40 "$out" 2>/dev/null || true + guide_fail "invoke timed out after ${TIMEOUT}s (headless-TTY hang -- the recipe likely needs a non-interactive/print flag)" + fi + return "$rc" +} + +# Read a value from an `export VAR=...` line in the connect --no-launch output. +# `unsloth start` writes each agent's session config off the user's ~ and points +# at it through a relocation env var (CODEX_HOME / OPENCODE_CONFIG / +# OPENCLAW_CONFIG_PATH), so the contract checks read the path from here. +raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) + local raw="$LOGS_DIR/connect-${AGENT}.txt" + local v; v="$(sed -n "s/^export $1=//p" "$raw" | tail -1)" + v="${v#\'}"; v="${v%\'}"; printf '%s' "$v" +} + +# ── 5-agent start.py path: parse env + command from --no-launch ───────── +# Populates globals CONNECT_ENV (export/unset lines) and CONNECT_CMD (the +# launch command on the last printed line), and runs start.py's config +# writers as a side effect (it writes each agent's relocated session config). +parse_connect() { + local raw="$LOGS_DIR/connect-${AGENT}.txt" + # CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their + # config (which now prompts by default), so the file-edit test opts into auto-approval + # here, the same intent as claude/codex's per-call bypass flags. + local yolo=() + [ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo) + if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + cat_redacted "$raw" + guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" + fi + echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw" + CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)" + # The launch command is the last non-export, non-status line. start.py + # prints "Studio · model " and "Updated ..." status lines first. + CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \ + | grep -E '[^[:space:]]' | tail -1)" + [ -n "$CONNECT_CMD" ] || guide_fail "could not parse a launch command from connect --no-launch output" + redact "$raw" +} + +# Cross-check the documented contract knobs so silent start.py changes +# (env-var rename, wire_api flip, attribution setting drop) also fail/flag. +crosscheck_contract() { + local raw="$LOGS_DIR/connect-${AGENT}.txt" + local cfg home + case "$AGENT" in + codex) + grep -q 'UNSLOTH_STUDIO_AUTH_TOKEN' "$raw" \ + || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (start.py _CODEX_ENV_KEY)" + home="$(raw_env CODEX_HOME)" + # An empty relocation var would make cfg "/config.toml" and silently + # skip the [ -f ] contract check below; fail loudly instead. + [ -n "$home" ] || guide_fail "CODEX_HOME missing from connect output (start.py codex())" + cfg="$home/config.toml" + if [ -f "$cfg" ]; then + grep -q 'wire_api = "responses"' "$cfg" \ + || guide_fail "Codex wire_api is no longer \"responses\" in \$CODEX_HOME/config.toml" + cp "$cfg" "$REDACTED_DIR/codex-config.toml" + fi + grep -q 'codex --oss --profile unsloth_api' "$raw" \ + || echo "::warning::Codex launch command changed from 'codex --oss --profile unsloth_api'" + ;; + claude) + grep -q 'ANTHROPIC_AUTH_TOKEN' "$raw" \ + || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (start.py claude())" + grep -q 'CLAUDE_CODE_ATTRIBUTION_HEADER' "$raw" \ + || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER no longer set for the session (start.py claude())" + ;; + hermes) + grep -q 'UNSLOTH_API_KEY' "$raw" \ + || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (start.py _HERMES_ENV_KEY)" + home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "HERMES_HOME missing from connect output (start.py hermes())" + cfg="$home/config.yaml" + [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/hermes-config.yaml" + ;; + openclaw) + cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + if [ -n "$cfg" ] && [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ + || echo "::warning::OpenClaw provider api is no longer 'openai-completions' (write_openclaw_config)" + cp "$cfg" "$REDACTED_DIR/openclaw.json" + fi + ;; + opencode) + cfg="$(raw_env OPENCODE_CONFIG)" + [ -n "$cfg" ] && [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/opencode.json" + ;; + pi) + # Pi has no config-dir env var; the session is HOME-relocated, and the + # provider config lives at $HOME/.pi/agent/models.json. + cfg="$(raw_env HOME)/.pi/agent/models.json" + if [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ + || echo "::warning::Pi provider api is no longer 'openai-completions' (write_pi_config)" + cp "$cfg" "$REDACTED_DIR/pi-models.json" + fi + ;; + esac + redact "$REDACTED_DIR"/* 2>/dev/null || true +} + +# Heavyweight agents (hermes, openclaw) bake a large system prompt + tool JSON +# schemas into every request, which a CPU runner cannot prefill before the invoke +# timeout. As with claude's --tools, we shrink the request from the agent's own +# config: zero tools for the connection probe collapses the prompt to a few +# hundred tokens, since both CLIs gate the bulk of their prompt on having tools. + +# Hermes: an explicit empty cli toolset disables all tools (and drops the +# tool-gated guidance blocks), so -z sends ~300 tokens instead of thousands. +# Hermes enables its default cli toolset when the session config does not pin one, +# so we must set platform_toolsets.cli explicitly to [] (not just append) to get +# zero tools. That needs a YAML parser, and the runner's bare python3 has no +# PyYAML -- but the venv that ships `unsloth` does (start.py imports yaml), so run +# the patch with that interpreter. We patch the relocated $HERMES_HOME/config.yaml +# that `unsloth start` printed, not the user's ~/.hermes. +# (-z reads platform_toolsets.cli; --ignore-rules is a no-op under -z.) +patch_hermes_tools() { # $1 = none|default + # Check the raw var BEFORE appending /config.yaml: the joined path is never + # empty, so the old guard could not fire and the patcher would die on + # "/config.yaml" with a bare traceback instead of this clear failure. + local home; home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "Hermes HERMES_HOME missing from connect output (start.py hermes())" + local cfg; cfg="$home/config.yaml" + # Find a python that can import yaml. The runner's bare python3 cannot, but the + # interpreter in the `unsloth` console-script shebang provably can (it runs + # start.py's write_hermes_config, which imports yaml). Try that first, then + # any python on PATH, then the venv sibling, picking the first with PyYAML. + local cand py="" shebang + shebang="$(head -1 "$(command -v unsloth)" 2>/dev/null | sed -n 's/^#![[:space:]]*//p' | awk '{print $1}')" + for cand in "$shebang" python3 python "$(dirname "$(command -v unsloth)")/python"; do + [ -n "$cand" ] || continue + { [ -x "$cand" ] || command -v "$cand" >/dev/null 2>&1; } || continue + if "$cand" -c 'import yaml' 2>/dev/null; then py="$cand"; break; fi + done + [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch the hermes session config" + echo "[hermes] patching $cfg with $py" + "$py" - "$1" "$cfg" <<'PY' +import os, sys +import yaml +mode = sys.argv[1] +p = sys.argv[2] +cfg = (yaml.safe_load(open(p)) or {}) if os.path.exists(p) else {} +ts = cfg.get("platform_toolsets") +if not isinstance(ts, dict): + ts = cfg["platform_toolsets"] = {} +if mode == "none": + ts["cli"] = [] # explicit empty list -> zero tools (not "defaults") +else: + ts.pop("cli", None) # file-edit needs real tools -> restore defaults +with open(p, "w") as fh: + yaml.safe_dump(cfg, fh, sort_keys=False) +print(f"[hermes] platform_toolsets.cli = {ts.get('cli', 'default')}") +PY +} + +# OpenClaw: 'openclaw agent' has no tool/prompt flags, so we define a 'ci' agent +# in openclaw.json. tools.deny ["*"] sends zero tool schemas (deny always wins) +# for the connection probe; contextInjection "never" + defaults.skipBootstrap +# drop the auto-injected AGENTS.md/SOUL.md bootstrap (the bulk of the prompt) for +# both modes. --agent must reference a defined agent, so write it before invoking. +patch_openclaw_agent() { # $1 = notools|tools + # OpenClaw reads its config from the relocated OPENCLAW_CONFIG_PATH that + # `unsloth start` printed, so patch THAT file (not the user's ~/.openclaw). + local cfg; cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + [ -n "$cfg" ] || guide_fail "OpenClaw OPENCLAW_CONFIG_PATH missing from connect output (start.py openclaw())" + python3 - "$1" "$cfg" <<'PY' +import os, sys, json +mode = sys.argv[1] +p = sys.argv[2] +cfg = json.load(open(p)) if os.path.exists(p) else {} +agents = cfg.setdefault("agents", {}) +agents.setdefault("defaults", {})["skipBootstrap"] = True +lst = [a for a in agents.get("list", []) if a.get("id") != "ci"] +agent = {"id": "ci", "contextInjection": "never"} +if mode == "notools": + agent["tools"] = {"deny": ["*"]} +lst.append(agent) +agents["list"] = lst +with open(p, "w") as fh: + json.dump(cfg, fh, indent=2) +print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}") +PY +} + +# Build an invoke script that applies start.py's env then runs the launch +# command (with extra args appended) under bash. We do NOT eval connect's env +# into this shell; we write it into a one-shot script so the export/unset +# semantics are exactly what start.py printed. The script path is absolute +# so it is valid even when the caller has cd'd into a scratch work dir. +invoke_via_connect() { # $1=outfile, rest=extra args appended to the command + local out="$1"; shift + local script="$LOGS_DIR/invoke-${AGENT}.sh" + local real; real="$(mktemp)" + # CONNECT_ENV_EXTRA / CONNECT_CMD_OVERRIDE let a caller (attribution-ab) flip a + # session knob without editing the user's config; empty -> use what start.py emitted. + local cmd="${CONNECT_CMD_OVERRIDE:-$CONNECT_CMD}" + { + echo "set -uo pipefail" + echo "$CONNECT_ENV" + [ -n "${CONNECT_ENV_EXTRA:-}" ] && echo "$CONNECT_ENV_EXTRA" + # Append extra args (the prompt / flags) to the launch command verbatim. + printf '%s' "$cmd" + local a + for a in "$@"; do printf ' %q' "$a"; done + printf '\n' + } > "$real" + # Upload a REDACTED copy of the script, but EXECUTE the un-redacted one from a + # temp path outside the artifact dir. Redacting the script we run would turn + # the real `export TOKEN=sk-...` line into `export TOKEN=`, which is + # invalid bash (the `<`/`>` are redirections) and silently breaks every agent. + # Writing the redacted copy up front keeps the key out of the artifact even if + # the run times out (run_timed exits before returning here). + cp "$real" "$script"; redact "$script" + # The connect one-liner now carries the key as an inline env assignment; scrub it on + # the way to the log (the executed $real keeps the live value). + echo "[$AGENT] invoking (timeout ${TIMEOUT}s): ${cmd//${UNSLOTH_API_KEY}/} $*" + run_timed "$out" bash "$real" + local rc=$? + rm -f "$real" + redact "$out" # the transcript can echo the token; scrub before upload + return "$rc" +} + +# ═════════════════════════════════════════════════════════════════════════ +case "$MODE" in + # ── connection: trivial prompt, assert a non-empty, error-free reply ──── + connection) + PROMPT='Reply with exactly the single word: pong' + OUT="$LOGS_DIR/${AGENT}-connection.txt" + parse_connect + crosscheck_contract + # claude/codex run in print mode via the flags start.py emits + # (claude -p / codex exec). For agents whose default subcommand prints + # to stdout we pass the prompt through ctx.args. + case "$AGENT" in + claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; + codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; + opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; + pi) invoke_via_connect "$OUT" -p "$PROMPT" ;; + hermes) patch_hermes_tools none + invoke_via_connect "$OUT" -z "$PROMPT" ;; + openclaw) patch_openclaw_agent notools + CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; + *) invoke_via_connect "$OUT" "$PROMPT" ;; + esac + # A non-zero exit from the documented launch command is drift even if it + # printed something: a benign-looking "command not found" / usage dump would + # otherwise slip past assert_reply (which only flags empty/error-keyword text). + rc=$? + [ "$rc" -eq 0 ] || guide_fail "the documented launch command exited non-zero (rc=$rc) -- see the transcript above" + assert_reply "$OUT" + echo "[$AGENT] connection OK" + ;; + + # ── file-edit: deterministic 2-turn hello.py test (Qwen3.5-2B) ────────── + file-edit) + WORK="$WORKDIR_BASE/${AGENT}" + rm -rf "$WORK"; mkdir -p "$WORK" + OUT1="$LOGS_DIR/${AGENT}-fileedit-turn1.txt" + OUT2="$LOGS_DIR/${AGENT}-fileedit-turn2.txt" + T1='Create a file named hello.py in the current directory whose entire contents are a single line: print("Hello"). Do not run it.' + T2='Run hello.py with python and show me the exact output.' + + # The start.py recipe writers + crosscheck must see the repo; run them + # from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw + # gate tool approval through their config (prompting by default), so file-edit + # opts them into auto-approval to run edits/commands headlessly. + case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac + parse_connect + crosscheck_contract + # File-edit needs real tools, so we cannot zero them as in connection. + # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md + # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work + # dir is empty, so no project context files are auto-loaded either. + case "$AGENT" in + hermes) patch_hermes_tools default ;; + openclaw) patch_openclaw_agent tools ;; + esac + + # Drive from inside the work dir so the agent edits files there. All log + # writes use absolute $LOGS_DIR, so cwd does not matter for them. + cd "$WORK" || guide_fail "could not enter work dir $WORK" + + invoke_turn() { # $1=outfile $2=continue? $3=prompt + local out="$1" cont="$2" prompt="$3" + case "$AGENT" in + pi) + # Pi continues the previous session with -c; provider/model come from + # the parsed `unsloth start pi` recipe (CONNECT_CMD), not hardcoded here. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" -p --continue "$prompt" + else + invoke_via_connect "$out" -p "$prompt" + fi ;; + claude) + # --dangerously-skip-permissions lets headless claude actually use the + # Write/Bash tools (otherwise it blocks on an approval prompt and emits + # nothing). IS_SANDBOX=1 (exported above) authorizes it. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" "${CLAUDE_EDIT_FLAGS[@]}" --dangerously-skip-permissions -p --continue "$prompt" + else + invoke_via_connect "$out" "${CLAUDE_EDIT_FLAGS[@]}" --dangerously-skip-permissions -p "$prompt" + fi ;; + codex) + # --dangerously-bypass-approvals-and-sandbox gives codex exec + # workspace-write (default is read-only -> cannot create hello.py) and + # skips the bubblewrap sandbox that the runner lacks. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" exec --dangerously-bypass-approvals-and-sandbox resume --last "$prompt" + else + invoke_via_connect "$out" exec --dangerously-bypass-approvals-and-sandbox "$prompt" + fi ;; + opencode) invoke_via_connect "$out" run "$prompt" ;; + hermes) invoke_via_connect "$out" -z "$prompt" ;; + openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;; + *) invoke_via_connect "$out" "$prompt" ;; + esac + } + + # Turn 1: create hello.py. + invoke_turn "$OUT1" fresh "$T1" + # Fail on a non-zero agent exit before trusting side effects: an agent can + # error out (API/tool failure) yet leave a plausible file/transcript behind, + # which would otherwise slip past the assertions below (mirrors connection). + rc=$? + [ "$rc" -eq 0 ] || { echo "[$AGENT] turn-1 transcript:"; tail -40 "$OUT1" 2>/dev/null || true; \ + guide_fail "turn 1 (create hello.py) exited non-zero (rc=$rc)"; } + + # Hard assertions on the side effect (the real test): file + content + run. + if [ ! -f hello.py ]; then + echo "[$AGENT] turn-1 transcript:"; tail -40 "$OUT1" 2>/dev/null || true + guide_fail "turn 1 did not create hello.py" + fi + grep -q 'Hello' hello.py || guide_fail "hello.py does not contain 'Hello'" + RUN_OUT="$(python3 hello.py 2>&1 || true)" + [ "$RUN_OUT" = "Hello" ] || guide_fail "python3 hello.py printed '$RUN_OUT', expected exactly 'Hello'" + echo "[$AGENT] turn 1 OK (file created, prints 'Hello')" + + # Turn 2: same cwd + session continuation; assert the agent's run output + # contains Hello. Narration drift is WARN-only, missing output is a hard fail. + invoke_turn "$OUT2" continue "$T2" + rc=$? + [ "$rc" -eq 0 ] || { echo "[$AGENT] turn-2 transcript:"; tail -60 "$OUT2" 2>/dev/null || true; \ + guide_fail "turn 2 (run hello.py) exited non-zero (rc=$rc)"; } + if grep -q 'Hello' "$OUT2"; then + echo "[$AGENT] turn 2 OK (run output contains 'Hello')" + else + echo "[$AGENT] turn-2 transcript:"; tail -60 "$OUT2" 2>/dev/null || true + guide_fail "turn 2 run/bash output did not contain 'Hello'" + fi + cd "$REPO_ROOT" || true + echo "[$AGENT] file-edit OK" + ;; + + # ── attribution-ab: Claude Code KV-cache HIT vs MISS ──────────────────── + attribution-ab) + [ "$AGENT" = "claude" ] || guide_fail "attribution-ab only applies to claude" + # The llama-server log filename uses the INTERNAL random llama.cpp port, + # not STUDIO_PORT, so we never glob by port: assert-prompt-cache.sh picks + # the newest llama-*.log and we slice it by a byte offset (`mark`) captured + # right before the measured turn, so an earlier turn's reuse can't leak in. + LLAMA_LOG_DIR="${UNSLOTH_LLAMA_LOG_DIR:-$HOME/.unsloth/studio/logs/llama-server}" + export LLAMA_LOG_DIR + parse_connect # prints session env + suppression flags (no ~/.claude write) + crosscheck_contract + PROMPT='Reply with exactly the single word: pong' + + # Phase A: the suppression start.py ships (CLAUDE_CODE_ATTRIBUTION_HEADER=0 + + # --exclude-dynamic-system-prompt-sections + --settings overlay) -> expect a + # HIT on the continued turn, since the system-prompt prefix is stable. + invoke_via_connect "$LOGS_DIR/claude-ab-hit-1.txt" -p "$PROMPT" # turn 1 primes + FROM_HIT="$(bash "$CACHE_HELPER" mark)" # offset before turn 2 + invoke_via_connect "$LOGS_DIR/claude-ab-hit-2.txt" -p --continue "$PROMPT again" + CACHE_LOG_FROM="$FROM_HIT" bash "$CACHE_HELPER" log HIT + + # Phase B: vanilla Claude with the header ENABLED -> expect a MISS. We flip + # the env var to 1 and strip the suppression flags from the launch command + # (without them the dynamic attribution line is included and changes every + # turn, so the shared prefix moves and the KV cache is invalidated, ~90% + # slower). This is session-only: nothing is written to ~/.claude. + CONNECT_ENV_EXTRA='export CLAUDE_CODE_ATTRIBUTION_HEADER=1' + CONNECT_CMD_OVERRIDE="$(printf '%s' "$CONNECT_CMD" \ + | sed -E "s/ --exclude-dynamic-system-prompt-sections//; s/ --settings '[^']*'//")" + invoke_via_connect "$LOGS_DIR/claude-ab-miss-1.txt" -p "$PROMPT" + FROM_MISS="$(bash "$CACHE_HELPER" mark)" + invoke_via_connect "$LOGS_DIR/claude-ab-miss-2.txt" -p --continue "$PROMPT again" + CACHE_LOG_FROM="$FROM_MISS" bash "$CACHE_HELPER" log MISS + unset CONNECT_ENV_EXTRA CONNECT_CMD_OVERRIDE + echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" + ;; + + # ── resume: does a launched agent's session survive exit and resume? ──── + # Unlike the other modes, this drives the real LAUNCH path (`unsloth start + # ...`, the interactive default), not the --no-launch recipe. That + # path relocates each agent's home to a throwaway temp dir wiped on exit, so + # a session cannot be resumed -- unless --persist routes it to the stable + # Unsloth agents dir instead. We run one headless turn per pass and check + # whether the turn left a session in a persistent store (deterministic, no + # reliance on the model recalling anything), for a baseline pass and a + # --persist pass, and assert the expected split for this agent. + resume) + CODEWORD="PLATYPUS7" + T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK." + T2="What codeword did I ask you to remember? Reply with just that word." + WORK="$WORKDIR_BASE/${AGENT}-resume" + + # STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to. + # Read it from a --no-launch probe (which also writes the agent's config + # there). codex/pi relocate their whole home/HOME here; opencode/claude keep + # their session data in a fixed user dir, so STABLE_HOME stays empty for them. + parse_connect + case "$AGENT" in + codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;; + pi) STABLE_HOME="$(raw_env HOME)" ;; + *) STABLE_HOME="" ;; + esac + + # The persistent stores a session would land in if it were NOT wiped. We + # count files here before/after each turn; a positive delta means the + # session persisted (is resumable), zero means it went to a wiped temp dir. + resume_tracked_dirs() { + case "$AGENT" in + codex) printf '%s\n' "$HOME/.codex" ;; + opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;; + claude) printf '%s\n' "$HOME/.claude" ;; + pi) printf '%s\n' "$HOME/.pi" ;; + *) : ;; + esac + [ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME" + } + count_session_files() { + local total=0 d n + while IFS= read -r d; do + [ -n "$d" ] && [ -d "$d" ] || continue + n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n)) + done < <(resume_tracked_dirs) + echo "$total" + } + + # The headless first-turn subcommand per agent (mirrors file-edit's map), + # forwarded verbatim through the launch path as passthrough args. + set_t1_cmd() { + case "$AGENT" in + claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;; + codex) T1_CMD=(exec "$T1") ;; + opencode) T1_CMD=(run "$T1") ;; + pi) T1_CMD=(-p "$T1") ;; + *) guide_fail "resume mode does not cover agent '$AGENT'" ;; + esac + } + + # Run one headless turn through the launch path. $1=outfile, $2="" or + # "--persist", rest = the agent subcommand. --yolo auto-approves so no tool + # prompt can hang; --api-key attaches to the already-served CI model. + launch_turn() { + local out="$1" rflag="$2"; shift 2 + local flag=(); [ -n "$rflag" ] && flag=("$rflag") + run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \ + --api-key "$UNSLOTH_API_KEY" "$@" + local rc=$? + redact "$out" + return "$rc" + } + + # One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED + # from the session-store delta. Runs in the main shell (not a command + # substitution) so a hang's guide_fail actually fails the job and the + # progress lines reach the CI log. $1 = "" (baseline) or "--persist". + RESULT="" + run_pass() { + local rflag="$1" label="baseline" + [ -n "$rflag" ] && label="resume" + rm -rf "$WORK"; mkdir -p "$WORK" + set_t1_cmd + local out="$LOGS_DIR/${AGENT}-resume-${label}.txt" + local before after rc + before="$(count_session_files)" + pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK" + launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$? + popd >/dev/null || true + after="$(count_session_files)" + echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})" + # The turn must succeed for the delta to mean anything: an agent that writes a + # session file then errors would otherwise be misread as PERSISTED. Mirror the + # file-edit mode and fail the pass on a non-zero launch (the flagship codex recall + # below stays WARN-only, driven by its own launch_turn calls). + [ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \ + guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; } + if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi + } + + run_pass ""; BASELINE="$RESULT" + # Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix. + # opencode/claude persist either way, so the baseline already proves it and a + # second full CPU turn only risks a timeout; skip it for them. + case "$AGENT" in + codex|pi) run_pass "--persist"; RESUME="$RESULT" ;; + *) RESUME="n/a (persists either way)" ;; + esac + + # Expected: codex/pi relocate their whole home to the temp dir, so a plain + # launch is WIPED and only --persist PERSISTS. opencode/claude keep their + # session data in a fixed user dir, so the baseline already PERSISTS. + case "$AGENT" in + codex|pi) EXPECT_BASELINE="WIPED" ;; + opencode|claude) EXPECT_BASELINE="PERSISTED" ;; + esac + + echo "──────────────────────────────────────────────" + echo "[$AGENT] RESUME EXPERIMENT" + echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})" + echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}" + echo "──────────────────────────────────────────────" + + [ "$BASELINE" = "$EXPECT_BASELINE" ] \ + || guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}" + case "$AGENT" in + codex|pi) + [ "$RESUME" = "PERSISTED" ] \ + || guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;; + esac + + # Flagship behavioral proof (codex only, WARN-only): after a --persist plant, + # resume the session and check the model actually recalls the codeword. A + # miss is not a failure (the CI model is small); the mechanism gate above is + # the real assertion. + if [ "$AGENT" = "codex" ]; then + rm -rf "$WORK"; mkdir -p "$WORK" + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true + if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then + echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}" + else + echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed" + fi + fi + echo "[$AGENT] resume OK" + ;; + + *) + echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2 + exit 2 + ;; +esac diff --git a/.github/scripts/agent-guides-install.sh b/.github/scripts/agent-guides-install.sh new file mode 100755 index 0000000..daf4bac --- /dev/null +++ b/.github/scripts/agent-guides-install.sh @@ -0,0 +1,108 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Install one coding-agent CLI for the Local Agent Guides CI. Isolated as +# failure class (b) "agent package install failed": npm/curl flakiness here +# is the single biggest source of false reds, so installs retry with +# backoff and the only ::error:: this script can emit is class (b). The +# install recipes mirror the install_hint strings in +# unsloth_cli/commands/start.py at HEAD. +# +# Usage: agent-guides-install.sh +# agent in: claude codex hermes openclaw opencode pi +set -uo pipefail + +AGENT="${1:?usage: agent-guides-install.sh }" +mkdir -p logs +LOG="logs/install-${AGENT}.log" + +install_fail() { + echo "::error::[agent install failed] agent=${AGENT}: $* (class (b): the agent CLI did not install; not a server or guide problem)." >&2 + echo "---- tail $LOG ----" >&2 + tail -60 "$LOG" 2>/dev/null || true + exit 1 +} + +# npm registry flakiness is common in CI; retry 3x with linear backoff. +# Extra npm flags may precede the package (e.g. npm_retry --ignore-scripts pkg). +npm_retry() { + local i + for i in 1 2 3; do + if npm install -g "$@" >> "$LOG" 2>&1; then + return 0 + fi + echo "[install] npm install -g $* attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + sleep "$((i * 10))" + done + return 1 +} + +# curl|bash installers, retried at the curl layer. We download to a temp file +# first and only execute on a fully successful fetch, so a truncated download +# (network hiccup mid-stream) can never run a half-written installer. +curl_bash() { + local url="$1"; shift + local i tmp + tmp="$(mktemp)" + for i in 1 2 3; do + if curl -fsSL --retry 3 --retry-delay 5 "$url" -o "$tmp" 2>>"$LOG" \ + && bash "$tmp" "$@" >> "$LOG" 2>&1; then + rm -f "$tmp" + return 0 + fi + echo "[install] curl|bash $url attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + sleep "$((i * 10))" + done + rm -f "$tmp" + return 1 +} + +echo "[install] agent=$AGENT (log=$LOG)" +case "$AGENT" in + claude) + # start.py install_hint: curl -fsSL https://claude.ai/install.sh | bash + curl_bash "https://claude.ai/install.sh" || install_fail "claude installer failed" + # The installer drops the binary under ~/.local/bin. + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + ;; + codex) + # start.py install_hint: npm install -g @openai/codex + npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed" + ;; + opencode) + # start.py install_hint: npm install -g opencode-ai + npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed" + ;; + openclaw) + # start.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash + # npm is the more deterministic path in CI and matches the agent's docs; + # fall back to the start.py curl installer if the npm tag is missing. + if ! npm_retry "openclaw@latest"; then + curl_bash "https://openclaw.ai/install.sh" || install_fail "openclaw install failed (npm + curl)" + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + fi + ;; + hermes) + # start.py install_hint: + # curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash + curl_bash "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh" \ + --non-interactive --skip-setup --skip-browser --no-skills \ + || install_fail "hermes installer failed" + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + ;; + pi) + # start.py install_hint: npm install -g --ignore-scripts @earendil-works/pi-coding-agent + # (--ignore-scripts matches Pi's documented recipe; exercising the exact hint + # catches guide drift). The CLI moved from the now-deprecated @mariozechner + # scope to @earendil-works (the old scope is frozen, so installing it would + # test a stale Pi against the API). + npm_retry --ignore-scripts "@earendil-works/pi-coding-agent" \ + || install_fail "npm install -g --ignore-scripts @earendil-works/pi-coding-agent failed" + ;; + *) + install_fail "unknown agent '$AGENT'" + ;; +esac + +echo "[install] OK for $AGENT" diff --git a/.github/scripts/assert-llama-loads.sh b/.github/scripts/assert-llama-loads.sh new file mode 100755 index 0000000..c2ffe27 --- /dev/null +++ b/.github/scripts/assert-llama-loads.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Assert Studio installed a llama.cpp that loads and runs on THIS macOS. Tests +# the contract that matters (binaries load and their minimum-OS is <= this host) +# instead of the old "did install.sh fall back to a source build?" grep, since a +# source build with a correct deployment target is a valid outcome. +set -uo pipefail + +UNSLOTH_HOME="${STUDIO_HOME:-$HOME/.unsloth}" +LLAMA_DIR="${LLAMA_CPP_DIR:-$UNSLOTH_HOME/llama.cpp}" +BIN_DIR="$LLAMA_DIR/build/bin" + +fail() { + echo "::error::$*" + if [ -f logs/install.log ]; then + echo "---- install.log (llama.cpp lines) ----" + grep -E "llama-prebuilt|llama\.cpp|macos prebuilt|falling back" logs/install.log | tail -80 || true + fi + exit 1 +} + +SERVER="$(find "$LLAMA_DIR" -type f -name 'llama-server' 2>/dev/null | head -1)" +QUANT="$(find "$LLAMA_DIR" -type f -name 'llama-quantize' 2>/dev/null | head -1)" +[ -n "$SERVER" ] || fail "llama-server not found under $LLAMA_DIR after install" +[ -n "$QUANT" ] || fail "llama-quantize not found under $LLAMA_DIR after install" + +HOST_VER="$(sw_vers -productVersion 2>/dev/null || echo '0')" +HOST_MAJOR="${HOST_VER%%.*}" + +# Static minimum-OS check on every Mach-O we ship. vtool ships with the Xcode +# command line tools, which GitHub macOS runners always have; if it is somehow +# missing we skip the static check and rely on the runtime launch below. +if command -v vtool >/dev/null 2>&1; then + while IFS= read -r macho; do + [ -n "$macho" ] || continue + minos="$(vtool -show-build "$macho" 2>/dev/null | awk '/minos/{print $2; exit}')" + [ -n "$minos" ] || continue + min_major="${minos%%.*}" + if [ "$min_major" -gt "$HOST_MAJOR" ] 2>/dev/null; then + fail "$(basename "$macho") is built for macOS $minos but this runner is macOS $HOST_VER (prebuilt is newer than the host)" + fi + done < <(find "$BIN_DIR" -type f \( -name '*.dylib' -o -name 'llama-server' -o -name 'llama-quantize' \) 2>/dev/null) +fi + +# Runtime launch: --version forces dyld to load every linked dylib (including +# libggml-metal.dylib). A missing Metal symbol or too-new binary fails here. +if ! "$SERVER" --version >/tmp/llama-server-version.txt 2>&1; then + echo "---- llama-server --version output ----" + cat /tmp/llama-server-version.txt || true + fail "llama-server failed to launch on macOS $HOST_VER (dyld load / symbol error)" +fi + +echo "llama.cpp load validation passed on macOS $HOST_VER" +echo " server: $SERVER" +sed -n '1,4p' /tmp/llama-server-version.txt 2>/dev/null || true diff --git a/.github/scripts/assert-prompt-cache.sh b/.github/scripts/assert-prompt-cache.sh new file mode 100755 index 0000000..f5b6b07 --- /dev/null +++ b/.github/scripts/assert-prompt-cache.sh @@ -0,0 +1,238 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Prompt-cache (KV-cache prefix reuse) detection, two strategies in one helper: +# +# mode=api A 2-turn /v1/chat/completions probe. Turn 2 prepends turn 1 + +# its reply, so the shared prefix must be served from llama.cpp's +# KV cache. Asserts usage.prompt_tokens_details.cached_tokens > 0 +# on turn 2. This is the OpenAI-dialect server cache sanity. +# WHY this works on chat completions: the chat path forwards +# llama-server's real cached_tokens through +# studio/backend/routes/inference.py:482-489 (_prompt_tokens_details) +# into prompt_tokens_details (inference.py:519). +# +# mode=log Read the llama-server log and decide HIT vs MISS from the +# prompt-reprocessing trace. WHY the log (not the API field): +# the Anthropic /v1/messages path builds AnthropicUsage( +# input_tokens=..., output_tokens=...) at inference.py:8787-8790 +# / :8829-8832 and NEVER sets cache_read_input_tokens, which +# therefore stays at its model default of 0 +# (studio/backend/models/inference.py:1655). So an Anthropic-path +# client (Claude Code, OpenClaw is openai-completions but Claude +# Code is the canonical Anthropic agent) can get a real KV-cache +# hit that the API usage field reports as 0. The only ground +# truth for the Anthropic path is the llama-server log. +# +# Log location (verified): studio/backend/core/inference/llama_cpp.py:4363-4365 +# _swa_cache_path().parent/"logs"/"llama-server"/llama-[label]-port-

[-try].log +# _swa_cache_path() => $UNSLOTH_STUDIO_HOME|$STUDIO_HOME or ~/.unsloth/studio +# (llama_cpp.py:337-340). So default: ~/.unsloth/studio/logs/llama-server/. +# +#

is the INTERNAL llama-server port (self._find_free_port(), +# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Studio port. So we must +# NOT filter the log glob by STUDIO_PORT (the brief's `port-` +# glob would never match). We pick the newest llama-*.log instead. +# +# Usage: +# assert-prompt-cache.sh api BASE_URL API_KEY +# assert-prompt-cache.sh log EXPECT # EXPECT = HIT | MISS +# # reads MARKER_BEFORE/MARKER_AFTER +# # byte offsets from env (see below) +# assert-prompt-cache.sh mark # print current log size to stdout +# # (use to bracket a turn) +# +# Env for mode=log: +# LLAMA_LOG_DIR override the log dir (default ~/.unsloth/studio/logs/llama-server) +# CACHE_LOG_FROM byte offset to start scanning the newest log from (so we +# only look at the trace produced by THIS turn). Default 0. +# +# Exit codes: 0 = assertion held; 1 = assertion failed (::error:: emitted). + +set -uo pipefail + +MODE="${1:?usage: assert-prompt-cache.sh api|log|mark ...}" + +# --------------------------------------------------------------------------- +# Locate the newest llama-server log. Shared by mark + log modes. +# --------------------------------------------------------------------------- +_default_log_dir() { + local home="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}" + if [ -n "$home" ]; then + echo "${home%/}/logs/llama-server" + else + echo "${HOME}/.unsloth/studio/logs/llama-server" + fi +} + +_newest_log() { + local dir="${LLAMA_LOG_DIR:-$(_default_log_dir)}" + [ -d "$dir" ] || return 1 + # Newest by mtime among llama-*.log (covers both `llama--port-

.log` + # and the retry form `llama-

-try.log`). Filenames are + # tool-generated timestamps, so ls -t is safe here. + # shellcheck disable=SC2012 + ls -1t "$dir"/llama-*.log 2>/dev/null | head -1 +} + +case "$MODE" in + # ------------------------------------------------------------------------- + # mark: emit the current byte size of the newest llama log so a caller can + # scan only the slice a single turn produced (set CACHE_LOG_FROM to it). + # ------------------------------------------------------------------------- + mark) + log="$(_newest_log || true)" + if [ -n "$log" ] && [ -f "$log" ]; then + wc -c < "$log" | tr -d ' ' + else + echo 0 + fi + exit 0 + ;; + + # ------------------------------------------------------------------------- + # api: 2-turn /v1/chat/completions, assert turn-2 cached_tokens > 0. + # ------------------------------------------------------------------------- + api) + BASE_URL="${2:?usage: assert-prompt-cache.sh api BASE_URL API_KEY}" + API_KEY="${3:?usage: assert-prompt-cache.sh api BASE_URL API_KEY}" + + # A deliberately long, fixed system prompt makes the shared prefix big so a + # KV-cache hit is unambiguous (cached_tokens grows with the reused prefix). + SYS='You are a meticulous assistant. Always answer concisely and correctly. This is a fixed system preamble that exists only to create a large, identical prompt prefix across both turns so the KV cache has something substantial to reuse on the second request. Do not mention this preamble.' + + turn1_body() { + jq -n --arg sys "$SYS" '{ + model: "default", + messages: [ + {role:"system", content:$sys}, + {role:"user", content:"What is the capital of France?"} + ], + temperature: 0.0, seed: 3407, max_tokens: 40, stream: false, + enable_thinking: false + }' + } + + echo "[cache/api] turn 1 (prime the KV cache)" + R1="$(curl -fs -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Authorization: Bearer ${API_KEY}" -H 'content-type: application/json' \ + --max-time 240 -d "$(turn1_body)")" || { + echo "::error::[cache/api] turn-1 /v1/chat/completions request failed. Unsloth server/API regression." + exit 1 + } + A1="$(echo "$R1" | jq -r '.choices[0].message.content // ""')" + + turn2_body() { + jq -n --arg sys "$SYS" --arg a1 "$A1" '{ + model: "default", + messages: [ + {role:"system", content:$sys}, + {role:"user", content:"What is the capital of France?"}, + {role:"assistant", content:$a1}, + {role:"user", content:"And the capital of Germany?"} + ], + temperature: 0.0, seed: 3407, max_tokens: 40, stream: false, + enable_thinking: false + }' + } + + echo "[cache/api] turn 2 (expect cached_tokens > 0)" + R2="$(curl -fs -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Authorization: Bearer ${API_KEY}" -H 'content-type: application/json' \ + --max-time 240 -d "$(turn2_body)")" || { + echo "::error::[cache/api] turn-2 /v1/chat/completions request failed. Unsloth server/API regression." + exit 1 + } + + CACHED="$(echo "$R2" | jq -r '.usage.prompt_tokens_details.cached_tokens // 0')" + PROMPT_TOK="$(echo "$R2" | jq -r '.usage.prompt_tokens // 0')" + echo "[cache/api] turn-2 usage: prompt_tokens=${PROMPT_TOK} cached_tokens=${CACHED}" + + if [ -z "$CACHED" ] || ! [ "$CACHED" -gt 0 ] 2>/dev/null; then + echo "::error::[cache/api] turn-2 usage.prompt_tokens_details.cached_tokens=${CACHED}, expected > 0. The server is not surfacing llama.cpp KV-cache hits on /v1/chat/completions. Check studio/backend/routes/inference.py:482-489 (_prompt_tokens_details) and :519. Full turn-2 usage:" + echo "$R2" | jq -c '.usage' 2>/dev/null || echo "$R2" + exit 1 + fi + echo "[cache/api] PASS server cache sanity (cached_tokens=${CACHED} > 0)" + exit 0 + ;; + + # ------------------------------------------------------------------------- + # log: classify the newest llama-server log (from CACHE_LOG_FROM bytes on) + # as HIT or MISS and compare to EXPECT. + # ------------------------------------------------------------------------- + log) + EXPECT="${2:?usage: assert-prompt-cache.sh log HIT|MISS}" + FROM="${CACHE_LOG_FROM:-0}" + + log="$(_newest_log || true)" + if [ -z "$log" ] || [ ! -f "$log" ]; then + echo "::error::[cache/log] no llama-server log under ${LLAMA_LOG_DIR:-$(_default_log_dir)}. Cannot read KV-cache trace. (Path contract: studio/backend/core/inference/llama_cpp.py:4363-4365.)" + exit 1 + fi + echo "[cache/log] reading $log from byte $FROM" + + # Scan only the slice produced after FROM. + slice="$(tail -c "+$((FROM + 1))" "$log" 2>/dev/null || cat "$log")" + + # ---- HIT detectors (most-specific first) ----------------------------- + # 1. Modern + legacy "re-used N tokens" / "reused N" (N>0). Primary signal + # per the design brief. + reused_n="$(printf '%s\n' "$slice" \ + | grep -aoiE 're-?used[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 2. "kv cache rm [START, end)" with START>0 => prefix [0,START) reused. + cache_rm_start="$(printf '%s\n' "$slice" \ + | grep -aoiE 'kv cache rm \[[0-9]+' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 3. "n_past = N" with N>0 after a prompt-processing line (prefix kept). + n_past_n="$(printf '%s\n' "$slice" \ + | grep -aoiE 'n_past[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 4. tokens_cached / tokens from cache (some builds). + tok_cached="$(printf '%s\n' "$slice" \ + | grep -aoiE 'tokens_cached[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + + # ---- MISS detectors -------------------------------------------------- + # Explicit forced full re-processing (SWA / recurrent) or kv cache rm [0,. + forced_full=0 + if printf '%s\n' "$slice" | grep -aqiE 'forcing full prompt re-?processing|kv cache rm \[0,'; then + forced_full=1 + fi + + HIT=0 + why="" + if [ -n "$reused_n" ] && [ "$reused_n" -gt 0 ] 2>/dev/null; then + HIT=1; why="re-used=$reused_n" + elif [ -n "$cache_rm_start" ] && [ "$cache_rm_start" -gt 0 ] 2>/dev/null; then + HIT=1; why="kv-cache-rm-start=$cache_rm_start" + elif [ -n "$tok_cached" ] && [ "$tok_cached" -gt 0 ] 2>/dev/null; then + HIT=1; why="tokens_cached=$tok_cached" + elif [ "$forced_full" = "0" ] && [ -n "$n_past_n" ] && [ "$n_past_n" -gt 0 ] 2>/dev/null; then + # n_past>0 is the weakest signal; only trust it if nothing forced a full + # reprocess. (On a cold slot n_past tracks total processed, so it is a + # last-resort fallback per the brief.) + HIT=1; why="n_past=$n_past_n(fallback)" + fi + [ "$HIT" = "1" ] || why="${why:-no-reuse-markers (forced_full=$forced_full)}" + + OBSERVED="MISS"; [ "$HIT" = "1" ] && OBSERVED="HIT" + echo "[cache/log] observed=$OBSERVED expected=$EXPECT ($why)" + + if [ "$OBSERVED" != "$EXPECT" ]; then + echo "::error::[cache/log] KV-cache observed=$OBSERVED but expected=$EXPECT ($why). See the attribution A/B note in the workflow." + echo "---- llama-server log slice (last 60 lines) ----" + printf '%s\n' "$slice" | tail -60 + exit 1 + fi + echo "[cache/log] PASS ($OBSERVED == $EXPECT)" + exit 0 + ;; + + *) + echo "::error::unknown mode '$MODE' (want api|log|mark)" + exit 1 + ;; +esac diff --git a/.github/scripts/ci-connect-prompt.txt b/.github/scripts/ci-connect-prompt.txt new file mode 100644 index 0000000..2d96f2b --- /dev/null +++ b/.github/scripts/ci-connect-prompt.txt @@ -0,0 +1 @@ +You are a helpful assistant in a CI connectivity check. Answer the user directly in plain text. Do not use any tools, do not take any actions, and do not explain. Just reply with the answer. diff --git a/.github/scripts/ci-min-system-prompt.txt b/.github/scripts/ci-min-system-prompt.txt new file mode 100644 index 0000000..d55b828 --- /dev/null +++ b/.github/scripts/ci-min-system-prompt.txt @@ -0,0 +1 @@ +You are a coding assistant running non-interactively in a CI smoke test. Use the available file-editing and shell tools to complete the user's request directly and concisely. Do not ask questions or explain; just do the task. diff --git a/.github/scripts/hf-download-with-retry.sh b/.github/scripts/hf-download-with-retry.sh new file mode 100755 index 0000000..013a459 --- /dev/null +++ b/.github/scripts/hf-download-with-retry.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# +# Download a single file from a Hugging Face repo with a stall-retry +# watchdog. Used by the Studio CI workflows so a hung hf-xet transfer +# kills + retries instead of silently consuming the job's timeout. +# +# Usage: hf-download-with-retry.sh REPO FILE LOCAL_DIR +# +# Why this exists +# --------------- +# huggingface_hub 1.15+ deprecated `hf_transfer` and routes every +# transfer through the `hf-xet` binary package. In CI we observed +# `hf download` on a 3 GB GGUF (gemma-4-E2B-it-UD-Q4_K_XL) progress +# to ~46% via Xet, then go completely silent for the remainder of +# the 30-min job timeout -- no progress bytes, no error, no exit. +# A sibling 940 MB mmproj on the same step downloaded in ~21s +# moments earlier, so the hang is per-file inside hf-xet rather +# than a network outage. The Xet env-vars below put hf-xet into +# its highest-throughput mode and force a 500 s client-read +# timeout; the watchdog loop ensures a stall does not eat the +# whole job: if the hf process has not exited after STALL_S +# seconds (default 180 = 3 min), we SIGTERM, then SIGKILL, then +# start a fresh attempt. Retries are unbounded -- the enclosing +# GitHub Actions job's `timeout-minutes` is the real bound. +# +# See https://huggingface.co/docs/huggingface_hub/package_reference/environment_variables +# for the HF_XET_* documentation, and npm/cli#7308's pattern (silent +# CI hang with no error) for prior art on this class of failure. + +set -uo pipefail + +REPO="${1:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}" +FILE="${2:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}" +# LOCAL_DIR is optional. If empty, hf falls back to HF_HUB_CACHE +# (~/.cache/huggingface/hub) which is the desired path for callers +# that populate HF_HOME for a downstream Studio model load. +LOCAL_DIR="${3:-}" + +# Stall threshold per attempt, in seconds. Override with +# HF_DOWNLOAD_STALL_SECONDS in the workflow env if 3 min is too tight +# for a specific runner / file. The script keeps retrying past this +# until the job timeout fires. +STALL_S="${HF_DOWNLOAD_STALL_SECONDS:-180}" + +# hf-xet tuning. HF_HUB_ENABLE_HF_TRANSFER is deliberately NOT set -- +# it is a no-op on huggingface_hub>=1.15 and only emits a deprecation +# FutureWarning. The five HF_XET_* knobs below mirror the settings +# Daniel asked for: max bandwidth + 64 parallel range gets, no chunk +# cache (download-once usage pattern), parallel disk writes (SSD/NVMe +# runners), and a generous 500 s read timeout so individual chunk +# requests fail loudly instead of stalling forever. +export HF_XET_HIGH_PERFORMANCE=1 +export HF_XET_CHUNK_CACHE_SIZE_BYTES=0 +export HF_XET_NUM_CONCURRENT_RANGE_GETS=64 +export HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY=0 +export HF_XET_CLIENT_READ_TIMEOUT=500 + +if [ -n "$LOCAL_DIR" ]; then + mkdir -p "$LOCAL_DIR" +fi + +attempt=1 +while : ; do + log="$(mktemp -t hf-download.XXXXXX)" + echo "[hf-download] $FILE attempt $attempt (stall threshold ${STALL_S}s, log=$log)" + + if [ -n "$LOCAL_DIR" ]; then + hf download "$REPO" "$FILE" --local-dir "$LOCAL_DIR" > "$log" 2>&1 & + else + hf download "$REPO" "$FILE" > "$log" 2>&1 & + fi + pid=$! + + elapsed=0 + while kill -0 "$pid" 2>/dev/null && [ "$elapsed" -lt "$STALL_S" ]; do + sleep 5 + elapsed=$((elapsed + 5)) + done + + if kill -0 "$pid" 2>/dev/null; then + echo "[hf-download] $FILE attempt $attempt exceeded ${STALL_S}s -- killing PID $pid and retrying" + kill -TERM "$pid" 2>/dev/null || true + sleep 2 + kill -KILL "$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null || true + echo "[hf-download] $FILE attempt $attempt log tail (last 40 lines):" + tail -40 "$log" || true + attempt=$((attempt + 1)) + continue + fi + + if wait "$pid"; then + rc=0 + else + rc=$? + fi + + if [ "$rc" -eq 0 ]; then + echo "[hf-download] $FILE attempt $attempt succeeded" + tail -20 "$log" || true + exit 0 + fi + + echo "[hf-download] $FILE attempt $attempt failed (exit $rc) -- retrying" + tail -40 "$log" || true + attempt=$((attempt + 1)) +done diff --git a/.github/scripts/serve-unsloth-run.sh b/.github/scripts/serve-unsloth-run.sh new file mode 100755 index 0000000..6ac98de --- /dev/null +++ b/.github/scripts/serve-unsloth-run.sh @@ -0,0 +1,172 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Boot `unsloth run --disable-tools` in the background, wait for it to be +# healthy, parse the minted API key from the banner, and resolve the +# /v1/models id. Exports everything downstream steps need into $GITHUB_ENV +# (or prints it when run outside Actions). Factored out of the workflow so +# the failure-isolation logic lives in one shellcheck-clean place. +# +# Usage: +# serve-unsloth-run.sh --model REPO --gguf-variant VAR --port PORT \ +# [--gguf-file PATH] [--extra "--seed 3407 --temp 0"] \ +# [--log-dir logs] [--health-timeout 300] +# +# Why a helper and not inline YAML +# -------------------------------- +# * Every `unsloth run` invocation here is the *Unsloth server* under test. +# A failure to come up healthy is class (a) "server/API regression" and +# must be reported with a distinct `::error::` BEFORE any agent runs. +# * The banner is the documented contract a human copies from. We parse the +# exact `API Key:` line printed by unsloth_cli/commands/studio.py +# (` API Key: ` non-silent, `API Key: ` silent) so a +# silent change to that line is also caught. +# * `unsloth run` re-execs into the studio venv ($STUDIO_HOME/unsloth_studio), +# so in CI after `install.sh --local` it runs the PR's repo code. +# +# Outputs written to $GITHUB_ENV (and echoed): +# UNSLOTH_API_KEY the sk-unsloth-* key minted on the banner +# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth start` +# finds THIS server, not the hardcoded :8888) +# UNSLOTH_BASE_URL same as UNSLOTH_STUDIO_URL (alias for clarity) +# UNSLOTH_MODEL_ID the canonical id reported by /v1/models +# UNSLOTH_SERVER_PID pid of the backgrounded `unsloth run` +# UNSLOTH_LLAMA_LOG_DIR ~/.unsloth/studio/logs/llama-server + +set -uo pipefail + +# ── arg parse ──────────────────────────────────────────────────────────── +MODEL="" +GGUF_VARIANT="" +GGUF_FILE="" +PORT="" +EXTRA="" +LOG_DIR="logs" +HEALTH_TIMEOUT="300" + +while [ "$#" -gt 0 ]; do + case "$1" in + --model) MODEL="$2"; shift 2 ;; + --gguf-variant) GGUF_VARIANT="$2"; shift 2 ;; + --gguf-file) GGUF_FILE="$2"; shift 2 ;; + --port) PORT="$2"; shift 2 ;; + --extra) EXTRA="$2"; shift 2 ;; + --log-dir) LOG_DIR="$2"; shift 2 ;; + --health-timeout) HEALTH_TIMEOUT="$2"; shift 2 ;; + *) echo "serve-unsloth-run.sh: unknown arg '$1'" >&2; exit 2 ;; + esac +done + +[ -n "$PORT" ] || { echo "serve-unsloth-run.sh: --port is required" >&2; exit 2; } +if [ -z "$MODEL" ] && [ -z "$GGUF_FILE" ]; then + echo "serve-unsloth-run.sh: one of --model or --gguf-file is required" >&2 + exit 2 +fi + +mkdir -p "$LOG_DIR" +SERVER_LOG="$LOG_DIR/unsloth-run-${PORT}.log" +BASE_URL="http://127.0.0.1:${PORT}" +STUDIO_HOME_DIR="${STUDIO_HOME:-$HOME/.unsloth/studio}" +LLAMA_LOG_DIR="${STUDIO_HOME_DIR}/logs/llama-server" + +# Emit a key=value pair to $GITHUB_ENV when set, always echo for local runs. +emit() { + echo "$1=$2" + if [ -n "${GITHUB_ENV:-}" ]; then + echo "$1=$2" >> "$GITHUB_ENV" + fi +} + +server_fail() { + echo "::error::Unsloth server/API regression: $*" >&2 + echo "---- last 200 lines of $SERVER_LOG ----" >&2 + tail -200 "$SERVER_LOG" 2>/dev/null || true + exit 1 +} + +# ── port collision guard ───────────────────────────────────────────────── +# A leftover listener (or a parallel matrix cell that wandered onto our port) +# would make us attach to the wrong server and mask a real regression. Fail +# fast instead. +if command -v ss >/dev/null 2>&1; then + if ss -tln 2>/dev/null | grep -q ":${PORT}\b"; then + server_fail "port ${PORT} already has a listener before we started (collision)" + fi +fi + +# ── build the command ──────────────────────────────────────────────────── +# `unsloth run` == alias of `unsloth studio run`. --disable-tools is REQUIRED +# (passthrough mode) so the agent's own tools relay instead of the server's. +# --no-cloudflare keeps us off the network (loopback bind, no tunnel attempt). +CMD=(unsloth run -H 127.0.0.1 -p "$PORT" --disable-tools --no-cloudflare) +if [ -n "$GGUF_FILE" ]; then + CMD+=(--model "$GGUF_FILE") +else + CMD+=(--model "$MODEL") + [ -n "$GGUF_VARIANT" ] && CMD+=(--gguf-variant "$GGUF_VARIANT") +fi +# Determinism knobs + any caller passthrough (e.g. --seed 3407 --temp 0). +# shellcheck disable=SC2206 # intentional word-split of caller-controlled flags +[ -n "$EXTRA" ] && CMD+=($EXTRA) + +echo "[serve] launching: ${CMD[*]}" +echo "[serve] server log: $SERVER_LOG" + +# Run detached, no controlling TTY (setsid avoids any TTY-prompt hang and +# detaches from this step's process group so the job's teardown is clean). +setsid "${CMD[@]}" > "$SERVER_LOG" 2>&1 < /dev/null & +SERVER_PID=$! +emit UNSLOTH_SERVER_PID "$SERVER_PID" + +# ── wait for /api/health == healthy ────────────────────────────────────── +HEALTHY=0 +for _ in $(seq 1 "$HEALTH_TIMEOUT"); do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + server_fail "process exited before becoming healthy (pid $SERVER_PID)" + fi + if curl -fs "${BASE_URL}/api/health" -o "$LOG_DIR/health-${PORT}.json" 2>/dev/null; then + if jq -e '.status == "healthy"' "$LOG_DIR/health-${PORT}.json" >/dev/null 2>&1; then + HEALTHY=1 + break + fi + fi + sleep 1 +done +[ "$HEALTHY" = "1" ] || server_fail "did not report /api/health healthy within ${HEALTH_TIMEOUT}s" +echo "[serve] /api/health healthy" + +# ── parse the API key from the banner ──────────────────────────────────── +# Match both the non-silent " API Key: " and silent "API Key: " +# forms. We do NOT trust a fixed column count; we take the sk-unsloth-* token. +API_KEY="" +for _ in $(seq 1 30); do + API_KEY="$(grep -aoE 'sk-unsloth-[A-Za-z0-9_-]+' "$SERVER_LOG" 2>/dev/null | head -1 || true)" + [ -n "$API_KEY" ] && break + sleep 1 +done +if [ -z "$API_KEY" ]; then + # Fallback: take whatever follows an "API Key:" label, in case the key + # prefix scheme changes. Still a parse-fragility guard, not silent. + API_KEY="$(grep -aE 'API Key:' "$SERVER_LOG" 2>/dev/null \ + | sed -E 's/.*API Key:[[:space:]]*//' | head -1 || true)" +fi +[ -n "$API_KEY" ] || server_fail "could not parse an API key from the banner (banner-parse fragility -- check the 'API Key:' line in unsloth_cli/commands/studio.py)" +echo "::add-mask::${API_KEY}" +emit UNSLOTH_API_KEY "$API_KEY" + +# ── resolve /v1/models id ──────────────────────────────────────────────── +if ! curl -fs "${BASE_URL}/v1/models" \ + -H "Authorization: Bearer ${API_KEY}" -o "$LOG_DIR/models-${PORT}.json" 2>/dev/null; then + server_fail "/v1/models did not respond (or rejected the banner key)" +fi +MODEL_ID="$(jq -r '.data[0].id // empty' "$LOG_DIR/models-${PORT}.json" 2>/dev/null || true)" +[ -n "$MODEL_ID" ] || server_fail "/v1/models returned no model id (model failed to load)" +echo "[serve] resolved model id: $MODEL_ID" + +emit UNSLOTH_MODEL_ID "$MODEL_ID" +emit UNSLOTH_STUDIO_URL "$BASE_URL" +emit UNSLOTH_BASE_URL "$BASE_URL" +emit UNSLOTH_LLAMA_LOG_DIR "$LLAMA_LOG_DIR" + +echo "[serve] server is up: ${BASE_URL} (model ${MODEL_ID})" diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml new file mode 100644 index 0000000..1bb4c2b --- /dev/null +++ b/.github/workflows/consolidated-tests-ci.yml @@ -0,0 +1,2356 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# One consolidated CPU-only job that runs every test_* function the existing +# CI does not already cover from this repo plus the full unsloth_zoo@main +# CPU test suite plus unsloth_zoo.compiler.test_apply_fused_lm_head. +# +# Why a separate workflow: +# - studio-backend-ci.yml's "Repo tests (CPU)" job already auto-discovers +# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 16 +# Bucket-A tests below live inside those --ignore dirs (CPU-runnable but +# historically excluded with their GPU siblings); pulling them out into +# a sibling job keeps the existing 760-passed baseline stable while we +# prove the new pieces are green. +# - unsloth_zoo has no CI on main today (.github/workflows/ is empty +# upstream as of HEAD 030e4ba). 106 of its 111 test_* functions are +# CPU-runnable; the 5 GPU/vLLM ones are deselected here. +# - test_apply_fused_lm_head lives at unsloth_zoo/compiler.py:1983, not +# under tests/, so it is not picked up by `pytest tests/`. It is a +# plain function with no fixtures: pure regex over transformers source +# strings, ~5-15 s wall, no GPU. +# +# Strict mode: every test step is gating (no `continue-on-error`). The +# upstream patch fixes that previously caused per-cell red have landed: +# - unslothai/unsloth#5319 (patch_fast_lora import, patch_sft_trainer +# Union, openenv OSError graceful skip). +# - unslothai/unsloth-zoo#628 (MoE coverage canary so old transformers +# skips legitimately while real discovery regressions still fail). +# After those merges every observed cell failure was one of these two +# things; if they regress we want a red cell, not a green-with-fail-prints +# cell. + +name: Core + +on: + pull_request: + paths: + - 'unsloth/**' + - 'unsloth_cli/**' + - 'studio/**' + - 'tests/**' + - 'pyproject.toml' + - '.github/workflows/consolidated-tests-ci.yml' + push: + branches: [main, pip] + workflow_dispatch: + inputs: + unsloth_zoo_ref: + description: 'unsloth_zoo git ref to test against (default main)' + required: false + default: 'main' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + consolidated: + # Matrix: three (transformers, TRL) combos cover the failure surface the + # PR cares about: + # 1. transformers==4.57.6 + TRL latest <1.0.0 (the just-before-5.x line) + # 2. transformers latest 5.x + TRL latest 1.x (the absolute upstream tip; + # currently 5.8.0 + 1.3.0, both BEYOND the unsloth/unsloth_zoo + # <=5.5.0 / <=0.24.0 caps -- the cell exists explicitly to surface + # drift signal) + # 3. transformers + TRL pinned by pyproject.toml's dependency entries + # (resolved dynamically at job time via tomllib) + # fail-fast: false so each cell runs independently and a transformers / + # TRL drift signal in one cell does not cancel the others. No + # job-level or per-step `continue-on-error` -- real test failures now + # fail the cell. Patches with legitimate CPU-runner preconditions + # (real CUDA dispatcher, runtime args) are explicitly skipped via + # NEEDS_PRECONDITION in the runtime check shim below. + strategy: + fail-fast: false + matrix: + combo: + - id: t4576-trl0latest + label: "HF=4.57.6 + TRL<1" + transformers_spec: "transformers==4.57.6" + trl_spec: "trl>=0.18.2,<1.0.0" + - id: tlatest5-trl1latest + label: "HF=latest + TRL=latest" + transformers_spec: "transformers>=5,<6" + trl_spec: "trl>=1,<2" + - id: pyproject + label: "HF=default + TRL=default" + transformers_spec: "__from_pyproject__" + trl_spec: "__from_pyproject__" + name: "Core (${{ matrix.combo.label }})" + runs-on: ubuntu-latest + timeout-minutes: 35 + # No job-level or per-step `continue-on-error`. Earlier iterations + # masked real test failures behind green check icons; that lie is + # gone. A failing test step fails the cell. NEEDS_PRECONDITION in + # the runtime check shim handles patches that legitimately cannot + # run on a CPU-only runner (real CUDA dispatcher, runtime args). + env: + UNSLOTH_ZOO_REF: ${{ inputs.unsloth_zoo_ref || 'main' }} + MATRIX_TRANSFORMERS_SPEC: ${{ matrix.combo.transformers_spec }} + MATRIX_TRL_SPEC: ${{ matrix.combo.trl_spec }} + MATRIX_COMBO_ID: ${{ matrix.combo.id }} + # Hoisted to job-level so every step (Sanity, Bucket-A, unsloth_zoo + # pytest, test_apply_fused_lm_head) inherits it. transformers' bundled + # *_pb2.py was generated against an older protoc; the C++ protobuf + # 4+/5+/6 implementation rejects them with "Descriptors cannot be + # created directly". The pure-Python parser bypasses the check; the + # speed cost is negligible for these tests. + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + PYTHONPATH: ${{ github.workspace }}/studio + UNSLOTH_COMPILE_DISABLE: '1' + # unsloth_zoo/__init__.py:314 raises ImportError unless UNSLOTH_IS_PRESENT + # is set — normally it is set by unsloth.__init__ when unsloth is imported + # first. In this job we sometimes import unsloth_zoo.* (e.g. + # unsloth_zoo.saving_utils, unsloth_zoo.temporary_patches) without going + # through `import unsloth` first; pin the env var to 1 so unsloth_zoo's + # bootstrap accepts it. Setting it has no effect on unsloth itself. + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + # Node 22 unblocks tests/studio/test_chat_preset_builtin_invariants.py's + # `node --experimental-strip-types` subprocess. Cheap to install; keeps + # the consolidated job self-sufficient even if studio-backend-ci.yml + # changes its node setup. + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - name: Install uv (some unsloth_zoo dev tooling expects it on PATH) + run: pip install uv + + - name: Resolve matrix specs (handle __from_pyproject__ sentinel) + # The pyproject cell uses a sentinel; resolve the real `transformers` + # and `trl` constraints from the project's pyproject.toml at job time. + # unsloth's pyproject puts the LLM stack pins in + # [project.optional-dependencies] under the `huggingfacenotorch` + # extra (top-level [project.dependencies] is just typer/pydantic/etc.), + # so we walk every optional extra and pick the first matching spec. + # Other cells pass their spec through unchanged. + run: | + set -euxo pipefail + python <<'PY' >> "$GITHUB_ENV" + import os, re, tomllib + spec_t = os.environ["MATRIX_TRANSFORMERS_SPEC"] + spec_r = os.environ["MATRIX_TRL_SPEC"] + + def _pkg_name(spec: str) -> str: + m = re.match(r"\s*([A-Za-z0-9_.-]+)", spec) + return (m.group(1).lower() if m else "") + + if spec_t == "__from_pyproject__" or spec_r == "__from_pyproject__": + with open("pyproject.toml", "rb") as f: + doc = tomllib.load(f) + proj = doc.get("project", {}) + # Try top-level deps first, then all optional extras. + all_deps: list[str] = list(proj.get("dependencies", [])) + for _name, dep_list in proj.get("optional-dependencies", {}).items(): + all_deps.extend(dep_list) + + if spec_t == "__from_pyproject__": + spec_t = next((x for x in all_deps if _pkg_name(x) == "transformers"), + "transformers") + if spec_r == "__from_pyproject__": + spec_r = next((x for x in all_deps if _pkg_name(x) == "trl"), + "trl") + print(f"RESOLVED_TRANSFORMERS_SPEC={spec_t}") + print(f"RESOLVED_TRL_SPEC={spec_r}") + PY + # Echo to logs so the matrix cell label maps cleanly to a spec. + grep RESOLVED_ "$GITHUB_ENV" || true + + - name: Install runtime deps (mirrors studio-backend-ci.yml + mlx-ci.yml) + # The shape matches studio-backend-ci.yml's "Repo tests (CPU)" install + # so we inherit the same CPU-spoof harness in tests/conftest.py and + # the same import-chain guarantees, plus the extra deps that the + # tests/saving + tests/utils Bucket-A files transitively need but + # which Repo tests (CPU) does not require because it --ignores + # those directories: + # - protobuf + sentencepiece: tests/saving/test_fix_sentencepiece_gguf_robustness.py + # does `from transformers.utils import sentencepiece_model_pb2`, + # which imports `google.protobuf`. Not pulled by transformers' + # base install. + # - triton: unsloth/_gpu_init.py:232 does an unconditional + # `import triton`. The triton PyPI wheel installs cleanly on + # Linux x86_64 even without CUDA (the import succeeds; runtime + # GPU work is what would fail, which we never do here). + # transformers + trl are matrix-parameterized. + run: | + set -euxo pipefail + python -m pip install --upgrade pip + pip install -r studio/backend/requirements/studio.txt + pip install \ + python-multipart aiofiles sqlalchemy cryptography \ + pyyaml jinja2 mammoth unpdf requests typer \ + 'numpy<3' pytest==9.0.3 pytest-asyncio httpx \ + protobuf sentencepiece triton \ + psutil packaging tqdm safetensors datasets \ + 'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \ + ipython + # torchvision: unsloth_zoo.vision_utils imports it at module scope. + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' + # transformers + trl from the matrix combo. + pip install "$RESOLVED_TRANSFORMERS_SPEC" + pip install "$RESOLVED_TRL_SPEC" + # bitsandbytes: hard import in unsloth/models/_utils.py. Recent + # versions ship a CPU build that imports cleanly on Linux. + pip install 'bitsandbytes>=0.45' + # unsloth itself, editable, no-deps so pip does not fight the + # explicit torch CPU-index install above. + pip install -e . --no-deps + echo "::group::Installed transformers + trl + torch + unsloth versions" + pip show transformers + pip show trl + pip show torch + pip show unsloth + echo "::endgroup::" + + - name: Clone unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} + # We need the repository tree (the wheel does not ship tests/), so + # clone shallow then editable-install so unsloth_zoo.* imports + # resolve to the cloned tree. We use `pip show` for the location + # check rather than `import unsloth_zoo` because the latter calls + # device_type.get_device_type() at module load and raises on a + # GPU-less runner; pytest steps below route through the existing + # tests/conftest.py spoof which handles that. + run: | + set -euxo pipefail + # github.com occasionally 500s on the git fetch; retry so a + # single upstream blip does not fail CI. + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 --branch="$UNSLOTH_ZOO_REF" \ + https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + pip install -e "$RUNNER_TEMP/unsloth-zoo" --no-deps + pip show unsloth_zoo + + - name: Sanity — collection only (both repos) + # Catches import-time breakage before we run the suite. Cheap; bails + # the job out fast if a transformers/torch resolution went sideways. + # Inherits PYTHONPATH / UNSLOTH_COMPILE_DISABLE / PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION + # from the job-level env block. + run: | + set -euxo pipefail + python -m pytest --collect-only -q \ + tests/saving/test_save_shell_injection.py \ + tests/saving/test_patch_saving_none_tokenizer.py \ + tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_compressed_export_schemes.py \ + tests/saving/test_export_api_surface.py \ + tests/saving/test_export_dispatch.py \ + tests/saving/test_imatrix_export.py \ + tests/utils/test_attention_masks.py \ + tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py + python -m pytest --collect-only -q "$RUNNER_TEMP/unsloth-zoo/tests/" + + - name: import_fixes drift detectors (18 tests, HARD GATE) + # One drift detector per fix_* / patch_* function in + # unsloth/import_fixes.py. The detectors assert the *healthy* + # upstream shape that the fix expects ABSENT the regression; + # ANY DRIFT DETECTED -> pytest.fail (NEVER skip) so the + # matrix cell goes red and the maintainer triages on the + # next PR, not in a downstream user's crash report. + # + # Pathologies covered by the suite (each maps to one fix + # function with the line range cited in the test docstring): + # * protobuf MessageFactory GetPrototype / GetMessageClass + # * datasets 4.4.x recursion range + # * TRL tuple-vs-bool _*_available caching + # * transformers PreTrainedModel.enable_input_require_grads + # source pattern flip + # * transformers torchcodec / causal_conv1d availability + # flags + # * transformers + accelerate is_wandb_available + # * peft.utils.transformers_weight_conversion importability + # + build_peft_weight_mapping signature + # * triton 3.6+ CompiledKernel num_ctas / cluster_dims + # * torch / torchvision pinned compatibility table + # * vllm guided_decoding_params / structured_outputs + + # aimv2 ovis config version + # * huggingface_hub is_offline_mode / HF_HUB_OFFLINE + # * torch.nn.init.trunc_normal_ presence (patch site for + # patch_trunc_normal_precision_issue) + # * xformers post-num_splits-key fix version + # HARD GATE: a red cell here is a real upstream regression + # without a corresponding zoo / unsloth-side workaround. + run: | + python -m pytest -v --tb=short tests/test_import_fixes_drift.py + + - name: public-api surface drift detectors (9 tests, HARD GATE) + # Companion to test_import_fixes_drift.py: that file catches + # third-party drift; this one catches drift in unsloth's OWN + # public surface (FastLanguageModel / FastVisionModel / + # FastModel + their classmethods + is_bf16_supported). A + # rename here would silently break the unslothai/notebooks tree + # one PR cycle later -- this gate catches it BEFORE the + # breakage reaches users. + run: | + python -m pytest -v --tb=short tests/test_public_api_surface.py + + - name: callback signature drift detector (HARD GATE) + # Catches the MLX-style bug from PR #5498: a producer in + # unsloth_zoo (or unsloth) grows a callback arg, but a consumer + # callback def still declares the old arity. The producer's + # try/except swallows the resulting TypeError and the symptom is + # "callback never fires" -- usually diagnosed downstream as a + # confusing assertion several seconds later. This static AST + # check fails fast at PR time. UNSLOTH_ZOO_SRC points at the + # freshly cloned main so the detector sees platform-specific + # submodules (e.g. unsloth_zoo/mlx/) that the released wheel + # may strip. + env: + UNSLOTH_ZOO_SRC: ${{ runner.temp }}/unsloth-zoo + run: | + python -m pytest -v --tb=short tests/test_callback_signature_drift.py + + - name: generation correctness guards (HARD GATE) + # Deterministic CPU guards, each validated to fail on its pre-fix code: + # leftpad = batched left-padded generation (#1066/#3699, fixed by + # #2216 + #4100; staging proof: unsloth-staging-2 PRs 170/172); + # rope_scaling_drift = config.rope_scaling dropped by replaced rotary + # classes (#2405). AST checks run first so import breakage cannot mask them. + run: | + python -m pytest -v --tb=short \ + tests/utils/test_prepare_inputs_leftpad.py \ + tests/utils/test_rope_scaling_drift.py + + - name: unsloth Bucket-A — CPU tests not in Repo tests (CPU) + # CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/ + # that Repo tests (CPU) --ignores. AST/protobuf/regex plus tiny CPU model + # loads; run cleanly here (transformers/torch installed). + run: | + python -m pytest -q --tb=short \ + tests/saving/test_save_shell_injection.py \ + tests/saving/test_patch_saving_none_tokenizer.py \ + tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_compressed_export_schemes.py \ + tests/saving/test_export_api_surface.py \ + tests/saving/test_export_dispatch.py \ + tests/saving/test_imatrix_export.py \ + tests/utils/test_attention_masks.py \ + tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py \ + tests/test_bad_mappings_redirect.py \ + tests/test_prefetch_snapshot_scope.py \ + tests/test_gemma_2b_mapper_key.py \ + --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' + # The deselected test monkeypatches flash_attn_varlen_func, which is + # only bound on the module when `flash_attn` is importable. flash_attn + # requires CUDA + dev toolchain, which the CPU-only ubuntu-latest + # runner does not have. The other Bucket-A tests pass cleanly. + + - name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU) + # 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip + # cases below auto-skip on a GPU-less runner; deselect them + # explicitly so the no-CUDA outcome is "deselected", not "skipped", + # making intent visible in the report. Env inherited from job block. + # + # test_get_peft_model_passes_finetune_last_n_layers_through is + # deselected because unsloth_zoo/mlx/loader.py at line 2972 calls + # model.trainable_parameters() on the fake-model fixture, which + # the test never stubbed; this fails on every platform regardless + # of CUDA. Tracked upstream as an unsloth_zoo bug; deselecting + # here unblocks unsloth CI until the loader fixture is fixed. + working-directory: ${{ runner.temp }}/unsloth-zoo + run: | + python -m pytest -q --tb=short tests/ \ + --deselect tests/test_unsloth_zoo_lora_merge.py::test_active_merge_device_returns_string_on_cuda_host \ + --deselect tests/test_unsloth_zoo_lora_merge.py::test_merge_lora_moves_cpu_inputs_to_active_device \ + --deselect tests/test_mlx_finetune_last_n_layers.py::test_get_peft_model_passes_finetune_last_n_layers_through + + - name: unsloth_zoo — test_apply_fused_lm_head (lives in compiler.py) + # `test_apply_fused_lm_head` lives at unsloth_zoo/compiler.py:1983, + # not under tests/, so pytest's default discovery does not pick it up. + # We route it through pytest by writing a one-shot shim test file + # inside the unsloth checkout's tests/ — pytest then walks UP and + # picks up tests/conftest.py, whose GPU-spoof harness (lines 84-141) + # patches torch.cuda.is_available, torch.cuda.memory.mem_get_info, + # torch.cuda.get_device_capability, and is_bf16_supported. That full + # spoof is required because unsloth_zoo/temporary_patches/gpt_oss.py + # at module load reads torch.cuda.memory.mem_get_info(0), which + # bare `is_available = True` doesn't cover. Env inherited. + run: | + set -euxo pipefail + cat > tests/_zoo_apply_fused_lm_head_shim.py <<'PY' + # Auto-generated by .github/workflows/consolidated-tests-ci.yml. + # Wraps unsloth_zoo.compiler.test_apply_fused_lm_head so that + # tests/conftest.py's GPU-spoof harness applies before the import. + # _zoo_aggressive_cuda_spoof extends conftest's harness with deeper + # patches (see tests/_zoo_aggressive_cuda_spoof.py). + import sys, pathlib + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + from unsloth_zoo.compiler import test_apply_fused_lm_head as _zoo_test + def test_zoo_apply_fused_lm_head_runs(): + _zoo_test() + PY + python -m pytest -q --tb=short tests/_zoo_apply_fused_lm_head_shim.py + rm -f tests/_zoo_apply_fused_lm_head_shim.py + + - name: Static checks — unsloth/trainer.py + unsloth/models/rl.py against latest pip TRL + # AST-only sanity: confirm both files parse and that every TRL symbol + # they reference still exists in the installed `trl`. Catches API + # drift (renamed / removed TRL classes) without running training. + # Pre-fetches latest pip transformers in case TRL pinned an older one. + run: | + set -euxo pipefail + # Use the matrix-resolved transformers + trl versions already + # installed by the runtime-deps step (don't upgrade here; that + # would defeat the matrix's purpose of testing against the + # specific (transformers, trl) combination the cell selected). + python <<'PY' + import ast, importlib, pathlib, sys + paths = [pathlib.Path("unsloth/trainer.py"), + pathlib.Path("unsloth/models/rl.py")] + for p in paths: + src = p.read_text() + tree = ast.parse(src, filename=str(p)) + # Collect every `from trl... import X` and `from trl... import (X, Y)` + missing = [] + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module and node.module.startswith("trl"): + mod = importlib.import_module(node.module) + for alias in node.names: + if alias.name == "*": + continue + if not hasattr(mod, alias.name): + missing.append(f"{node.module}.{alias.name}") + print(f"{p}: TRL symbols referenced and resolved -> {'OK' if not missing else 'MISSING ' + ', '.join(missing)}") + if missing: + sys.exit(1) + PY + + - name: Static checks — unsloth_zoo/tiled_mlp.py against latest pip transformers + # AST parse + transformers symbol-resolution. The user flagged tiled + # MLP patching as the path that breaks first when transformers ships + # an MLP class rename; this step is the canary against whatever + # transformers version the matrix cell selected. + working-directory: ${{ runner.temp }}/unsloth-zoo + run: | + set -euxo pipefail + python <<'PY' + import ast, importlib, pathlib, sys + p = pathlib.Path("unsloth_zoo/tiled_mlp.py") + src = p.read_text() + tree = ast.parse(src, filename=str(p)) + missing = [] + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module and node.module.startswith("transformers"): + try: + mod = importlib.import_module(node.module) + except Exception as e: + missing.append(f"{node.module} (import failed: {type(e).__name__})") + continue + for alias in node.names: + if alias.name == "*": + continue + if not hasattr(mod, alias.name): + missing.append(f"{node.module}.{alias.name}") + print(f"{p}: transformers symbols referenced -> {'OK' if not missing else 'MISSING ' + ', '.join(missing)}") + if missing: + sys.exit(1) + PY + + - name: Static checks — unsloth_zoo/hf_utils.py syntax + import-graph + working-directory: ${{ runner.temp }}/unsloth-zoo + run: | + set -euxo pipefail + python <<'PY' + import ast, pathlib + p = pathlib.Path("unsloth_zoo/hf_utils.py") + tree = ast.parse(p.read_text(), filename=str(p)) + # Surface every public function + class so the PR check log shows + # what's covered, not just OK/FAIL. + public = [] + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and not node.name.startswith("_"): + public.append(f"{type(node).__name__.replace('Def','').lower()}:{node.name}") + print(f"hf_utils.py public surface ({len(public)}): " + ", ".join(public)) + PY + + - name: Runtime checks — invoke every zero-arg patch_* across both repos (via pytest shim) + # Routed through pytest so tests/conftest.py's GPU-spoof harness + # applies before any unsloth_zoo.temporary_patches.* import. + # Locally validated 50/51 zero-arg patches succeed; the lone failure + # surfaces a real bug (unsloth.models._utils.patch_fast_lora raises + # NameError: name 'fast_lora_forward' is not defined). The shim + # reports the full ledger but only fails when one of the two + # `required` helpers is absent. + run: | + set -euxo pipefail + cat > tests/_runtime_patch_check_shim.py <<'PY' + # Auto-generated by .github/workflows/consolidated-tests-ci.yml. + # Wraps the runtime patch_* validation into a pytest test so the + # tests/conftest.py GPU-spoof harness applies. continue-on-error + # at the workflow level catches per-patch failures; this shim only + # asserts that the two `required` helpers are reachable. + import sys, pathlib + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + import importlib, inspect + + MODULES = [ + "unsloth.models._utils", "unsloth.models.rl", "unsloth.import_fixes", + "unsloth.kernels.cross_entropy_loss", "unsloth.kernels.rms_layernorm", + "unsloth.tokenizer_utils", "unsloth.save", + "unsloth_zoo.patching_utils", "unsloth_zoo.gradient_checkpointing", + "unsloth_zoo.loss_utils", "unsloth_zoo.tokenizer_utils", + "unsloth_zoo.tiled_mlp", "unsloth_zoo.dataset_utils", + "unsloth_zoo.patch_torch_functions", + "unsloth_zoo.temporary_patches.gemma", + "unsloth_zoo.temporary_patches.ministral", + "unsloth_zoo.temporary_patches.pixtral", + "unsloth_zoo.temporary_patches.deepseek_v3_moe", + "unsloth_zoo.temporary_patches.qwen3_5_moe", + "unsloth_zoo.temporary_patches.mxfp4", + "unsloth_zoo.temporary_patches.bitsandbytes", + "unsloth_zoo.temporary_patches.flex_attention_bwd", + ] + REQUIRED = { + "patch_unsloth_smart_gradient_checkpointing", + "patch_gradient_accumulation_fix", + } + # Patches whose signature looks zero-arg (`()` or all-defaulted) + # but which actually require either runtime args or real CUDA. + # Calling these in isolation is meaningless, so skip the + # invocation. Symbol presence (REQUIRED above) is still verified. + # patch_linear_scaling / patch_llama_rope_scaling: defaults are + # None placeholders; the bodies start with + # `assert is not None`. + # patch_unsloth_smart_gradient_checkpointing: legitimately + # allocates CUDA tensors via aten::empty.memory_format inside + # initialize_unsloth_gradient_checkpointing(); the + # torch.cuda.* spoof can't intercept that at the dispatcher + # level. + NEEDS_PRECONDITION = { + "patch_linear_scaling", + "patch_llama_rope_scaling", + "patch_unsloth_smart_gradient_checkpointing", + } + + def test_zero_arg_patch_invocations(): + ok, fail, args, skipped, miss_imports = 0, [], [], [], {} + seen_required = set() + for mod_name in MODULES: + try: + mod = importlib.import_module(mod_name) + except Exception as e: + miss_imports[mod_name] = f"{type(e).__name__}: {e}" + continue + for name in sorted(dir(mod)): + if not name.startswith("patch_"): continue + fn = getattr(mod, name, None) + if not callable(fn): continue + if name in REQUIRED: seen_required.add(name) + try: + sig = inspect.signature(fn) + need = [p.name for p in sig.parameters.values() + if p.default is inspect.Parameter.empty + and p.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.POSITIONAL_ONLY)] + except (TypeError, ValueError): + need = [] + if need: + args.append((mod_name, name, need)); continue + if name in NEEDS_PRECONDITION: + skipped.append(f"{mod_name}.{name}") + print(f" SKIP {mod_name}.{name} (needs precondition / CUDA)") + continue + try: + fn() + ok += 1 + print(f" OK {mod_name}.{name}") + except Exception as e: + fail.append((mod_name, name, type(e).__name__, str(e)[:200])) + print(f" FAIL {mod_name}.{name} -> {type(e).__name__}: {str(e)[:200]}") + print(f"\nzero-arg patch_*: ok={ok} fail={len(fail)} skipped={len(skipped)}") + print(f"arg-required patch_* (skipped, listed for review): {len(args)}") + for m, n, r in args: + print(f" needs={r}: {m}.{n}") + if skipped: + print(f"explicitly skipped (needs precondition / CUDA): {skipped}") + if miss_imports: + print("\nmodules failed to import (skipped):") + for k, v in miss_imports.items(): + print(f" {k}: {v}") + print(f"required patch_* helpers seen: {sorted(seen_required)}") + missing = REQUIRED - seen_required + assert not missing, f"required patch_* helpers MISSING: {sorted(missing)}" + # Strict: any zero-arg patch that raises is a real + # regression now that #5319 has landed (the three previously + # known-broken patches are fixed; legitimate + # CPU-precondition skips are recorded in NEEDS_PRECONDITION + # above, not in `fail`). Print all failures and re-raise + # them as one assertion message. + if fail: + raise AssertionError( + f"zero-arg patch_* invocation failures (ok={ok}, " + f"fail={len(fail)}, skipped={len(skipped)}):\n " + + "\n ".join( + f"{m}.{n} -> {ec}: {msg}" for m, n, ec, msg in fail + ) + ) + PY + python -m pytest -q --tb=short tests/_runtime_patch_check_shim.py -s + rm -f tests/_runtime_patch_check_shim.py + + - name: Runtime checks — patch_tiled_mlp on a synthetic MLP module (via pytest shim) + # Same shim pattern: pytest picks up tests/conftest.py before importing + # unsloth_zoo.tiled_mlp, so the GPU-spoof harness covers + # unsloth_zoo.temporary_patches.gpt_oss's mem_get_info call. + run: | + set -euxo pipefail + cat > tests/_tiled_mlp_check_shim.py <<'PY' + # Auto-generated by .github/workflows/consolidated-tests-ci.yml. + import sys, pathlib + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + import torch + import torch.nn as nn + from unsloth_zoo.tiled_mlp import patch_tiled_mlp, patch_mlp + + class _MLP(nn.Module): + def __init__(self, hidden=64, intermediate=128): + super().__init__() + self.gate_proj = nn.Linear(hidden, intermediate, bias=False) + self.up_proj = nn.Linear(hidden, intermediate, bias=False) + self.down_proj = nn.Linear(intermediate, hidden, bias=False) + self.act_fn = nn.SiLU() + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + class _FakeModel(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([nn.ModuleDict({"mlp": _MLP()}) for _ in range(2)]) + def forward(self, x): + for layer in self.layers: + x = x + layer["mlp"](x) + return x + + def test_patch_tiled_mlp_numerical_equivalence(): + # `patch_mlp(target_arctic=True)` sets `chunk_size = max(1, H)` + # and shards the SEQUENCE dim with `n_shards = max(1, S // + # chunk_size)`. Pick S > H so the tiled path actually runs + # multi-shard (n_shards = 192 // 64 = 3, plus a remainder + # shard) rather than degenerating to n_shards = 1 which is + # bit-exact and only confirms patching installed something. + # If the tiled implementation is correct, multi-shard output + # must still match the un-tiled reference within FP32 noise. + torch.manual_seed(0) + m = _FakeModel().eval() + hidden = 64 + # 192 = 3 * hidden, so divmod(192, 64) = (3, 0) -> 3 shards, + # no remainder; gives a clean multi-shard verification. + x = torch.randn(2, 192, hidden) + with torch.no_grad(): + y_before = m(x).clone() + patch_mlp(m.layers[0]["mlp"]) + patch_tiled_mlp(m) + # Sanity-check we are actually exercising the multi-shard + # path: poke chunk_size by re-deriving it the same way + # `tiled_forward_arctic_size` does. + S = x.shape[1] + chunk = max(1, hidden) + n_shards_expected = max(1, S // chunk) + assert n_shards_expected > 1, ( + "tiled MLP shim is not exercising multi-shard: " + f"S={S}, chunk={chunk}, n_shards={n_shards_expected}" + ) + with torch.no_grad(): + y_after = m(x).clone() + err = (y_before - y_after).abs().max().item() + print( + f"patch_tiled_mlp multi-shard (n_shards={n_shards_expected}) " + f"output diff = {err:.3e}" + ) + assert err < 1e-3, f"tiled MLP output drifted: {err}" + PY + python -m pytest -q --tb=short tests/_tiled_mlp_check_shim.py -s + rm -f tests/_tiled_mlp_check_shim.py + + - name: Compiler cache hygiene + source-rewriter invariants (synthetic inputs) + # Lightweight pipeline coverage for unsloth_zoo.compiler. Pure regex + # / tokenize / ast paths driven by tiny synthetic source strings: + # - higher_precision_softmax (basic + idempotent) + # - fix_rotary_embedding_dtype (no-op + active under + # UNSLOTH_FORCE_CUSTOM_DTYPE) + # - fix_attention_dtype_consistency (insert + idempotent) + # - convert_attention_masks_to_bool (rewrite + no-op) + # - create_new_function happy-path (versioning block, license + # header, AST parse, importlib re-import) + # - create_new_function **kwargs collision (exercises + # _rewrite_kwargs_param + _insert_kwargs_alias) + # - UNSLOTH_COMPILE_OVERWRITE=0 forced-recompile on transformers + # version mismatch (compiler.py:947-963) + # - matching short-circuit when versions are equal + # No real transformers modeling module is loaded; complements the + # heavier real-class round-trip step below. Wall-time ~10-25s. + run: | + set -euxo pipefail + cat > tests/_compiler_cache_invariants_shim.py <<'PY' + # Auto-generated by .github/workflows/consolidated-tests-ci.yml. + # Cache-hygiene + source-rewriter invariants for unsloth_zoo.compiler. + import sys, pathlib, os, ast, importlib, importlib.util, time + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + import pytest + import torch # noqa: F401 (compiler.py imports torch at module load) + + + def _isolate_cache(tmp_path, monkeypatch): + """Point UNSLOTH_COMPILE_LOCATION at tmp_path and reset module + globals. The compiler.py global is captured at module load + (line 75/179), so we delete + reimport per test.""" + monkeypatch.setenv("UNSLOTH_COMPILE_LOCATION", str(tmp_path)) + if "unsloth_zoo.compiler" in sys.modules: + del sys.modules["unsloth_zoo.compiler"] + import unsloth_zoo.compiler as compiler + compiler.UNSLOTH_COMPILE_LOCATION = str(tmp_path) + compiler.UNSLOTH_COMPILE_USE_TEMP = False + return compiler + + + def test_higher_precision_softmax_basic_and_idempotent(tmp_path, monkeypatch): + c = _isolate_cache(tmp_path, monkeypatch) + src = ( + "y = nn.functional.softmax(x, dim=-1)\n" + "z = F.softmax(a, dim=1, dtype=torch.bfloat16)\n" + ) + out = c.higher_precision_softmax(src) + assert "dtype = torch.float32).to(x.dtype)" in out + assert "dtype = torch.float32).to(a.dtype)" in out + # Idempotency landed in unslothai/unsloth-zoo#631 + # (negative-lookahead on `.to(.dtype)` so a second + # pass does not append another cast). + assert c.higher_precision_softmax(out) == out + + + def test_fix_rotary_dtype_no_op_without_env(tmp_path, monkeypatch): + c = _isolate_cache(tmp_path, monkeypatch) + monkeypatch.delenv("UNSLOTH_FORCE_CUSTOM_DTYPE", raising=False) + src = "out = cos.to(dtype=x.dtype) + sin.to(dtype=x.dtype)\n" + assert c.fix_rotary_embedding_dtype(src) == src + + + def test_fix_rotary_dtype_active(tmp_path, monkeypatch): + c = _isolate_cache(tmp_path, monkeypatch) + monkeypatch.setenv( + "UNSLOTH_FORCE_CUSTOM_DTYPE", + "float16;torch.float32;torch.bfloat16;torch.float16;pass", + ) + monkeypatch.setenv("UNSLOTH_FORCE_FLOAT32", "1") + src = "out = cos.to(dtype=x.dtype) + sin.to(dtype=x.dtype)\n" + out = c.fix_rotary_embedding_dtype(src) + # Active form rewrites cos.to / sin.to. Either the conditional + # form or the cast form is acceptable -- different transformers + # versions surface slightly different outputs from the rewriter. + assert "cos.to(dtype=x.dtype)" not in out + assert "sin.to(dtype=x.dtype)" not in out + + + def test_fix_attention_dtype_consistency_insert_then_idempotent(tmp_path, monkeypatch): + c = _isolate_cache(tmp_path, monkeypatch) + src = ( + " query_states, key_states = apply_rotary_pos_emb(" + "query_states, key_states, cos, sin)\n" + " attn = q @ k.T\n" + ) + out = c.fix_attention_dtype_consistency(src) + assert out.count("value_states = value_states.to(query_states.dtype)") == 1 + assert c.fix_attention_dtype_consistency(out) == out + + + def test_convert_attention_masks_to_bool_rewrites(tmp_path, monkeypatch): + c = _isolate_cache(tmp_path, monkeypatch) + src = ( + "def make_mask(x):\n" + " out = torch.finfo(x.dtype).min * x\n" + " return out\n" + ) + out = c.convert_attention_masks_to_bool("make_mask", src) + # Loose match: rewriter inserts a `!=torch.finfo(...).min` check + # somewhere on the return path. Tightening to an exact + # last-line match is brittle across transformers versions. + assert "!=torch.finfo" in out + + + def test_convert_attention_masks_to_bool_no_op(tmp_path, monkeypatch): + c = _isolate_cache(tmp_path, monkeypatch) + src = "def make_mask(x):\n return x\n" + assert c.convert_attention_masks_to_bool("make_mask", src) == src + + + def _versioning_lines(file_text): + """Extract the four version strings from the versioning block.""" + assert file_text.startswith('"""\n'), "missing opening triple-quote" + head = file_text.split("__UNSLOTH_VERSIONING__", 1)[0] + lines = [ln for ln in head.splitlines() if ln and ln != '"""'] + return lines + + + def test_create_new_function_happy_path(tmp_path, monkeypatch): + c = _isolate_cache(tmp_path, monkeypatch) + src = "def f(x):\n return nn.functional.softmax(x, dim=-1)\n" + c.create_new_function( + name="f_happy", new_source=src, model_location="builtins", + functions=[], overwrite=True, + ) + cached = tmp_path / "f_happy.py" + assert cached.exists() + text = cached.read_text(encoding="utf-8") + versions = _versioning_lines(text) + assert len(versions) == 4, versions + assert text.count(c._full_license_header) == 1 + ast.parse(text) + spec = importlib.util.spec_from_file_location("f_happy_reimport", cached) + m2 = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m2) + assert callable(m2.f) + import inspect as _inspect + # higher_precision_softmax should have promoted to float32. + assert "dtype = torch.float32" in _inspect.getsource(m2.f) + + + def test_create_new_function_overwrite_zero_recompiles_on_version_mismatch( + tmp_path, monkeypatch, + ): + c = _isolate_cache(tmp_path, monkeypatch) + name = "vmismatch" + cached = tmp_path / f"{name}.py" + stub = ( + '"""\n0.0.0\n0.0.0\n0.0.0-stub\n0.0.0\n__UNSLOTH_VERSIONING__\n"""\n' + + c._full_license_header + + "def vmismatch(x):\n return x\n" + ) + cached.write_text(stub, encoding="utf-8") + monkeypatch.setenv("UNSLOTH_COMPILE_OVERWRITE", "0") + src = "def vmismatch(x):\n return x + 1\n" + c.create_new_function( + name=name, new_source=src, model_location="builtins", + functions=[], overwrite=False, + ) + text = cached.read_text(encoding="utf-8") + assert "0.0.0-stub" not in text, ( + "OVERWRITE=0 + transformers-version-mismatch did NOT recompile" + ) + versions = _versioning_lines(text) + import importlib.metadata as _md + assert versions[2] == _md.version("transformers") + + + def test_create_new_function_overwrite_zero_short_circuits_when_versions_match( + tmp_path, monkeypatch, + ): + c = _isolate_cache(tmp_path, monkeypatch) + name = "vmatch" + src = "def vmatch(x):\n return x\n" + c.create_new_function( + name=name, new_source=src, model_location="builtins", + functions=[], overwrite=True, + ) + cached = tmp_path / f"{name}.py" + mtime_before = cached.stat().st_mtime_ns + time.sleep(0.05) + monkeypatch.setenv("UNSLOTH_COMPILE_OVERWRITE", "0") + c.create_new_function( + name=name, new_source=src, model_location="builtins", + functions=[], overwrite=False, + ) + assert cached.stat().st_mtime_ns == mtime_before, ( + "OVERWRITE=0 + matching versions should NOT rewrite the file" + ) + PY + python -m pytest -q --tb=short tests/_compiler_cache_invariants_shim.py + rm -f tests/_compiler_cache_invariants_shim.py + + - name: Compiler full-model-sweep (every transformers.models.*) + SFT trainer round-trip + # Calls `unsloth_compile_transformers(model_type=...)` against EVERY + # `transformers.models.` package the matrix's transformers ships + # (pkgutil.iter_modules walk -- 383 packages on 4.57.6, similar on + # latest), then ast.parse / importlib-load / introspect the + # generated unsloth_compiled_cache/*.py file per model. Catches + # regex / source-rewriter drift across the matrix's (transformers, + # trl) combination -- the dominant failure mode of + # `unsloth_compile_transformers` after a transformers point release. + # + # 21 model_types currently break the compiler (verified locally on + # transformers 4.57.6). They are listed in KNOWN_BROKEN below with + # their failure mode so the sweep stays green and any NEW breakage + # surfaces as red. Each entry is tracked for an individual fix + # PR on unsloth-zoo. The list is split by failure category so + # follow-up PRs can target one bug at a time. + # + # Hermetic cache dir per pytest invocation; we override the + # job-level UNSLOTH_COMPILE_DISABLE=1 inside the shim so + # compilation actually runs here. Wall-time estimate ~2-3 min + # warm (mean ~0.3s/model, 383 models = ~110s on the runner). + run: | + set -euxo pipefail + cat > tests/_zoo_compiler_cache_shim.py <<'PY' + # Auto-generated by .github/workflows/consolidated-tests-ci.yml. + import os, sys, ast, pathlib, importlib.util, tempfile + _HERE = pathlib.Path(__file__).parent + sys.path.insert(0, str(_HERE)) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + + # Hermetic cache dir + force compile path. The compiler's + # globals (UNSLOTH_COMPILE_LOCATION, UNSLOTH_COMPILE_USE_TEMP) + # are captured at module load; an earlier conftest `import + # unsloth` may have already imported unsloth_zoo.compiler with + # the default "unsloth_compiled_cache" path. Mutate the live + # module globals after import so this shim is robust to that + # ordering. Otherwise the compiler silently writes to the + # default cache and the per-model file assertion fails. + _CACHE = pathlib.Path(tempfile.mkdtemp(prefix="unsloth_cache_")) + os.environ["UNSLOTH_COMPILE_LOCATION"] = str(_CACHE) + os.environ["UNSLOTH_COMPILE_OVERWRITE"] = "1" + os.environ.pop("UNSLOTH_COMPILE_DISABLE", None) + + import pytest + import unsloth_zoo.compiler as _zoo_compiler + _zoo_compiler.UNSLOTH_COMPILE_LOCATION = str(_CACHE) + _zoo_compiler.UNSLOTH_COMPILE_USE_TEMP = False + from unsloth_zoo.compiler import unsloth_compile_transformers + + + def _verify_file(path: pathlib.Path, must_expose): + assert path.exists(), f"compiler did not write {path}" + src = path.read_text(encoding="utf-8") + ast.parse(src, filename=str(path)) + spec = importlib.util.spec_from_file_location(path.stem, path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + for name in must_expose: + assert hasattr(mod, name), ( + f"{path.name} missing expected attr {name!r}; " + f"found: {sorted(n for n in dir(mod) if not n.startswith('_'))[:25]}" + ) + + + # ---------- Full transformers.models.* compile sweep ---------- + # Track the model_types that currently break the compiler on + # transformers >=5,<6. After unsloth-zoo#632 landed, transformers + # 4.57.6 has zero failures across all model_types; the 27 entries + # below are the residual failures on the tf 5.x line. New breakage + # on any OTHER model_type fails the cell. Each entry is a + # tracking item for a follow-up unsloth-zoo PR. + KNOWN_BROKEN_COMPILE = { + # Category A: `string index out of range` in source rewriter. + "colpali": "string index out of range", + "colqwen2": "string index out of range", + "colmodernvbert": "string index out of range", + "dpr": "string index out of range", + "gemma4_assistant":"string index out of range", + "rag": "string index out of range", + "shieldgemma2": "string index out of range", + "timm_backbone": "string index out of range", + # Category B: rewriter emits invalid Python source. + "clvp": "emitted file: unexpected indent", + "falcon_mamba": "emitted file: unexpected indent", + "gpt2": "emitted file: unexpected indent", + "imagegpt": "emitted file: unexpected indent", + "mamba": "emitted file: unexpected indent", + "tapas": "emitted file: expected ':'", + "xlstm": "emitted file: unexpected indent", + # Category B-2: emit unterminated string literal (latest tf). + "audioflamingo3": "emitted file: unterminated string literal", + "musicflamingo": "emitted file: unterminated string literal", + "voxtral": "emitted file: unterminated string literal", + "voxtral_realtime":"emitted file: unterminated string literal", + # Category C: rewriter emits unclosed paren. + "kosmos2": "emitted file: '(' was never closed", + "kosmos2_5": "emitted file: '(' was never closed", + # Category D: imports list builder picks up a non-exported name. + "auto": "module has no attribute _BaseModelWithGenerate", + "bit": "module has no attribute Linear", + "regnet": "module has no attribute Linear", + "resnet": "module has no attribute Linear", + # Category E: undefined name in emitted file. + "perceiver": "name 'AbstractPreprocessor' is not defined", + "sam3_lite_text": "name 'Sam3LiteTextLayerScaledResidual' is not defined", + # Category F: compile exceeds 60s budget on the runner. + # First seen on transformers >=5,<6; each represents a slow + # or recursive source-rewriter path the zoo can address. + "beit": "TimeoutError: compile exceeds per-model budget", + "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", + "sam": "TimeoutError: compile exceeds per-model budget", + "sam_hq": "TimeoutError: compile exceeds per-model budget", + "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", + } + + + def _all_model_types(): + import pkgutil, transformers.models as tm + return sorted(s.name for s in pkgutil.iter_modules(tm.__path__) if s.ispkg) + + + def test_compile_every_transformers_model_type(): + """Run unsloth_compile_transformers across every model_type + the matrix's transformers ships. Allowed outcomes: + ok -> compile emitted a parseable, importable cache file + skipped -> no `modeling_.py` file (expected for some + umbrella packages like `auto`, `deprecated`) + known -> in KNOWN_BROKEN_COMPILE; tracked for follow-up. + Any uncaught failure fails the cell. + + Per-model SIGALRM cap so one infinite-looping model_type + cannot wedge the whole sweep + nuke the job timeout + (observed on transformers >=5,<6 -- 30+ min hang before + this guard landed).""" + import importlib as _il + import signal + ok = 0 + skipped = [] + known = [] + new_failures = [] + models = _all_model_types() + def _on_timeout(signum, frame): + raise TimeoutError("compile exceeded per-model budget") + prev_handler = signal.signal(signal.SIGALRM, _on_timeout) + try: + for i, model_type in enumerate(models): + if i % 25 == 0: + print(f" sweep progress: {i}/{len(models)} -> {model_type}", flush=True) + modeling_path = f"transformers.models.{model_type}.modeling_{model_type}" + try: + _il.import_module(modeling_path) + except (ModuleNotFoundError, ImportError): + skipped.append((model_type, "no modeling file")) + continue + signal.alarm(60) + try: + unsloth_compile_transformers( + model_type=model_type, fast_lora_forwards=False, + ) + except Exception as e: + signal.alarm(0) + msg = f"{type(e).__name__}: {str(e)[:200]}" + if model_type in KNOWN_BROKEN_COMPILE: + known.append((model_type, msg)) + else: + new_failures.append((model_type, msg)) + continue + signal.alarm(0) + if model_type in KNOWN_BROKEN_COMPILE: + # Came back green unexpectedly -- that's GOOD news, + # the bug was fixed. Surface it so we can drop the + # entry from KNOWN_BROKEN_COMPILE. + print( + f" UNEXPECTED-OK {model_type}: was in " + "KNOWN_BROKEN_COMPILE, now compiles cleanly. " + "Drop the entry." + ) + ok += 1 + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, prev_handler) + print(f"\nCompile sweep: ok={ok} skipped={len(skipped)} " + f"known-broken={len(known)} new-failures={len(new_failures)}") + for m, r in known: + print(f" KNOWN {m}: {r}") + for m, r in new_failures[:30]: + print(f" NEW {m}: {r}") + if len(new_failures) > 30: + print(f" ...and {len(new_failures)-30} more new failures") + assert not new_failures, ( + f"unsloth_compile_transformers introduced new failures on " + f"{len(new_failures)} model_types not in the known-broken " + f"list: {[m for m, _ in new_failures]}" + ) + # Sanity floor: at least 200 model_types should compile cleanly + # (we observed 362 ok / 383 total on transformers 4.57.6). + assert ok >= 200, ( + f"only {ok} model_types compiled cleanly; expected >=200. " + "Possible transformers-version-induced regression." + ) + + + @pytest.mark.parametrize("model_type,rms_class", [ + ("llama", "LlamaRMSNorm"), + ("qwen3", "Qwen3RMSNorm"), + ("gemma3", "Gemma3RMSNorm"), + ]) + def test_compile_real_modeling_module(model_type, rms_class): + """Spot-check on the three production-relevant families that + the compile_every sweep also covers; this case verifies the + emitted cache file has the model-specific RMSNorm class + attribute, not just that the file parses + imports. + + ``unsloth_compile_transformers`` is not idempotent in- + process: calling it twice on the same modeling module + after rewriting class attributes corrupts the inspect + source/line cache and the second emitted file is malformed + Python. The sweep above already produced a valid cache + file for every non-KNOWN_BROKEN model_type, so just verify + that artefact here. Trigger a compile only when running + this test in isolation (no sweep preceded).""" + import importlib as _il + try: + modeling = _il.import_module( + f"transformers.models.{model_type}.modeling_{model_type}" + ) + except ModuleNotFoundError: + pytest.skip( + f"transformers build lacks model_type={model_type}" + ) + combined = _CACHE / f"unsloth_compiled_module_{model_type}.py" + if not combined.exists(): + unsloth_compile_transformers( + model_type=model_type, fast_lora_forwards=False, + ) + modeling = _il.import_module( + f"transformers.models.{model_type}.modeling_{model_type}" + ) + assert getattr(modeling, "__UNSLOTH_PATCHED__", False) is True + _verify_file(combined, must_expose=[rms_class]) + + + def test_compile_disable_writes_nothing(): + """Negative control: when UNSLOTH_COMPILE_DISABLE=1 the + compile path must early-return without producing new files.""" + os.environ["UNSLOTH_COMPILE_DISABLE"] = "1" + try: + before = set(_CACHE.iterdir()) + # Pick a model_type that still resolves on this transformers. + for mt in ("llama", "mistral", "qwen2"): + try: + import importlib as _il + _il.import_module( + f"transformers.models.{mt}.modeling_{mt}" + ) + break + except ModuleNotFoundError: + continue + else: + pytest.skip("no probe model_type available") + unsloth_compile_transformers( + model_type=mt, fast_lora_forwards=False, + ) + after = set(_CACHE.iterdir()) + assert after == before, ( + f"DISABLE=1 still wrote: {[p.name for p in after - before]}" + ) + finally: + os.environ.pop("UNSLOTH_COMPILE_DISABLE", None) + + + def test_compile_sft_trainer_patch(): + """Round-trip TRL's SFTTrainer through the rl.py patch path + and verify the generated UnslothSFTTrainer.py.""" + pytest.importorskip("trl") + try: + from unsloth.models.rl import _patch_trl_rl_trainers + except ImportError: + pytest.skip("unsloth.models.rl._patch_trl_rl_trainers absent") + try: + _patch_trl_rl_trainers("sft_trainer") + except Exception as e: + # TRL 1.x renames break the patch helper internally; we + # accept that here and skip rather than fail the cell. + pytest.skip(f"_patch_trl_rl_trainers raised: {type(e).__name__}: {e}") + sft = _CACHE / "UnslothSFTTrainer.py" + if not sft.exists(): + pytest.skip( + "_patch_trl_rl_trainers ran but did not emit " + "UnslothSFTTrainer.py on this TRL version." + ) + _verify_file(sft, must_expose=["UnslothSFTTrainer"]) + PY + python -m pytest -q --tb=short tests/_zoo_compiler_cache_shim.py + rm -f tests/_zoo_compiler_cache_shim.py + + - name: TRL trainer + Config auto-discovery + dynamic patch coverage + # Mirror unsloth/models/rl.py:patch_trl_rl_trainers AND verify the + # dynamic per-version patch surface: + # 1. AST-parse every *_trainer / *_config submodule. + # 2. Apply the same *Trainer / *Config discovery rules + # _patch_trl_rl_trainers uses (rl.py:553-620). + # 3. Orphan check: every _trainer must have a sibling + # _config OR an inline *Config. + # 4. Dynamic count: enumerate every canonical trainer that + # imports cleanly, run patch_trl_rl_trainers(), assert + # every one ends up Unsloth-prefixed in-place. Floor matches + # the cohort sizes from the version sweep: + # TRL 0.22-0.23 -> 18 canonical trainers + # TRL 0.24-0.28 -> 15 canonical trainers + # TRL 0.29-1.x -> 6 canonical (rest are experimental + # thin-wrappers; covered next) + # 5. Experimental coverage (TRL 0.29+): walk trl.experimental.*, + # find every *Trainer class, verify the umbrella patch + # reaches them via the thin-wrapper MRO walk in + # _patch_trl_rl_trainers (rl.py:677-702). + # Per-cell wall-time ~30-60s. + run: | + set -euxo pipefail + cat > tests/_trl_trainer_discovery_shim.py <<'PY' + # Auto-generated by .github/workflows/consolidated-tests-ci.yml. + # Walks every *_trainer / *_config module in trl.trainer and + # validates that unsloth's auto-discovery rules in + # unsloth/models/rl.py:_patch_trl_rl_trainers (lines 542-620, + # 1934-1949) still pick out exactly one *Trainer and one + # *Config per module on the matrix's TRL version. + import sys, pathlib, importlib, importlib.util, ast, inspect + + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + + import pytest + pytest.importorskip("trl") + import trl # noqa: F401 (forces lazy-module init) + import trl.trainer + + + def _is_real_submodule(qual_name: str) -> bool: + """True iff `qual_name` resolves to an importable submodule + with a file on disk (i.e. has a non-None find_spec().origin). + + TRL re-exports utility FUNCTIONS into `trl.trainer.__init__` + whose names happen to end with `_config` (e.g. + `get_peft_config`, `get_quantization_config`). Without this + filter the `endswith` check below picks them up as if they + were submodules and the AST stage fails on `no spec`. The + same trap exists for `_trainer` (none today, but defensive). + """ + try: + spec = importlib.util.find_spec(qual_name) + except (ImportError, ValueError): + return False + return spec is not None and bool(getattr(spec, "origin", None)) + + + # Replicate rl.py:1939-1943 verbatim, then filter to actual + # submodules so re-exported utility functions (e.g. + # `get_peft_config`) do not pollute the AST sweep. + def _trainer_files(): + return [ + x for x in dir(trl.trainer) + if x.islower() + and x.endswith("_trainer") + and x != "base_trainer" + and _is_real_submodule(f"trl.trainer.{x}") + ] + + + def _config_files(): + return [ + x for x in dir(trl.trainer) + if x.islower() + and x.endswith("_config") + and _is_real_submodule(f"trl.trainer.{x}") + ] + + + def _ast_parse_module_via_spec(qual_name: str): + """AST-parse a module's source on disk WITHOUT importing it. + `trl.trainer` uses _LazyModule so `find_spec` resolves the + file path without firing the module-level `__init__`. This + dodges optional-dep ImportErrors (e.g. grpo_trainer's vllm + import) and still surfaces real syntax drift in the file.""" + spec = importlib.util.find_spec(qual_name) + if spec is None or not spec.origin: + return None, "no spec" + path = pathlib.Path(spec.origin) + if not path.is_file(): + return None, f"spec.origin not a file: {path}" + src = path.read_text(encoding="utf-8") + ast.parse(src, filename=str(path)) + return path, None + + + def test_every_trl_trainer_and_config_module_ast_parses(): + """Stage 1: pure file-on-disk AST parse. Catches a TRL + source-level syntax issue on any matrix cell without + triggering optional-dep imports.""" + fail = [] + ok = 0 + for name in _trainer_files() + _config_files(): + qual = f"trl.trainer.{name}" + try: + path, err = _ast_parse_module_via_spec(qual) + if err: + fail.append((qual, err)) + else: + ok += 1 + except SyntaxError as e: + fail.append((qual, f"SyntaxError: {e}")) + except Exception as e: + fail.append((qual, f"{type(e).__name__}: {e}")) + print(f"AST-parsed {ok} TRL trainer+config modules; failed={len(fail)}") + for q, e in fail: + print(f" AST FAIL {q}: {e}") + assert not fail, f"AST parse failed for {len(fail)} TRL modules" + + + def _apply_unsloth_discovery_rules(mod, trainer_file): + """Replicate the four endswith filters in + rl.py:553-569 verbatim.""" + prefix = trainer_file.split("_")[0] + names = [ + x for x in dir(mod) + if x.endswith("Trainer") and x != "Trainer" + and not x.startswith("_") and prefix in x.lower() + ] + configs = [ + x for x in dir(mod) + if x.endswith("Config") and x != "Config" + and not x.startswith("_") and prefix in x.lower() + ] + return names, configs + + + def _resolve_config_via_fallbacks(trainer_file, name_list, mod): + """Replicate rl.py:575-615: try the sibling *_config.py + module, then the MRO walk fallback. Returns the resolved + config-name list (length 0 or 1).""" + # Fallback 1: _config.py module sibling. + cfg_module_name = trainer_file.replace("_trainer", "_config") + try: + cfg_mod = getattr(trl.trainer, cfg_module_name) + except Exception: + cfg_mod = None + if cfg_mod is not None: + prefix = trainer_file.split("_")[0] + hits = [ + x for x in dir(cfg_mod) + if x.endswith("Config") and x != "Config" + and not x.startswith("_") and prefix in x.lower() + ] + if len(hits) == 1: + return hits + # Fallback 2: MRO walk into experimental parent module. + if len(name_list) != 1: + return [] + try: + trainer_cls = getattr(mod, name_list[0]) + except Exception: + return [] + prefix = trainer_file.split("_")[0] + for parent in trainer_cls.__mro__[1:]: + if parent is object: + continue + parent_mod = inspect.getmodule(parent) + if parent_mod is None: + continue + if parent_mod.__name__ == f"trl.trainer.{trainer_file}": + continue + hits = [ + x for x in dir(parent_mod) + if x.endswith("Config") and x != "Config" + and not x.startswith("_") and prefix in x.lower() + ] + if len(hits) == 1: + return hits + return [] + + + def test_unsloth_auto_discovery_finds_trainer_and_config_per_module(): + """Stage 2: drive the same unsloth rules over every trainer + file. import-failures (optional deps) are recorded as + `import-skipped`, mirroring rl.py:1944-1948 try/except.""" + ok = 0 + import_skipped = [] + discovery_skipped = [] + fail = [] + for trainer_file in _trainer_files(): + qual = f"trl.trainer.{trainer_file}" + try: + mod = getattr(trl.trainer, trainer_file) + except Exception as e: + import_skipped.append((qual, f"{type(e).__name__}: {e}")) + continue + trainers, configs = _apply_unsloth_discovery_rules( + mod, trainer_file, + ) + if len(trainers) != 1: + discovery_skipped.append( + (qual, f"trainers={trainers}") + ) + continue + if len(configs) != 1: + configs = _resolve_config_via_fallbacks( + trainer_file, trainers, mod, + ) + if len(configs) != 1: + fail.append( + (qual, + f"trainer={trainers[0]} but config not found " + "(checked module, *_config sibling, and MRO)") + ) + continue + ok += 1 + print(f" OK {qual}: trainer={trainers[0]}, config={configs[0]}") + print( + f"\nDiscovery: ok={ok} import_skipped={len(import_skipped)} " + f"discovery_skipped={len(discovery_skipped)} fail={len(fail)}" + ) + for q, r in import_skipped: + print(f" IMPORT-SKIP {q}: {r}") + for q, r in discovery_skipped: + print(f" DISC-SKIP {q}: {r}") + for q, r in fail: + print(f" FAIL {q}: {r}") + # Hard contract: every TRAINER that imports cleanly AND has + # exactly one *Trainer must also resolve exactly one *Config + # via one of the three rules. import-skipped + discovery- + # skipped (no/multiple *Trainer) are tolerated. + assert not fail, ( + f"unsloth discovery rules failed for {len(fail)} trainers" + ) + # Sanity: at least 3 trainers should fully discover on any + # matrix cell (sft + reward + dpo are the historical core). + assert ok >= 3, ( + f"only {ok} trainers fully discovered; expected >=3 " + "(sft/reward/dpo). Possible TRL surface regression." + ) + + + def test_orphan_trainer_modules_do_not_exist(): + """Stage 3: every _trainer module should have a sibling + _config (TRL 0.26+ convention) OR an inline *Config. An + ORPHAN _trainer with neither is a TRL refactor we want + to know about: it would silently break unsloth's + auto-discovery without raising.""" + orphans = [] + for trainer_file in _trainer_files(): + cfg_module_name = trainer_file.replace("_trainer", "_config") + has_sibling_cfg = ( + importlib.util.find_spec( + f"trl.trainer.{cfg_module_name}" + ) is not None + ) + if has_sibling_cfg: + continue + # No sibling -> require an inline *Config in the + # trainer module itself (resolved via discovery rules). + try: + mod = getattr(trl.trainer, trainer_file) + except Exception: + # Optional-dep failure -> skip; the AST-parse stage + # already covered the file. + continue + _, configs = _apply_unsloth_discovery_rules( + mod, trainer_file, + ) + if not configs: + orphans.append(trainer_file) + assert not orphans, ( + "Orphan TRL trainer modules with neither sibling " + f"_config.py nor an inline *Config: {orphans}. " + "unsloth auto-discovery would silently skip these." + ) + + + # ---- Dynamic patch coverage: count + verify Unsloth-prefixed ---- + + def _enumerate_canonical_trainer_classes(): + """Walk trl.trainer/*_trainer.py on disk (the source of + truth for what `dir(trl.trainer)` should expose) and return + [(trainer_file, TrainerClass), ...] for every entry that + imports + has exactly-one resolvable *Trainer per the + unsloth rules. Skips optional-dep ImportErrors.""" + out = [] + for trainer_file in _trainer_files(): + try: + mod = getattr(trl.trainer, trainer_file) + except Exception: + continue + trainers, _ = _apply_unsloth_discovery_rules(mod, trainer_file) + if len(trainers) != 1: + continue + try: + cls = getattr(mod, trainers[0]) + except Exception: + continue + out.append((trainer_file, cls)) + return out + + + def _enumerate_experimental_trainer_packages(): + """TRL 0.29+ moved many trainers (bco, cpo, gkd, nash_md, + online_dpo, orpo, ppo, prm, xpo, ...) to `trl.experimental.`, + re-exposing them via thin-wrapper deprecation shims in + `trl.trainer._trainer`. List every `trl.experimental.` + that defines at least one *Trainer class, parsed by AST so we + do NOT trigger the optional-dep imports on the package init.""" + spec = importlib.util.find_spec("trl.experimental") + if spec is None or not spec.submodule_search_locations: + return [] + import re as _re + hits = [] + for root in spec.submodule_search_locations: + rp = pathlib.Path(root) + for sub in sorted(rp.iterdir()): + if not sub.is_dir() or sub.name.startswith("_"): + continue + classes = [] + for py in sub.rglob("*.py"): + try: + src = py.read_text(encoding="utf-8") + except Exception: + continue + for m in _re.finditer( + r"^class\s+([A-Za-z0-9_]+Trainer)\b", src, _re.M, + ): + classes.append(m.group(1)) + if classes: + hits.append((sub.name, sorted(set(classes)))) + return hits + + + def _is_unsloth_patched(cls) -> bool: + return getattr(cls, "__name__", "").startswith("Unsloth") + + + def test_unsloth_patches_every_canonical_trainer_in_this_trl_version(): + """Verify the count + identity of canonically-patched trainers + matches the trainer surface this TRL version actually ships. + + For TRL 0.22.x-0.23.x: ~18 canonical trainers expected. + For TRL 0.24.x-0.28.x: ~15 canonical trainers expected. + For TRL 0.29.x-1.x: 6 canonical (rest are experimental + thin-wrappers; covered by the next test).""" + from unsloth.models.rl import patch_trl_rl_trainers + before = _enumerate_canonical_trainer_classes() + before_count = len(before) + before_unpatched = [ + (tf, cls.__name__) for tf, cls in before + if not _is_unsloth_patched(cls) + ] + # Apply unsloth's umbrella patch. + patch_trl_rl_trainers() + # Re-enumerate (some classes may have been replaced in-module). + after = _enumerate_canonical_trainer_classes() + after_count = len(after) + patched = [(tf, cls.__name__) for tf, cls in after + if _is_unsloth_patched(cls)] + unpatched = [(tf, cls.__name__) for tf, cls in after + if not _is_unsloth_patched(cls)] + print( + f"\nCanonical trainer surface for TRL {trl.__version__}: " + f"discoverable_before={before_count} " + f"discoverable_after={after_count} " + f"patched={len(patched)} unpatched={len(unpatched)}" + ) + for tf, n in patched: + print(f" PATCHED {tf}: {n}") + for tf, n in unpatched: + print(f" UNPATCHED {tf}: {n}") + # Hard contract: every canonical trainer that imports + # cleanly must end up Unsloth-prefixed after the umbrella + # patch. If a trainer was discoverable BEFORE the patch but + # is missing from `after`, that is a separate (rare) issue + # we surface as failure. + assert before_count == after_count, ( + f"trainer-class set changed across patching: " + f"before={[n for _, n in before_unpatched]} " + f"after={[n for _, n in unpatched]}" + ) + assert not unpatched, ( + "unsloth.models.rl.patch_trl_rl_trainers did NOT patch: " + + ", ".join(f"{tf}:{n}" for tf, n in unpatched) + ) + # Floor matches the cohort sizes from the TRL version sweep: + # 18 (0.22-0.23), 15 (0.24-0.28), 6 (0.29+ canonical only). + assert len(patched) >= 6, ( + f"only {len(patched)} canonical trainers patched; " + "expected >= 6 (the smallest production cohort)." + ) + + + def test_unsloth_patches_experimental_trainers_via_thin_wrappers(): + """TRL 0.29+ ships canonical-`trl.trainer._trainer` modules + for many trainers as deprecation thin-wrappers that forward + to `trl.experimental.`. unsloth's + `_patch_trl_rl_trainers` (rl.py:677-702) detects + `trl.experimental` in the trainer source and resolves to + the parent class -- so patching the canonical entry should + also Unsloth-prefix the experimental class via in-module + setattr. + + Verify by walking trl.experimental.* AST for every *Trainer + class, then checking whether it (or any class with the same + name in the experimental package) carries the Unsloth + prefix after the umbrella patch.""" + from unsloth.models.rl import patch_trl_rl_trainers + patch_trl_rl_trainers() + experimental_pkgs = _enumerate_experimental_trainer_packages() + if not experimental_pkgs: + pytest.skip( + f"TRL {trl.__version__} has no trl.experimental.* " + "trainer surface (pre-0.29 cohort). The canonical " + "test above already covers patching here." + ) + found = [] + missing = [] + for pkg_name, class_names in experimental_pkgs: + qual = f"trl.experimental.{pkg_name}" + try: + pkg_mod = importlib.import_module(qual) + except Exception as e: + # Optional-dep ImportError: experimental package + # could not be loaded. Match unsloth's runtime + # tolerance: this would also be silently skipped + # by `_patch_trl_rl_trainers`. Record but do not + # fail. + print( + f" IMPORT-SKIP {qual}: " + f"{type(e).__name__}: {str(e)[:120]}" + ) + continue + for cls_name in class_names: + cls = getattr(pkg_mod, cls_name, None) + if cls is None: + # Class is defined inside the package but not + # re-exported on the package init. Walk + # submodules to find it. + import pkgutil as _pku + for sub in _pku.walk_packages( + pkg_mod.__path__, prefix=qual + "." + ): + try: + sub_mod = importlib.import_module(sub.name) + except Exception: + continue + cls = getattr(sub_mod, cls_name, None) + if cls is not None: + break + if cls is None: + missing.append((pkg_name, cls_name)) + continue + if _is_unsloth_patched(cls): + found.append((pkg_name, cls_name)) + print(f" PATCHED trl.experimental.{pkg_name}.{cls_name}") + else: + # Not Unsloth-prefixed: either unsloth chose + # not to patch this surface (e.g. the canonical + # thin-wrapper module did not exist) or the + # patch silently failed. Record both + # outcomes; the assertion below tolerates the + # gap as informational, not failure -- the + # canonical test enforces the hard contract. + print( + f" NOT-PATCHED trl.experimental.{pkg_name}." + f"{cls_name} (no Unsloth-prefix on the " + "experimental surface)" + ) + total_experimental = sum(len(cs) for _, cs in experimental_pkgs) + print( + f"\nExperimental trainer surface (TRL {trl.__version__}): " + f"{len(experimental_pkgs)} packages, " + f"{total_experimental} *Trainer classes; " + f"unsloth-patched={len(found)} class-missing={len(missing)}" + ) + # Hard contract: a *Trainer class declared in a python + # source file must be locatable in its package after import. + # If we saw the class definition but cannot find the symbol + # at runtime, the package's public surface drifted. + assert not missing, ( + "experimental *Trainer classes declared in source but " + f"not importable: {missing}" + ) + PY + python -m pytest -q --tb=short -s tests/_trl_trainer_discovery_shim.py + rm -f tests/_trl_trainer_discovery_shim.py + + - name: MoE per-family coverage + GRPO patches + grouped_gemm AST + # Catches the recurring class of bugs that PR #624 (gemma4 missing + # extractor), PR #612 (gemma4 GRPO patch silently dropped), PR #607 + # (gate_up LoRA dropped from grad graph), PR #601 (qwen MoE shape + # mismatch), unsloth#4934 (TRL disable_gradient_checkpointing + # corrupts unsloth GC), and unsloth#3598 (gradient_accumulation + # double-scale on accepts_loss_kwargs=False) targeted. Coverage: + # + # 1. Per-MoE-family side-effect contract: for every patch_*_moe + # function in unsloth_zoo.temporary_patches, if its target + # transformers class is importable on this matrix cell, the + # patch must mark the class with `_unsloth_already_patched=True` + # after running. This is exactly what unsloth_zoo's existing + # test_moe_lora_extractor_coverage walks at the registration + # level; here we tie each patch fn to its declared target so a + # silent early-return (PR #612 style) surfaces as red rather + # than a coverage skip. + # + # 2. PR #4934 (GRPO + TRL 1.0): patch_trl_disable_gradient_checkpointing + # must rebind trl.models.utils.disable_gradient_checkpointing to + # the unsloth no-op AND propagate the rebinding to every trl.* + # module that imported the symbol by reference. + # + # 3. PR #3598 (gradient_accumulation): patch_gradient_accumulation_fix + # must run cleanly on a synthetic Trainer whose training_step + # signature carries `num_items_in_batch`. The original bug was + # that `accepts_loss_kwargs=False` (Qwen3VL, Gemma3 in t-4.57) + # caused double loss-scaling; here we verify the rewrite path + # itself does not raise on a CPU-resolvable shape. + # + # 4. unsloth/kernels/moe/grouped_gemm AST smoke: the Triton kernels + # are GPU-only at runtime, but a SyntaxError or stray + # string-literal in the source still surfaces as a test-time + # ImportError on every install. ast.parse the .py files without + # executing. + # + # Wall-time per cell ~30-60s. Routed through pytest for the spoof + # harness so unsloth_zoo.temporary_patches imports are clean. + run: | + set -euxo pipefail + cat > tests/_moe_coverage_shim.py <<'PY' + # Auto-generated by .github/workflows/consolidated-tests-ci.yml. + import sys, pathlib, ast, importlib, importlib.util, contextlib, os + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + + import pytest + + # Map each MoE patch function to the transformers classes it is + # contractually responsible for marking with _unsloth_already_patched + # after a successful run. Sourced from + # unsloth_zoo/temporary_patches/_moe.py: + # - qwen3_moe.py:382-398 patches Qwen3MoeExperts (new path) or + # Qwen3MoeSparseMoeBlock (old path). + # - qwen3_5_moe.py + qwen3_next_moe.py + qwen3_vl_moe.py register + # extractors on Qwen3_5MoeExperts / Qwen3NextExperts / + # Qwen3VLMoeTextExperts respectively. + # - gemma4_moe.py marks Gemma4TextExperts (current) or + # Gemma4TextMoEBlock (legacy). + # - glm4_moe.py marks Glm4MoeLiteNaiveMoe. + # - deepseek_v3_moe.py marks DeepseekV3NaiveMoe. + # - gpt_oss.py:patch_gpt_oss_moe_for_lora marks GptOssExperts. + # Each cell skips a target if the transformers version lacks it + # (legitimate version-skew); only patches with at least one + # importable target are exercised. + # Each entry = ((patch_module, patch_fn), targets, env_setup, + # version_gate). env_setup runs before the patch fn (e.g. set + # UNSLOTH_MODEL_NAME for gpt_oss). version_gate is a callable + # returning True when the patch SHOULD run on this transformers; + # if False, the test skips with a documented reason. + def _v5_or_later(): + try: + import transformers + major = int(transformers.__version__.split(".")[0]) + return major >= 5 + except Exception: + return False + + MOE_PATCHES = [ + { + "module": "unsloth_zoo.temporary_patches.qwen3_moe", + "fn": "patch_qwen3_moe", + "targets": [ + ("transformers.models.qwen3_moe.modeling_qwen3_moe", "Qwen3MoeExperts"), + ("transformers.models.qwen3_moe.modeling_qwen3_moe", "Qwen3MoeSparseMoeBlock"), + ], + "env": {}, + "gate": lambda: True, + "gate_reason": "", + }, + { + "module": "unsloth_zoo.temporary_patches.qwen3_5_moe", + "fn": "patch_qwen3_5_moe", + "targets": [ + ("transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", "Qwen3_5MoeExperts"), + ], + "env": {}, "gate": lambda: True, "gate_reason": "", + }, + { + "module": "unsloth_zoo.temporary_patches.qwen3_next_moe", + "fn": "patch_qwen3_next_moe", + "targets": [ + ("transformers.models.qwen3_next.modeling_qwen3_next", "Qwen3NextExperts"), + ], + "env": {}, "gate": lambda: True, "gate_reason": "", + }, + { + "module": "unsloth_zoo.temporary_patches.qwen3_vl_moe", + "fn": "patch_qwen3_vl_moe", + "targets": [ + ("transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe", "Qwen3VLMoeTextExperts"), + ], + "env": {}, "gate": lambda: True, "gate_reason": "", + }, + { + "module": "unsloth_zoo.temporary_patches.gemma4_moe", + "fn": "patch_gemma4_moe", + "targets": [ + ("transformers.models.gemma4.modeling_gemma4", "Gemma4TextExperts"), + ], + "env": {}, "gate": lambda: True, "gate_reason": "", + }, + { + "module": "unsloth_zoo.temporary_patches.glm4_moe", + "fn": "patch_glm4_moe", + "targets": [ + ("transformers.models.glm4_moe.modeling_glm4_moe", "Glm4MoeLiteNaiveMoe"), + ], + "env": {}, "gate": lambda: True, "gate_reason": "", + }, + { + "module": "unsloth_zoo.temporary_patches.deepseek_v3_moe", + "fn": "patch_deepseek_v3_moe", + "targets": [ + ("transformers.models.deepseek_v3.modeling_deepseek_v3", "DeepseekV3NaiveMoe"), + ], + "env": {}, "gate": lambda: True, "gate_reason": "", + }, + { + "module": "unsloth_zoo.temporary_patches.gpt_oss", + "fn": "patch_gpt_oss_moe_for_lora", + "targets": [ + ("transformers.models.gpt_oss.modeling_gpt_oss", "GptOssExperts"), + ], + # The patch reads UNSLOTH_MODEL_NAME and only runs when + # "gpt_oss" is in the normalized form. Set it explicitly + # so the gate at gpt_oss.py:1387 passes; otherwise the + # patch silently early-returns and the test would + # spuriously fail. + "env": {"UNSLOTH_MODEL_NAME": "gpt_oss"}, + # Additionally only runs on transformers >= 5 + # (gpt_oss.py:1392 `_is_transformers_v5()` gate). + "gate": _v5_or_later, + "gate_reason": ( + "patch_gpt_oss_moe_for_lora gates on " + "transformers >= 5 (split-LoRA grouped_mm path)" + ), + }, + ] + + + def _resolve_target_classes(targets): + """Return [(qual, cls), ...] for every importable target.""" + out = [] + for mod_path, cls_name in targets: + try: + mod = importlib.import_module(mod_path) + except Exception: + continue + cls = getattr(mod, cls_name, None) + if cls is None: + continue + out.append((f"{mod_path}.{cls_name}", cls)) + return out + + + @pytest.mark.parametrize( + "spec", + MOE_PATCHES, + ids=lambda s: s["fn"], + ) + def test_moe_patch_marks_its_target_when_class_present(spec, monkeypatch): + """If at least one target class is importable AND the + version gate passes, run the patch fn and assert at least + one target is marked patched afterwards. Skips when the + transformers version lacks every target or when the + version gate blocks the patch (legitimate). Fails on + silent patch-fn early-returns (PR #612 class of bug).""" + targets = spec["targets"] + patch_module = spec["module"] + patch_name = spec["fn"] + importable = _resolve_target_classes(targets) + if not importable: + pytest.skip( + f"{patch_name}: no target class importable on this " + f"transformers (looked for {[c for _, c in targets]})." + ) + if not spec["gate"](): + pytest.skip( + f"{patch_name}: version gate blocks this cell. " + f"Reason: {spec['gate_reason']}" + ) + for k, v in spec["env"].items(): + monkeypatch.setenv(k, v) + try: + pmod = importlib.import_module(patch_module) + except Exception as e: + pytest.skip( + f"{patch_module} import failed (likely optional dep): " + f"{type(e).__name__}: {e}" + ) + fn = getattr(pmod, patch_name, None) + if fn is None or not callable(fn): + pytest.skip(f"{patch_module} has no callable {patch_name}") + try: + fn() + except Exception as e: + raise AssertionError( + f"{patch_name}() raised on a transformers that " + f"DOES ship at least one target class ({importable}). " + f"This is the silent-failure mode PR #612 fixed: " + f"{type(e).__name__}: {e}" + ) + # At least one importable target must now carry SOME marker + # showing unsloth touched it. Accepted signals (each is set + # by a different patch flow in unsloth_zoo): + # - `_unsloth_already_patched=True` (gemma4, deepseek_v3, glm4) + # - `_unsloth_lora_patched=True` (gpt_oss_moe_for_lora) + # - `_unsloth_lora_extractor_fn` is callable (qwen3_*, glm4_moe) + # - `_original___forward` attr + # (set by patch_function: qwen3_moe SparseMoeBlock, etc.) + # - `_original_forward` attribute (gpt_oss in-place patch) + # Accept any one as "patched". + def _is_patched(cls) -> bool: + if getattr(cls, "_unsloth_already_patched", False) is True: + return True + if getattr(cls, "_unsloth_lora_patched", False) is True: + return True + if callable(getattr(cls, "_unsloth_lora_extractor_fn", None)): + return True + if "_original_forward" in dir(cls): + return True + cls_name = cls.__name__ + for attr in dir(cls): + if attr.startswith("_original_") and attr.endswith( + f"_{cls_name}_forward" + ): + return True + return False + + after = _resolve_target_classes(targets) + marked = [qual for qual, cls in after if _is_patched(cls)] + if not marked: + raise AssertionError( + f"{patch_name}() ran without exception but no target " + f"in {importable} carries any of the unsloth markers " + "(_unsloth_already_patched / _unsloth_lora_patched / " + "_unsloth_lora_extractor_fn / _original_*_forward). " + "Patch silently no-op'd (PR #612 class of bug)." + ) + print(f" {patch_name}: marked {marked}") + + + # ---- PR #4934 (TRL 1.0+ GRPO disable_gradient_checkpointing) ---- + + def test_patch_trl_disable_gradient_checkpointing(): + """unsloth/models/rl.py:patch_trl_disable_gradient_checkpointing + must rebind trl.models.utils.disable_gradient_checkpointing to + the unsloth no-op when TRL >= 1.0. Pre-1.0 TRL has no such + symbol -> the patch returns early.""" + try: + import trl.models.utils as _tmu + except ImportError: + pytest.skip("trl not installed") + had_symbol = hasattr(_tmu, "disable_gradient_checkpointing") + try: + from unsloth.models.rl import patch_trl_disable_gradient_checkpointing + except ImportError: + pytest.skip( + "unsloth.models.rl.patch_trl_disable_gradient_checkpointing " + "absent (older unsloth than #4934)" + ) + patch_trl_disable_gradient_checkpointing() + if not had_symbol: + # Pre-1.0 TRL: patch is a no-op early-return. Verify + # nothing broke. + pytest.skip( + "TRL pre-1.0 has no disable_gradient_checkpointing; " + "patch correctly early-returned." + ) + fn = getattr(_tmu, "disable_gradient_checkpointing", None) + assert fn is not None, ( + "trl.models.utils.disable_gradient_checkpointing missing " + "after patch -- patch removed the symbol entirely?" + ) + assert getattr(fn, "_unsloth_noop_patched", False) is True, ( + "trl.models.utils.disable_gradient_checkpointing was NOT " + "rebound to the unsloth no-op. PR #4934 regression." + ) + # PR #4934 also walks sys.modules to rebind trl.* modules + # that imported the symbol by reference. Verify at least the + # canonical trainer modules picked up the rebinding when + # they re-export it. + import sys + checked = 0 + missed = [] + for mod_name, mod in list(sys.modules.items()): + if not mod_name.startswith("trl."): + continue + bound = getattr(mod, "disable_gradient_checkpointing", None) + if bound is None: + continue + checked += 1 + if not getattr(bound, "_unsloth_noop_patched", False): + missed.append(mod_name) + print(f" rebound disable_gradient_checkpointing in {checked} trl.* modules") + assert not missed, ( + "trl.* modules that imported disable_gradient_checkpointing " + f"by reference but did not get rebound: {missed}" + ) + + + # ---- PR #3598 (gradient_accumulation loss-scaling rewrite) ---- + + def test_patch_gradient_accumulation_fix_runs_on_synthetic_trainer(): + """patch_gradient_accumulation_fix rewrites a Trainer's + `training_step` source via inspect+exec when the signature + carries `num_items_in_batch`. PR #3598 fixed the rewrite + path to not double-scale for trainers with + `accepts_loss_kwargs=False`. Verify the patch fn runs + without raising on a synthetic Trainer carrying that + signature.""" + try: + from unsloth.models._utils import patch_gradient_accumulation_fix + except ImportError: + pytest.skip( + "unsloth.models._utils.patch_gradient_accumulation_fix absent" + ) + try: + from transformers import Trainer + except ImportError: + pytest.skip("transformers.Trainer absent") + # The patch reads the live Trainer.training_step source. We + # exercise the standard transformers.Trainer here -- if the + # bug is reintroduced in the source rewriter (e.g. broken + # exec, missing import injection), the patch fn raises. + try: + patch_gradient_accumulation_fix(Trainer) + except Exception as e: + raise AssertionError( + "patch_gradient_accumulation_fix raised on a vanilla " + f"transformers.Trainer: {type(e).__name__}: {e}" + ) + # Idempotency: second call must not raise either (the rewrite + # adds `_unsloth_training_step` marker so the second call + # short-circuits per _utils.py:1692-1693). + patch_gradient_accumulation_fix(Trainer) + + + # ---- unsloth/kernels/moe/grouped_gemm AST smoke ---- + + def _walk_py_files(root: pathlib.Path): + for p in root.rglob("*.py"): + if "__pycache__" in p.parts: + continue + yield p + + + def test_unsloth_kernels_moe_grouped_gemm_ast_parses(): + """unsloth/kernels/moe/grouped_gemm hosts the Triton MoE + kernels (GPU-only at runtime). A SyntaxError or stray token + at the SOURCE level still surfaces as ImportError on every + install, so AST-parse the .py files without executing.""" + # Locate `unsloth/kernels/moe/grouped_gemm` via the installed + # `unsloth` package. + import unsloth as _unsloth + kernel_root = ( + pathlib.Path(_unsloth.__file__).parent + / "kernels" / "moe" / "grouped_gemm" + ) + if not kernel_root.exists(): + pytest.skip( + f"{kernel_root} not present in this unsloth checkout." + ) + fail = [] + ok = 0 + for p in _walk_py_files(kernel_root): + try: + ast.parse(p.read_text(encoding="utf-8"), filename=str(p)) + ok += 1 + except SyntaxError as e: + fail.append((str(p), f"SyntaxError: {e}")) + except Exception as e: + fail.append((str(p), f"{type(e).__name__}: {e}")) + print(f"AST-parsed {ok} grouped_gemm files; failed={len(fail)}") + for path, err in fail: + print(f" AST FAIL {path}: {err}") + assert not fail, ( + f"AST parse failed for {len(fail)} grouped_gemm files" + ) + # Sanity: the directory MUST contain at least the interface + # + kernels + reference subtrees as documented. + expected = [ + "interface.py", + "kernels/forward.py", + "kernels/backward.py", + "reference/moe_block.py", + "reference/moe_ops.py", + ] + missing = [e for e in expected if not (kernel_root / e).is_file()] + assert not missing, ( + "grouped_gemm directory layout regressed; missing: " + f"{missing}" + ) + PY + python -m pytest -q --tb=short -s tests/_moe_coverage_shim.py + rm -f tests/_moe_coverage_shim.py + + - name: Summary + if: always() + run: | + echo "::group::Versions" + python -c "import sys, platform; print(sys.version); print(platform.platform())" + python -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())" + python -c "import transformers; print('transformers', transformers.__version__)" + # `pip show` instead of `import unsloth_zoo` — its __init__ raises + # without an accelerator and the spoof harness only kicks in under + # pytest. Cheap and accurate. + pip show unsloth_zoo + echo "::endgroup::" + echo "Consolidated job done. Coverage:" + echo " - 16 unsloth Bucket-A tests under tests/saving/ + tests/utils/" + echo " - unsloth_zoo @ ${UNSLOTH_ZOO_REF} pytest tests/ (5 GPU cases deselected)" + echo " - unsloth_zoo.compiler.test_apply_fused_lm_head" + + llama-cpp-smoke: + # Standalone llama.cpp build + smoke. Earlier this lived inside every + # consolidated matrix cell and re-cmake'd llama.cpp ~5 min per cell -- + # 3 cells x 275 s = ~14 min of duplicated CPU on every PR for an + # artefact that has nothing to do with the (transformers, TRL) combo. + # `install_llama_cpp` clones ggml-org/llama.cpp at a pinned commit and + # builds the LLAMA_CPP_TARGETS list; the result is independent of the + # HF stack version. Run once, gate the PR. + name: llama.cpp build + smoke + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + UNSLOTH_ZOO_REF: ${{ inputs.unsloth_zoo_ref || 'main' }} + # Same env contract the matrix cells use: protobuf python parser + # (transformers' bundled *_pb2.py needs it), studio on PYTHONPATH, + # compile-disable + UNSLOTH_IS_PRESENT so unsloth_zoo's __init__ + # bootstrap accepts a pure-import. + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + PYTHONPATH: ${{ github.workspace }}/studio + UNSLOTH_COMPILE_DISABLE: '1' + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Install runtime deps for unsloth_zoo.llama_cpp + # unsloth_zoo's `__init__` imports `temporary_patches`, which + # in turn pulls per-architecture submodules (gemma3n, gemma4, + # qwen3_*_moe, glm4_moe, deepseek_v3_moe, pixtral, ministral, + # mxfp4, bitsandbytes, flex_attention_bwd) -- many of those + # transitively touch transformers and peft / accelerate. Mirror + # the matrix job's install minus the heavy bits that have no + # bearing on `install_llama_cpp` itself: studio.txt's FastAPI + # stack, bitsandbytes (CUDA-only build dependency), triton, + # mammoth/unpdf (PDF tools), datasets, sqlalchemy/cryptography, + # pytest (we run no tests). The remaining pin shape matches + # studio-backend-ci.yml's "Repo tests (CPU)" baseline. + run: | + set -euxo pipefail + python -m pip install --upgrade pip + # Match the matrix job's torch path so unsloth_zoo's + # `import torch` resolves to the same CPU build. + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' + pip install \ + 'numpy<3' protobuf sentencepiece \ + requests tqdm psutil packaging safetensors \ + 'peft>=0.18,<0.20' 'accelerate>=0.34,<2' + # transformers + trl come from pyproject.toml's pinned line + # so this job stays in sync with whatever the consolidated + # `__from_pyproject__` matrix cell is using. + pip install transformers trl + pip install -e . --no-deps + + - name: Clone unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} + # Same shallow clone as the matrix job; we install editable so + # `unsloth_zoo.llama_cpp` resolves to the cloned tree (and any + # main-branch fixes flow into the smoke without a release). + run: | + set -euxo pipefail + # github.com occasionally 500s on the git fetch; retry so a + # single upstream blip does not fail CI. + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 --branch="$UNSLOTH_ZOO_REF" \ + https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + pip install -e "$RUNNER_TEMP/unsloth-zoo" --no-deps + pip show unsloth_zoo + + - name: llama.cpp install via unsloth_zoo.llama_cpp + CLI `--help` smoke + # Exercise the canonical `unsloth_zoo.llama_cpp.install_llama_cpp` + # flow that GGUF export uses at runtime: clone ggml-org/llama.cpp + # into ~/.unsloth/llama.cpp, build the LLAMA_CPP_TARGETS list + # (llama-quantize, llama-cli, llama-mtmd-cli, llama-gguf-split, + # llama-server) via cmake, then run `--help` on whichever CLI + # inference binary the build actually produced. + # + # This replaces the previous "download upstream prebuilt zip" + # approach, which silently exited 0 with the message + # "no ubuntu-x64 prebuilt asset" when ggml-org's release-asset + # naming drifted (the regex `bin-ubuntu-x64.*\.zip$` no longer + # matched their current asset names). The build path is the same + # one Unsloth users hit in production via `model.save_pretrained_gguf`. + # + # We do NOT hard-require `llama-cli` specifically: upstream + # ggml-org/llama.cpp moved the cli/server/ui targets behind the + # `LLAMA_BUILD_SERVER` cmake option (tools/CMakeLists.txt) and the + # set of binaries that survive a given checkout drifts over time + # (e.g. a recent build root shipped llama-server + llama-quantize + # + llama-diffusion-cli but no llama-cli). The durable contract is + # "install_llama_cpp produced a working CLI inference binary AND a + # working quantizer", so we --help-probe the first of + # llama-cli / llama-mtmd-cli / llama-server that exists. If a + # future llama.cpp restores llama-cli it is first in the list and + # is preferred, so this stays backwards compatible. + # + # Wall-time budget: ~3-5 min cold, dominated by cmake build of + # 5 targets on the runner's 4 cores. Apt-package install is + # handled by `install_llama_cpp` itself via its + # `check_build_requirements` -> `install_package` chain. + run: | + set -euxo pipefail + # libssl-dev / libcurl4-openssl-dev are needed by llama.cpp's + # cmake build for HTTPS support; install up-front so the + # `install_llama_cpp` requirement-check is a no-op. + sudo apt-get update -qq + sudo apt-get install -y -qq build-essential cmake git curl \ + libgomp1 libssl-dev libcurl4-openssl-dev + python <<'PY' + import os, shutil, subprocess, sys, pathlib + # Apply the same CPU spoof the pytest shims use BEFORE any + # unsloth_zoo import: unsloth_zoo/__init__.py calls + # device_type.get_device_type() at module load and raises + # `NotImplementedError: Unsloth cannot find any torch + # accelerator` on a GPU-less runner. The spoof flips + # torch.cuda.is_available() to True so the device probe takes + # the cuda branch; we never actually run CUDA tensor ops in + # this step (just clone+cmake+--help on the binaries). + sys.path.insert(0, str(pathlib.Path("tests").resolve())) + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + from unsloth_zoo.llama_cpp import ( + install_llama_cpp, + LLAMA_CPP_DEFAULT_DIR, + LLAMA_CPP_TARGETS, + ) + print(f"Unsloth llama.cpp default dir: {LLAMA_CPP_DEFAULT_DIR}") + print(f"Build targets: {LLAMA_CPP_TARGETS}") + # install_llama_cpp returns (quantizer_path, converter_script_path). + # The quantizer's directory is the `llama.cpp` install root, which + # also holds the CLI inference binaries after build/bin/llama-* gets + # copied up (llama_cpp.py:1450-1454; on Windows they stay in + # build/bin/Release/). + quantizer, converter = install_llama_cpp(print_output=True) + assert quantizer and os.path.exists(quantizer), ( + f"install_llama_cpp returned quantizer={quantizer!r} but file missing" + ) + assert converter and os.path.isfile(converter), ( + f"install_llama_cpp returned converter={converter!r} but missing" + ) + install_root = os.path.dirname(quantizer) + is_windows = sys.platform == "win32" + exe = ".exe" if is_windows else "" + # Search both the copied-up root and the Windows build/bin/Release/ + # location the quantizer might already live in. + search_dirs = [install_root] + win_release = os.path.join(install_root, "build", "bin", "Release") + if win_release not in search_dirs: + search_dirs.append(win_release) + # Any of these proves a working llama.cpp CLI inference binary was + # built. Order = preference: llama-cli is canonical (restored first + # if upstream brings it back), then the multimodal CLI, then the + # server (always built whenever cli would be, behind LLAMA_BUILD_SERVER). + cli_names = [f"llama-cli{exe}", f"llama-mtmd-cli{exe}", f"llama-server{exe}"] + cli = None + cli_name = None + for name in cli_names: + for d in search_dirs: + candidate = os.path.join(d, name) + if os.path.exists(candidate) and (is_windows or os.access(candidate, os.X_OK)): + cli, cli_name = candidate, name + break + if cli is not None: + break + if cli is None: + found = [] + for d in search_dirs: + if os.path.isdir(d): + found += [p for p in os.listdir(d) if p.startswith("llama-")] + raise AssertionError( + f"No CLI inference binary ({', '.join(cli_names)}) found after " + f"build in {search_dirs}. Build root contents: {sorted(set(found))[:20]}" + ) + print(f"Using CLI inference binary: {cli_name} -> {cli}") + # `--help` exits non-zero on some builds; the contract is that + # recognizable help text appears on stdout/stderr. llama-server + # exposes a different flag set than llama-cli, so accept its + # tokens too (e.g. --host / --port / "server"). + proc = subprocess.run( + [cli, "--help"], capture_output=True, text=True, timeout=30, + ) + combined = (proc.stdout or "") + (proc.stderr or "") + print(f"--- {cli_name} --help (first 30 lines) ---") + print("\n".join(combined.splitlines()[:30])) + assert any( + tok in combined.lower() + for tok in ("usage", "--help", "--model", "-m,", "--host", "--port", "server") + ), ( + f"{cli_name} --help produced no recognizable help text. " + f"exit={proc.returncode}\nstdout: {proc.stdout[:400]!r}\n" + f"stderr: {proc.stderr[:400]!r}" + ) + # Also exercise the quantizer the way GGUF export does: --help + # round-trip on the binary that does the actual heavy lifting. + q = subprocess.run( + [quantizer, "--help"], capture_output=True, text=True, timeout=15, + ) + q_combined = (q.stdout or "") + (q.stderr or "") + assert "usage" in q_combined.lower() or "type" in q_combined.lower(), ( + f"llama-quantize --help produced no help text. " + f"exit={q.returncode}\nstdout: {q.stdout[:400]!r}\n" + f"stderr: {q.stderr[:400]!r}" + ) + print( + f"\nOK: install_llama_cpp produced a working {cli_name} at {cli} " + f"and llama-quantize at {quantizer}." + ) + PY diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml new file mode 100644 index 0000000..bb7dcbf --- /dev/null +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs installer parity and autostart opt-out tests on Windows and macOS. +# +# Why: that test is the guard that install.sh and install.ps1 stay in +# sync, but today it only runs on ubuntu-latest (auto-discovered by +# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both +# installer scripts, and on Windows Path.read_text() defaults to the +# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already +# contains a U+274C) raises UnicodeDecodeError there even though Linux and +# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in +# #6166; this job keeps that from silently regressing by exercising the +# test on the platforms it claims parity for. Pure pytest, no GPU, +# sub-second, so the matrix is cheap. + +name: Cross-platform parity + +on: + pull_request: + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/test_installer_skip_autostart.py' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + push: + branches: [main] + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/test_installer_skip_autostart.py' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + parity: + name: parity (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 10 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - run: python -m pip install -U pip pytest + - name: Cross-platform parity tests + env: + UNSLOTH_NO_TORCH: '1' + run: >- + python -m pytest + tests/python/test_cross_platform_parity.py + tests/test_installer_skip_autostart.py + -q diff --git a/.github/workflows/lint-ci.yml b/.github/workflows/lint-ci.yml new file mode 100644 index 0000000..bd859a6 --- /dev/null +++ b/.github/workflows/lint-ci.yml @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Whole-repo, multi-language source-lint gate. Runs on every PR +# (no path filter) because each step is sub-second to a few seconds +# and together they catch a class of breakage the focused build +# workflows would miss: +# +# - Python syntax + ruff + leftover debugger calls (across 350+ +# committed .py files, not just studio/backend). +# - Shell `bash -n` parse for every committed *.sh. +# - `yaml.safe_load` and `json.loads` round-trip for every +# committed YAML / JSON config. +# +# TypeScript and Rust are NOT duplicated here on purpose: +# - Studio Frontend CI runs `npm run typecheck` (= `tsc --noEmit`) +# and `npm run build` (vite/swc) on every studio/frontend/** +# change, which is a full TS AST + type check. +# - Studio Tauri CI runs `tauri build --debug --no-bundle` on +# every studio/src-tauri/** or studio/frontend/** change, which +# compiles the Rust crate (= cargo check + cargo build). +# Each is a stricter check than a parse-only step would be, so a +# fast-fail duplicate here would only burn cache; the dedicated +# workflows already block merges on Rust / TS regressions. + +name: Lint CI + +on: + pull_request: + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + source-lint: + name: Source lint (Python + shell + YAML + JSON + safety nets) + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + # Pin ruff to match .pre-commit-config.yaml so a CI-only ruff + # bump cannot disagree with what pre-commit accepted. + # codespell is pinned for the same reason: a reviewer should + # never see a typo report appear and disappear depending on + # which codespell version the runner happened to install. + - run: pip install 'ruff==0.15.12' 'pyyaml>=6' 'codespell>=2.3,<3' + + - name: Linux deps for shellcheck + run: sudo apt-get update -qq && sudo apt-get install -y --no-install-recommends shellcheck + + - name: Python AST/syntax check (every committed .py must compile) + # python -m compileall uses the same parser the interpreter + # uses, so anything broken here would also crash at + # `import X` on a user's machine. Sub-second across 350+ + # files. Hard gate. + run: | + python -m compileall -q -j 0 \ + unsloth unsloth_cli studio tests cli.py unsloth-cli.py + + - name: Python ruff check (whole repo) + # The narrow rule set in pyproject.toml [tool.ruff.lint] + # selects E9 / F63 / F7 / F82 -- syntax errors, broken + # comparisons, undefined names. The whole repo passes today, + # so this is a hard gate. + run: | + ruff check unsloth unsloth_cli studio tests cli.py unsloth-cli.py + + - name: Import-hoist verifier self-test + # scripts/verify_import_hoist.py is a scope-aware (LEGB) AST + # resolver that gates import-hoisting / alias-rename refactors + # against two bugs ruff and pyflakes both miss: + # 1. dangling alias -- `from a import b as _b` hoisted to + # `from a import b` but a leftover `_b` reference now + # resolves to nothing (or to some other module-level `_b`). + # 2. rename clash -- `_b -> b` silently re-points at a + # different object already named `b` in that scope. + # This step runs the tool's 8 negative-control cases so a + # regression in the verifier itself fails before we trust it on + # a diff. Hermetic, stdlib-only, sub-second. Hard gate. + run: | + python scripts/verify_import_hoist.py --self-test + + - name: Import-hoist / alias-rename safety (changed Python files) + # Runs the verifier in compare mode on every in-place-modified + # .py in the PR: parses each file BEFORE (base branch) and AFTER + # (this diff), resolves every name load, and fails on a BLOCKER + # (dangling alias / rename clash / re-pointed import). INFO + # findings (a helper relocated to another file) do not fail. + # + # --diff-filter=M (in-place edits only) is deliberate: that is + # exactly where a hoist refactor lives, and it skips brand-new + # files whose re-export imports would otherwise look "unused". + # + # Diff against the true merge-base, not the base tip. A two-dot + # diff against the tip re-lints every file the base branch + # changed after the PR branched, comparing newer base code + # (BEFORE) against the PR's older snapshot (AFTER) - a + # time-reversed comparison that flags the base branch's own + # refactors as blockers on PRs that never touched those files. + # The compare API returns the merge-base without needing local + # history, and fetching that single commit by SHA keeps the + # shallow (fetch-depth: 1) clone. + if: github.event_name == 'pull_request' + env: + GH_TOKEN: ${{ github.token }} + run: | + MERGE_BASE=$(gh api \ + "repos/${{ github.repository }}/compare/${{ github.event.pull_request.base.sha }}...${{ github.event.pull_request.head.sha }}" \ + --jq .merge_base_commit.sha) + git fetch --no-tags --depth=1 origin "$MERGE_BASE" + mapfile -t CHANGED < <( + git diff --name-only --diff-filter=M \ + "$MERGE_BASE" HEAD -- '*.py' \ + | grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true + ) + if [ "${#CHANGED[@]}" -eq 0 ]; then + echo "no in-place-modified Python files to check" + exit 0 + fi + printf 'merge base: %s\n' "$MERGE_BASE" + printf 'checking %d file(s):\n' "${#CHANGED[@]}" + printf ' %s\n' "${CHANGED[@]}" + python scripts/verify_import_hoist.py \ + --before "$MERGE_BASE" --after HEAD "${CHANGED[@]}" + + - name: No leftover debugger / pdb / breakpoint calls + # Catches the "I'll just stick a breakpoint() here" mistake + # before it ships. AST-based so commented-out debugger + # markers don't false-positive (a bare grep would; there + # are three commented `# breakpoint()` markers in + # unsloth/models/rl* today). Sub-second. + run: | + python <<'PY' + import ast, pathlib, sys + + SKIP_PARTS = {".venv", "venv", "build", "dist", ".git", + "unsloth_compiled_cache", "node_modules", + "unsloth.egg-info"} + + bad = [] + scanned = 0 + for path in sorted(pathlib.Path(".").rglob("*.py")): + if any(part in SKIP_PARTS for part in path.parts): + continue + scanned += 1 + try: + tree = ast.parse(path.read_text(encoding="utf-8", errors="replace")) + except SyntaxError: + continue # compileall step above already failed this + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + fn = node.func + if isinstance(fn, ast.Name) and fn.id == "breakpoint": + bad.append((path, node.lineno, "breakpoint()")) + elif (isinstance(fn, ast.Attribute) and fn.attr == "set_trace" + and isinstance(fn.value, ast.Name) + and fn.value.id in {"pdb", "ipdb"}): + bad.append((path, node.lineno, f"{fn.value.id}.set_trace()")) + + if bad: + for path, lineno, what in bad: + print(f"::error file={path},line={lineno}::leftover {what} -- remove before merging") + sys.exit(1) + print(f"no leftover debugger calls (scanned {scanned} files)") + PY + + - name: License-header drift (informational; whole repo) + # Three header families are accepted across the repo: + # 1. SPDX one-liner: `# SPDX-License-Identifier: ...` + # Used across studio/ (AGPL-3.0-only) and a few new + # files elsewhere. + # 2. Apache-2.0 long form, marker phrase + # "Licensed under the Apache License". Used across + # unsloth/ and unsloth_cli/. + # 3. GNU long form, marker phrase "General Public License". + # That single substring covers GPL, LGPL ("GNU Lesser + # General Public License") and AGPL ("GNU Affero + # General Public License") preambles, all three of + # which appear in unsloth/kernels/* (LGPL/AGPL) without + # the SPDX line. + # Empty files (mainly empty __init__.py) are skipped. + # Surfaced as a warning; cleaning up the actual misses is a + # follow-up PR, not a CI fix. + continue-on-error: true + run: | + python <<'PY' + import pathlib + + ACCEPTED = ( + "SPDX-License-Identifier", # any SPDX line + "Licensed under the Apache License", # Apache-2.0 long form + "General Public License", # GPL / LGPL / AGPL long form + ) + SKIP_PARTS = {".venv", "venv", "build", "dist", ".git", + "unsloth_compiled_cache", "node_modules", + "unsloth.egg-info"} + + studio_missing = [] + other_missing = [] + for path in sorted(pathlib.Path(".").rglob("*.py")): + if any(part in SKIP_PARTS for part in path.parts): + continue + text = path.read_text(encoding="utf-8", errors="replace") + if not text.strip(): + continue # empty __init__.py etc. + head = "\n".join(text.splitlines()[:25]) + if any(marker in head for marker in ACCEPTED): + continue + if "studio" in path.parts: + studio_missing.append(path) + else: + other_missing.append(path) + + total = len(studio_missing) + len(other_missing) + if total == 0: + print("every committed .py has a recognised license header") + else: + print(f"::warning::{total} Python files have no recognised license " + f"header (SPDX / Apache-2.0 / GNU long form): " + f"studio={len(studio_missing)}, other={len(other_missing)}") + for path in (studio_missing + other_missing)[:30]: + print(f" {path}") + if total > 30: + print(f" ... and {total - 30} more") + PY + + - name: Shell scripts parse cleanly (`bash -n`) + # Same idea as Python's compileall: parse-only check that + # every committed *.sh would not blow up at `bash script.sh` + # invocation time on a release box. tests/sh/ is the largest + # cluster (the install.sh shape tests). + run: | + shopt -s globstar + fail=0 + for f in $(git ls-files '*.sh'); do + if ! bash -n "$f"; then + echo "::error file=$f::shell parse error" + fail=1 + fi + done + if [ "$fail" -ne 0 ]; then + exit 1 + fi + n=$(git ls-files '*.sh' | wc -l) + echo "$n shell scripts parse cleanly" + + - name: YAML files parse cleanly (yaml.safe_load) + # Catches truncated workflow files, broken indents in + # dependabot.yml / pre-commit configs, etc. Includes + # .github/workflows/*.yml so a typo in the file we just + # added shows up immediately. + run: | + python <<'PY' + import pathlib, sys, yaml + + SKIP_PARTS = {".venv", "venv", "build", "dist", ".git", + "node_modules", "unsloth_compiled_cache", + "unsloth.egg-info"} + + bad = [] + scanned = 0 + for path in sorted(list(pathlib.Path(".").rglob("*.yml")) + + list(pathlib.Path(".").rglob("*.yaml"))): + if any(part in SKIP_PARTS for part in path.parts): + continue + scanned += 1 + try: + with path.open("r", encoding="utf-8") as fh: + list(yaml.safe_load_all(fh)) + except Exception as exc: + bad.append((path, exc)) + + if bad: + for path, exc in bad: + print(f"::error file={path}::YAML parse failed: {exc}") + sys.exit(1) + print(f"{scanned} YAML files parse cleanly") + PY + + - name: JSON files parse cleanly (json.loads) + # Catches malformed package.json, biome.json, etc. Skips: + # - huge npm/bun lockfiles (machine-generated, slow to + # parse, no value). + # - tsconfig*.json: TypeScript convention is JSONC (JSON + # with `/* ... */` comments), which standard json.loads + # rejects. Strip-and-validate would need json5 or a + # hand-rolled comment scrubber for marginal value, since + # `tsc --noEmit` already validates these in Frontend CI. + run: | + python <<'PY' + import fnmatch, json, pathlib, sys + + SKIP_PARTS = {".venv", "venv", "build", "dist", ".git", + "node_modules", "unsloth_compiled_cache", + "unsloth.egg-info"} + SKIP_NAMES = {"package-lock.json", "bun.lock"} + SKIP_PATTERNS = ("tsconfig*.json",) + + bad = [] + scanned = 0 + for path in sorted(pathlib.Path(".").rglob("*.json")): + if any(part in SKIP_PARTS for part in path.parts): + continue + if path.name in SKIP_NAMES: + continue + if any(fnmatch.fnmatch(path.name, pat) for pat in SKIP_PATTERNS): + continue + scanned += 1 + try: + json.loads(path.read_text(encoding="utf-8")) + except Exception as exc: + bad.append((path, exc)) + + if bad: + for path, exc in bad: + print(f"::error file={path}::JSON parse failed: {exc}") + sys.exit(1) + print(f"{scanned} JSON files parse cleanly") + PY + + - name: codespell typo check (informational) + # Catches typos in code, comments, and docs across the repo. + # Skips lockfiles, generated assets, binary artefacts, and + # the LICENSE files (US/UK spelling drift in legal text is + # not ours to second-guess). The ignore-words-list pulls + # out short identifiers + valid technical terms that + # codespell's default dictionary would otherwise flag + # (e.g. `ans` as a math-quiz variable name in + # tests/utils/aime_eval.py, `parm`/`parms` in PyTorch + # nn.Module idioms). Non-blocking until the surfaced typos + # are fixed; drop continue-on-error after the cleanup. + continue-on-error: true + run: | + codespell \ + --skip='*.lock,*.lockb,*.json,*.svg,*.png,*.jpg,*.jpeg,*.gif,*.ico,*.woff*,*.ttf,*.eot,*.zip,*.gz,*.gguf,*.safetensors,*.bin,node_modules,.git,build,dist,unsloth_compiled_cache,unsloth.egg-info,target,studio/frontend/dist,*.pyc,*-licenses.txt,LICENSE*' \ + --ignore-words-list='ans,bu,hel,fo,te,ot,hist,ned,sav,recurser,datas,nin,parm,parms,checkin,nd,fr,inout,donot,uint' \ + --quiet-level=2 + + - name: shellcheck on committed *.sh (informational) + # Goes beyond `bash -n` (which only parses): catches subtle + # shell bugs like unquoted variable expansions, useless + # `cat`, command substitutions inside `[[`, etc. The + # install/setup scripts are critical-path so the signal is + # worth surfacing. Non-blocking until install.sh's + # hand-rolled patterns get cleaned up; drop continue-on-error + # afterwards. + continue-on-error: true + run: | + # Exclude SC1090 ("source not followable") -- legitimate + # for installer scripts that source files at runtime + # paths shellcheck cannot resolve statically. + # SC2034 ("variable assigned but never used") fires on + # the export-only assignment idiom we use in install.sh. + shellcheck -e SC1090,SC2034 $(git ls-files '*.sh') + + - name: ruff format drift (informational) + # The canonical formatter is scripts/run_ruff_format.py + # = ruff format + scripts/enforce_kwargs_spacing.py, so plain + # `ruff format --check` reports the kwarg-spacing diff as + # drift. Surface the count for visibility but keep + # non-blocking until the custom pipeline is wired in here. + continue-on-error: true + run: | + ruff format --check unsloth unsloth_cli studio tests cli.py unsloth-cli.py diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml new file mode 100644 index 0000000..25796bd --- /dev/null +++ b/.github/workflows/local-agent-guides-ci.yml @@ -0,0 +1,787 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Local Agent Guides CI +# ===================== +# Detects when our local-agent setup recipes drift out of sync with +# `unsloth run`. Boots a real `unsloth run --disable-tools` server and +# drives the coding agents end to end through the *exact* recipes defined +# in unsloth_cli/commands/start.py (the in-repo source of truth -- there +# is no docs/ tree). Wherever start.py has a recipe we drive the agent +# via `unsloth start --no-launch` and execute what it prints, so +# the test self-updates against start.py and catches silent recipe drift. +# +# Source-of-truth files this workflow guards: +# unsloth_cli/commands/start.py the `unsloth start ` recipes +# unsloth_cli/commands/studio.py the `unsloth run` banner (API Key line) +# +# Failure taxonomy (each surfaced with a distinct ::error:: + the agent name +# + the start.py location, so a red X is immediately triageable): +# (a) Unsloth server/API regression -- the dialect HTTP preflight fails +# BEFORE the agent runs (or the server never becomes healthy). +# (b) Agent package install failed -- npm/curl install of the CLI failed. +# (c) Guide drift -- preflight passed + install ok, but +# the documented `unsloth start` flow produced no/garbled output. +# +# Agents covered (6): claude, codex, hermes, openclaw, opencode, pi. +# - All six have a `unsloth start ` recipe, so each cell obtains its +# env + command from `unsloth start --no-launch` and runs THAT +# (self-updating: a recipe change is exercised automatically). + +name: Local Agent Guides CI + +on: + # Off-peak weekly, deliberately a NON-:00 minute to dodge the top-of-hour + # GitHub-hosted-runner stampede. + schedule: + - cron: '37 7 * * 1' + workflow_dispatch: + pull_request: + paths: + - 'unsloth_cli/**' + - 'studio/backend/routes/**' + # Contracts this workflow asserts that live outside routes/**: the + # /api/health endpoint, the llama-server KV-cache log behavior, and the + # request/response schemas the agent dialects depend on. + - 'studio/backend/main.py' + - 'studio/backend/core/inference/llama_cpp.py' + - 'studio/backend/models/**' + - 'install.sh' + - '.github/workflows/local-agent-guides-ci.yml' + - '.github/scripts/serve-unsloth-run.sh' + - '.github/scripts/assert-prompt-cache.sh' + - '.github/scripts/agent-guides-install.sh' + - '.github/scripts/agent-guides-drive.sh' + - '.github/scripts/ci-connect-prompt.txt' + - '.github/scripts/ci-min-system-prompt.txt' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +# Secret handling on pull_request: these jobs check out and run PR-controlled code +# (install.sh, .github/scripts/**), so HF_TOKEN (an external HF credential) is gated +# off pull_request at each step below -- public GGUF repos still download anonymously. +# GH_TOKEN (GITHUB_TOKEN) is kept: it is the job-scoped contents:read token and +# install_llama_prebuilt.py needs it for the GitHub releases API (else 403s). + +env: + # Determinism precedent (studio-inference-smoke.yml): temp 0 + fixed seed. + UNSLOTH_SEED: '3407' + # A single invoke must never hang the runner on a headless TTY prompt. With + # prefill-shrinking flags (minimal system prompt + restricted tools) a turn on + # a 4B model finishes in a couple of minutes on CPU; this also caps how long a + # still-large-prompt agent burns before failing. Well under the 6h job cap. + AGENT_INVOKE_TIMEOUT: '600' + +jobs: + # ═════════════════════════════════════════════════════════════════════ + # Job 1: connection + # Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect, + # install the agent, run `unsloth start --no-launch`, execute + # the emitted recipe with a trivial prompt, assert a non-empty reply. + # Runs on PR + weekly + dispatch. Each matrix cell is its own runner so + # it serves exactly one model on its own port. + # ═════════════════════════════════════════════════════════════════════ + connection: + name: connection (${{ matrix.agent }}) + runs-on: ubuntu-latest + timeout-minutes: 40 + strategy: + fail-fast: false + matrix: + agent: [claude, codex, hermes, openclaw, opencode, pi] + include: + # OpenClaw needs Node 24; everything else is happy on 22. + - agent: openclaw + node: '24' + env: + # gemma-4-E4B (128K context, capable enough to drive every agent for a + # trivial reply; the 270m model produced empty/failed responses for + # codex/openclaw). Hermes' 64K context floor no longer constrains the model + # choice: write_hermes_config claims the floor for smaller windows and + # scales compaction back to the real window. Served as a flat + # GGUF file (the -MTP- repo ships no separate draft, so this is plain 4B). + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18901' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: ${{ matrix.node || '22' }} + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + # ── boot the server under test (factored helper) ────────────────── + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + # ── (a) server/API preflight: prove the dialect works BEFORE the agent ─ + # Distinct error class. If this step fails it is a SERVER regression, + # not the agent's or the guide's fault, and the agent steps never run. + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect; this is class (a), not guide drift). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + # Anthropic Messages dialect. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + # Codex always streams /v1/responses. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). + # OpenClaw's start.py recipe writes an "openai-completions" + # provider (write_openclaw_config), so it uses this path, not + # /v1/messages. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + # ── (b) install the agent CLI (hardened npm/curl, retried) ───────── + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + # ── (c) drive the agent via start.py and assert a reply ────────── + # For the 5 agents with a start.py recipe we run + # `unsloth start --no-launch`, eval its env/unset exports, + # then run the printed command with a hard timeout (no headless-TTY + # hang). Pi has no connect recipe, so it is driven by hand and the + # cell asserts that absence is the (known) reason. + - name: Drive ${{ matrix.agent }} via unsloth start (class-c isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: connection-${{ matrix.agent }}-log + path: | + logs/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job 2: file-edit + # The deterministic 2-turn hello.py test on Qwen3.5-4B (smaller models + # can't reliably drive the heavyweight agents' edit flows). Weekly + + # dispatch only -- it is the slow, model-heavy job and must not gate PRs. + # ═════════════════════════════════════════════════════════════════════ + file-edit: + name: file-edit (${{ matrix.agent }}) + if: github.event_name != 'pull_request' + runs-on: ubuntu-latest + timeout-minutes: 60 + # hermes and openclaw drive a multi-turn tool loop that a CPU-only runner + # cannot finish in time (e.g. openclaw holds its 300s session-write-lock past + # expiry; each turn re-prefills the tool prompt at ~16 tok/s). Their endpoint + # wiring + generation are already hard-gated by the connection job, so the + # file-edit cell is best-effort here -- it still runs and uploads logs, but a + # timeout does not fail the workflow. Drop best_effort (or move e2e to a GPU + # runner) to make it blocking again. + continue-on-error: ${{ matrix.best_effort || false }} + strategy: + fail-fast: false + matrix: + agent: [claude, codex, hermes, openclaw, opencode, pi] + include: + - agent: openclaw + node: '24' + best_effort: true + - agent: hermes + best_effort: true + env: + # gemma-4-E4B served as a flat GGUF file (cache size tracks the .gguf 1:1, + # no xet-chunk inflation; the -MTP- repo ships no separate draft file). + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18902' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: ${{ matrix.node || '22' }} + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect; this is class (a), not guide drift). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + # Probe the same dialect the agent will use, so a streaming/messages + # regression in the weekly run is reported as class (a) here instead of + # surfacing later as guide drift (mirrors the connection job). + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: 2-turn hello.py test (class-c isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh file-edit "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: file-edit-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job: resume + # Does a conversation started with `unsloth start ` survive exit + # and resume? This drives the REAL launch path (not the --no-launch + # recipe the other jobs use). A plain launch relocates the agent home to + # a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the + # session to the stable Unsloth agents dir so it persists. opencode/claude + # keep their session data in a fixed user dir, so they persist either way. + # Dispatch-only: it is an end-to-end experiment, not a PR gate. + # ═════════════════════════════════════════════════════════════════════ + resume: + name: resume (${{ matrix.agent }}) + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + # codex/pi relocate their whole home (resume broken without --persist); + # opencode/claude keep session data in a fixed dir (resume already works). + # One agent from each class proves the split end to end; openclaw/hermes + # share codex's relocation mechanism and are covered by the unit tests. + agent: [codex, opencode, claude, pi] + env: + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18904' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: Resume experiment (launch path) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: resume-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job 3: prompt-cache + # (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0 + # (server prompt-cache sanity). + # (b) Claude Code attribution A/B: with CLAUDE_CODE_ATTRIBUTION_HEADER=0 + # expect a llama-server KV-cache HIT on turn 2; without it expect a + # MISS. If it inverts, the guide flag is stale. + # PR + weekly + dispatch (cheap, gemma-3-270m). + # ═════════════════════════════════════════════════════════════════════ + prompt-cache: + name: prompt-cache (gemma-3-270m) + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18903' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-3-270m) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --model "$GGUF_REPO" --gguf-variant "$GGUF_VARIANT" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" + + # (a) server prompt-cache sanity on the OpenAI chat path. The helper runs + # the 2-turn probe internally (turn 2 reuses turn 1's prefix) and asserts + # turn-2 usage.prompt_tokens_details.cached_tokens > 0. This is the hard + # gate -- it proves llama.cpp KV reuse is surfaced on /v1/chat/completions. + - name: Server prompt-cache sanity (cached_tokens > 0) + run: bash .github/scripts/assert-prompt-cache.sh api "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" + + - name: Install Claude Code (class-b isolation) + env: + AGENT: claude + run: bash .github/scripts/agent-guides-install.sh claude + + # (b) Claude attribution A/B against the llama-server log. This is the most + # environment-sensitive check (it depends on the bundled llama.cpp's + # slot-reuse log wording and on claude --continue reusing the prefix), so + # it is non-blocking until calibrated on the first scheduled run; the + # server cache sanity above is the hard gate. The step still prints the + # observed HIT/MISS so drift is visible in the log + artifacts. + - name: Claude attribution A/B (HIT with header=0, MISS without) + continue-on-error: true + run: bash .github/scripts/agent-guides-drive.sh attribution-ab claude + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: prompt-cache-log + path: | + logs/ + redacted-configs/ + retention-days: 7 diff --git a/.github/workflows/lockfile-audit.yml b/.github/workflows/lockfile-audit.yml new file mode 100644 index 0000000..aaf258d --- /dev/null +++ b/.github/workflows/lockfile-audit.yml @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Fast, focused supply-chain audit of every checked-in lockfile. +# +# Runs scripts/lockfile_supply_chain_audit.py on PRs that touch any +# npm or cargo lockfile, on push to main, and on a daily schedule so +# newly-published IOCs surface even when no PR opens. +# +# Default behavior is "advisory": only public indicator-of-compromise +# strings, known-malicious pinned versions, and structurally broken +# lockfiles fail the build. Structural anomalies (missing integrity, +# non-default registry, etc.) are emitted as GitHub Actions warnings +# but do not block merges. This deliberately keeps the noise floor +# low while still failing the moment a checked-in lockfile starts +# pointing at known-bad bytes. +# +# This workflow is intentionally separate from security-audit.yml: +# - security-audit.yml is the umbrella job (pip-audit + npm audit + +# cargo audit + OSV + Semgrep + secret scanning + SBOM + ...); +# it takes ~25 minutes and runs only when dep manifests change. +# - lockfile-audit.yml is a ~30 second pure-Python parse + grep on +# the lockfiles themselves; it runs on every PR that even nudges +# a lockfile so reviewers always see the audit result inline. + +name: Lockfile supply-chain audit + +on: + pull_request: + paths: + - 'studio/frontend/package-lock.json' + - 'studio/backend/core/data_recipe/oxc-validator/package-lock.json' + - 'studio/package-lock.json' + - 'studio/src-tauri/Cargo.lock' + - 'scripts/lockfile_supply_chain_audit.py' + - '.github/workflows/lockfile-audit.yml' + push: + branches: [main] + paths: + - 'studio/frontend/package-lock.json' + - 'studio/backend/core/data_recipe/oxc-validator/package-lock.json' + - 'studio/package-lock.json' + - 'studio/src-tauri/Cargo.lock' + - 'scripts/lockfile_supply_chain_audit.py' + - '.github/workflows/lockfile-audit.yml' + schedule: + - cron: '37 5 * * *' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + audit: + name: lockfile supply-chain audit + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: '3.12' + + - name: Verify audit script parses + run: python3 -c "import ast; ast.parse(open('scripts/lockfile_supply_chain_audit.py').read())" + + - name: Run lockfile supply-chain audit + # Default mode: only known-malicious pinned versions, known IOC + # strings, and structurally broken lockfiles fail the build. + # Missing-integrity and other structural anomalies are emitted + # as ::warning:: annotations and do not gate merges. + run: python3 scripts/lockfile_supply_chain_audit.py diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml new file mode 100644 index 0000000..a2f716a --- /dev/null +++ b/.github/workflows/mlx-ci.yml @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Focused PR gate for the MLX dispatch surface, running on a real +# Apple Silicon runner. +# +# Runner: macos-14 (M1, 3 vCPU / 7 GB / Apple Silicon standard runner +# -- FREE for public repositories per the GitHub Actions billing +# reference; larger variants like macos-14-large/-xlarge are paid so +# we deliberately avoid those). +# +# Why a single Mac job (no Linux+spoof leg): the dispatch tests are +# 100% spoofed monkeypatches and run identically on any host, so the +# Linux leg was duplicating the matrix tests already covered on Mac +# while missing everything Apple-specific. The Mac job runs the SAME +# spoofed matrix PLUS three things only a real Apple Silicon host +# can prove: +# +# 1. unsloth._IS_MLX flips True on Darwin+arm64 with mlx genuinely +# installed (no spoof). +# 2. Every PR-A MLX-only unsloth_zoo module (mlx_loader, mlx_trainer, +# mlx_compile, mlx_utils, mlx_cce, gated_delta_vjp) imports +# against the real `mlx` + `mlx-lm` + `mlx-vlm` PyPI wheels -- +# each does `import mlx.core as mx` at module top level, so this +# catches a future change that breaks the real wheels without +# needing a Mac developer in the loop. +# 3. The hardware-dispatch spoofs do not collide with the real +# environment (the test fixture installs a MetaPathFinder that +# blocks `import mlx.core` for "no-mlx" profiles, faithfully +# simulating a Mac without mlx even when mlx IS installed). +# 4. End-to-end MLX training + inference smoke test: +# run_real_mlx_smoke.py trains unsloth/gemma-3-270m-it for 7 +# deterministic LoRA steps on a single repeated text row, then +# verifies the trained model can complete the prompt and that +# losses + grad norms are finite and well-behaved. This is the +# only place in CI that exercises a real MLX backward pass + +# optimizer step + inference call. +# +# Three dispatch test files documented in tests/studio/README.md: +# - test_hardware_dispatch_matrix.py parametrized 7-profile matrix +# + 2 dispatch-priority canaries +# - test_is_mlx_dispatch_gate.py AST + runtime guard on +# unsloth._IS_MLX +# - test_mlx_training_worker_behaviors.py AST contract checks on +# studio/backend/core/training/worker.py +# +# Surfaces a single PR check ("MLX CI on Mac M1 / dispatch"). +# +# Security audit footprint: every package this workflow installs is +# already covered by .github/workflows/security-audit.yml -- the deps +# come from studio/backend/requirements/studio.txt and unsloth-zoo's +# pyproject (resolved transitively). The git+ install of unsloth-zoo +# is intentionally skipped by the audit (pip-audit cannot resolve a +# git URL through PyPI metadata; the audit comment in security-audit.yml +# documents this). No new package is introduced solely by MLX CI. + +name: MLX CI on Mac M1 + +on: + pull_request: + paths: + - 'unsloth/__init__.py' + - 'unsloth/_gpu_init.py' + - 'studio/backend/utils/hardware/**' + - 'studio/backend/core/training/worker.py' + - 'studio/backend/core/inference/mlx_inference.py' + - 'tests/studio/test_hardware_dispatch_matrix.py' + - 'tests/studio/test_is_mlx_dispatch_gate.py' + - 'tests/studio/test_mlx_training_worker_behaviors.py' + - 'tests/studio/run_real_mlx_smoke.py' + - 'tests/conftest.py' + - '.github/workflows/mlx-ci.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + dispatch: + name: dispatch + runs-on: macos-14 + # 25 min: dispatch + spoofed matrix + 7-step real LoRA training is + # under 2 min; GGUF export builds llama.cpp via cmake on Apple + # Silicon (~5-7 min), so we budget headroom. + timeout-minutes: 25 + steps: + # harden-runner audit mode: macOS runners cannot use blocking mode + # today (eBPF egress enforcement is Linux-only), but audit mode is + # supported cross-platform and surfaces the egress destinations in + # the runner log. This produces the data needed to graduate this + # job to a block-mode allowlist once macOS support lands. + - name: Harden runner (audit) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: audit + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + # macOS install ladder, validated locally against a Linux + # mac-sim venv (platform spoofed + mlx_simulation shim + real + # datasets/transformers/structlog). + # + # 1. studio/backend/requirements/studio.txt brings structlog, + # fastapi, etc. The hardware probe imports structlog at + # module top level. + # 2. Same pytest / numpy / httpx stack the rest of the repo CI + # uses. + # 3. torch is explicitly installed: unsloth-zoo's pyproject + # deliberately excludes torch on darwin+arm64 (mlx replaces + # it for runtime use), but the dispatch tests spoof + # torch.cuda / torch.xpu / torch.backends.mps via monkeypatch + # and so the test process needs torch importable. We pull + # from the PyTorch CPU index so Apple Silicon gets the + # explicit cpu+MPS arm64 wheel rather than something the + # default PyPI resolver might pick up. The CPU index hosts + # macosx_*_arm64 wheels alongside the Linux x86_64 ones. + # 4. unsloth-zoo from git main (NOT PyPI), WITH deps. PR-A's + # MLX support landed after the most recent unsloth-zoo PyPI + # release; the wheel still raises NotImplementedError on + # Apple Silicon when device_type.get_device_type() runs + # unguarded. Studio's own install.sh overlays unsloth-zoo + # from git main for the same reason. Pulling deps lets pip + # resolve the platform-conditional MLX-only wheels (mlx, + # mlx-lm, mlx-vlm gated on darwin+arm64 in unsloth-zoo's + # pyproject) AND the shared deps (datasets, transformers, + # sentencepiece, ...) that unsloth's MLX branch loads via + # dataprep/raw_text.py. + # 5. unsloth -e . --no-deps so the editable install does not + # fight the unsloth-zoo dep set. + # + # All explicit pip installs are version-pinned to a single + # released version (the latest as of 2026-05-07 within each + # project's existing constraint range). bump alongside the rest + # of the security audit when a new release lands. + - name: Install deps + run: | + python -m pip install --upgrade pip + pip install -r studio/backend/requirements/studio.txt + pip install \ + 'python-multipart==0.0.27' \ + 'aiofiles==25.1.0' \ + 'sqlalchemy==2.0.49' \ + 'cryptography==48.0.0' \ + 'pyyaml==6.0.3' \ + 'jinja2==3.1.6' \ + 'mammoth==1.12.0' \ + 'unpdf==1.0.0' \ + 'requests==2.33.1' \ + 'typer==0.25.1' \ + 'numpy==2.4.4' \ + 'pytest==9.0.3' \ + 'pytest-asyncio==1.3.0' \ + 'httpx==0.28.1' + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch==2.10.0' + # github.com occasionally 500s on the git fetch; retry the + # zoo install so a single upstream blip does not fail CI. + for attempt in 1 2 3; do + if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::pip install unsloth_zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::unsloth_zoo install failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + pip install -e . --no-deps + + # Real Apple Silicon sanity: confirm _IS_MLX activates on real + # hardware with no platform spoof. + - name: Verify _IS_MLX flips True on real Apple Silicon + run: | + python -c " + import platform + assert platform.system() == 'Darwin', platform.system() + assert platform.machine() == 'arm64', platform.machine() + import unsloth + assert unsloth._IS_MLX is True, f'expected _IS_MLX=True on real Apple Silicon, got {unsloth._IS_MLX}' + print('OK: _IS_MLX activated on real Apple Silicon') + " + + # Real Apple Silicon sanity: confirm every PR-A MLX-only module + # loads against real mlx + mlx-lm + mlx-vlm wheels. + - name: Smoke-import every MLX-only unsloth_zoo module + run: | + python -c " + import importlib + for name in [ + 'unsloth_zoo.mlx_loader', + 'unsloth_zoo.mlx_trainer', + 'unsloth_zoo.mlx_compile', + 'unsloth_zoo.mlx_utils', + 'unsloth_zoo.mlx_cce', + 'unsloth_zoo.gated_delta_vjp', + ]: + importlib.import_module(name) + print('OK:', name) + from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig + assert hasattr(FastMLXModel, 'from_pretrained') + print('OK: FastMLXModel + MLXTrainer surface present') + " + + # Spoofed dispatch matrix. Runs on the real Mac too -- the + # test fixture installs a MetaPathFinder that blocks + # `import mlx.core` for "no-mlx" profiles, so the spoofs + # faithfully simulate every supported hardware combo regardless + # of whether mlx is installed for real. + - name: MLX dispatch tests (3 files, 36 tests) + env: + PYTHONPATH: ${{ github.workspace }}/studio + UNSLOTH_COMPILE_DISABLE: '1' + run: | + python -m pytest -v --tb=short \ + tests/studio/test_hardware_dispatch_matrix.py \ + tests/studio/test_is_mlx_dispatch_gate.py \ + tests/studio/test_mlx_training_worker_behaviors.py + + # Real MLX training + inference smoke test. Trains + # unsloth/gemma-3-270m-it for 7 deterministic LoRA steps + # (batch_size=2, gradient_accumulation_steps=3) on a single + # repeated row ("<> My name is Unsloth!"), then saves + # the trained model in 3 export formats. The `train` subcommand + # captures per-phase timing + peak GPU + peak RSS into + # train_metrics.json so we can detect regressions across CI runs. + - name: MLX export round-trip — TRAIN + SAVE 3 formats + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + UNSLOTH_COMPILE_DISABLE: '1' + run: | + mkdir -p mlx_workdir + # Authenticate llama.cpp's release-API lookup (anonymous 403s on rate-limit); + # read-only GITHUB_TOKEN scoped here only, never to steps that run binaries. + GH_TOKEN="${{ secrets.GITHUB_TOKEN }}" GITHUB_TOKEN="${{ secrets.GITHUB_TOKEN }}" \ + python tests/studio/run_real_mlx_smoke.py train \ + --workdir "$PWD/mlx_workdir" + + # Each reload step runs in a FRESH Python process to confirm + # the cold-start path users would hit in production also works + # (not just the in-memory continuation of a still-running + # trainer). FastMLXModel.from_pretrained gets called from + # scratch; mx.random is re-seeded; per-step timing + peak + # memory are emitted to {format}_reload_metrics.json next to + # the saved dir. + - name: MLX export round-trip — RELOAD LoRA (fresh process) + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + UNSLOTH_COMPILE_DISABLE: '1' + run: | + python tests/studio/run_real_mlx_smoke.py reload \ + --format lora \ + --dir "$PWD/mlx_workdir/lora" + + - name: MLX export round-trip — RELOAD merged_16bit (fresh process) + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + UNSLOTH_COMPILE_DISABLE: '1' + run: | + python tests/studio/run_real_mlx_smoke.py reload \ + --format merged \ + --dir "$PWD/mlx_workdir/merged_16bit" + + # GGUF reload uses the llama-cli binary that save_pretrained_gguf + # built. If save_pretrained_gguf was skipped during train (e.g. + # llama.cpp's convert_hf_to_gguf asserts on the model's tokenizer + # vocab -- a downstream llama.cpp limitation, not an unsloth_zoo + # bug), this step emits a workflow warning and exits 0 so the + # LoRA + merged_16bit assertions remain the gating signal. + - name: MLX export round-trip — RELOAD GGUF via llama-cli (fresh process) + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then + python tests/studio/run_real_mlx_smoke.py reload \ + --format gguf \ + --dir "$PWD/mlx_workdir/gguf" + else + REASON=$(python -c "import json; m=json.load(open('mlx_workdir/train_metrics.json')); print(m.get('gguf_skip_reason') or 'unknown')") + echo "::warning title=GGUF round-trip skipped::${REASON}" + echo "GGUF export was skipped during the train phase. Reason:" + echo " ${REASON}" + echo "Continuing without failing the job; the LoRA + merged_16bit" + echo "reload assertions are still gating this PR." + fi + + # Print all metrics JSON files so regressions are visible in the + # job log. always() so we get telemetry even if a reload step + # asserted gibberish. + - name: MLX export round-trip — aggregate metrics + if: always() + run: | + for f in mlx_workdir/train_metrics.json \ + mlx_workdir/lora_reload_metrics.json \ + mlx_workdir/merged_reload_metrics.json \ + mlx_workdir/gguf_reload_metrics.json; do + echo "=== $f ===" + cat "$f" 2>/dev/null || echo "(missing)" + echo + done + + # Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the + # unslothai/llama.cpp fork's latest release, download a small public GGUF, and + # check llama-server /completion end to end. Split and placed last so the + # untrusted binary runs only in the final smoke step, after every HF_TOKEN step, + # leaving no token-bearing step or shared workspace for a tampered prebuilt to + # corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch. + - name: Studio prebuilt llama.cpp install + GGUF download (Mac M1) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -euo pipefail + INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" + rm -rf "$INSTALL_DIR" + # Download only -- no llama-quantize / llama-server launch in this step. + python studio/install_llama_prebuilt.py \ + --install-dir "$INSTALL_DIR" \ + --published-repo unslothai/llama.cpp + mkdir -p /tmp/ggufs + bash .github/scripts/hf-download-with-retry.sh \ + 'unsloth/gemma-3-270m-it-GGUF' \ + 'gemma-3-270m-it-Q4_K_M.gguf' \ + /tmp/ggufs + + # Final step: runs the downloaded binaries with no secrets present, and clears + # the GitHub Actions command files so a tampered prebuilt cannot influence the job. + - name: Studio prebuilt llama.cpp GGUF inference smoke (Mac M1) + run: | + set -euo pipefail + unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY + INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" + # Studio bundles only llama-server + llama-quantize (not llama-cli); + # inference goes through llama-server's HTTP /completion endpoint. + LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server" + LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize" + [ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; } + [ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; } + echo "llama-server : $LLAMA_SERVER" + echo "llama-quantize: $LLAMA_QUANT" + "$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK" + + PORT=18080 + echo "=== starting llama-server on 127.0.0.1:$PORT ===" + "$LLAMA_SERVER" \ + -m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \ + --host 127.0.0.1 \ + --port "$PORT" \ + -c 256 \ + -n 16 \ + --no-warmup \ + > /tmp/llama-server.log 2>&1 & + SERVER_PID=$! + trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT + + # Wait for /health to come up + for i in $(seq 1 30); do + if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + echo " server up after ${i}s" + break + fi + sleep 1 + done + if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + echo "::error::llama-server never became healthy" + tail -40 /tmp/llama-server.log + exit 1 + fi + + PROMPT="Hello, my name is" + echo "=== POST /completion ===" + RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \ + -H 'Content-Type: application/json' \ + -d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}") + echo "raw response (head): $(echo "$RESP" | head -c 600)" + CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))") + echo "completion content: $CONTENT" + + if [ -z "$CONTENT" ]; then + echo "::error::llama-server /completion returned empty content" + tail -40 /tmp/llama-server.log + exit 1 + fi + echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works" diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml new file mode 100644 index 0000000..0e0b35d --- /dev/null +++ b/.github/workflows/notebooks-ci.yml @@ -0,0 +1,448 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Cross-repo notebook validator. Lives in unslothai/unsloth (this repo) +# and inspects every notebook in unslothai/notebooks at HEAD (or the +# ref dispatched in via repository_dispatch). +# +# Catches the bug classes that landed in: +# - unslothai/notebooks#258 Colab torchao 0.10 vs peft 0.19 floor +# - unslothai/notebooks#260 DONT_UPDATE_EXCEPTIONS coverage drift +# - unslothai/notebooks#261 torch/torchcodec ABI; --no-deps tokenizers +# - unslothai/notebooks#264 --no-deps transformers + Colab tokenizers drift +# - unslothai/notebooks#221 git+ HEAD installs in install cells +# - unslothai/notebooks commit 51b1462 template/notebook drift +# +# CPU-only by design. Layer 2 (api-introspect) reuses the existing +# tests/_zoo_aggressive_cuda_spoof.py harness so `import unsloth` +# succeeds on a GPU-less ubuntu-latest runner. + +name: Notebooks CI + +on: + pull_request: + paths: + - 'unsloth/**' + - 'scripts/notebook_validator.py' + - 'scripts/notebook_to_python.py' + - 'scripts/data/colab_pip_freeze.gpu.txt' + - 'scripts/data/colab_to_cpu_pin.json' + - 'tests/notebooks/**' + - 'tests/_zoo_aggressive_cuda_spoof.py' + - '.github/workflows/notebooks-ci.yml' + schedule: + # Daily 06:17 UTC. Catches Colab preinstall bumps (the upstream image + # is rebuilt roughly weekly) without us waiting on a PR. Off the + # :00/:30 fleet-collision spots. + - cron: '17 6 * * *' + workflow_dispatch: + inputs: + notebooks_ref: + description: 'unslothai/notebooks ref to lint (branch / SHA / tag)' + default: 'main' + include_smoke: + description: 'Also run the install-cell smoke matrix (longer)' + type: boolean + default: false + repository_dispatch: + # Fired by a tiny companion workflow on unslothai/notebooks. + types: [notebooks_pr_opened, notebooks_main_pushed] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +env: + NOTEBOOKS_REF: >- + ${{ github.event.inputs.notebooks_ref || + github.event.client_payload.ref || + 'main' }} + +jobs: + static: + name: static (drift + lint + exceptions) + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + # Validate the dispatched ref before it reaches actions/checkout's `ref:` + # input. Reading via env (NOT direct ${{ ... }} interpolation in the + # regex test) closes the GitHub-Actions-injection class where a + # client_payload.ref like `main"; rm -rf / #` would be embedded into the + # shell command. NOTEBOOKS_REF defaults to 'main' on non-dispatch + # events, but only repository_dispatch can supply attacker-controlled + # values, so we gate this check on that event type. + - name: Validate client_payload.ref shape + if: github.event_name == 'repository_dispatch' + env: + NOTEBOOKS_REF: ${{ github.event.client_payload.ref }} + run: | + if ! printf '%s' "$NOTEBOOKS_REF" | grep -Eq '^[A-Za-z0-9._/-]+$'; then + echo "::error::client_payload.ref contains disallowed characters" >&2 + exit 1 + fi + + - name: Checkout unsloth (this PR) + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + path: unsloth + persist-credentials: false + + - name: Checkout unslothai/notebooks @ ${{ env.NOTEBOOKS_REF }} + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + repository: unslothai/notebooks + ref: ${{ env.NOTEBOOKS_REF }} + path: notebooks + fetch-depth: 0 # drift check needs git status / diff + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Install validator deps + run: | + python -m pip install --upgrade pip + # nbformat + nbconvert come from the converter's requirements; + # spellchecker + huggingface_hub are imported at module top of + # update_all_notebooks.py. + pip install \ + 'nbformat>=5.10' 'nbconvert>=7.16' 'pyspellchecker>=0.8' \ + 'huggingface_hub>=0.34' 'tqdm>=4.66' + + - name: Refresh Colab pip-freeze (best-effort; falls back to snapshot) + run: | + python unsloth/scripts/notebook_validator.py refresh-colab \ + --out unsloth/scripts/data/colab_pip_freeze.gpu.txt \ + || echo "::warning::refresh-colab failed; using committed snapshot" + + - name: Diff Colab oracle vs committed snapshots (advisory) + # Pulls pip-freeze.gpu.txt + apt-list-gpu.txt + os-info-gpu.txt + # from googlecolab/backend-info and prints NEW / REMOVED / + # CHANGED entries against scripts/data/colab_*.txt. Non-blocking + # on PRs; the daily cron job below runs the same step with + # --strict so upstream rotations surface within ~24h. + continue-on-error: true + working-directory: ${{ github.workspace }} + run: | + python unsloth/scripts/notebook_validator.py colab-diff \ + --snapshot-dir unsloth/scripts/data + + - name: Drift check (re-run update_all_notebooks.py + git diff) + working-directory: ${{ github.workspace }} + # Reported as non-blocking until the upstream `unslothai/notebooks` + # tree is regenerated. The first run on @main surfaces ~463 files + # of drift (7359 / 9634 line delta), which is a real backlog the + # notebooks-side maintainers need to clear in their own repo -- + # this PR's role is to surface the count, not auto-fix it. + continue-on-error: true + run: | + python unsloth/scripts/notebook_validator.py drift \ + --notebooks-dir notebooks + + - name: Convert sanity (every nb / kaggle / original_template -> .py) + # Same rationale as Drift: a handful of upstream notebooks fail + # the converter (custom magics, malformed JSON, etc). Surface + # the count without blocking; the team triages in unslothai/notebooks. + continue-on-error: true + run: | + python unsloth/scripts/notebook_validator.py convert \ + --notebooks-dir notebooks \ + --out _converted + + - name: Lint (install cells + AST scan, env-scoped) + # Reported as non-blocking (continue-on-error: true) until the + # backlog of pre-existing findings on unslothai/notebooks@main is + # cleared. Same pattern PR #5298 used for biome:check on the + # frontend. As of this commit the live tree surfaces 27 errors + + # 6 warnings, all real (peft/torchao floor missing in 6 nb/ + # notebooks, 14 git+ HEAD installs in hand-tuned exception + # notebooks, 6 torch/torchcodec ABI mismatches, 1 + # transformers/tokenizers --no-deps drift). The count surfaces + # in the PR check UI. Drop continue-on-error once it hits zero. + continue-on-error: true + run: | + python unsloth/scripts/notebook_validator.py lint \ + --notebooks-dir notebooks \ + --colab-pin unsloth/scripts/data/colab_pip_freeze.gpu.txt \ + --no-pypi + # --no-pypi skips R-INST-002 (transitive resolve via PyPI metadata). + # Layer 1 keeps PR-time wall-clock predictable; the daily cron run + # below drops --no-pypi and refreshes the cache. + + - name: DONT_UPDATE_EXCEPTIONS coverage + run: | + python unsloth/scripts/notebook_validator.py exceptions \ + --notebooks-dir notebooks + + static-with-pypi: + name: static + transitive resolve (cron / dispatch only) + if: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }} + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + # See `static.Validate client_payload.ref shape` for rationale. This + # job's `if:` excludes repository_dispatch today, so the validation + # step is a defence-in-depth no-op until that gate ever relaxes. + - name: Validate client_payload.ref shape + if: github.event_name == 'repository_dispatch' + env: + NOTEBOOKS_REF: ${{ github.event.client_payload.ref }} + run: | + if ! printf '%s' "$NOTEBOOKS_REF" | grep -Eq '^[A-Za-z0-9._/-]+$'; then + echo "::error::client_payload.ref contains disallowed characters" >&2 + exit 1 + fi + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + repository: unslothai/notebooks + ref: ${{ env.NOTEBOOKS_REF }} + path: notebooks + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: { python-version: '3.12', cache: 'pip' } + - name: Install + run: pip install -U pip + - name: Refresh Colab oracle + run: | + python unsloth/scripts/notebook_validator.py refresh-colab \ + --out unsloth/scripts/data/colab_pip_freeze.gpu.txt + - name: Diff Colab oracle vs committed snapshots (--strict on cron) + # Cron-only escalation of the advisory PR-time check. Fails if + # any of pip-freeze.gpu.txt / apt-list-gpu.txt / os-info-gpu.txt + # has drifted from scripts/data/colab_*.txt; refresh the + # snapshots in this repo to acknowledge. + run: | + python unsloth/scripts/notebook_validator.py colab-diff \ + --snapshot-dir unsloth/scripts/data --strict + - name: Lint with live PyPI metadata + run: | + python unsloth/scripts/notebook_validator.py lint \ + --notebooks-dir notebooks \ + --colab-pin unsloth/scripts/data/colab_pip_freeze.gpu.txt + + api-introspect: + name: api surface (under CUDA spoof) + runs-on: ubuntu-latest + timeout-minutes: 12 + steps: + - name: Validate client_payload.ref shape + if: github.event_name == 'repository_dispatch' + env: + NOTEBOOKS_REF: ${{ github.event.client_payload.ref }} + run: | + if ! printf '%s' "$NOTEBOOKS_REF" | grep -Eq '^[A-Za-z0-9._/-]+$'; then + echo "::error::client_payload.ref contains disallowed characters" >&2 + exit 1 + fi + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + repository: unslothai/notebooks + ref: ${{ env.NOTEBOOKS_REF }} + path: notebooks + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: { python-version: '3.12', cache: 'pip' } + + - name: Install CPU torch + pinned unsloth + trl + converter deps + run: | + python -m pip install --upgrade pip + # CPU torch + torchvision. torchvision is required because + # unsloth_zoo.vision_utils imports PIL at module top, and the + # easiest way to get a torch-compatible PIL on a CPU runner is + # to let torchvision pull the right Pillow version. + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.8,<2.11' 'torchvision<0.26' + # Pin to the same versions update_all_notebooks.py installs in + # generated notebooks. Keep these in lockstep with PIN_TRL / + # PIN_TRANSFORMERS in unslothai/notebooks/update_all_notebooks.py. + # `triton` is added because unsloth/_gpu_init.py:232 does an + # unconditional `import triton`; the PyPI wheel installs cleanly + # on Linux x86_64 even without CUDA (same rationale as + # consolidated-tests-ci.yml line 192-205). + # Pillow is listed explicitly as a defensive belt-and-braces + # next to torchvision (vision_utils crashes ModuleNotFoundError + # if torchvision skipped its Pillow dep for any reason). + pip install 'transformers>=4.56,<5.6' 'trl>=0.22,<0.26' 'accelerate>=1.0' \ + 'datasets>=3.4,<5' 'peft>=0.15,<0.20' \ + 'bitsandbytes>=0.43' 'sentencepiece' 'protobuf' triton \ + Pillow safetensors tqdm packaging psutil + # Converter deps (nbformat for notebook_to_python.py). + pip install 'nbformat>=5.10' 'nbconvert>=7.16' + # Install unsloth from the LOCAL checkout (the PR head), not PyPI. + # The PR-time CI must validate the code in this PR; PyPI unsloth + # may lag the in-repo CPU-torch fallback in unsloth/kernels/utils.py + # (lines 162-170) that handles missing torch._C._cuda_getCurrentRawStream. + # unsloth_zoo from git main mirrors every other CI (Core / MLX / + # install.sh) so PR-time validation sees the same zoo HEAD. + for attempt in 1 2 3; do + if pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done + pip install --no-deps -e ./unsloth + + - name: Convert notebooks for AST scan + # Same upstream-conversion-error tolerance as the static job. + continue-on-error: true + run: | + python unsloth/scripts/notebook_validator.py convert \ + --notebooks-dir notebooks --out _converted + + - name: Dump unsloth + trl API surface (under CUDA spoof) + run: | + PYTHONPATH=unsloth/tests python -u - <<'PY' + import sys, json, inspect + import _zoo_aggressive_cuda_spoof as _spoof + _spoof.apply() + import unsloth + import trl + surface = {} + for cls_name in ("FastLanguageModel", "FastVisionModel", "FastModel"): + cls = getattr(unsloth, cls_name, None) + if cls is None: + continue + surface[cls_name] = sorted(n for n in dir(cls) if not n.startswith("_")) + surface["SFTConfig_kwargs"] = sorted(inspect.signature(trl.SFTConfig.__init__).parameters) + json.dump(surface, open("_api_surface.json", "w"), indent=2) + print("dumped surface for:", list(surface)) + PY + + - name: Run API rule against converted notebooks + run: | + python unsloth/scripts/notebook_validator.py api \ + --converted-dir _converted \ + --surface _api_surface.json + + smoke-install: + name: smoke install (Colab-shaped venv, opt-in) + if: ${{ github.event.inputs.include_smoke == 'true' || github.event_name == 'schedule' }} + runs-on: ubuntu-latest + timeout-minutes: 25 + strategy: + fail-fast: false + matrix: + # One representative notebook per installation_*_content template. + # Add rows when a new install template lands in update_all_notebooks.py. + notebook: + - 'nb/Llama3.1_(8B)-Alpaca.ipynb' # installation_content + - 'nb/Gemma3_(4B)-Vision.ipynb' # installation_content + vision + - 'nb/Llama3.1_(8B)-GRPO.ipynb' # installation_extra_grpo_content + - 'nb/gpt-oss-(20B)-Fine-tuning.ipynb' # installation_gpt_oss_content + - 'nb/Qwen3_5_(4B)_Vision.ipynb' # installation_qwen3_5_content + - 'nb/Nemotron-3-Nano-30B-A3B_A100.ipynb' # installation_nemotron_nano_content + - 'nb/Whisper.ipynb' # installation_whisper_content + - 'nb/Synthetic_Data_Hackathon.ipynb' # installation_synthetic_data_content + steps: + - name: Validate client_payload.ref shape + if: github.event_name == 'repository_dispatch' + env: + NOTEBOOKS_REF: ${{ github.event.client_payload.ref }} + run: | + if ! printf '%s' "$NOTEBOOKS_REF" | grep -Eq '^[A-Za-z0-9._/-]+$'; then + echo "::error::client_payload.ref contains disallowed characters" >&2 + exit 1 + fi + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + repository: unslothai/notebooks + ref: ${{ env.NOTEBOOKS_REF }} + path: notebooks + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: { python-version: '3.12' } + + - name: Seed Colab-shaped venv from pip-freeze (CPU-mapped) + run: | + # Strip cu128 local versions, route torch/torchvision to the CPU + # wheel index, drop CUDA-specific deps the runner can't use. + python -u - <<'PY' > /tmp/seed_pins.txt + import json, re + mapping = json.load(open("unsloth/scripts/data/colab_to_cpu_pin.json")) + rewrite = mapping["rewrite"] + skip = set(mapping["skip"]) + spoof = set(mapping["module_spoof"]) + out = [] + for line in open("unsloth/scripts/data/colab_pip_freeze.gpu.txt"): + line = line.strip() + if not line or line.startswith("#"): + continue + m = re.match(r"^([A-Za-z0-9._-]+)\s*==\s*(.+)$", line) + if not m: + continue + name, ver = m.group(1).lower(), m.group(2) + if name in skip: + continue + if name in spoof: + continue + if name in rewrite: + ver = re.sub(r"[+\-].+$", "", ver) + out.append(f"{name}=={ver}") + else: + ver = re.sub(r"[+\-].+$", "", ver) + out.append(f"{name}=={ver}") + print("\n".join(out)) + PY + head -5 /tmp/seed_pins.txt + wc -l /tmp/seed_pins.txt + + - name: Install Colab-shaped venv + run: | + python -m pip install --upgrade pip + # Best-effort: any single line that fails to resolve on CPU is + # tolerated; the smoke contract is "the install cell + the unsloth + # import works", not "the entire Colab venv reproduces." + while IFS= read -r spec; do + pip install "$spec" --index-url https://download.pytorch.org/whl/cpu \ + --extra-index-url https://pypi.org/simple || \ + echo "::warning::pin failed: $spec" + done < /tmp/seed_pins.txt + + - name: Run install cell + run: | + python unsloth/scripts/notebook_validator.py convert \ + --notebooks-dir notebooks --out _converted + # Take the converted .py and run the install cell only. + BASE="$(basename '${{ matrix.notebook }}' .ipynb | tr -d '()' | tr -c '[:alnum:]_' _)" + PY="_converted/${BASE}.py" + [ -f "$PY" ] || { echo "::error::$PY not found"; ls _converted | head; exit 1; } + # Truncate at the first `from unsloth import` so we run install + + # core imports only. + awk '/^from unsloth import/ { print "import sys; sys.exit(0)"; exit } { print }' "$PY" > _smoke.py + PYTHONPATH=unsloth/tests python -u - <<'PY' + import _zoo_aggressive_cuda_spoof as _s; _s.apply() + # Stub torchcodec for cells that import it — no CPU wheel exists. + import sys, types + if "torchcodec" not in sys.modules: + sys.modules["torchcodec"] = types.ModuleType("torchcodec") + exec(open("_smoke.py").read(), {"__name__": "__main__"}) + PY + + - name: Verify imports under spoof + run: | + PYTHONPATH=unsloth/tests python -u - <<'PY' + import sys, types + if "torchcodec" not in sys.modules: + sys.modules["torchcodec"] = types.ModuleType("torchcodec") + import _zoo_aggressive_cuda_spoof as _s; _s.apply() + import unsloth, peft, torch, torchao, transformers, tokenizers + print("OK: imports pass under CUDA spoof") + PY diff --git a/.github/workflows/ossf.yml b/.github/workflows/ossf.yml new file mode 100644 index 0000000..f9a2705 --- /dev/null +++ b/.github/workflows/ossf.yml @@ -0,0 +1,78 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: '21 20 * * 0' + push: + branches: [ "main" ] + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + # `publish_results: true` only works when run from the default branch. conditional can be removed if disabled. + if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request' + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore + # file_mode: git + + # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF + # format to the repository Actions tab. + - name: "Upload artifact" + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: SARIF file + path: results.sarif + retention-days: 5 + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@v3 + with: + sarif_file: results.sarif diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml new file mode 100644 index 0000000..188e078 --- /dev/null +++ b/.github/workflows/release-desktop.yml @@ -0,0 +1,995 @@ +name: Release Desktop App + +on: + workflow_dispatch: + inputs: + studio_version: + description: 'Studio version tag to release (for example, v0.1.39-beta)' + type: string + required: true + pypi_version: + description: 'Exact PyPI unsloth version just published/stamped (for example, 2026.5.3); leave blank to use MIN_DESKTOP_BACKEND_VERSION' + type: string + required: false + draft: + description: 'Create as draft release; draft runs do not advance desktop-latest updater channel' + type: boolean + default: true + +permissions: + contents: read + +concurrency: + group: release-desktop-${{ github.repository }} + cancel-in-progress: false + +jobs: + prepare-version: + name: Prepare release versions + runs-on: ubuntu-latest + outputs: + studio_version: ${{ steps.prepare.outputs.studio_version }} + app_version: ${{ steps.prepare.outputs.app_version }} + desktop_release_tag: ${{ steps.prepare.outputs.desktop_release_tag }} + prerelease: ${{ steps.prepare.outputs.prerelease }} + pypi_version: ${{ steps.prepare.outputs.pypi_version }} + + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd + with: + persist-credentials: false + + - name: Validate release versions + id: prepare + shell: bash + env: + INPUT_STUDIO_VERSION: ${{ inputs.studio_version }} + INPUT_PYPI_VERSION: ${{ inputs.pypi_version }} + run: | + python3 <<'PY' + import os + import pathlib + import re + import sys + + studio_version = os.environ['INPUT_STUDIO_VERSION'].strip() + if not studio_version: + sys.exit('studio_version is required, for example v0.1.39-beta') + if re.fullmatch(r'v?20\d{2}\.\d+\.\d+(?:[-+][0-9A-Za-z.-]+)?', studio_version): + sys.exit(f'studio_version must be a Studio SemVer tag, not a date-style backend version: {studio_version}') + + semver_tag = re.compile( + r'^v(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)' + r'(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$' + ) + if not semver_tag.fullmatch(studio_version): + sys.exit(f'studio_version must be a SemVer tag with leading v, for example v0.1.39-beta: {studio_version}') + + app_version = studio_version.removeprefix('v') + desktop_release_tag = f'desktop-v{app_version}' + prerelease = 'true' if '-' in app_version.split('+', 1)[0] else 'false' + + def parse_backend_version(version): + match = re.fullmatch( + r'(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)' + r'(?:([a-zA-Z]|\.dev|dev|\.rc|rc|\.post|post)(\d*))?' + r'(?:[-+]([0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?', + version, + ) + if not match: + return None + major, minor, patch, suffix_name, suffix_number, suffix_text = match.groups() + if suffix_name: + normalized = suffix_name.lower().lstrip('.') + order = {'dev': 0, 'a': 1, 'b': 2, 'rc': 3, 'post': 5}.get(normalized) + if order is None: + return None + number = int(suffix_number or '0') + elif suffix_text: + order = 3 if version[version.find(suffix_text) - 1] == '-' else 4 + number = 0 + else: + order = 4 + number = 0 + return (int(major), int(minor), int(patch), order, number) + + preflight = pathlib.Path('studio/src-tauri/src/preflight/version.rs').read_text() + match = re.search(r'MIN_DESKTOP_BACKEND_VERSION:\s*&str\s*=\s*"([^"]+)"', preflight) + if not match: + sys.exit('Could not read MIN_DESKTOP_BACKEND_VERSION') + min_backend_version = match.group(1) + + input_pypi_version = os.environ.get('INPUT_PYPI_VERSION', '').strip() + parsed_min_backend = parse_backend_version(min_backend_version) + if parsed_min_backend is None: + sys.exit(f'MIN_DESKTOP_BACKEND_VERSION is not a supported backend package version: {min_backend_version}') + + pypi_version = input_pypi_version or min_backend_version + parsed_pypi = parse_backend_version(pypi_version) + if parsed_pypi is None: + sys.exit(f'pypi_version is not a supported backend package version: {pypi_version}') + if parsed_pypi < parsed_min_backend: + sys.exit( + f'pypi_version {pypi_version} is lower than desktop minimum ' + f'MIN_DESKTOP_BACKEND_VERSION {min_backend_version}' + ) + + if input_pypi_version: + print( + 'Using exact PyPI unsloth version from pypi_version input: ' + f'{pypi_version} (desktop minimum: {min_backend_version})' + ) + else: + print( + 'Using exact PyPI unsloth version from MIN_DESKTOP_BACKEND_VERSION: ' + f'{pypi_version}' + ) + + with open(os.environ['GITHUB_OUTPUT'], 'a', encoding='utf-8') as output: + print(f'studio_version={studio_version}', file=output) + print(f'app_version={app_version}', file=output) + print(f'desktop_release_tag={desktop_release_tag}', file=output) + print(f'prerelease={prerelease}', file=output) + print(f'pypi_version={pypi_version}', file=output) + PY + + - name: Verify PyPI package and Studio stamp + shell: bash + env: + STUDIO_VERSION: ${{ steps.prepare.outputs.studio_version }} + PYPI_VERSION: ${{ steps.prepare.outputs.pypi_version }} + run: | + set -euo pipefail + python3 <<'PY' + import json + import os + import pathlib + import sys + import time + import urllib.error + import urllib.request + + pypi_version = os.environ['PYPI_VERSION'] + dist_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'pypi-unsloth-dist') + dist_dir.mkdir(parents=True, exist_ok=True) + metadata_url = f'https://pypi.org/pypi/unsloth/{pypi_version}/json' + + last_error = None + for attempt in range(1, 6): + try: + with urllib.request.urlopen(metadata_url, timeout=30) as response: + metadata = json.load(response) + break + except Exception as exc: + last_error = exc + if attempt < 5: + time.sleep(10 * attempt) + else: + sys.exit(f'Publish unsloth=={pypi_version} to PyPI before the desktop release ({last_error})') + + files = metadata.get('urls') or [] + if not files: + sys.exit(f'PyPI returned no distribution files for unsloth=={pypi_version}') + + for file_info in files: + filename = file_info.get('filename') + url = file_info.get('url') + if not filename or '/' in filename or not url: + sys.exit(f'Unexpected PyPI file entry for unsloth=={pypi_version}: {file_info!r}') + target = dist_dir / filename + for attempt in range(1, 4): + try: + with urllib.request.urlopen(url, timeout=60) as response: + target.write_bytes(response.read()) + break + except Exception as exc: + last_error = exc + if attempt < 3: + time.sleep(5 * attempt) + else: + sys.exit(f'Could not download {filename} from PyPI ({last_error})') + PY + + if [ -f scripts/stamp_studio_release.py ]; then + mapfile -t dists < <(find "$RUNNER_TEMP/pypi-unsloth-dist" -type f \( -name '*.whl' -o -name '*.tar.gz' \) | sort) + if [ "${#dists[@]}" -eq 0 ]; then + echo "No PyPI wheel/sdist artifacts downloaded for unsloth==$PYPI_VERSION" >&2 + exit 1 + fi + python3 scripts/stamp_studio_release.py --verify-dist "$RUNNER_TEMP/pypi-unsloth-dist" --expected "$STUDIO_VERSION" + else + echo "scripts/stamp_studio_release.py not found; release-desktop requires #5308 to verify the PyPI Studio stamp." >&2 + exit 1 + fi + + - name: Guard public updater channel version + if: ${{ !inputs.draft }} + shell: bash + env: + GH_REPO: ${{ github.repository }} + GH_TOKEN: ${{ github.token }} + APP_VERSION: ${{ steps.prepare.outputs.app_version }} + run: | + set -euo pipefail + mkdir -p "$RUNNER_TEMP/desktop-current" + if ! gh release download desktop-latest --pattern latest.json --dir "$RUNNER_TEMP/desktop-current" --clobber 2>/dev/null; then + echo "No existing desktop-latest latest.json found; allowing first channel publish." + exit 0 + fi + python3 <<'PY' + import json + import os + import pathlib + import re + import sys + + def parse(value: str): + value = value.removeprefix('v') + match = re.fullmatch( + r'(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)' + r'(?:-([0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?' + r'(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?', + value, + ) + if not match: + sys.exit(f'desktop-latest latest.json has invalid version: {value}') + major, minor, patch, prerelease = match.groups() + return (int(major), int(minor), int(patch), prerelease) + + def numeric_tail(identifier: str) -> tuple[str, int] | None: + match = re.fullmatch(r'([A-Za-z-]+)(\d+)', identifier) + if not match: + return None + return (match.group(1).lower(), int(match.group(2))) + + def compare_identifier(left: str, right: str) -> int: + left_num = left.isdigit() + right_num = right.isdigit() + if left_num and right_num: + return (int(left) > int(right)) - (int(left) < int(right)) + if left_num: + return -1 + if right_num: + return 1 + + left_tail = numeric_tail(left) + right_tail = numeric_tail(right) + if left_tail and right_tail and left_tail[0] == right_tail[0]: + return (left_tail[1] > right_tail[1]) - (left_tail[1] < right_tail[1]) + + return (left > right) - (left < right) + + def compare_prerelease(left: str | None, right: str | None) -> int: + if left == right: + return 0 + if left is None: + return 1 + if right is None: + return -1 + left_parts = left.split('.') + right_parts = right.split('.') + for left_part, right_part in zip(left_parts, right_parts): + order = compare_identifier(left_part, right_part) + if order: + return order + return (len(left_parts) > len(right_parts)) - (len(left_parts) < len(right_parts)) + + def compare(left: str, right: str) -> int: + left_major, left_minor, left_patch, left_pre = parse(left) + right_major, right_minor, right_patch, right_pre = parse(right) + left_core = (left_major, left_minor, left_patch) + right_core = (right_major, right_minor, right_patch) + if left_core != right_core: + return (left_core > right_core) - (left_core < right_core) + return compare_prerelease(left_pre, right_pre) + + current_path = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-current', 'latest.json') + current = json.loads(current_path.read_text()).get('version') + next_version = os.environ['APP_VERSION'] + if not isinstance(current, str): + sys.exit('desktop-latest latest.json has missing version') + if compare(next_version, current) < 0: + sys.exit( + f'Refusing to publish {next_version}; desktop-latest currently points at newer version {current}.' + ) + PY + + build: + # TODO: split into a "build (no secrets)" + "publish (secrets)" job pair + # with actions/upload-artifact handoff so the matrix build cannot + # publish a Release on its own. The current matrix runs across + # Linux/macOS/Windows in a single job, so the split needs artefact + # collection across the OS matrix and is out of scope for this + # hardening pass. + permissions: + contents: write # tauri-apps/tauri-action creates / uploads a GitHub Release + strategy: + fail-fast: false + max-parallel: 1 + 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 }} + needs: prepare-version + runs-on: ${{ matrix.platform }} + + env: + FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true + APP_VERSION: ${{ needs.prepare-version.outputs.app_version }} + STUDIO_VERSION: ${{ needs.prepare-version.outputs.studio_version }} + DESKTOP_RELEASE_TAG: ${{ needs.prepare-version.outputs.desktop_release_tag }} + DESKTOP_PRERELEASE: ${{ needs.prepare-version.outputs.prerelease }} + + steps: + # harden-runner in audit mode: surfaces every egress destination in + # the runner log so the allowlist for a future `egress-policy: block` + # promotion can be derived from observed traffic. Audit mode is + # cross-platform (Linux / macOS / Windows runners); blocking mode is + # currently Linux-only, so we deliberately stay in audit until the + # macOS + Windows codesign paths have been observed. + - name: Harden runner (audit) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: audit + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd + with: + persist-credentials: false + + # ── 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 libappindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf + + # ── Node.js ── + - name: Setup Node.js + uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e + with: + node-version: 24 + + - name: Install pinned Tauri CLI + # Lifecycle scripts (esbuild native-binary postinstall, etc.) are + # required for `vite build`. The pre-install lockfile structural + # audit (lockfile_supply_chain_audit.py) is the practical defence + # against the npm postinstall-dropper class -- it fires BEFORE any + # tarball runs, on the injection pattern itself rather than an + # advisory-DB lookup. + run: npm install --save-dev --prefix studio @tauri-apps/cli@2.10.1 --no-fund --no-audit + + - 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: Verify desktop updater and Linux package config + shell: bash + run: | + node <<'JS' + const { readFileSync } = require('node:fs'); + + const expected = 'https://github.com/unslothai/unsloth/releases/download/desktop-latest/latest.json'; + const config = JSON.parse(readFileSync('studio/src-tauri/tauri.conf.json', 'utf8')); + const endpoints = config.plugins?.updater?.endpoints; + if (!Array.isArray(endpoints) || endpoints.length !== 1) { + throw new Error('Expected exactly one desktop updater endpoint'); + } + if (endpoints[0] !== expected) { + throw new Error('Desktop updater endpoint must be ' + expected + ', got ' + endpoints[0]); + } + if (endpoints.some((endpoint) => endpoint.includes('/releases/latest/'))) { + throw new Error('Desktop updater endpoint must not use repo-wide /releases/latest/'); + } + + const targets = config.bundle?.targets; + if (Array.isArray(targets) && targets.some((target) => String(target).toLowerCase() === 'rpm')) { + throw new Error('Desktop release must not target RPM packages'); + } + if (config.bundle?.linux?.rpm) { + throw new Error('bundle.linux.rpm must not be configured'); + } + if (config.bundle?.linux?.appimage?.bundleMediaFramework !== false) { + throw new Error('Linux AppImage bundleMediaFramework must stay false'); + } + + const workflow = readFileSync('.github/workflows/release-desktop.yml', 'utf8'); + const lines = workflow.split(/\r?\n/); + const linuxInstallLines = lines.filter((line) => line.includes('sudo apt-get install')); + const ayatanaPackage = ['libayatana', 'appindicator3-dev'].join('-'); + if (linuxInstallLines.some((line) => line.includes(ayatanaPackage))) { + throw new Error('Desktop Linux release must not install the Ayatana appindicator dev package'); + } + if (!linuxInstallLines.some((line) => line.includes('libappindicator3-dev'))) { + throw new Error('Desktop Linux release must install libappindicator3-dev'); + } + const linuxdeployLines = lines.filter((line) => line.includes('github.com/linuxdeploy/linuxdeploy/releases/download')); + if (!linuxdeployLines.some((line) => line.includes('1-alpha-20250213-2/linuxdeploy-x86_64.AppImage'))) { + throw new Error('Desktop Linux release must pin linuxdeploy 1-alpha-20250213-2'); + } + // A pinned version/path is reproducibility, not integrity: the asset + // can be replaced after upload. Require the immutable SHA-256 digest + // to be pinned AND verified before chmod +x. Scope every check to the + // real "Pin linuxdeploy for AppImage" step so this guard cannot + // satisfy itself; a file-wide scan would match the guard's own code. + const expectedLinuxdeployDigest = '4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a'; + const isComment = (line) => { + const trimmed = line.trim(); + return trimmed.startsWith('#') || trimmed.startsWith('//'); + }; + const stepStart = lines.findIndex((line) => /^\s*- name: Pin linuxdeploy for AppImage\s*$/.test(line)); + if (stepStart === -1) { + throw new Error('Desktop Linux release must keep the "Pin linuxdeploy for AppImage" step'); + } + const stepIndent = lines[stepStart].search(/\S/); + let stepEnd = lines.length; + for (let i = stepStart + 1; i < lines.length; i += 1) { + const line = lines[i]; + if (line.trim() === '') continue; + const indent = line.search(/\S/); + // The next sibling step ('- ...') at the same indent, or any dedent + // below the step, ends this step's block. + if (indent < stepIndent || (indent === stepIndent && /^\s*-\s/.test(line))) { + stepEnd = i; + break; + } + } + const stepLines = lines.slice(stepStart, stepEnd); + const digestEnvRe = /^\s*LINUXDEPLOY_SHA256:\s*["']([0-9a-f]{64})["']\s*$/; + const digestEnvLine = stepLines.find((line) => digestEnvRe.test(line)); + if (!digestEnvLine || digestEnvLine.match(digestEnvRe)[1] !== expectedLinuxdeployDigest) { + throw new Error('Desktop Linux release must pin the linuxdeploy SHA-256 digest in the LINUXDEPLOY_SHA256 env'); + } + const sha256Idx = stepLines.findIndex((line) => !isComment(line) && line.includes('sha256sum -c')); + if (sha256Idx === -1) { + throw new Error('Desktop Linux release must verify the linuxdeploy digest with sha256sum -c before use'); + } + const chmodIdx = stepLines.findIndex((line) => !isComment(line) && /chmod\s+\+x/.test(line)); + if (chmodIdx !== -1 && sha256Idx > chmodIdx) { + throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x'); + } + const releaseBodies = []; + for (let i = 0; i < lines.length; i += 1) { + const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/); + if (!match) continue; + const baseIndent = match[1].length; + const bodyLines = []; + i += 1; + for (; i < lines.length; i += 1) { + const line = lines[i]; + if (line.trim() === '') { + bodyLines.push(''); + continue; + } + const indent = line.match(/^\s*/)[0].length; + if (indent <= baseIndent) { + i -= 1; + break; + } + bodyLines.push(line.slice(baseIndent + 2)); + } + releaseBodies.push(bodyLines.join('\n')); + } + if (releaseBodies.length === 0) { + throw new Error('Expected at least one desktop release body'); + } + for (const body of releaseBodies) { + if (/\brpm\b|\.rpm/i.test(body)) { + throw new Error('Desktop release body must not advertise RPM packages'); + } + if (/AppImage.*universal|universal.*AppImage/i.test(body)) { + throw new Error('Desktop release body must not advertise AppImage as universal'); + } + if (!/AppImage.*experimental/i.test(body)) { + throw new Error('Desktop release body must mark AppImage as experimental'); + } + } + JS + + - name: Install frontend dependencies + working-directory: studio/frontend + # Lifecycle scripts (esbuild native-binary postinstall, etc.) are + # required for `vite build`. The pre-install lockfile structural + # audit (lockfile_supply_chain_audit.py) is the practical defence + # against the npm postinstall-dropper class -- it fires BEFORE any + # tarball runs, on the injection pattern itself rather than an + # advisory-DB lookup. + run: npm install --no-fund --no-audit + + # ── Rust ── + - name: Install Rust stable + uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 + with: + targets: ${{ matrix.platform == 'macos-latest' && 'aarch64-apple-darwin,x86_64-apple-darwin' || '' }} + + - name: Patch desktop app version + shell: bash + working-directory: studio/src-tauri + run: | + set -euo pipefail + if command -v python3 >/dev/null 2>&1; then + PYTHON=python3 + else + PYTHON=python + fi + "$PYTHON" <<'PY' + import os + import pathlib + import re + import sys + + app_version = os.environ['APP_VERSION'] + if not app_version: + sys.exit('APP_VERSION is required') + + cargo_toml = pathlib.Path('Cargo.toml') + lines = cargo_toml.read_text().splitlines(keepends=True) + in_package = False + patched = False + for index, line in enumerate(lines): + stripped = line.strip() + if stripped == '[package]': + in_package = True + continue + if stripped.startswith('[') and stripped.endswith(']'): + in_package = False + if in_package and re.fullmatch(r'version\s*=\s*"[^"]+"\s*', stripped): + lines[index] = f'version = "{app_version}"\n' + patched = True + break + if not patched: + sys.exit('Could not patch [package] version in Cargo.toml') + cargo_toml.write_text(''.join(lines)) + + cargo_lock = pathlib.Path('Cargo.lock') + lock_text = cargo_lock.read_text() + lock_text, count = re.subn( + r'(?m)(^\[\[package\]\]\nname = "unsloth-studio"\nversion = ")[^"]+(")', + lambda match: f'{match.group(1)}{app_version}{match.group(2)}', + lock_text, + ) + if count != 1: + sys.exit(f'Could not patch unsloth-studio version in Cargo.lock (matches={count})') + cargo_lock.write_text(lock_text) + PY + + cargo metadata --locked --no-deps --format-version 1 > "$RUNNER_TEMP/cargo-metadata.json" + "$PYTHON" <<'PY' + import json + import os + import pathlib + import sys + + app_version = os.environ['APP_VERSION'] + metadata = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'cargo-metadata.json').read_text()) + versions = [package['version'] for package in metadata.get('packages', []) if package.get('name') == 'unsloth-studio'] + if versions != [app_version]: + sys.exit(f'cargo metadata unsloth-studio version mismatch: expected {app_version}, got {versions}') + PY + + git diff -- Cargo.toml Cargo.lock + + - name: Rust cache + uses: swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 + 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.10.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: pin AppImage packaging toolchain ── + - name: Pin linuxdeploy for AppImage + if: matrix.platform == 'ubuntu-22.04' + shell: bash + env: + # Pinning the versioned release path is reproducibility, not + # integrity: a GitHub release asset can be replaced (or its delivery + # path compromised) after upload. The SHA-256 below is the immutable + # digest of this exact asset and is the integrity gate. If linuxdeploy + # publishes a new build under this tag, this run fails closed and the + # digest must be re-pinned deliberately. + LINUXDEPLOY_URL: "https://github.com/linuxdeploy/linuxdeploy/releases/download/1-alpha-20250213-2/linuxdeploy-x86_64.AppImage" + LINUXDEPLOY_SHA256: "4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a" + run: | + set -euo pipefail + tools_dir="$RUNNER_TEMP/tauri-tools-cache/tauri" + mkdir -p "$tools_dir" + dest="$tools_dir/linuxdeploy-x86_64.AppImage" + curl -fsSL "$LINUXDEPLOY_URL" -o "$dest" + # Verify the digest BEFORE the binary is ever marked executable. The + # next step builds the AppImage with the Tauri signing key and a + # contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy + # that ran here could exfiltrate signing material or tamper with + # published release artifacts. Fail closed on any mismatch. + echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c - + chmod +x "$dest" + + # ── 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 }} + XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache + with: + projectPath: studio + tauriScript: npx --prefix . tauri + tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }} + releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' + releaseBody: | + Desktop app for Unsloth Studio. + + **macOS**: Download the Apple Silicon `.dmg`. + **Windows**: Download the `-setup.exe` installer. + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. + + > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. + > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. + releaseDraft: ${{ inputs.draft }} + prerelease: ${{ needs.prepare-version.outputs.prerelease }} + args: -v ${{ matrix.args }} + + # ── macOS: build + sign + notarize + upload ── + - 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: ${{ needs.prepare-version.outputs.desktop_release_tag }} + releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' + releaseBody: | + Desktop app for Unsloth Studio. + + **macOS**: Download the Apple Silicon `.dmg`. + **Windows**: Download the `-setup.exe` installer. + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. + + > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. + > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. + releaseDraft: ${{ inputs.draft }} + prerelease: ${{ needs.prepare-version.outputs.prerelease }} + args: -v ${{ matrix.args }} + + # ── Windows: build + sign + upload ── + - 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: ${{ needs.prepare-version.outputs.desktop_release_tag }} + releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' + releaseBody: | + Desktop app for Unsloth Studio. + + **macOS**: Download the Apple Silicon `.dmg`. + **Windows**: Download the `-setup.exe` installer. + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. + + > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. + > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. + releaseDraft: ${{ inputs.draft }} + prerelease: ${{ needs.prepare-version.outputs.prerelease }} + args: -v ${{ matrix.args }} + + # Release process note: only non-draft workflow runs advance the public + # desktop-latest updater channel. Draft builds are for private review; if a + # draft is manually published later, this channel intentionally remains + # unchanged until a narrow manual channel-publish flow is added or a public + # desktop release is created by running this workflow with draft=false. + publish-updater-channel: + name: Publish desktop updater channel + needs: [prepare-version, build] + if: ${{ !inputs.draft }} + runs-on: ubuntu-latest + permissions: + contents: write + env: + GH_REPO: ${{ github.repository }} + APP_VERSION: ${{ needs.prepare-version.outputs.app_version }} + STUDIO_VERSION: ${{ needs.prepare-version.outputs.studio_version }} + DESKTOP_RELEASE_TAG: ${{ needs.prepare-version.outputs.desktop_release_tag }} + DESKTOP_PRERELEASE: ${{ needs.prepare-version.outputs.prerelease }} + + steps: + - name: Download versioned updater metadata + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + mkdir -p "$RUNNER_TEMP/desktop-updater" + gh api "repos/${GITHUB_REPOSITORY}/releases/tags/${DESKTOP_RELEASE_TAG}" > "$RUNNER_TEMP/source-release.json" + python3 <<'PY' + import json + import os + import pathlib + import sys + + source = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'source-release.json').read_text()) + expected_tag = os.environ['DESKTOP_RELEASE_TAG'] + if source.get('tag_name') != expected_tag: + sys.exit(f'Expected source release {expected_tag}, got {source.get("tag_name")}') + if source.get('draft'): + sys.exit(f'Source desktop release {expected_tag} is draft; refusing to publish public updater channel') + PY + gh release download "$DESKTOP_RELEASE_TAG" --pattern latest.json --dir "$RUNNER_TEMP/desktop-updater" --clobber + test -s "$RUNNER_TEMP/desktop-updater/latest.json" + + - name: Validate versioned updater metadata + shell: bash + run: | + python3 <<'PY' + import json + import os + import pathlib + import re + import sys + + app_version = os.environ['APP_VERSION'] + release_tag = os.environ['DESKTOP_RELEASE_TAG'] + latest_path = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-updater', 'latest.json') + data = json.loads(latest_path.read_text()) + if not isinstance(data, dict): + sys.exit('latest.json must be a JSON object') + + version = data.get('version') + if not isinstance(version, str) or not version: + sys.exit('latest.json missing version') + if not re.fullmatch(r'v?\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.-]+)?', version): + sys.exit(f'latest.json version is not SemVer-like: {version}') + if version.removeprefix('v') != app_version: + sys.exit(f'latest.json version {version} does not match desktop app version {app_version}') + + platforms = data.get('platforms') + if not isinstance(platforms, dict) or not platforms: + sys.exit('latest.json missing platforms') + + required_families = { + 'darwin-aarch64': False, + 'linux-x86_64': False, + 'windows-x86_64': False, + } + expected_prefix = f'https://github.com/unslothai/unsloth/releases/download/{release_tag}/' + forbidden_fragments = ('/releases/latest/', '/releases/download/desktop-latest/') + + for platform, entry in platforms.items(): + if not isinstance(entry, dict): + sys.exit(f'Platform {platform} must be an object') + url = entry.get('url') + signature = entry.get('signature') + if not isinstance(url, str) or not url.strip(): + sys.exit(f'Platform {platform} missing url') + if not isinstance(signature, str) or not signature.strip(): + sys.exit(f'Platform {platform} missing signature') + if any(fragment in url for fragment in forbidden_fragments): + sys.exit(f'Platform {platform} points at a moving updater channel: {url}') + if not url.startswith(expected_prefix): + sys.exit(f'Platform {platform} URL must point at {release_tag}: {url}') + for family in required_families: + if platform == family or platform.startswith(family + '-'): + required_families[family] = True + + missing = [family for family, found in required_families.items() if not found] + if missing: + sys.exit('latest.json missing required platform families: ' + ', '.join(missing)) + PY + + - name: Ensure desktop updater channel release + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + channel_json="$RUNNER_TEMP/desktop-latest-release.json" + if ! gh api "repos/${GITHUB_REPOSITORY}/releases/tags/desktop-latest" > "$channel_json" 2>/dev/null; then + gh release create desktop-latest \ + --title "Unsloth Studio Desktop updater channel" \ + --notes "Machine-managed desktop updater channel; latest.json is replaced by release-desktop.yml." \ + --prerelease \ + --latest=false \ + --target "$GITHUB_SHA" + gh api "repos/${GITHUB_REPOSITORY}/releases/tags/desktop-latest" > "$channel_json" + fi + + python3 <<'PY' + import json + import os + import pathlib + import sys + + channel = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-latest-release.json').read_text()) + if channel.get('draft'): + sys.exit('desktop-latest release is draft; refusing to publish updater channel') + if channel.get('immutable'): + sys.exit('desktop-latest release is immutable; cannot replace latest.json') + if not channel.get('prerelease'): + sys.exit('desktop-latest release must be a prerelease so it cannot compete with repo-wide latest') + PY + + - name: Prevent updater channel downgrade + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + mkdir -p "$RUNNER_TEMP/desktop-current" + if ! gh release download desktop-latest --pattern latest.json --dir "$RUNNER_TEMP/desktop-current" --clobber 2>/dev/null; then + echo "No existing desktop-latest latest.json found; allowing first channel publish." + exit 0 + fi + python3 <<'PY' + import json + import os + import pathlib + import re + import sys + + def parse(value: str): + value = value.removeprefix('v') + match = re.fullmatch( + r'(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)' + r'(?:-([0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?' + r'(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?', + value, + ) + if not match: + sys.exit(f'desktop-latest latest.json has invalid version: {value}') + major, minor, patch, prerelease = match.groups() + return (int(major), int(minor), int(patch), prerelease) + + def numeric_tail(identifier: str) -> tuple[str, int] | None: + match = re.fullmatch(r'([A-Za-z-]+)(\d+)', identifier) + if not match: + return None + return (match.group(1).lower(), int(match.group(2))) + + def compare_identifier(left: str, right: str) -> int: + left_num = left.isdigit() + right_num = right.isdigit() + if left_num and right_num: + return (int(left) > int(right)) - (int(left) < int(right)) + if left_num: + return -1 + if right_num: + return 1 + + left_tail = numeric_tail(left) + right_tail = numeric_tail(right) + if left_tail and right_tail and left_tail[0] == right_tail[0]: + return (left_tail[1] > right_tail[1]) - (left_tail[1] < right_tail[1]) + + return (left > right) - (left < right) + + def compare_prerelease(left: str | None, right: str | None) -> int: + if left == right: + return 0 + if left is None: + return 1 + if right is None: + return -1 + left_parts = left.split('.') + right_parts = right.split('.') + for left_part, right_part in zip(left_parts, right_parts): + order = compare_identifier(left_part, right_part) + if order: + return order + return (len(left_parts) > len(right_parts)) - (len(left_parts) < len(right_parts)) + + def compare(left: str, right: str) -> int: + left_major, left_minor, left_patch, left_pre = parse(left) + right_major, right_minor, right_patch, right_pre = parse(right) + left_core = (left_major, left_minor, left_patch) + right_core = (right_major, right_minor, right_patch) + if left_core != right_core: + return (left_core > right_core) - (left_core < right_core) + return compare_prerelease(left_pre, right_pre) + + current_path = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-current', 'latest.json') + next_path = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-updater', 'latest.json') + current = json.loads(current_path.read_text()).get('version') + next_version = json.loads(next_path.read_text()).get('version') + if not isinstance(current, str) or not isinstance(next_version, str): + sys.exit('Could not compare desktop-latest channel versions') + if compare(next_version, current) < 0: + sys.exit( + f'Refusing to move desktop-latest from {current} to older version {next_version}.' + ) + PY + + - name: Publish desktop updater channel metadata + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + gh release upload desktop-latest "$RUNNER_TEMP/desktop-updater/latest.json" --clobber + gh api "repos/${GITHUB_REPOSITORY}/releases/tags/desktop-latest" > "$RUNNER_TEMP/desktop-latest-release.json" + python3 <<'PY' + import json + import os + import pathlib + import sys + + channel = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-latest-release.json').read_text()) + assets = [asset for asset in channel.get('assets', []) if asset.get('name') == 'latest.json'] + if len(assets) != 1: + sys.exit(f'Expected exactly one desktop-latest latest.json asset, found {len(assets)}') + expected_url = f'https://github.com/{os.environ["GITHUB_REPOSITORY"]}/releases/download/desktop-latest/latest.json' + actual_url = assets[0].get('browser_download_url') + if actual_url != expected_url: + sys.exit(f'desktop-latest latest.json URL mismatch: expected {expected_url}, got {actual_url}') + PY diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml new file mode 100644 index 0000000..1275d12 --- /dev/null +++ b/.github/workflows/security-audit.yml @@ -0,0 +1,1246 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Multi-language supply-chain audit. Triggers: +# - PRs touching any dependency manifest (Python / npm / Cargo), a +# scanner or its allowlist baseline, or this workflow file, +# - push to main / pip, +# - nightly @ 04:13 UTC so newly-published advisories surface even +# when no PR opens, +# - workflow_dispatch for ad-hoc invocations. +# +# Two jobs: +# - advisory-audit: one runner that runs pip-audit + npm audit + +# cargo audit back-to-back. All three are +# advisory-DB lookups -- fast, lockfile-driven, +# no archive download. Setting up the python / +# node / rust toolchains on one runner and +# running the three commands serially is +# cheaper than spinning up three runners. +# - pip-scan-packages: 3-shard matrix that downloads + pattern-scans +# every PyPI archive in the transitive closure. +# This is the expensive job (~6 min/shard, +# running in parallel) and it must stay +# independent so a CVE-DB hit in advisory-audit +# does not block the supply-chain pattern scan +# (or vice versa). +# +# All steps are non-blocking initially. The default branch already +# carries a known-vuln backlog (the dependabot banner shows 17 today, +# pip-audit catches 2 more, npm/cargo will catch their own); a hard +# gate now would block every PR on a baseline we have not triaged. +# As each baseline closes, drop continue-on-error per step. +# +# Dependency coverage: +# - unsloth core (pyproject.toml [project.dependencies]) +# - unsloth `huggingfacenotorch` extras (the canonical install path +# for fine-tuning users; pulls transformers / peft / accelerate / +# trl / datasets / diffusers / sentence-transformers / etc.) +# - all six Studio backend requirements files +# - Studio frontend (npm) and Tauri shell (cargo) +# Each Python step builds a filtered dep list from pyproject.toml + +# requirements/*.txt before auditing. We do NOT install any of these +# -- pip-audit resolves through PyPI metadata, scan_packages.py +# downloads sdist/wheel archives and inspects them without running +# install hooks, so an attacker who has compromised a transitive dep +# cannot execute code in this workflow. + +name: Security audit + +on: + pull_request: + paths: + - 'studio/backend/requirements/**' + - 'studio/frontend/package.json' + - 'studio/frontend/package-lock.json' + - 'studio/src-tauri/Cargo.toml' + - 'studio/src-tauri/Cargo.lock' + - 'pyproject.toml' + - 'scripts/scan_packages.py' + - 'scripts/scan_packages_baseline.json' + - 'scripts/scan_npm_packages.py' + - 'scripts/scan_npm_packages_baseline.json' + - '.github/workflows/security-audit.yml' + push: + branches: [main, pip] + schedule: + - cron: '13 4 * * *' # 04:13 UTC daily, off the cron rush + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +# ────────────────────────────────────────────────────────────────────── +# Network-resilience knobs, applied to every job/step. These add retries +# and backoff ONLY; they do not relax a single integrity check. cargo +# still resolves against Cargo.lock (--locked), pip still verifies the +# wheels it downloads, npm still enforces package-lock integrity, the +# harden-runner egress allowlists below are unchanged, and every action +# stays SHA-pinned. The advisory-audit run on 2026-05-29 red-failed when +# one crates.io tarball fetch hit "Recv failure: Connection reset by +# peer" (curl 56); cargo's default of 3 retries over an HTTP/2-multiplexed +# connection did not recover. The settings below make that class of +# transient fault self-heal instead of failing the whole run. +env: + # pip: raise the built-in retry count and per-connection timeout. + PIP_RETRIES: "10" + PIP_DEFAULT_TIMEOUT: "60" + # cargo: retry network ops and disable HTTP/2 multiplexing -- the + # documented mitigation for the curl-56 connection resets above. + CARGO_NET_RETRY: "10" + CARGO_HTTP_MULTIPLEXING: "false" + CARGO_NET_GIT_FETCH_WITH_CLI: "true" + # npm: retry registry fetches with capped exponential backoff. + NPM_CONFIG_FETCH_RETRIES: "5" + NPM_CONFIG_FETCH_RETRY_MINTIMEOUT: "2000" + NPM_CONFIG_FETCH_RETRY_MAXTIMEOUT: "60000" + +jobs: + # ───────────────────────────────────────────────────────────────────── + # Combined advisory-DB audit: pip-audit + npm audit + cargo audit + # all on one runner. Each step is continue-on-error so a finding in + # one toolchain does not suppress the others. + # ───────────────────────────────────────────────────────────────────── + advisory-audit: + name: advisory audit (pip + npm + cargo) + runs-on: ubuntu-latest + timeout-minutes: 25 + steps: + # step-security/harden-runner installs an eBPF-based egress + # firewall on the runner. In `audit` mode it logs every outbound + # connection without blocking; in `block` mode it rejects + # anything outside `allowed-endpoints`. We run audit-only + # initially: the next time this job hits a real PyPI advisory or + # an attacker-funded archive in pip-scan-packages, the audit log + # tells us exactly which hosts were dialed and we promote the + # allowlist to block. Would have *contained* the litellm exfil + # even if scan_packages had missed the .pth payload. + # SHA-pinned (not @v2): the litellm 1.82.7 attack chain hijacked + # mutable tags on aquasecurity/trivy-action and would have hit + # anyone using @v0 / @v2 / @latest references. Pinning to a 40- + # char SHA freezes this action at known-good code; Dependabot's + # github-actions ecosystem will auto-bump the SHA. + # v2.19.1 commit: + # Per-job allowlist: advisory-audit hits PyPI, npm registry, + # crates.io advisories, GitHub release artefacts (osv-scanner + # binary), Semgrep registry, and TruffleHog's own GitHub action. + - name: Harden runner (egress block) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: block + disable-sudo: true + allowed-endpoints: > + api.github.com:443 + github.com:443 + codeload.github.com:443 + objects.githubusercontent.com:443 + raw.githubusercontent.com:443 + release-assets.githubusercontent.com:443 + registry.npmjs.org:443 + pypi.org:443 + files.pythonhosted.org:443 + static.rust-lang.org:443 + index.crates.io:443 + static.crates.io:443 + crates.io:443 + semgrep.dev:443 + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + # Full history so TruffleHog can diff base..head; without + # this it sees only the latest commit and reports nothing. + fetch-depth: 0 + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 + + - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + with: + workspaces: studio/src-tauri -> target + + - name: Install pip-audit + cargo-audit + # cargo-audit pulls advisories from the RustSec advisory-db on + # first run and caches them under ~/.cargo/advisory-db. Pin + # --locked so the version we install matches Cargo.lock + # determinism. cargo-audit 0.22 supports the CVSS 4.0 schema + # used in 2026 advisories (e.g. RUSTSEC-2026-0073); 0.21 + # crashes with a TOML parse error on that file. + # npm audit is bundled with the node toolchain, no install. + run: | + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + retry 5 python -m pip install --upgrade pip 'pip-audit>=2.7' + # --locked keeps the resolved tree identical to Cargo.lock; the + # CARGO_NET_* env above plus this outer loop survive transient + # crates.io connection resets without weakening that guarantee. + retry 5 cargo install --locked --version '^0.22' cargo-audit + + # ───────────────────────────────────────────────────────────── + # Python: pip-audit + # ───────────────────────────────────────────────────────────── + - name: Build filtered Python requirements set + # Two transforms: + # (1) Generate audit-reqs/unsloth-deps.txt from pyproject.toml + # so pip-audit sees the unsloth pip package's own dep set + # (core + huggingfacenotorch extras: transformers / peft / + # accelerate / trl / datasets / diffusers / + # sentence-transformers / huggingface_hub / hf_transfer / + # etc.). + # (2) Copy each studio/backend/requirements/*.txt into + # audit-reqs/ with `git+` lines stripped. pip-audit's `-r` + # mode does a dry-run resolve against PyPI metadata; a + # `git+https://...` spec forces it to clone, which is + # both slow and outside the threat model (we audit + # PyPI-served archives; a git ref is whatever HEAD says + # on the runner). A comment line is left in place so the + # skipped specs are obvious in the artifact. + # The `huggingface` extra is `huggingfacenotorch` plus torch / + # torchvision / triton, deliberately skipped: Studio backend + # already pins a torch and the +cu* / +cpu local-version tags + # trip up the PyPI resolver in `-r` mode. + run: | + mkdir -p audit-reqs + python <<'PY' > audit-reqs/unsloth-deps.txt + import tomllib + with open("pyproject.toml", "rb") as f: + d = tomllib.load(f) + core = d["project"]["dependencies"] + extras = d["project"]["optional-dependencies"]["huggingfacenotorch"] + print("# Auto-generated from pyproject.toml by security-audit.yml.") + print("# core deps + huggingfacenotorch extras.") + for spec in core + extras: + print(spec) + PY + for f in studio.txt extras.txt extras-no-deps.txt \ + no-torch-runtime.txt overrides.txt triton-kernels.txt; do + python < "audit-reqs/$f" + src = "studio/backend/requirements/$f" + with open(src) as fh: + for line in fh: + stripped = line.strip() + before_comment = stripped.split("#", 1)[0] + if "git+" in before_comment: + print(f"# [security-audit] skipped git+ spec: {stripped}") + continue + print(line.rstrip("\n")) + PY + done + + - name: pip-audit (declared Python deps, no install) + # `-r requirements.txt` resolves the requirements through pip's + # dependency resolver against PyPI metadata and audits the + # resolved tree without ever executing setup.py / install + # hooks. Way faster than installing the full Studio runtime + # and -- critically -- safer: an attacker who has compromised + # a transitive dep cannot run code in this job. + # + # extras.txt + extras-no-deps.txt have legacy setup.py + # packages (notably openai-whisper) whose setup.py imports + # `pkg_resources`, which the isolated build env's current + # setuptools no longer ships. PIP_CONSTRAINT pins an older + # setuptools into the build env so those builds resolve. + # Per-file loop so one bad file doesn't take out the whole + # audit. + continue-on-error: true + env: + PIP_CONSTRAINT: ${{ github.workspace }}/audit-reqs/build-constraints.txt + run: | + set +e + cat > audit-reqs/build-constraints.txt <<'CONSTRAINTS' + setuptools<78 + wheel + CONSTRAINTS + : > logs-pip-audit.txt + for f in unsloth-deps studio extras extras-no-deps \ + no-torch-runtime overrides triton-kernels; do + if ! grep -qE '^[^#[:space:]]' "audit-reqs/$f.txt"; then + echo "[security-audit] $f.txt has no PyPI specs after git+ filter, skipping" \ + | tee -a logs-pip-audit.txt + continue + fi + echo "::group::pip-audit -r audit-reqs/$f.txt" + { + echo + echo "=== $f ===" + pip-audit -r "audit-reqs/$f.txt" --format=columns + echo "=== end $f (rc=$?) ===" + } 2>&1 | tee -a logs-pip-audit.txt + echo "::endgroup::" + done + { + echo "## pip-audit (Python)" + echo + echo '### Coverage' + echo '- unsloth core + `huggingfacenotorch` extras (pyproject.toml)' + echo '- studio/backend/requirements/{studio,extras,extras-no-deps,no-torch-runtime,overrides,triton-kernels}.txt' + echo '- `git+` specs are stripped before audit (out of scope: we audit PyPI archives)' + echo + echo '### Findings' + echo '```' + cat logs-pip-audit.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # Pre-install lockfile supply-chain audit (npm + cargo). + # Catches structural anomalies (non-registry resolved URLs, + # missing integrity hashes, known IOC strings) BEFORE `npm + # audit` or OSV-Scanner consult the advisory DB. The advisory + # path is reactive -- there is a window between a malicious + # publication and the GHSA landing. This step fires on the + # injection pattern itself so it catches the same class of + # attack the moment the lockfile shape becomes wrong. + # ───────────────────────────────────────────────────────────── + - name: Lockfile supply-chain audit (pre-install scan) + run: | + python3 scripts/lockfile_supply_chain_audit.py + { + echo "## Lockfile supply-chain audit" + echo + echo "Scanned: studio/frontend/package-lock.json + studio/src-tauri/Cargo.lock" + echo + echo "No structural anomalies or known IOC strings." + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # npm: Studio frontend + # ───────────────────────────────────────────────────────────── + - name: npm audit (Studio frontend) + # `npm audit` resolves the lockfile through the npmjs.com + # advisory DB. `--audit-level=high` filters the noise floor + # to only HIGH and CRITICAL. We do NOT pass --omit=dev: a + # malicious dev-only dep can still steal secrets from a CI + # runner, so dev deps need to be in the audit surface. + continue-on-error: true + working-directory: studio/frontend + run: | + set +e + npm audit --audit-level=high | tee ../../logs-npm-audit.txt + # Always also write the full JSON for grep-ability. + npm audit --json > ../../logs-npm-audit.json || true + { + echo "## npm audit (Studio frontend)" + echo + echo '```' + tail -200 ../../logs-npm-audit.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # cargo: Studio Tauri shell + # ───────────────────────────────────────────────────────────── + - name: cargo audit (Studio Tauri) + # `--deny warnings` would make the job fail on any advisory. + # Keep non-blocking initially; drop continue-on-error after + # the baseline closes. + continue-on-error: true + working-directory: studio/src-tauri + run: | + set +e + cargo audit | tee ../../logs-cargo-audit.txt + { + echo "## cargo audit (Studio Tauri)" + echo + echo '```' + tail -200 ../../logs-cargo-audit.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # OSV-Scanner: cross-ecosystem advisory DB (PyPI + npm + cargo) + # ───────────────────────────────────────────────────────────── + - name: Download + verify OSV-Scanner + # Split out from the scan below so binary integrity is a HARD gate: + # a checksum mismatch (swapped release asset, the Trivy-style pivot + # this workflow refuses) fails the job instead of being swallowed by + # the scan step's continue-on-error. A download still failing after + # retries is transient, so we skip the scan rather than red-fail. + # SHA-256 verified BEFORE chmod +x / exec. Bump OSV_SHA256 in lockstep + # with OSV_VERSION (value from the release's osv-scanner_SHA256SUMS). + run: | + set -euo pipefail + OSV_VERSION="v2.0.2" + OSV_SHA256="3abcfd7126c453a00421487e721b296e0cb68085bd431d6cef60872774170fc8" + if ! curl --proto '=https' --tlsv1.2 -fsSL \ + --retry 5 --retry-delay 3 --retry-connrefused --retry-all-errors \ + -o /tmp/osv-scanner \ + "https://github.com/google/osv-scanner/releases/download/${OSV_VERSION}/osv-scanner_linux_amd64"; then + echo "::warning::osv-scanner download failed after retries; skipping scan" >&2 + rm -f /tmp/osv-scanner + exit 0 # transient availability: do not red-fail the job + fi + if ! echo "${OSV_SHA256} /tmp/osv-scanner" | sha256sum -c -; then + echo "::error::osv-scanner checksum mismatch; refusing to execute" >&2 + rm -f /tmp/osv-scanner + exit 1 # integrity failure: hard-fail + fi + chmod +x /tmp/osv-scanner + /tmp/osv-scanner --version + + - name: OSV-Scanner (PyPI + npm + cargo, cross-ecosystem advisories) + # OSV's advisory feed is a superset of GitHub-Advisory + RustSec + # + npm advisories; running it alongside the per-ecosystem audit + # tools catches CVEs that haven't propagated to the per-ecosystem + # DBs yet (e.g. langchain-core CVE-2025-68664 was on OSV before + # GitHub Advisory). Single binary, one transitive resolver, all + # three lockfile types in one pass. Binary is checksum-verified in + # the step above; only the advisory scan stays non-blocking until + # baselines close. + continue-on-error: true + run: | + set +e + if [ ! -x /tmp/osv-scanner ]; then + echo "osv-scanner unavailable this run; skipping scan" | tee logs-osv-scanner.txt + else + /tmp/osv-scanner scan source \ + --lockfile=studio/frontend/package-lock.json \ + --lockfile=studio/src-tauri/Cargo.lock \ + --lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \ + --lockfile=requirements.txt:audit-reqs/studio.txt \ + --lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \ + --lockfile=requirements.txt:audit-reqs/overrides.txt \ + --lockfile=requirements.txt:audit-reqs/extras.txt \ + --lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \ + --format=table 2>&1 | tee logs-osv-scanner.txt + fi + { + echo "## OSV-Scanner (cross-ecosystem)" + echo + echo '```' + tail -200 logs-osv-scanner.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # Semgrep: design-flaw detection (catches what regex-pattern + # scanning of malicious authors cannot, e.g. first-party logic bugs + # like langchain-core CVE-2025-68664 dumps/dumpd injection, + # n8n CVE-2025-68668 _pyodide.eval_code sandbox escape, marimo + # CVE-2026-39987 unauth WebSocket). + # ───────────────────────────────────────────────────────────── + - name: Semgrep (supply-chain + python rule packs) + continue-on-error: true + run: | + set +e + python -m pip install --quiet 'semgrep>=1.95' + semgrep --version + semgrep scan \ + --config p/supply-chain \ + --config p/python \ + --config p/javascript \ + --config p/security-audit \ + --severity ERROR --severity WARNING \ + --metrics off \ + --timeout 120 \ + studio/backend unsloth scripts \ + 2>&1 | tee logs-semgrep.txt + { + echo "## Semgrep (supply-chain + python + javascript rules)" + echo + echo '```' + tail -200 logs-semgrep.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # Lockfile pin verifier. The litellm 1.82.7 attack window was + # ~40 minutes; anyone resolving with `>=` got the malicious + # version automatically. Flag every spec in the requirements + # files that does not pin to an exact `==` (or `@` for git + # refs, or `===` for arbitrary equality). Warning-only for now; + # graduate to blocking once the baseline is clean. + # ───────────────────────────────────────────────────────────── + - name: Lockfile pin verifier (Python requirements) + continue-on-error: true + run: | + python <<'PY' | tee logs-pin-verifier.txt + import re + from pathlib import Path + + # Specs that look like `pkg==1.2.3` or `pkg @ git+...` or + # bare comments / -r lines are pinned-or-not-applicable. + PINNED = re.compile(r"^\s*[A-Za-z0-9_.\-]+\s*(?:===|==)\s*[^,;]+\s*$") + GIT_OR_URL = re.compile(r"^\s*[A-Za-z0-9_.\-]+\s*@\s*(?:git\+|https?://)") + + unpinned = [] + for f in sorted(Path("studio/backend/requirements").glob("*.txt")): + for i, raw in enumerate(f.read_text().splitlines(), 1): + line = raw.strip() + if not line or line.startswith("#") or line.startswith("-"): + continue + spec = line.split("#", 1)[0].strip().split(";", 1)[0].strip() + if not spec: + continue + if "git+" in spec or PINNED.match(spec) or GIT_OR_URL.match(spec): + continue + unpinned.append((str(f), i, line)) + + print(f"::group::Lockfile pin status") + if unpinned: + print(f"WARN: {len(unpinned)} non-`==` specs across requirements/*.txt") + print("(litellm 1.82.7 wave hit anyone on `>=`; tighten when feasible.)") + for f, i, line in unpinned[:80]: + print(f" {f}:{i}: {line}") + if len(unpinned) > 80: + print(f" ... and {len(unpinned) - 80} more") + else: + print("OK: every spec is exact-pinned.") + print("::endgroup::") + PY + { + echo "## Lockfile pin verifier" + echo + echo '```' + cat logs-pin-verifier.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # Trivy is deliberately NOT installed here. Trivy was the entry + # point for the litellm 1.82.7 supply-chain compromise (March + # 2026): attackers force-rewrote 76 of 77 tags in + # aquasecurity/trivy-action to point at malicious commits; + # anyone running the action with a tag ref auto-pulled a + # credential-harvesting payload. By design a security scanner + # has broad read access to runner secrets, which is exactly + # what made it the ideal pivot. We pick up Trivy's CVE coverage + # from OSV-Scanner (NVD + GHSA + GitLab) and its secret + # detection from TruffleHog. IaC misconfig detection (Trivy's + # one unique value-add) is unfilled for now -- revisit with + # checkov / kics when we ship a Dockerfile or k8s manifests. + # See https://docs.litellm.ai/blog/security-update-march-2026 + # and the Microsoft / Trend Micro / Snyk incident write-ups. + # ───────────────────────────────────────────────────────────── + + # ───────────────────────────────────────────────────────────── + # TruffleHog secret-leak scan on the PR diff. Catches API keys + # / tokens / cred files committed accidentally. --only-verified + # filters out probabilistic findings, so we only flag tokens + # that the source provider confirmed are live. On push to main + # / pip we scan the full repo; on PR we scan base..head. + # SHA-pinned for the same reason as harden-runner above. + # v3.95.2 commit: + # ───────────────────────────────────────────────────────────── + - name: TruffleHog (secrets in diff) + continue-on-error: true + uses: trufflesecurity/trufflehog@37b77001d0174ebec2fcca2bd83ff83a6d45a3ab # v3.95.3 + with: + path: ./ + base: ${{ github.event.pull_request.base.sha || '' }} + head: ${{ github.event.pull_request.head.sha || github.sha }} + # The action passes --no-update internally; passing it here + # too triggers `flag 'no-update' cannot be repeated`. Stick + # with --only-verified so we only flag tokens the source + # provider confirmed are live (no probabilistic findings). + extra_args: --only-verified + + # ───────────────────────────────────────────────────────────── + # CycloneDX SBOM. Lets downstream consumers audit what's + # actually shipped in unsloth wheels and the Studio backend + # runtime. Generates one JSON file per requirements input plus + # a combined SBOM keyed off pyproject.toml; uploads as a build + # artifact (and a future step can attest it via SLSA). + # ───────────────────────────────────────────────────────────── + - name: Generate CycloneDX SBOM + continue-on-error: true + run: | + set +e + python -m pip install --quiet 'cyclonedx-bom>=4.6' + mkdir -p sbom + # Per-requirements-file SBOM (the audit-reqs/ files are the + # filtered, git+-stripped views built earlier in this job). + # cyclonedx-py 4.x uses `--sv` for spec version and `-o` for + # the output file; the older `--schema-version`/`--outfile` + # spellings are not accepted. + for f in audit-reqs/*.txt; do + base=$(basename "$f" .txt) + if grep -qE '^[^#[:space:]]' "$f"; then + cyclonedx-py requirements "$f" \ + --sv 1.6 \ + --of JSON \ + -o "sbom/sbom-$base.json" 2>&1 | tail -5 || true + fi + done + # Project-level SBOM from pyproject.toml. + cyclonedx-py environment \ + --sv 1.6 \ + --of JSON \ + -o sbom/sbom-environment.json 2>&1 | tail -5 || true + ls -la sbom/ + { + echo "## CycloneDX SBOM" + echo + echo "Generated SBOM files:" + ls sbom/ | sed 's/^/- sbom\//' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # GitHub Actions pinning verifier. tj-actions/changed-files + # was compromised in March 2025; anyone using `@v4` (a mutable + # ref) auto-shipped the malicious version. Catch every + # non-SHA-pinned `uses:` across the workflows tree. Warn-only + # initially so the existing baseline doesn't block PRs. + # ───────────────────────────────────────────────────────────── + - name: GitHub Actions pinning verifier + continue-on-error: true + run: | + python <<'PY' | tee logs-actions-pinning.txt + import re + from pathlib import Path + # SHA pin = 40 hex chars after @ + SHA_PIN = re.compile(r"@[0-9a-f]{40}\b") + # First-party / GitHub-published actions get a softer pass + # (still recommended to pin; not a security gate). + FIRST_PARTY = re.compile(r"^\s*-\s*uses:\s*(actions|github)/[^@]+@") + USES = re.compile(r"^\s*-\s*uses:\s*([^@\s]+)@(\S+)") + unpinned_third = [] + unpinned_first = [] + for f in sorted(Path(".github/workflows").glob("*.yml")): + for i, line in enumerate(f.read_text().splitlines(), 1): + m = USES.match(line) + if not m: + continue + name, ref = m.group(1), m.group(2) + if SHA_PIN.search(line): + continue + bucket = unpinned_first if FIRST_PARTY.match(line) else unpinned_third + bucket.append((str(f), i, name, ref)) + print("::group::Action pinning status") + print(f"third-party actions on mutable refs: {len(unpinned_third)}") + for f, i, n, r in unpinned_third: + print(f" HIGH {f}:{i}: {n}@{r}") + print() + print(f"first-party (actions/* | github/*) on mutable refs: {len(unpinned_first)}") + for f, i, n, r in unpinned_first[:30]: + print(f" WARN {f}:{i}: {n}@{r}") + if len(unpinned_first) > 30: + print(f" ... and {len(unpinned_first) - 30} more") + print() + print("Recommendation: pin third-party actions to a 40-char SHA.") + print("Dependabot's github-actions ecosystem will auto-bump them.") + print("::endgroup::") + PY + { + echo "## GitHub Actions pinning verifier" + echo + echo '```' + cat logs-actions-pinning.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + # ───────────────────────────────────────────────────────────── + # Hash-pin verifier. `==` pinning protects against version + # drift but not against a re-uploaded malicious wheel at the + # same version (PyPI lets a yanked release be re-published with + # different bytes for ~5 minutes via `--filename` collision). + # `pip install --require-hashes` rejects any download whose + # SHA-256 doesn't match. Inspector step that reports how many + # specs would gain from a hash pin -- conversion is a roadmap + # item (needs pip-tools / uv pip compile --generate-hashes). + # ───────────────────────────────────────────────────────────── + - name: Hash-pin verifier (Python requirements) + continue-on-error: true + run: | + python <<'PY' | tee logs-hash-verifier.txt + import re + from pathlib import Path + PINNED = re.compile(r"^\s*[A-Za-z0-9_.\-]+\s*==\s*[^,;]+\s*$") + HASH_LINE = re.compile(r"--hash=sha256:[0-9a-f]{64}") + total_pinned = 0 + with_hash = 0 + for f in sorted(Path("studio/backend/requirements").glob("*.txt")): + text = f.read_text() + for raw in text.splitlines(): + line = raw.strip() + if not line or line.startswith("#") or line.startswith("-"): + continue + spec = line.split("#", 1)[0].strip().split(";", 1)[0] + if PINNED.match(spec): + total_pinned += 1 + if HASH_LINE.search(raw): + with_hash += 1 + print(f"::group::Hash-pin status") + print(f" exact == pins: {total_pinned}") + print(f" with --hash=sha256: {with_hash}") + print(f" without --hash: {total_pinned - with_hash}") + print() + print("Roadmap: convert to hash-locked installs via") + print("`uv pip compile --generate-hashes` and `pip install --require-hashes`.") + print("Hash-locked installs would have refused a republished") + print("malicious litellm 1.82.7 wheel even at the same version.") + print("::endgroup::") + PY + { + echo "## Hash-pin verifier" + echo + echo '```' + cat logs-hash-verifier.txt + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: always() + with: + name: advisory-audit-logs + path: | + logs-pip-audit.txt + logs-npm-audit.txt + logs-npm-audit.json + logs-cargo-audit.txt + logs-osv-scanner.txt + logs-semgrep.txt + logs-pin-verifier.txt + logs-actions-pinning.txt + logs-hash-verifier.txt + audit-reqs/ + sbom/ + retention-days: 30 + + # ───────────────────────────────────────────────────────────────────── + # Python: pre-install package scan (no install, no execution) + # ───────────────────────────────────────────────────────────────────── + pip-scan-packages: + # Downloads each declared dep WITHOUT installing it and inspects + # the archive contents for known malicious patterns: weaponized + # .pth files, credential stealers, obfuscated payloads, + # install-time droppers, suspicious subprocess / network / + # base64-blob combinations. + # + # This is the kind of check that would have caught: + # - litellm 1.82.7 / 1.82.8 (March 2026, supply-chain compromise) + # - the typo-squat campaign against PyTorch Lightning + # before either landed in the install path. pip-audit only knows + # about CVE-published vulnerabilities, so it does NOT see novel + # malicious uploads. scan_packages.py runs deterministic regex + # pattern matching, no LLM calls. + # + # `--with-deps` makes the scan transitive: every package the + # declared set resolves to gets fetched and pattern-scanned, not + # just the top-level pins. Resolving the full transitive closure + # of the unsloth + Studio dep tree downloads several hundred + # archives, hence the longer timeout. + # + # Sharded across runners for wall-clock parallelism. Each shard + # runs scan_packages.py once with --with-deps so its own slice + # benefits from pip's deduped transitive resolve. Shard + # composition tries to balance load: + # - hf-stack: pyproject extras + no-torch-runtime + # (~150 archives, transformers/peft/accelerate/...) + # - studio: FastAPI/Studio backend + overrides + extras-no-deps + # (~150 archives, smaller scientific stack) + # - extras: the heavy openai-whisper / scikit-learn / librosa + # stack (~250 archives, dominant cost) + # triton-kernels.txt is git+-only, fully skipped. + name: ${{ matrix.shard.name }} + runs-on: ubuntu-latest + timeout-minutes: 25 + strategy: + fail-fast: false + matrix: + shard: + - name: 'pip scan-packages :: hf-stack' + id: hf-stack + files: 'unsloth-deps no-torch-runtime' + - name: 'pip scan-packages :: studio' + id: studio + files: 'studio overrides extras-no-deps' + - name: 'pip scan-packages :: extras' + id: extras + files: 'extras' + steps: + # Egress block on every shard. Each shard pulls hundreds of + # PyPI archives -- if a malicious wheel ever phones home from + # within the scanner sandbox (it shouldn't; we never execute + # the archive), harden-runner now rejects the connect outright. + # Per-job allowlist: pip-scan-packages only fetches PyPI archives + # via scan_packages.py + pip download. No npm or cargo traffic. + - name: Harden runner (egress block) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: block + disable-sudo: true + allowed-endpoints: > + api.github.com:443 + github.com:443 + codeload.github.com:443 + objects.githubusercontent.com:443 + pypi.org:443 + files.pythonhosted.org:443 + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Install scan_packages.py runtime deps + # scan_packages.py imports requests + packaging at runtime to + # talk to PyPI's JSON API and to parse version specifiers. We + # do not install the packages it scans -- those are downloaded + # raw and inspected without ever touching `pip install`. + run: python -m pip install --upgrade pip requests packaging + + - name: Build filtered requirements set + # Mirrors the advisory-audit job's input transform: pyproject.toml + # extraction + git+ stripping. scan_packages.py downloads + # PyPI archives without building, so it tolerates legacy + # setup.py packages (no resolver dry-run); but `--with-deps` + # delegates resolution to a single `pip download` call that + # cannot satisfy `git+` specs without git operations, so we + # strip them here too. + run: | + mkdir -p audit-reqs + python <<'PY' > audit-reqs/unsloth-deps.txt + import tomllib + with open("pyproject.toml", "rb") as f: + d = tomllib.load(f) + core = d["project"]["dependencies"] + extras = d["project"]["optional-dependencies"]["huggingfacenotorch"] + print("# Auto-generated from pyproject.toml by security-audit.yml.") + print("# core deps + huggingfacenotorch extras.") + for spec in core + extras: + print(spec) + PY + for f in studio.txt extras.txt extras-no-deps.txt \ + no-torch-runtime.txt overrides.txt triton-kernels.txt; do + python < "audit-reqs/$f" + src = "studio/backend/requirements/$f" + with open(src) as fh: + for line in fh: + stripped = line.strip() + before_comment = stripped.split("#", 1)[0] + if "git+" in before_comment: + print(f"# [security-audit] skipped git+ spec: {stripped}") + continue + print(line.rstrip("\n")) + PY + done + + - name: Sanity-check scan_packages.py + # The scanner lives at scripts/scan_packages.py in this repo + # so we don't depend on a network fetch at job time. + run: | + test -f scripts/scan_packages.py + head -3 scripts/scan_packages.py + grep -q "Standalone pre-install package scanner" scripts/scan_packages.py + + - name: Scan declared + transitive Python deps + # scan_packages.py exits 1 on NON-baselined CRITICAL/HIGH + # findings, 0 otherwise. It scans code-only (docstrings and + # comments are blanked first) and suppresses reviewed + # known-good findings via scripts/scan_packages_baseline.json, + # so legitimate-library noise no longer red-fails the gate. + # The step stays advisory until SCAN_ENFORCE=1 (see env below); + # then PIPESTATUS propagates the scanner's exit code. + # + # `--with-deps` walks PyPI metadata to enumerate every + # transitive dep the declared set would install, then scans + # them all. Without this flag, we'd only catch a malicious + # *direct* dep -- and supply-chain attacks usually land + # several hops down (litellm 1.82.7 was a dep of a dep for + # most users). + # + # This step runs once per matrix shard. Within a shard, every + # -r file is fed to a single `pip download` call so pip + # intersects version constraints and yields a deduped + # transitive set (no point fetching the same transformers + # wheel five times). Across shards we accept some redundant + # downloads in exchange for wall-clock parallelism. + env: + SHARD_FILES: ${{ matrix.shard.files }} + # Enforcement switch. "1" = blocking: a non-baselined CRITICAL/HIGH + # fails the build. scan_packages.py scans code-only (docstrings/comments + # stripped), fetches sdist-only packages directly from PyPI (no build) + # so every shard resolves, and honors the reviewed allowlist at + # scripts/scan_packages_baseline.json, so only NON-baselined + # CRITICAL/HIGH cause its exit 1. The committed baseline makes all three + # shards exit 0 today; set this back to "0" to return to advisory. + SCAN_ENFORCE: "1" + run: | + set +e + mkdir -p logs + LOG="logs-scan-packages-${{ matrix.shard.id }}.txt" + echo "::group::shard ${{ matrix.shard.id }} input files" + REQ_ARGS=() + for f in $SHARD_FILES; do + if grep -qE '^[^#[:space:]]' "audit-reqs/$f.txt"; then + echo " + audit-reqs/$f.txt" + REQ_ARGS+=( -r "audit-reqs/$f.txt" ) + else + echo " - audit-reqs/$f.txt (empty after git+ filter, skipping)" + fi + done + echo "::endgroup::" + rc=0 + if [ ${#REQ_ARGS[@]} -eq 0 ]; then + echo "[security-audit] shard ${{ matrix.shard.id }}: no PyPI specs, nothing to scan" \ + | tee "$LOG" + else + python scripts/scan_packages.py --with-deps "${REQ_ARGS[@]}" \ + 2>&1 | tee "$LOG" + rc=${PIPESTATUS[0]} + fi + { + echo "## scan_packages :: shard ${{ matrix.shard.id }}" + echo + echo "### Files in this shard" + for f in $SHARD_FILES; do echo "- audit-reqs/$f.txt"; done + echo + echo "scan_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)" + echo + echo '### Findings (tail)' + echo '```' + tail -200 "$LOG" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + # Advisory by default; blocking once SCAN_ENFORCE=1 and the baseline + # is committed. PIPESTATUS is captured above so `tee` does not mask the + # scanner's exit code. + if [ "$SCAN_ENFORCE" = "1" ]; then + exit "$rc" + fi + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: always() + with: + name: scan-packages-log-${{ matrix.shard.id }} + path: | + logs-scan-packages-${{ matrix.shard.id }}.txt + audit-reqs/ + retention-days: 30 + + # ───────────────────────────────────────────────────────────────────── + # npm: pre-install tarball content scan. + # ───────────────────────────────────────────────────────────────────── + npm-scan-packages: + # Counterpart to pip-scan-packages for the npm side. Reads + # studio/frontend/package-lock.json, downloads each resolved + # tarball DIRECTLY from registry.npmjs.org (never via `npm + # install` -- no lifecycle scripts ever run), verifies the + # lockfile integrity hash, unpacks each tarball into a sandboxed + # temp dir behind size / count / path-escape / symlink guards, + # and pattern-scans the extracted file contents for the + # signatures common to npm supply-chain attacks: + # + # - lifecycle (preinstall / install / postinstall / prepare) + # scripts in any package.json that fetch + execute external + # code, + # - C2 / exfiltration hosts (getsession.org, AWS IMDS, + # Kubernetes ServiceAccount token paths, GitHub Actions OIDC, + # HashiCorp Vault endpoints), + # - credential-stealing references (.npmrc, .aws/credentials, + # GITHUB_TOKEN / NPM_TOKEN in JS sources), + # - known IOC filenames (router_init.js, tanstack_runner.js, + # router_runtime.js), + # - obfuscation shapes (Function/eval against base64 blobs). + # + # Threat model: every tarball is hostile. Safety guarantees are + # documented at scripts/scan_npm_packages.py top-of-file. The + # script is stdlib-only so adding it does not increase the + # transitive supply-chain surface. + name: npm scan-packages (Studio frontend tarballs) + runs-on: ubuntu-latest + timeout-minutes: 30 + needs: [] + steps: + # Per-job allowlist: npm-scan-packages only fetches tarballs from + # registry.npmjs.org. GitHub endpoints retained for checkout + + # setup-python action machinery. + - name: Harden runner (egress block) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: block + disable-sudo: true + allowed-endpoints: > + api.github.com:443 + github.com:443 + codeload.github.com:443 + objects.githubusercontent.com:443 + registry.npmjs.org:443 + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Sanity-check scan_npm_packages.py + run: | + test -f scripts/scan_npm_packages.py + python3 -c "import ast; ast.parse(open('scripts/scan_npm_packages.py').read())" + + - name: Scan npm tarballs (declared + transitive, no install) + # scan_npm_packages.py exits 1 on NON-baselined HIGH/CRITICAL + # findings, 0 otherwise. It scans code-only (JS/TS comments are + # blanked first) and honors a reviewed allowlist at + # scripts/scan_npm_packages_baseline.json. It never runs + # `npm install`, never executes anything from a downloaded + # tarball, and only fetches from registry.npmjs.org. The npm + # corpus is clean (the baseline is empty), so the gate is + # enforcing (SCAN_ENFORCE=1) and any new finding fails the build. + env: + SCAN_ENFORCE: "1" + run: | + set +e + LOG=logs-scan-npm.txt + python3 scripts/scan_npm_packages.py 2>&1 | tee "$LOG" + rc=${PIPESTATUS[0]} + { + echo "## scan_npm_packages" + echo + echo "scan_npm_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)" + echo + echo '### Findings (tail)' + echo '```' + tail -300 "$LOG" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + # Blocking: the npm corpus is clean, so any non-baselined + # HIGH/CRITICAL is new and should fail the build. PIPESTATUS is + # captured above so `tee` does not mask the scanner's exit code. + if [ "$SCAN_ENFORCE" = "1" ]; then + exit "$rc" + fi + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: always() + with: + name: scan-npm-packages-log + path: logs-scan-npm.txt + retention-days: 30 + + # ───────────────────────────────────────────────────────────────────── + # Workflow-trigger lint. Refuses two patterns that together powered the + # TanStack GHSA-g7cv-rxg3-hmpx supply-chain compromise: + # + # 1. `pull_request_target` -- runs a fork's workflow YAML against + # the base repository's secrets. There is no safe use of this + # trigger for a public open-source project. + # + # 2. Shared cache keys between PR-triggered workflows and the + # publish workflow. A fork PR can poison the cache; the publish + # workflow then restores the poisoned cache on next run. + # + # Cheap pure-Python lint, runs in seconds. Fail-closed. + # ───────────────────────────────────────────────────────────────────── + workflow-trigger-lint: + name: workflow-trigger lint (pull_request_target / cache-poisoning) + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - name: Harden runner (egress block) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: block + disable-sudo: true + allowed-endpoints: > + api.github.com:443 + github.com:443 + codeload.github.com:443 + objects.githubusercontent.com:443 + pypi.org:443 + files.pythonhosted.org:443 + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Install PyYAML + run: pip install pyyaml + + - name: Lint workflow triggers + cache keys + run: python3 scripts/lint_workflow_triggers.py + + # ───────────────────────────────────────────────────────────────────── + # Regression tests: pin scanner IOC tables and pre-install fixtures. + # Hard gate (no continue-on-error) so future drift in the IOC tables + # or scanner exit semantics fails this PR at review time. + # ───────────────────────────────────────────────────────────────────── + tests-security: + name: pytest tests/security + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Harden runner (egress block) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: block + disable-sudo: true + allowed-endpoints: > + api.github.com:443 + github.com:443 + codeload.github.com:443 + objects.githubusercontent.com:443 + pypi.org:443 + files.pythonhosted.org:443 + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Install pytest + PyYAML + # PyYAML is imported by scripts/lint_workflow_triggers.py, which the + # `tests/security/test_lint_workflow_triggers.py` regression suite + # exercises as a subprocess. Without it the lint script bails with + # `ERROR: PyYAML is required` (exit 2) and the 5 lint regression + # tests fail. Pinned the same way pytest is pinned. + run: pip install pytest==9.0.3 pyyaml==6.0.2 + + - name: Run security regression tests + run: python3 -m pytest tests/security -v + + # ───────────────────────────────────────────────────────────────────── + # npm provenance + new install-script diff. Catches the two npm + # supply-chain levers we don't yet gate on: + # + # 1. `npm audit signatures` validates the registry-signed + # provenance of every tarball laid down in node_modules. Pulled + # from the public npm transparency log; surfaces unsigned or + # mis-signed deps. Informational for now (continue-on-error) + # while the baseline settles. + # + # 2. `check_new_install_scripts.py` diffs the PR's lockfile + # against the base ref and refuses any newly-added dep that + # ships a postinstall hook. Every recent npm supply-chain + # compromise leveraged a postinstall as the execution lever, so + # blocking new ones at PR time is a small, high-signal gate. + # ───────────────────────────────────────────────────────────────────── + npm-provenance-and-install-scripts: + name: npm provenance + new install-script diff + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Harden runner (egress block) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: audit + disable-sudo: true + allowed-endpoints: > + api.github.com:443 + github.com:443 + codeload.github.com:443 + objects.githubusercontent.com:443 + registry.npmjs.org:443 + + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + # Need the base commit accessible for `git show + # :studio/frontend/package-lock.json` below. + fetch-depth: 0 + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Install Studio frontend deps (--ignore-scripts) + # `npm audit signatures` requires node_modules to be populated. + # `--ignore-scripts` is mandatory: this is exactly the lever the + # new-install-script gate below protects against, and we must + # not run any third-party hook to set up the audit. + working-directory: studio/frontend + run: | + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + # --ignore-scripts is mandatory here (no third-party hook runs); + # the retry only re-attempts the registry fetch, it never relaxes + # that flag or the package-lock integrity check npm ci enforces. + retry 5 npm ci --ignore-scripts + + - name: npm audit signatures (informational) + # Surfaces unsigned / mis-signed packages from the npm + # transparency log. continue-on-error during baseline-build + # phase; promote to hard gate once the lockfile is fully + # signed (most major maintainers signed by mid-2025). + working-directory: studio/frontend + continue-on-error: true + run: | + set -o pipefail + LOG=logs-audit-signatures.txt + npm audit signatures 2>&1 | tee "$LOG" + { + echo "## npm audit signatures" + echo + echo '```' + tail -200 "$LOG" + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + + - name: Extract base-ref lockfile (PR triggers only) + if: github.event_name == 'pull_request' + run: | + set -e + BASE_SHA="${{ github.event.pull_request.base.sha }}" + git show "$BASE_SHA:studio/frontend/package-lock.json" \ + > /tmp/base-package-lock.json + + - name: Diff for newly-added install-script deps + if: github.event_name == 'pull_request' + run: | + python3 scripts/check_new_install_scripts.py \ + --base /tmp/base-package-lock.json \ + --head studio/frontend/package-lock.json + + - name: Skip install-script diff (non-PR trigger) + if: github.event_name != 'pull_request' + run: | + echo "Not a pull_request event; install-script diff requires a base ref." + echo "This step is intentionally a no-op outside PR triggers." + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: always() + with: + name: npm-audit-signatures-log + path: studio/frontend/logs-audit-signatures.txt + if-no-files-found: ignore + retention-days: 30 diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml new file mode 100644 index 0000000..1a4cf84 --- /dev/null +++ b/.github/workflows/stale.yml @@ -0,0 +1,37 @@ +name: 'Inactive Issue Pinger' + +on: + schedule: + - cron: '30 5 * * *' # Runs at 5:30 UTC every day + +jobs: + stale: + runs-on: ubuntu-latest + permissions: + issues: write + + steps: + - uses: actions/stale@b5d41d4e1d5dceea10e7104786b73624c18a190f # v10.2.0 + with: + # The message to post on stale issues. + # This message will ping the issue author. + # Note: The stale bot action does not currently support a direct placeholder for the last commenter. + # As a workaround, this message encourages any participant to reply. + stale-issue-message: > + Is this issue still important to you? + Apologies in advance we might have missed this issue as well. + For faster response times, please post on our Reddit server - https://www.reddit.com/r/unsloth or our Discord - https://discord.com/invite/unsloth + + # The number of days of inactivity before an issue is considered stale. + days-before-issue-stale: 9999 + + # Set to -1 to never close stale issues. + days-before-issue-close: -1 + + # A label to apply to stale issues. + stale-issue-label: 'inactive' + + # The number of operations to perform per run to avoid rate limiting. + operations-per-run: 500 + + enable-statistics: false diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml new file mode 100644 index 0000000..15efee3 --- /dev/null +++ b/.github/workflows/studio-api-smoke.yml @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Studio API & Auth Tests -- HTTP-level integration tests for the +# FastAPI surface. No Playwright, no model UI; tests/studio/test_studio_api_smoke.py +# runs ~30 s and asserts: +# - CORS hardening (no wildcard + credentials, no bootstrap leak) +# - /api/system + /api/system/hardware require auth +# - Auth state machine + JWT expiry +# - API key lifecycle E2E (create / list / use / delete / reject) +# - Auth file-mode hardening (Linux only) +# - Inference lifecycle (force reload, bogus variant, /v1/models, /v1/embeddings, /v1/responses) +# - Endpoint-by-endpoint auth audit +# +# Reuses the GGUF cache key from studio-ui-smoke.yml so the model +# download is one cache-hit on the second job. + +name: Studio API CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.sh' + - 'pyproject.toml' + - 'tests/studio/**' + - '.github/workflows/studio-api-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + api-smoke: + name: Studio API & Auth Tests + runs-on: ubuntu-latest + timeout-minutes: 12 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18893' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + # Same key as studio-ui-smoke.yml so the two jobs share a + # single GGUF download across CI. + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Install pyjwt for the JWT-expiry forge test + run: pip install 'pyjwt>=2.6' + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + + - name: Pass bootstrap password + rotated targets to the test + # The test does its own bootstrap-login + rotation to exercise + # the auth state machine; we just pre-mint two random rotated + # passwords for it. Mask them so the log is clean. + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="ApiSmoke-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + NEW2="ApiSmoke-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "::add-mask::$NEW2" + echo "STUDIO_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" + echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" + + - name: Run Studio API & Auth tests + # The script is named WITHOUT a `test_` prefix so it isn't + # auto-collected by pytest in Backend CI's `tests/` walk + # (which doesn't set BASE_URL and would crash at import). + env: + BASE_URL: http://127.0.0.1:18893 + STUDIO_AUTH_DIR: /home/runner/.unsloth/studio/auth + run: python tests/studio/studio_api_smoke.py + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + + - name: Upload API smoke logs + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: studio-api-smoke-log + path: | + logs/install.log + logs/studio.log + retention-days: 7 diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml new file mode 100644 index 0000000..3022127 --- /dev/null +++ b/.github/workflows/studio-backend-ci.yml @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs the existing studio/backend/tests/ suite (~860 tests, all CPU-friendly) +# on every PR that touches the backend or unsloth library. Until this lands, +# none of those tests run automatically. Verified locally on Python 3.13 with +# the surgical exclusions below: 861 pass, 4 skipped. +# +# Exclusions: +# - tests/test_studio_api.py: end-to-end against a live model + GGUF download, +# too heavy for free runners. Run separately when GPU CI is available. +# - -k 'not llama_cpp_load_progress_live': spawns a real llama.cpp process, +# not appropriate for CPU-only runners. +# +# Two jobs: +# - pytest matrix (3.10/3.11/3.12/3.13) over studio/backend/tests +# - repo-cpu-tests: auto-discovered tests/ + state-isolated spoof files +# +# Whole-repo Python lint (syntax + ruff + debugger-leftover scan) +# moved to the dedicated `Lint CI` workflow (.github/workflows/lint-ci.yml) +# so it fires on every PR rather than only on studio/unsloth/tests +# path changes. + +name: Backend CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'tests/**' + - 'pyproject.toml' + - '.github/workflows/studio-backend-ci.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + pytest: + name: (Python ${{ matrix.python }}) + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + python: ['3.10', '3.11', '3.12', '3.13'] + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '${{ matrix.python }}' + cache: 'pip' + + - name: Install backend test dependencies (CPU only) + run: | + python -m pip install --upgrade pip + # Studio's declared backend deps: + pip install -r studio/backend/requirements/studio.txt + # Extras that studio.txt does not list but the import chain needs + # (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography + # for the auth DB, yaml/jinja2 for utils.models.model_config, psutil for + # the orphan-cleanup process scan, etc.): + pip install \ + python-multipart aiofiles sqlalchemy cryptography psutil \ + pyyaml jinja2 mammoth unpdf requests \ + 'numpy<3' pytest pytest-asyncio httpx + # Torch CPU + transformers are required by a chunk of the backend test + # suite (gpu_selection, kv_cache_estimation, utils). CPU-only torch + # keeps the install ~250 MB / ~1 min on a clean runner. + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple 'torch>=2.4,<2.11' + pip install 'transformers>=4.51,<5.5' + + - name: Backend tests + working-directory: studio/backend + # Locally validated against this dep set: 831 passed, 5 skipped, 35 deselected. + # Deselections (all environment-specific, would never pass on a GPU-less + # `ubuntu-latest` runner regardless of code correctness): + # - llama_cpp_load_progress_live: spawns a real llama.cpp process + # - TestGpuAutoSelection / TestPreSpawnGpuResolution / TestPerGpuFitGuardAllCounts: + # require live transformers config introspection on real GPUs + # - TestTransformersIntrospection: same + # - test_returns_cuda_when_cuda_available / test_calls_cuda_cache_when_cuda: + # assume CUDA-capable GPU + run: | + python -m pytest tests/ -q --tb=short \ + --ignore=tests/test_studio_api.py \ + -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' + + repo-cpu-tests: + # Auto-discover everything under tests/ that is not GPU-bound by + # design. New tests added in covered directories are picked up + # without a workflow edit. Locally validated: 760 passed, 1 skipped, + # 23 deselected. tests/conftest.py (mirroring unsloth-zoo PR #624) + # pre-loads unsloth_zoo.device_type and unsloth.device_type under a + # mocked torch.cuda.is_available so the unsloth import chain + # succeeds on CPU. + name: Repo tests (CPU) + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + # node + uv unlock ~60 tests that previously skipped on CI: + # - 9 tests in test_chat_preset_builtin_invariants.py need node to + # compile a tiny TS harness against the frontend chat sources. + # - tests/python/* spawn fresh `uv venv`s to verify the no-torch + # install path; they self-skip when uv is missing. + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - name: Install uv (for tests/python/* sandboxed venvs) + run: pip install uv + + - name: Install deps (shared shape with backend pytest job) + run: | + python -m pip install --upgrade pip + pip install -r studio/backend/requirements/studio.txt + pip install \ + python-multipart aiofiles sqlalchemy cryptography psutil \ + pyyaml jinja2 mammoth unpdf requests typer \ + 'numpy<3' pytest pytest-asyncio httpx + # torchvision: unsloth_zoo.vision_utils imports it at module scope. + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' + pip install 'transformers>=4.51,<5.5' + # bitsandbytes: hard import in unsloth/models/_utils.py. Recent + # versions ship a CPU build that imports cleanly on Linux. + pip install 'bitsandbytes>=0.45' + # unsloth.device_type imports unsloth_zoo.utils.Version at module + # scope, so the conftest preload needs unsloth_zoo. Pull from + # git main so this job sees the same zoo HEAD as Core / MLX / + # install.sh do (otherwise a fix on zoo main hides until release). + # No --no-deps: matches prior `pip install 'unsloth_zoo>=2026.5.1'` + # behaviour so triton etc. still come in for the Repo tests CPU + # collection imports. + for attempt in 1 2 3; do + if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done + pip install -e . --no-deps + + - name: Repo tests (CPU, auto-discovered) + env: + # tests/python/* import install_python_stack from studio/. + PYTHONPATH: ${{ github.workspace }}/studio + # Skip lazy compilation work the unsloth import chain wants to + # do at import time on a real GPU. + UNSLOTH_COMPILE_DISABLE: '1' + # --ignore: GPU-bound directories (qlora/saving need real weights; + # tests/sh is the shell suite the next step handles; tests/utils + # is a helpers folder); tests/vllm_compat + tests/version_compat + # are dedicated multi-version drift canaries with their own job + # in version-compat-ci.yml that installs the heavier dep set + # (torchcodec, full transformers/peft/bnb pins) those tests need. + # State-sensitive hardware-spoofing files run in isolation in the + # next step because they mutate hardware.py module globals. + # -m: honour markers from tests/python/conftest.py (`server` = + # needs studio venv, `e2e` = needs network). + # --deselect: + # - test_model_registration / test_all_model_registration: + # hit huggingface_hub for live model existence checks. + # - test_autoconfig_works_with_no_torch_runtime / test_autoconfig_succeeds: + # fail because no-torch-runtime.txt does not pin tokenizers + # and the latest tokenizers (0.23.1) is incompatible with the + # transformers it resolves to. Tracked separately; this is a + # real bug in the no-torch install path, not a CI issue. + run: | + python -m pytest tests/ -q --tb=short \ + --ignore=tests/qlora \ + --ignore=tests/saving \ + --ignore=tests/utils \ + --ignore=tests/sh \ + --ignore=tests/studio/test_hardware_dispatch_matrix.py \ + --ignore=tests/studio/test_is_mlx_dispatch_gate.py \ + --ignore=tests/vllm_compat \ + --ignore=tests/version_compat \ + -m 'not server and not e2e' \ + --deselect tests/test_model_registry.py::test_model_registration \ + --deselect tests/test_model_registry.py::test_all_model_registration \ + --deselect 'tests/python/test_tokenizers_and_torch_constraint.py::TestE2ETokenizersFix::test_autoconfig_works_with_no_torch_runtime' \ + --deselect 'tests/python/test_tokenizers_and_torch_constraint.py::TestE2EFullNoTorchSandbox::test_autoconfig_succeeds' + + - name: Hardware-spoof tests (state-sensitive, run in isolation) + env: + PYTHONPATH: ${{ github.workspace }}/studio + UNSLOTH_COMPILE_DISABLE: '1' + # These two files mutate hardware.py module globals at runtime + # via the spoof fixtures, which leaks state into any other test + # that imports hardware. Run them in their own pytest invocation + # so the leak does not cross file boundaries. + run: | + python -m pytest -q --tb=short \ + tests/studio/test_hardware_dispatch_matrix.py \ + tests/studio/test_is_mlx_dispatch_gate.py + + - name: Shell installer tests + # Subset that does not depend on a writable / pristine install.sh + # tree; test_install_host_defaults.sh checks install.ps1 layout + # which has drifted (separate followup). + run: | + set -e + for s in \ + tests/sh/test_get_torch_index_url.sh \ + tests/sh/test_mac_intel_compat.sh \ + tests/sh/test_node_decision.sh \ + tests/sh/test_studio_home_node_dir.sh \ + tests/sh/test_system_node_readonly.sh \ + tests/sh/test_nvcc_meets_llama_minimum.sh \ + tests/sh/test_resolve_cuda_archs.sh \ + tests/sh/test_tauri_install_exit_order.sh \ + tests/sh/test_torch_constraint.sh \ + tests/sh/test_torch_flavor.sh \ + tests/sh/test_with_llama_cpp_dir_flag.sh \ + tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do + echo "::group::$s" + bash "$s" + echo "::endgroup::" + done + diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml new file mode 100644 index 0000000..1ee6489 --- /dev/null +++ b/.github/workflows/studio-export-capability-ci.yml @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs studio/backend/tests/test_export_capability.py on Linux, Windows and macOS. +# +# export_capability() is per-OS (is_apple_silicon() and the PyTorch-import probe differ per +# platform) and the export backend must import without PyTorch, so this confirms the gating and +# import-safety on hosted Windows/macOS. Hosted runners have no GPU/MLX, so a real accelerator +# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block +# torch/unsloth, so the job installs only a CPU PyTorch plus import deps. + +name: Studio export capability + +on: + pull_request: + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + push: + branches: [main] + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + capability: + name: capability (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + env: + # No accelerator on hosted runners; keep detection on the CPU path. + CUDA_VISIBLE_DEVICES: "" + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Upgrade pip + run: python -m pip install --upgrade pip + - name: Install CPU PyTorch + # CPU wheel index so every OS gets a CPU build; keep PyPI as an extra index so torch's + # transitive deps still resolve (matching the other workflows in this repo). + run: python -m pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple "torch>=2.4,<2.13" + - name: Install backend import deps + # Enough to import utils.hardware and core.export.export; NOT unsloth (needs a GPU, and + # the import-safety test blocks it) or triton/llama.cpp (Linux-only / native builds). + run: python -m pip install + transformers peft accelerate safetensors huggingface_hub datasets + sentencepiece protobuf fastapi starlette structlog psutil + python-multipart pydantic httpx "numpy<3" pytest + - name: Export capability + import-safety tests + working-directory: studio/backend + run: python -m pytest tests/test_export_capability.py -q diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml new file mode 100644 index 0000000..b42086f --- /dev/null +++ b/.github/workflows/studio-frontend-ci.yml @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Frontend PR gate: lockfile freshness, typecheck, build, and a bundle grep +# that catches the 2026.5.1 chat-history regression at the JS level. +# +# biome runs as non-blocking for now: the codebase currently has accumulated +# ~470 errors and ~1650 warnings against the existing biome config. Surfacing +# the count in CI lets us drive it down without forcing a fleet-wide cleanup +# in the same PR. Drop `continue-on-error` once that number is zero. + +name: Frontend CI + +on: + pull_request: + paths: + - 'studio/frontend/**' + - 'scripts/check_frontend_dep_removal.py' + - 'tests/studio/test_frontend_dep_removal.py' + - 'scripts/sync_allow_scripts_pins.py' + - 'tests/studio/test_sync_allow_scripts_pins.py' + - '.github/workflows/studio-frontend-ci.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + build: + name: Frontend build + bundle sanity + runs-on: ubuntu-latest + timeout-minutes: 10 + defaults: + run: + working-directory: studio/frontend + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + # FIXME: drop this step once @assistant-ui/* and assistant-stream + # leave 0.x -- on 1.x, caret ranges are conventional. Until then, + # every 0.minor on this surface is a SemVer-major (this is exactly + # how 2026.5.1 shipped a broken chat runtime: ^0.12.19 quietly + # resolved to 0.12.28). + - name: '@assistant-ui must be pinned exactly (no caret/tilde)' + working-directory: ${{ github.workspace }} + run: | + set -e + if grep -nE '"(@assistant-ui/[a-z-]+|assistant-stream)":[[:space:]]*"[\^~]' studio/frontend/package.json; then + echo "::error file=studio/frontend/package.json::These packages must be pinned to exact versions until they leave 0.x. Drop the leading ^ or ~." + exit 1 + fi + echo "All assistant-ui packages are pinned exactly." + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + # node 22 bundles npm 10.x, which predates allowScripts. Move to the + # 11.x line and fail loudly if the gate is still missing, so the + # strict flag below can never silently degrade into a warning. + - name: Upgrade npm to 11.x (allowScripts enforcement) + working-directory: ${{ github.workspace }} + run: | + npm install -g npm@^11 --no-fund --no-audit + V=$(npm -v) + case "$V" in + 11.1[6-9].*|11.[2-9][0-9].*|1[2-9].*) echo "npm $V has allowScripts" ;; + *) echo "::error::npm $V lacks allowScripts (need >=11.16)"; exit 1 ;; + esac + + # Run the structural lockfile scan BEFORE npm ci. A compromised + # tarball runs its `prepare` / `postinstall` during `npm ci`, + # so any catch has to fire upstream of that. The scanner is + # pure-Python read-only; safe to call ahead of every install. + - name: Lockfile supply-chain audit (pre-install scan) + working-directory: ${{ github.workspace }} + run: python3 scripts/lockfile_supply_chain_audit.py + + # Dependency bumps strand the version-pinned allowScripts entries. + # The paired pre-commit hook auto-fixes PRs; this is the backstop. + - name: allowScripts pins must match the lockfile + working-directory: ${{ github.workspace }} + run: | + python3 tests/studio/test_sync_allow_scripts_pins.py + python3 scripts/sync_allow_scripts_pins.py --check + + - name: Lockfile must agree with package.json (npm ci is strict) + # The vite 8 chain (rolldown, lightningcss, tailwind oxide) ships napi + # binaries with no install scripts. The only script-bearing deps are + # covered by `allowScripts` in package.json (npm >=11.16, default in + # npm 12). The pre-install lockfile audit above stays the first line + # of defence -- it fires before any tarball can run code. + # --strict-allow-scripts: any unreviewed install script hard-fails + # the job; the sync hook keeps the pins fresh after bumps. + run: npm ci --strict-allow-scripts --no-fund --no-audit + + - name: npm ci must not have modified the working tree + working-directory: ${{ github.workspace }} + run: | + if ! git diff --quiet -- studio/frontend; then + echo "::error::npm ci modified files; commit the updated lockfile" + git status -- studio/frontend + exit 1 + fi + + # Catch the common foot-gun: a dep dropped from package.json that is + # still imported somewhere. The script walks the lockfile dep graph + # from the new top-level deps and only counts top-level node_modules + # paths as valid resolution targets for bare src/ imports. + # + # actions/checkout uses fetch-depth: 1 by default, so the base branch + # is not available locally. Fetch the single base commit with an + # explicit refspec so origin/ is reliably created (a bare + # `git fetch origin ` only updates FETCH_HEAD in some configs). + - name: Dependency removal safety check + if: github.event_name == 'pull_request' + working-directory: ${{ github.workspace }} + run: | + git fetch --no-tags --depth=1 origin \ + "${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}" + python3 scripts/check_frontend_dep_removal.py \ + --base "origin/${{ github.base_ref }}" \ + --enumerate-dead + python3 tests/studio/test_frontend_dep_removal.py + + - name: Typecheck + run: npm run typecheck + + - name: Build + run: npm run build + + - name: Built bundle must not contain Studio's unstable_Provider call site + run: | + set -e + JS=$(ls dist/assets/index-*.js | head -1) + HITS=$(grep -c 'unstable_Provider:' "$JS" || echo 0) + echo "main bundle: $JS" + echo "unstable_Provider: hits=$HITS (assistant-ui internals contribute up to 3)" + if [ "$HITS" -gt 3 ]; then + echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Studio bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead." + exit 1 + fi + + - name: Bundle size budget (75 MB) + run: | + SIZE=$(du -sb dist | cut -f1) + BUDGET=$((75 * 1024 * 1024)) + echo "dist size: $SIZE bytes ($((SIZE/1024/1024)) MB), budget: $BUDGET bytes (75 MB)" + if [ "$SIZE" -gt "$BUDGET" ]; then + echo "::error::studio/frontend/dist/ exceeded the 75 MB budget. Drop dead deps (e.g. the unused next dep) or split chunks." + exit 1 + fi + + - name: Biome (non-blocking until accumulated drift is cleared) + continue-on-error: true + run: npm run biome:check + + - name: Upload built dist + # Always upload so a green run is reviewable too -- the dist + # output catches "tests passed but bundle changed unexpectedly" + # regressions that would be invisible if we only kept artifacts + # on failure. + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: studio-frontend-dist + path: studio/frontend/dist + retention-days: 3 diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml new file mode 100644 index 0000000..f540c11 --- /dev/null +++ b/.github/workflows/studio-inference-smoke.yml @@ -0,0 +1,1128 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Three end-to-end smoke jobs that boot a freshly-installed Studio and +# exercise the surfaces real users hit through the OpenAI / Anthropic +# SDKs and curl. Each job picks the smallest model that exercises the +# behaviour under test, primes HF_HOME via actions/cache, and shares +# the install.sh --local --no-torch bootstrap. +# +# 1. OpenAI, Anthropic API tests +# gemma-3-270m-it UD-Q4_K_XL (~254 MiB). +# Password rotation via /api/auth/change-password (old fails, +# new works), then OpenAI + Anthropic Python SDKs against /v1/* +# with temperature=0 and a fixed seed. Asserts the four-turn +# conversation is deterministic across two runs. +# +# 2. Tool calling Tests +# Qwen3.5-2B UD-IQ3_XXS (~890 MiB). OpenAI function calling, +# server-side tools (python, terminal, web_search) via +# enable_tools / enabled_tools, and enable_thinking on/off. +# +# 3. JSON, images +# Qwen3-VL-2B-Instruct UD-Q4_K_XL (~1.1 GiB) + mmproj-F16 (~780 MiB). +# response_format JSON-schema decoding and OpenAI image_url +# (data URI) plus Anthropic source/base64 image inputs. +# +# All three jobs run in parallel. Total wall time is dominated by job 3 +# on a cold cache; warm cache cuts that to ~3 min. + +name: Studio GGUF CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.sh' + - 'pyproject.toml' + - '.github/workflows/studio-inference-smoke.yml' + push: + branches: [main, pip] + # Manual trigger for pre-warming HF_HOME caches on main, or re-running + # against an arbitrary branch without pushing a no-op commit. + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + # ───────────────────────────────────────────────────────────────────── + # Job 1: OpenAI, Anthropic API tests + # ───────────────────────────────────────────────────────────────────── + openai-anthropic: + name: OpenAI, Anthropic API tests + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18888' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Install OpenAI + Anthropic Python SDKs + run: pip install 'openai>=1.50' 'anthropic>=0.40' + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json + exit 0 + fi + sleep 1 + done + echo "Studio did not become healthy in 180s" + tail -200 logs/studio.log + exit 1 + + - name: Password rotation (old must fail, new must work) + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIRotated-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + # 1. Login with the bootstrap password. + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + [ -n "$OLD_TOKEN" ] && [ "$OLD_TOKEN" != "null" ] || { echo "bootstrap login failed"; exit 1; } + # 2. Rotate to a fresh random password. + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + # 3. Old password must now be rejected (HTTP 401). + OLD_STATUS=$(curl -s -o /dev/null -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}") + if [ "$OLD_STATUS" != "401" ]; then + echo "::error::Login with old password returned $OLD_STATUS, expected 401" + exit 1 + fi + # 4. New password must succeed; capture the JWT for downstream steps. + NEW_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + [ -n "$NEW_TOKEN" ] && [ "$NEW_TOKEN" != "null" ] || { echo "new login failed"; exit 1; } + echo "TOKEN=$NEW_TOKEN" >> "$GITHUB_ENV" + echo "password rotation OK (old=401, new=200)" + + - name: Load the GGUF (HF repo + variant, served from HF_HOME cache) + run: | + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}" \ + | jq '{status, display_name, is_gguf, context_length}' + + - name: Multi-turn determinism via OpenAI + Anthropic SDKs + env: + BASE_URL: http://127.0.0.1:18888 + run: | + python - <<'PY' + import json + import os + from openai import OpenAI + from anthropic import Anthropic + + BASE = os.environ["BASE_URL"] + KEY = os.environ["TOKEN"] # JWT also accepted as Bearer on /v1/* + SEED = 3407 + + # Four-turn conversation: the second and fourth turns can only be + # answered correctly if the model sees the prior turns, so this + # also exercises the conversation-history wiring. + PROMPTS = [ + "What is 1+1?", + "What did I ask before?", + "What is the capital of France?", + "Repeat the city name", + ] + + def run_openai(): + client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY) + history, replies = [], [] + for prompt in PROMPTS: + history.append({"role": "user", "content": prompt}) + resp = client.chat.completions.create( + model = "default", + messages = history, + temperature = 0.0, + max_tokens = 80, + seed = SEED, + extra_body = {"enable_thinking": False}, + ) + text = resp.choices[0].message.content or "" + replies.append(text) + history.append({"role": "assistant", "content": text}) + return replies + + def run_anthropic(): + # Two SDK quirks vs. Studio: + # 1. base_url must NOT include /v1 -- the SDK appends + # /v1/messages itself; otherwise the request hits + # /v1/v1/messages and 405s. + # 2. The SDK sends `x-api-key` by default, but Studio's + # auth layer is HTTPBearer-only. Override via + # default_headers so Authorization: Bearer ... is + # sent instead. + client = Anthropic( + base_url = BASE, + api_key = "unused", + default_headers = {"Authorization": f"Bearer {KEY}"}, + ) + history, replies = [], [] + for prompt in PROMPTS: + history.append({"role": "user", "content": prompt}) + msg = client.messages.create( + model = "default", + max_tokens = 80, + messages = history, + temperature = 0.0, + extra_body = {"seed": SEED, "enable_thinking": False}, + ) + text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text") + replies.append(text) + history.append({"role": "assistant", "content": text}) + return replies + + for label, runner in (("openai", run_openai), ("anthropic", run_anthropic)): + first = runner() + second = runner() + determinism_failures = [] + for i, (a, b) in enumerate(zip(first, second), start = 1): + print(f"[{label} turn {i}] {a!r}") + # Both runs must be non-empty; small-quant drift + # across runs is WARN-only (grounding asserts below + # are the stronger signal). + assert a, f"{label}: empty turn {i} response in first run" + assert b, f"{label}: empty turn {i} response in second run" + if a.strip() != b.strip(): + determinism_failures.append( + f"turn {i}: run1={a!r} run2={b!r}" + ) + if determinism_failures: + print( + f"[{label}] WARN non-determinism at temperature=0.0 across " + f"{len(determinism_failures)} of {len(first)} turn(s); " + f"small-quant model drift, not a Studio regression. " + f"Details: " + " | ".join(determinism_failures) + ) + # Sanity: turn-2 reply should mention the earlier question, and + # turn-4 reply should mention Paris (model echoes the city it + # produced for turn 3). Lower-cased substring checks keep the + # assertion robust to formatting jitter. + joined = " ".join(first).lower() + assert "1" in first[0], f"{label}: turn-1 answer should contain '1', got {first[0]!r}" + assert "paris" in joined, f"{label}: expected 'paris' somewhere in the four-turn transcript: {first}" + status_word = "PASS" if not determinism_failures else "PASS (with drift)" + print(f"[{label}] {status_word} -- 4 turns, history grounded ('paris' present)") + PY + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + ss -tln | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + # Always upload so green runs are still reviewable. + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: openai-anthropic-log + path: | + logs/studio.log + logs/install.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job 2: Tool calling Tests + # ───────────────────────────────────────────────────────────────────── + tool-calling: + name: Tool calling Tests + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + # Tool calling is the highest-volume GGUF in this workflow + # (Qwen3.5-2B at Q4_K_XL = ~1.28 GiB). Caching HF_HOME would + # store xet chunks + blobs + snapshots = ~4 GiB compressed -- + # 4-5x file-size inflation, dominated by xet chunks. Use main's + # `--local-dir gguf-cache` pattern to cache the flat .gguf only. + # Studio's /api/inference/load accepts either a HF repo (which + # uses HF_HOME) or an absolute file path; passing the absolute + # path keeps the test off HF_HOME entirely so the cache size + # tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images + # jobs still cover the gguf_variant resolution path. + # Q4_K_XL, not IQ3_XXS: at IQ3_XXS this model emits malformed + # tool calls that llama-server's peg-native parser rejects with a + # 500. Mac/Windows already use Q4_K_XL for the same reason. + GGUF_REPO: unsloth/Qwen3.5-2B-GGUF + GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf + STUDIO_PORT: '18889' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Reset auth + boot Studio (API-only, default tool policy) + # We deliberately use the API-only mode rather than + # `unsloth studio run` because the latter calls + # `set_tool_policy(...)` with a resolved bool: on loopback the + # default resolves to True, which forces every request through + # the server-side agentic loop and breaks the standard + # function-calling test below. API-only mode leaves + # tool_policy=None so each request's `enable_tools` field is + # honoured. + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, change password, load model + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CITool-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + GGUF_PATH="$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" + ls -lh "$GGUF_PATH" + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_PATH\",\"is_lora\":false,\"max_seq_length\":2048}" \ + | jq '{status, display_name}' + + - name: Tool calling, server-side tools, thinking on/off + env: + BASE_URL: http://127.0.0.1:18889 + run: | + python - <<'PY' + import json + import os + import time + import urllib.error + import urllib.request + + BASE = os.environ["BASE_URL"] + KEY = os.environ["API_KEY"] + SEED = 3407 + + def post(path, body, *, timeout = 240): + """Plain JSON POST. For requests that don't go through + the server-side agentic loop, the response is one JSON + object.""" + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{BASE}{path}", + data = data, + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) + + def post_sse(path, body, *, timeout = 600): + """POST a streaming request and accumulate the assistant + text deltas. The server-side agentic loop ALWAYS returns + SSE regardless of the request's `stream` field, so any + call with enable_tools=true must use this helper. + + Returns (content, raw_payloads): + content -- concatenated assistant delta.content + raw_payloads -- list of every raw "data: ..." event + payload (JSON strings). Callers asserting + that a server-side tool actually ran (and + not just that the model emitted some + text) should grep raw_payloads for tool + invocation markers / tool output, since + `delta.content` alone is not evidence + that the tool path executed. + """ + body = {**body, "stream": True} + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{BASE}{path}", + data = data, + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + parts = [] + events = [] + with urllib.request.urlopen(req, timeout = timeout) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data: "): + continue + payload = line[6:] + if payload == "[DONE]": + break + events.append(payload) + try: + chunk = json.loads(payload) + except json.JSONDecodeError: + continue + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + if delta.get("content"): + parts.append(delta["content"]) + return "".join(parts), events + + _STUDIO_TOOL_TYPES = { + "tool_start", "tool_end", "tool_use", "tool_result", + } + + def _tool_invoked(events): + """Structural check: True iff some SSE payload is a real + tool envelope (Studio tool_start/tool_end, Anthropic + tool_use/tool_result, OpenAI non-empty delta.tool_calls / + message.tool_calls / finish_reason='tool_calls' / + role:'tool' / function_call). tool_status is NOT + evidence: Studio emits empty tool_status events on + iteration boundaries even when no tool ran. + """ + for raw in events: + try: + ev = json.loads(raw) + except (json.JSONDecodeError, TypeError): + continue + if not isinstance(ev, dict): + continue + if ev.get("type") in _STUDIO_TOOL_TYPES: + return True + for choice in ev.get("choices", []) or []: + if not isinstance(choice, dict): + continue + if choice.get("finish_reason") == "tool_calls": + return True + for src_key in ("delta", "message"): + src = choice.get(src_key) or {} + if not isinstance(src, dict): + continue + tc = src.get("tool_calls") + if isinstance(tc, list) and tc: + return True + if src.get("function_call"): + return True + if src.get("role") == "tool": + return True + for item in ev.get("output", []) or []: + if isinstance(item, dict) and item.get("type") in { + "tool_call", "function_call", "tool_use", + }: + return True + content = ev.get("content") + if isinstance(content, list): + for blk in content: + if isinstance(blk, dict) and blk.get("type") in { + "tool_use", "tool_result", + }: + return True + return False + + def _tool_output_contains(events, *needles): + """True iff any tool_end.result / tool_result.content / + tool-role message content contains a needle. Inspects + the tool's own output, not the model's narration.""" + for raw in events: + try: + ev = json.loads(raw) + except (json.JSONDecodeError, TypeError): + continue + if not isinstance(ev, dict): + continue + if ev.get("type") == "tool_end": + result = ev.get("result") + if isinstance(result, str) and any(n in result for n in needles if n): + return True + if ev.get("type") == "tool_result": + content = ev.get("content") + if isinstance(content, str) and any(n in content for n in needles if n): + return True + if isinstance(content, list): + for blk in content: + if isinstance(blk, dict): + text = blk.get("text") or blk.get("content") + if isinstance(text, str) and any(n in text for n in needles if n): + return True + for choice in ev.get("choices", []) or []: + delta = (choice or {}).get("delta") or {} + msg = (choice or {}).get("message") or {} + for src in (delta, msg): + if src.get("role") == "tool": + content = src.get("content") or "" + if isinstance(content, str) and any(n in content for n in needles if n): + return True + return False + + # ── 1. Standard OpenAI function calling ────────────────────── + weather_tool = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + + status, data = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "What is the weather in Paris?"}], + "tools": [weather_tool], + "tool_choice": "required", + "stream": False, + "temperature": 0.0, + "seed": SEED, + "max_tokens": 120, + }) + assert status == 200, f"tool call status {status}: {data}" + choice = data["choices"][0] + assert choice["finish_reason"] == "tool_calls", f"finish_reason={choice['finish_reason']!r}" + tc = choice["message"]["tool_calls"][0] + assert tc["function"]["name"] == "get_weather" + args = json.loads(tc["function"]["arguments"]) + assert args.get("city"), f"missing city arg: {args}" + print(f"[tools] PASS function calling -> {tc['function']['name']}({args})") + + # T=0 = deterministic argmax in llama.cpp; T>0 lets seed + # rotation explore distinct trajectories on retry. + TOOL_PROBE_TEMP = 0.4 + + def _run_tool_probe(*, label, prompt, enabled, session, needles, + max_attempts = 4): + """Drive a server-side tool with retries. Hard FAIL if no + attempt has structural invocation evidence. WARN (not + FAIL) if invoked but no attempt produces the expected + literal in tool_end.result -- small-quant Qwen3.5-2B can + emit OpenAI tool_calls deltas without Studio's GGUF + agentic loop intercepting them, and that GGUF-vs-OpenAI + format mismatch is out of scope for #5642. + """ + attempts_log = [] + best = None + for attempt_i in range(max_attempts): + attempt_seed = SEED + attempt_i + content, events = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": prompt}], + "enable_tools": True, + "enabled_tools": enabled, + "session_id": f"{session}-att{attempt_i}", + "temperature": TOOL_PROBE_TEMP, + "seed": attempt_seed, + "max_tokens": 600, + }) + invoked = _tool_invoked(events) + produced = _tool_output_contains(events, *needles) + attempts_log.append({ + "attempt": attempt_i, "seed": attempt_seed, + "n_events": len(events), + "tool_invoked": invoked, "tool_output_contains": produced, + "content_len": len(content), + }) + if invoked and produced: + print(f"[tools] PASS {label} attempt {attempt_i}") + return content, events, attempts_log + if invoked and best is None: + best = (content, events) + print(f"[tools] retry {label} attempt {attempt_i}: invoked={invoked} output_ok={produced} events={len(events)}") + if best is not None: + print(f"[tools] WARN {label}: invoked but no tool_end.result match (small-quant flake). Attempts: {attempts_log}") + content, events = best + return content, events, attempts_log + raise AssertionError( + f"{label}: no structural tool-invocation evidence across " + f"{max_attempts} attempts. enable_tools may be silently " + f"ignored. Attempts: {attempts_log}" + ) + + # ── 2. Server-side python tool ─────────────────────────────── + content, events, _attempts = _run_tool_probe( + label = "python tool", + prompt = "What is 123 * 456? Use the python tool to compute it and tell me the number.", + enabled = ["python"], + session = "ci-tool-calling-py", + needles = ("56088", "56,088"), + ) + if "56088" in content or "56,088" in content: + print(f"[tools] python tool narration OK") + else: + print(f"[tools] python tool narration drifted -- content={content!r}") + + # ── 3. Server-side bash (terminal) tool ────────────────────── + content, events, _attempts = _run_tool_probe( + label = "bash/terminal tool", + prompt = "Use the terminal tool to run `echo hello-bash-tool` and tell me the exact output.", + enabled = ["terminal"], + session = "ci-tool-calling-bash", + needles = ("hello-bash-tool",), + ) + if "hello-bash-tool" in content: + print(f"[tools] bash/terminal narration OK") + else: + print(f"[tools] bash/terminal narration dropped literal -- content={content!r}") + + # ── 4. Server-side web_search tool ─────────────────────────── + # DuckDuckGo is flaky from CI runners and small Qwen3.5-2B + # may not actually search. Only assert that the SSE stream + # opens and yields any data; HTTP / parser failures already + # raise above. Tool-invocation strictness is relaxed here + # because (a) the search may legitimately return no results, + # and (b) DuckDuckGo upstream blocks GHA IP ranges often + # enough that requiring a tool_call marker would create + # red-herring failures from infra rather than from Studio. + try: + content, events = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], + "enable_tools": True, + "enabled_tools": ["web_search"], + "session_id": "ci-tool-calling-web", + "temperature": 0.0, + "seed": SEED, + "max_tokens": 400, + }) + print( + f"[tools] PASS web_search stream ({len(content)} chars in content, " + f"{len(events)} raw events)" + ) + except Exception as exc: + print(f"[tools] WARN web_search probe failed (non-blocking): {exc}") + + # ── 5. Thinking on / off ───────────────────────────────────── + # Studio strips think blocks from message.content for tools-mode + # responses, so we toggle plain chat (no enable_tools) and look + # at the surfaced reasoning_content / message.thinking field. + def thinking_call(enable): + status, data = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Briefly: is 17 prime?"}], + "stream": False, + "enable_thinking": enable, + "temperature": 0.0, + "seed": SEED, + "max_tokens": 300, + }) + assert status == 200 + msg = data["choices"][0]["message"] + # Studio surfaces thinking via reasoning_content (OpenAI + # extension). Fall back to inline markers for + # robustness across template versions. + raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "") + return raw + + on_text = thinking_call(True) + off_text = thinking_call(False) + had_think_on = ("" in on_text) or len(on_text) > 80 + had_think_off = ("" in off_text) and len(off_text) > 0 + assert had_think_on, ( + f"enable_thinking=True produced no thinking signal: {on_text!r}" + ) + # Off-mode should not contain the literal marker. + assert "" not in off_text, ( + f"enable_thinking=False but still present: {off_text!r}" + ) + print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)") + PY + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + ss -tln | grep ":${STUDIO_PORT}" || true + # Capture backend + llama-server logs so a 500 has a server-side traceback. + mkdir -p logs/server-logs + cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true + + - name: Upload logs + # Always upload so green runs are still reviewable. + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: tool-calling-log + path: | + logs/studio.log + logs/install.log + logs/server-logs/ + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job 3: JSON, images + # ───────────────────────────────────────────────────────────────────── + json-images: + name: JSON, images + runs-on: ubuntu-latest + timeout-minutes: 30 + env: + GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF + # UD-Q4_K_XL, not UD-IQ2_XXS: at 2-bit the temp-0 answer to the JSON + # step's capital-of-France probe flips with the host's SIMD kernels + # (GitHub runners deterministically answered France while other CPUs + # answer Paris; seeds do not rescue it, 1/5 Paris at temp 0.7). The + # Q4 quant answered Paris 13/13 across temps and seeds on the same + # runners, so the hard Paris assertion below stays reliable. + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: Qwen3-VL-2B-Instruct-UD-Q4_K_XL.gguf + MMPROJ_FILE: mmproj-F16.gguf + STUDIO_PORT: '18890' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj) + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 + + - name: Prime HF_HOME with the GGUF + mmproj + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj) + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Install OpenAI + Anthropic Python SDKs + run: pip install 'openai>=1.50' 'anthropic>=0.40' + + - name: Reset auth + boot Studio (API-only) + # See Job 2's comment: API-only mode keeps tool_policy=None so + # response_format requests aren't routed through the agentic + # tool loop. + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, change password, load model + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIJson-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + # Retry: llama-server startup can race process teardown after a + # failed attempt. Keep curl out of a pipe so HTTP failures are not + # masked by jq. + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 900 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP; response:" + cat /tmp/load.json || true + sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_vision}' /tmp/load.json + + - name: JSON schema decoding + image input + env: + BASE_URL: http://127.0.0.1:18890 + run: | + python - <<'PY' + import base64 + import json + import os + import time + import urllib.error + import urllib.request + from openai import OpenAI + from anthropic import Anthropic + + BASE = os.environ["BASE_URL"] + KEY = os.environ["API_KEY"] + SEED = 3407 + + def post(path, body, *, timeout = 240): + req = urllib.request.Request( + f"{BASE}{path}", + data = json.dumps(body).encode(), + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) + + # ── 1. response_format = json_object (JSON mode) ───────────── + # llama.cpp's HTTP server supports OpenAI-compatible JSON + # mode: `response_format: {"type": "json_object"}` constrains + # the model to emit syntactically-valid JSON. We use raw HTTP + # rather than the OpenAI SDK so that the field shape Studio + # forwards to llama-server is unambiguous (the SDK rewrites + # response_format depending on which variant it recognises). + # We deliberately do NOT pass a strict JSON schema -- on + # small Gemma-4 quants the GBNF-from-schema path occasionally + # produces empty output, and JSON mode is the surface we care + # about exposing through Studio. + status, data = post("/v1/chat/completions", { + "model": "default", + "messages": [ + {"role": "system", "content": 'Reply with a single JSON object of the form {"city": "...", "country": "..."}. Output ONLY the JSON, nothing else.'}, + {"role": "user", "content": "What is the capital of France?"}, + ], + "temperature": 0.0, + "max_tokens": 200, + "seed": SEED, + "stream": False, + "enable_thinking": False, + "response_format": {"type": "json_object"}, + }, timeout = 600) + assert status == 200, f"json status {status}: {data}" + content = (data["choices"][0]["message"].get("content") or "").strip() + # Some chat templates wrap JSON in ```json fences even in JSON + # mode -- strip those before parsing. + if content.startswith("```"): + content = content.split("```", 2)[1] + if content.startswith("json"): + content = content[4:] + content = content.strip("`\n ") + parsed = json.loads(content) + assert "paris" in str(parsed.get("city", "")).lower(), ( + f"city != Paris: {parsed}" + ) + print(f"[json] PASS json_object -> {parsed}") + + # ── 2. OpenAI image_url (data URI base64) ─────────────────── + # 64x64 solid-red PNG. stb_image (used by Studio's image + # normaliser at routes/inference.py:3410) rejects 4x4 or + # smaller PNGs as truncated, so we go up to 64x64 -- still + # tiny in token cost. The assertion is loose: any non-empty + # response from the vision path proves multimodal end-to-end + # wiring; small VL quants are weak at colour identification. + PNG_64X64_RED_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAYklEQVR4nO3PMQ0AIADAMEAI/k" + "UhBhEcDcmqYJtn7/GzpQNeNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA" + "1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaBdCJ0BmMJ25zMAAAAASUVORK5CYII=" + ) + data_uri = f"data:image/png;base64,{PNG_64X64_RED_B64}" + + client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY) + openai_resp = client.chat.completions.create( + model = "default", + temperature = 0.0, + max_tokens = 80, + seed = SEED, + messages = [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_uri}}, + {"type": "text", "text": "What colour dominates this image? Reply in one word."}, + ], + }], + ) + openai_text = (openai_resp.choices[0].message.content or "").lower() + print(f"[image/openai] reply: {openai_text!r}") + assert openai_text, "OpenAI image_url returned empty content" + # We do not strictly require 'red' -- some quants of small VL + # models are weak at colour names. Just require a non-empty + # answer; the vision path is the part under test. + print("[image/openai] PASS image_url accepted, non-empty response") + + # ── 3. Anthropic source/base64 image ──────────────────────── + # Two SDK quirks vs. Studio: base_url must NOT include /v1 + # (the SDK appends it itself; otherwise /v1/v1/messages -> 405), + # and Studio's auth is HTTPBearer-only so the SDK's default + # x-api-key header is ignored -- send Authorization: Bearer + # via default_headers. + anthropic = Anthropic( + base_url = BASE, + api_key = "unused", + default_headers = {"Authorization": f"Bearer {KEY}"}, + ) + a_msg = anthropic.messages.create( + model = "default", + max_tokens = 80, + temperature = 0.0, + extra_body = {"seed": SEED}, + messages = [{ + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": PNG_64X64_RED_B64, + }, + }, + {"type": "text", "text": "Describe this image briefly."}, + ], + }], + ) + a_text = "".join(b.text for b in a_msg.content if getattr(b, "type", None) == "text") + print(f"[image/anthropic] reply: {a_text!r}") + assert a_text, "Anthropic source/base64 returned empty content" + print("[image/anthropic] PASS source/base64 accepted, non-empty response") + PY + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + ss -tln | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + # Always upload so green runs are still reviewable. + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: json-images-log + path: | + logs/studio.log + logs/install.log + retention-days: 7 diff --git a/.github/workflows/studio-load-orchestrator-ci.yml b/.github/workflows/studio-load-orchestrator-ci.yml new file mode 100644 index 0000000..93d1a77 --- /dev/null +++ b/.github/workflows/studio-load-orchestrator-ci.yml @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Event-loop regression test for the Studio model-load orchestrator. +# Pins down issue #5642 (Win10 UI freeze on model load): the /load +# route calls LlamaCppBackend.detect_audio_type synchronously, blocking +# the FastAPI event loop on a chain of sync httpx.Client.post() probes. +# +# The suite stands up a stdlib fake llama-server + a tiny FastAPI app +# via uvicorn and asserts that detect_audio_type runs via +# asyncio.to_thread so concurrent /api/inference/load-progress polling +# stays responsive. CPU-only, no torch, no real llama.cpp binary, no +# GPU -- the matching cross-OS staging proof lives on +# danielhanchen/unsloth-staging-2 (Ubuntu / macOS / Windows all +# green at PR time). + +name: Studio load-orchestrator CI + +on: + pull_request: + paths: + - 'studio/backend/routes/inference.py' + - 'studio/backend/core/inference/llama_cpp.py' + - 'tests/studio/load_freeze/**' + - '.github/workflows/studio-load-orchestrator-ci.yml' + push: + branches: [main] + paths: + - 'studio/backend/routes/inference.py' + - 'studio/backend/core/inference/llama_cpp.py' + - 'tests/studio/load_freeze/**' + - '.github/workflows/studio-load-orchestrator-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install minimal deps (no torch, no unsloth) + # The test stubs `loggers` and `structlog`, imports + # core.inference.llama_cpp directly, and drives a small + # FastAPI app. Nothing here pulls torch or any GPU code, + # so the entire job typically completes in well under 60 s. + run: | + python -m pip install --upgrade pip + python -m pip install \ + 'pytest>=8' \ + 'httpx>=0.27,<1' \ + 'fastapi>=0.110,<1' \ + 'uvicorn>=0.30,<1' \ + 'anyio>=4' + - name: Run load-orchestrator tests + run: python -m pytest -v --tb=short tests/studio/load_freeze/ diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml new file mode 100644 index 0000000..617ce18 --- /dev/null +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Mac counterpart to studio-api-smoke.yml. Same tests/studio/ +# studio_api_smoke.py exercise (CORS hardening, auth state machine, +# JWT expiry, API key lifecycle, /v1/models / /v1/embeddings / +# /v1/responses, endpoint-by-endpoint auth audit) but on a real +# Apple Silicon (macos-14, M1) runner. Drops the apt-get block; +# GitHub-hosted macos-14 ships curl + jq. + +name: Mac Studio API CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.sh' + - 'pyproject.toml' + - 'tests/studio/**' + - '.github/workflows/studio-mac-api-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + api-smoke: + name: Studio API & Auth Tests + runs-on: macos-14 + timeout-minutes: 25 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18895' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Install pyjwt for the JWT-expiry forge test + run: pip install 'pyjwt>=2.6' + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + + - name: Pass bootstrap password + rotated targets to the test + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="ApiSmoke-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + NEW2="ApiSmoke-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "::add-mask::$NEW2" + echo "STUDIO_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" + echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" + + - name: Run Studio API & Auth tests + env: + BASE_URL: http://127.0.0.1:18895 + STUDIO_AUTH_DIR: /Users/runner/.unsloth/studio/auth + run: python tests/studio/studio_api_smoke.py + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + + - name: Upload API smoke logs + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: mac-studio-api-smoke-log + path: | + logs/install.log + logs/studio.log + retention-days: 7 diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml new file mode 100644 index 0000000..03c0a85 --- /dev/null +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -0,0 +1,1086 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Three end-to-end smoke jobs that boot a freshly-installed Studio and +# exercise the surfaces real users hit through the OpenAI / Anthropic +# SDKs and curl. Each job picks the smallest model that exercises the +# behaviour under test, primes a model cache via actions/cache, and +# shares the install.sh --local --no-torch bootstrap. +# +# 1. OpenAI, Anthropic API tests +# gemma-3-270m-it UD-Q4_K_XL (~254 MiB). +# Password rotation via /api/auth/change-password (old fails, +# new works), then OpenAI + Anthropic Python SDKs against /v1/* +# with temperature=0 and a fixed seed. Asserts the four-turn +# conversation is deterministic across two runs. +# +# 2. Tool calling Tests +# Qwen3.5-2B UD-IQ3_XXS (~890 MiB). OpenAI function calling, +# server-side tools (python, terminal, web_search) via +# enable_tools / enabled_tools, and enable_thinking on/off. +# +# 3. JSON, images +# gemma-4-E2B-it UD-IQ3_XXS (~2.4 GiB) + mmproj-F16 (~986 MiB). +# response_format JSON-schema decoding and OpenAI image_url +# (data URI) plus Anthropic source/base64 image inputs. +# +# All three jobs run in parallel. Total wall time is dominated by job 3 +# on a cold cache; warm cache cuts that to ~3 min. + +name: Mac Studio GGUF CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.sh' + - 'pyproject.toml' + - '.github/workflows/studio-mac-inference-smoke.yml' + push: + branches: [main, pip] + # Manual trigger for pre-warming model caches on main, or re-running + # against an arbitrary branch without pushing a no-op commit. + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + # ───────────────────────────────────────────────────────────────────── + # Job 1: OpenAI, Anthropic API tests + # ───────────────────────────────────────────────────────────────────── + openai-anthropic: + name: OpenAI, Anthropic API tests + runs-on: macos-14 + timeout-minutes: 25 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18888' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + # Save partial caches on cancel/timeout -- hf download resumes by + # content hash. `outcome != skipped` keeps cache-hit a no-op. + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome != 'skipped' && hashFiles('hf-cache/**/*.gguf') != '' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Install OpenAI + Anthropic Python SDKs + run: pip install 'openai>=1.50' 'anthropic>=0.40' + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json + exit 0 + fi + sleep 1 + done + echo "Studio did not become healthy in 180s" + tail -200 logs/studio.log + exit 1 + + - name: Password rotation (old must fail, new must work) + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIRotated-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + # 1. Login with the bootstrap password. + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + [ -n "$OLD_TOKEN" ] && [ "$OLD_TOKEN" != "null" ] || { echo "bootstrap login failed"; exit 1; } + # 2. Rotate to a fresh random password. + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + # 3. Old password must now be rejected (HTTP 401). + OLD_STATUS=$(curl -s -o /dev/null -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}") + if [ "$OLD_STATUS" != "401" ]; then + echo "::error::Login with old password returned $OLD_STATUS, expected 401" + exit 1 + fi + # 4. New password must succeed; capture the JWT for downstream steps. + NEW_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + [ -n "$NEW_TOKEN" ] && [ "$NEW_TOKEN" != "null" ] || { echo "new login failed"; exit 1; } + echo "TOKEN=$NEW_TOKEN" >> "$GITHUB_ENV" + echo "password rotation OK (old=401, new=200)" + + - name: Load the GGUF (HF repo + variant, served from HF_HOME cache) + run: | + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}" \ + | jq '{status, display_name, is_gguf, context_length}' + + - name: Multi-turn determinism via OpenAI + Anthropic SDKs + env: + BASE_URL: http://127.0.0.1:18888 + run: | + python - <<'PY' + import json + import os + from openai import OpenAI + from anthropic import Anthropic + + BASE = os.environ["BASE_URL"] + KEY = os.environ["TOKEN"] # JWT also accepted as Bearer on /v1/* + SEED = 3407 + + # Four-turn conversation: the second and fourth turns can only be + # answered correctly if the model sees the prior turns, so this + # also exercises the conversation-history wiring. + PROMPTS = [ + "What is 1+1?", + "What did I ask before?", + "What is the capital of France?", + "Repeat the city name", + ] + + def run_openai(): + client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY) + history, replies = [], [] + for prompt in PROMPTS: + history.append({"role": "user", "content": prompt}) + resp = client.chat.completions.create( + model = "default", + messages = history, + temperature = 0.0, + max_tokens = 80, + seed = SEED, + extra_body = {"enable_thinking": False}, + ) + text = resp.choices[0].message.content or "" + replies.append(text) + history.append({"role": "assistant", "content": text}) + return replies + + def run_anthropic(): + # Two SDK quirks vs. Studio: + # 1. base_url must NOT include /v1 -- the SDK appends + # /v1/messages itself; otherwise the request hits + # /v1/v1/messages and 405s. + # 2. The SDK sends `x-api-key` by default, but Studio's + # auth layer is HTTPBearer-only. Override via + # default_headers so Authorization: Bearer ... is + # sent instead. + client = Anthropic( + base_url = BASE, + api_key = "unused", + default_headers = {"Authorization": f"Bearer {KEY}"}, + ) + history, replies = [], [] + for prompt in PROMPTS: + history.append({"role": "user", "content": prompt}) + msg = client.messages.create( + model = "default", + max_tokens = 80, + messages = history, + temperature = 0.0, + extra_body = {"seed": SEED, "enable_thinking": False}, + ) + text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text") + replies.append(text) + history.append({"role": "assistant", "content": text}) + return replies + + for label, runner in (("openai", run_openai), ("anthropic", run_anthropic)): + first = runner() + second = runner() + for i, (a, b) in enumerate(zip(first, second), start = 1): + print(f"[{label} turn {i}] {a!r}") + assert a, f"{label}: empty turn {i} response" + # Compare on stripped content: llama-server can vary + # trailing whitespace (specifically a final '\n') between + # otherwise-identical greedy runs depending on the + # batch-flush boundary at which the stream is closed. The + # generated tokens are identical; only the trailing + # whitespace differs. Keep the raw repr in the failure + # message so a real divergence is still legible. + assert a.strip() == b.strip(), ( + f"{label} non-deterministic at turn {i} with temperature=0.0:\n" + f" run1: {a!r}\n run2: {b!r}" + ) + # Sanity: turn-2 reply should mention the earlier question, and + # turn-4 reply should mention Paris (model echoes the city it + # produced for turn 3). Lower-cased substring checks keep the + # assertion robust to formatting jitter. + joined = " ".join(first).lower() + assert "1" in first[0], f"{label}: turn-1 answer should contain '1', got {first[0]!r}" + assert "paris" in joined, f"{label}: expected 'paris' somewhere in the four-turn transcript: {first}" + print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded") + PY + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + ss -tln | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + # Always upload so green runs are still reviewable. + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: openai-anthropic-log + path: | + logs/studio.log + logs/install.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job 2: Tool calling Tests + # ───────────────────────────────────────────────────────────────────── + tool-calling: + name: Tool calling Tests + runs-on: macos-14 + timeout-minutes: 25 + env: + # Tool calling is the highest-volume GGUF in this workflow + # (Qwen3.5-2B at Q4_K_XL = ~1.28 GiB on Mac, where IQ3_XXS + # collapses for tool-call grammar under Metal at temperature=0). + # Caching HF_HOME stores xet chunks + blobs + snapshots = ~4.6 + # GiB compressed -- 3.6x file-size inflation. Use main's + # `--local-dir gguf-cache` pattern to cache the flat .gguf only. + # The OpenAI/Anth and JSON+images jobs still cover the + # gguf_variant resolution path. + GGUF_REPO: unsloth/Qwen3.5-2B-GGUF + GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf + STUDIO_PORT: '18898' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + # Save partial caches on cancel; next run resumes via content hash. + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome != 'skipped' && hashFiles('gguf-cache/**/*.gguf') != '' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Reset auth + boot Studio (API-only, default tool policy) + # We deliberately use the API-only mode rather than + # `unsloth studio run` because the latter calls + # `set_tool_policy(...)` with a resolved bool: on loopback the + # default resolves to True, which forces every request through + # the server-side agentic loop and breaks the standard + # function-calling test below. API-only mode leaves + # tool_policy=None so each request's `enable_tools` field is + # honoured. + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, change password, load model + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CITool-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + GGUF_PATH="$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" + ls -lh "$GGUF_PATH" + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_PATH\",\"is_lora\":false,\"max_seq_length\":2048}" \ + | jq '{status, display_name}' + + - name: Tool calling, server-side tools, thinking on/off + env: + BASE_URL: http://127.0.0.1:18898 + run: | + python - <<'PY' + import json + import os + import time + import urllib.error + import urllib.request + + BASE = os.environ["BASE_URL"] + KEY = os.environ["API_KEY"] + SEED = 3407 + + def post(path, body, *, timeout = 240): + """Plain JSON POST. For requests that don't go through + the server-side agentic loop, the response is one JSON + object.""" + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{BASE}{path}", + data = data, + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) + + def post_sse(path, body, *, timeout = 600): + """POST a streaming request and accumulate the assistant + text deltas. The server-side agentic loop ALWAYS returns + SSE regardless of the request's `stream` field, so any + call with enable_tools=true must use this helper.""" + body = {**body, "stream": True} + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{BASE}{path}", + data = data, + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + parts = [] + with urllib.request.urlopen(req, timeout = timeout) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data: "): + continue + payload = line[6:] + if payload == "[DONE]": + break + try: + chunk = json.loads(payload) + except json.JSONDecodeError: + continue + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + if delta.get("content"): + parts.append(delta["content"]) + return "".join(parts) + + # ── 1. Standard OpenAI function calling ────────────────────── + weather_tool = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + + # Mac Metal at temperature=0 is pathological for these small + # quants (Qwen3.5-2B emits ',,,,,,...' or 'The The The...'), + # gemma-4-E2B emits '' tokens). The Linux CPU + # backend hides the issue. Use a small non-zero temperature + # with a fixed seed so we stay deterministic but escape the + # degenerate sampling trap. + TEMP = 0.2 + + status, data = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "What is the weather in Paris?"}], + "tools": [weather_tool], + "tool_choice": "required", + "stream": False, + "temperature": TEMP, + "seed": SEED, + # tool_choice='required' constrains the grammar so the + # model emits a tool_call quickly when it works at all; + # 128 tokens is enough for `{"city":"Paris"}` plus the + # JSON envelope. + "max_tokens": 128, + }, timeout = 180) + assert status == 200, f"tool call status {status}: {data}" + choice = data["choices"][0] + tool_calls = (choice.get("message") or {}).get("tool_calls") or [] + # Studio's contract: when tool_choice='required', llama.cpp's + # grammar should force a tool_calls payload. On Mac that + # contract is sometimes broken by the underlying quant; the + # PASS path is "tool_calls present + correct schema", the + # WARN path documents Studio still returned 200 with a + # well-formed choices[] envelope. + if tool_calls: + tc = tool_calls[0] + assert tc["function"]["name"] == "get_weather", ( + f"unexpected tool name: {tc['function']['name']!r}" + ) + args = json.loads(tc["function"]["arguments"]) + assert args.get("city"), f"missing city arg: {args}" + print(f"[tools] PASS function calling -> {tc['function']['name']}({args}) finish={choice.get('finish_reason')!r}") + else: + # Infrastructure path is correct; model output drifted. + print( + f"[tools] WARN function calling: no tool_calls (finish_reason=" + f"{choice.get('finish_reason')!r}); HTTP path OK, this is a " + f"Mac Metal quant degeneracy." + ) + + # ── 2. Server-side python tool ─────────────────────────────── + # 123 * 456 = 56088. The agentic loop streams SSE; we + # accumulate the assistant text and look for the answer. On + # Mac the model often loses the tool calling contract before + # producing the answer; accept either the answer OR a + # non-empty SSE stream as proof the path completes. + # macos-14 free runner is ~10 tok/s on Qwen3.5-2B Q4_K_XL; + # cap max_tokens tightly so each SSE round stays under ~30s + # even when the model stalls in a degenerate output state. + content = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], + "enable_tools": True, + "enabled_tools": ["python"], + "session_id": "ci-tool-calling-py", + "temperature": TEMP, + "seed": SEED, + "max_tokens": 128, + }, timeout = 180) + if "56088" in content or "56,088" in content: + print(f"[tools] PASS python tool ({len(content)} chars, found 56088)") + else: + # Empty stream is a known Mac-quant degeneracy too; log + # but do not fail. + print( + f"[tools] WARN python tool: SSE OK ({len(content)} chars) but " + f"model didn't return 56088 -- Mac quant drift" + ) + + # NOTE: the dedicated "Server-side bash (terminal) tool" axis + # was dropped in favour of the python axis above. Both share + # the SAME server-side agentic loop wiring (only the registry + # entry differs); the python axis is the canonical proof. On + # macos-14 the duplicated SSE round was the dominant cost in + # this step, so collapsing the two saves ~30-60 s wallclock + # without losing distinct coverage. + + # ── 3. Server-side web_search tool ─────────────────────────── + # DuckDuckGo is flaky from CI runners and small Qwen3.5-2B + # may not actually search. Only assert that the SSE stream + # opens and yields any data; HTTP / parser failures already + # raise above. + try: + content = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], + "enable_tools": True, + "enabled_tools": ["web_search"], + "session_id": "ci-tool-calling-web", + "temperature": TEMP, + "seed": SEED, + "max_tokens": 96, + }, timeout = 180) + print(f"[tools] PASS web_search stream ({len(content)} chars)") + except Exception as exc: + print(f"[tools] WARN web_search probe failed (non-blocking): {exc}") + + # ── 4. Thinking on / off ───────────────────────────────────── + # Studio strips think blocks from message.content for tools-mode + # responses, so we toggle plain chat (no enable_tools) and look + # at the surfaced reasoning_content / message.thinking field. + def thinking_call(enable): + status, data = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Briefly: is 17 prime?"}], + "stream": False, + "enable_thinking": enable, + "temperature": TEMP, + "seed": SEED, + # 80 tokens lands within the 25-minute job timeout + # on the macos-14 free runner. 17 is small; this is + # plenty of room for either "Yes" + brief reasoning + # or a degenerate empty completion. + "max_tokens": 80, + }, timeout = 180) + assert status == 200 + msg = data["choices"][0]["message"] + # Studio surfaces thinking via reasoning_content (OpenAI + # extension). Fall back to inline markers for + # robustness across template versions. + raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "") + return raw + + on_text = thinking_call(True) + off_text = thinking_call(False) + # Mac quant drift: the model may produce empty / degenerate + # output regardless of enable_thinking. Assert ONLY that the + # endpoint returned 200 (already enforced inside thinking_call) + # and that toggling the flag doesn't surface a hard + # marker when off. + had_think_on = ("" in on_text) or len(on_text) > 80 + if not had_think_on: + print( + f"[tools] WARN enable_thinking=True produced no thinking signal: " + f"{on_text[:200]!r} -- Mac quant drift" + ) + # Off-mode should not contain the literal marker. + assert "" not in off_text, ( + f"enable_thinking=False but still present: {off_text!r}" + ) + print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)") + PY + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + ss -tln | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + # Always upload so green runs are still reviewable. + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: tool-calling-log + path: | + logs/studio.log + logs/install.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job 3: JSON, images + # ───────────────────────────────────────────────────────────────────── + json-images: + name: JSON, images + runs-on: macos-14 + timeout-minutes: 30 + env: + GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF + # Linux smoke uses UD-IQ3_XXS, but on Mac Metal that gemma-4 + # quant emits sentinel tokens () for any prompt at + # temperature=0 -- inference path is fine, the quant itself is + # broken on Metal. UD-Q4_K_XL is the smallest published variant + # that generates real text on M1. + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-4-E2B-it-UD-Q4_K_XL.gguf + MMPROJ_FILE: mmproj-F16.gguf + STUDIO_PORT: '18899' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + # Cache flat .gguf + mmproj (Job 2's pattern). HF_HOME inflates + # ~3.6x via xet/blobs/snapshots, which made macOS saves never land. + # mmproj is auto-detected as a sibling via detect_mmproj_file + # (studio/backend/utils/models/model_config.py). + - name: Restore GGUF + mmproj files + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2 + + - name: Verify cache contains BOTH gguf + mmproj + id: verify-cache + if: steps.cache-gguf.outputs.cache-hit == 'true' + run: | + if [[ -f "gguf-cache/$GGUF_FILE" && -f "gguf-cache/$MMPROJ_FILE" ]]; then + echo "ok=true" >> "$GITHUB_OUTPUT" + else + echo "Partial cache hit -- forcing re-download." + echo "ok=false" >> "$GITHUB_OUTPUT" + fi + + - name: Download GGUF + mmproj if cache miss or partial + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.verify-cache.outputs.ok != 'true' + # Authenticated + parallel: shared macos-14 NAT egress stalls + # multi-GB anonymous downloads. + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache & + MODEL_PID=$! + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" gguf-cache & + MMPROJ_PID=$! + wait "$MODEL_PID" + wait "$MMPROJ_PID" + # Fail loud on a partial download instead of in the next step. + ls -lh "gguf-cache/$GGUF_FILE" "gguf-cache/$MMPROJ_FILE" + + # Save partial caches on cancel. hashFiles guard avoids a hard + # save failure when the download step exits with no files. The + # additional mmproj-presence check stops a partial save from + # poisoning the cache for the next run. + - name: Save GGUF + mmproj files + if: always() && steps.download-gguf.outcome != 'skipped' && hashFiles('gguf-cache/**/*.gguf') != '' && hashFiles(format('gguf-cache/{0}', env.MMPROJ_FILE)) != '' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Install OpenAI + Anthropic Python SDKs + run: pip install 'openai>=1.50' 'anthropic>=0.40' + + - name: Reset auth + boot Studio (API-only) + # See Job 2's comment: API-only mode keeps tool_policy=None so + # response_format requests aren't routed through the agentic + # tool loop. + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, change password, load model + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIJson-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + # Load via local file path; mmproj sibling auto-detected by + # detect_mmproj_file (model_config.py). gguf_variant omitted + # -- it routes through _find_local_gguf_by_variant which + # expects a directory, not a file path. + GGUF_PATH="$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" + MMPROJ_PATH="$GITHUB_WORKSPACE/gguf-cache/${MMPROJ_FILE}" + ls -lh "$GGUF_PATH" "$MMPROJ_PATH" + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 900 \ + -d "{\"model_path\":\"$GGUF_PATH\",\"is_lora\":false,\"max_seq_length\":2048}" \ + | jq '{status, display_name, is_vision}' + + - name: JSON schema decoding + image input + env: + BASE_URL: http://127.0.0.1:18899 + run: | + python - <<'PY' + import base64 + import json + import os + import time + import urllib.error + import urllib.request + from openai import OpenAI + from anthropic import Anthropic + + BASE = os.environ["BASE_URL"] + KEY = os.environ["API_KEY"] + SEED = 3407 + # Mac Metal degenerates these gemma-4 quants at temperature=0 + # (any prompt yields '...' padding tokens). Use a + # small non-zero temperature with the same seed so we stay + # deterministic-enough but escape the trap. + TEMP = 0.2 + + def post(path, body, *, timeout = 240): + req = urllib.request.Request( + f"{BASE}{path}", + data = json.dumps(body).encode(), + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) + + # ── 1. response_format = json_object (JSON mode) ───────────── + # llama.cpp's HTTP server supports OpenAI-compatible JSON + # mode: `response_format: {"type": "json_object"}` constrains + # the model to emit syntactically-valid JSON. We use raw HTTP + # rather than the OpenAI SDK so that the field shape Studio + # forwards to llama-server is unambiguous (the SDK rewrites + # response_format depending on which variant it recognises). + # We deliberately do NOT pass a strict JSON schema -- on + # small Gemma-4 quants the GBNF-from-schema path occasionally + # produces empty output, and JSON mode is the surface we care + # about exposing through Studio. + status, data = post("/v1/chat/completions", { + "model": "default", + "messages": [ + {"role": "system", "content": 'Reply with a single JSON object of the form {"city": "...", "country": "..."}. Output ONLY the JSON, nothing else.'}, + {"role": "user", "content": "What is the capital of France?"}, + ], + "temperature": TEMP, + # Trimmed for Mac runner timeout budget; json_object + # grammar terminates quickly when working. + "max_tokens": 200, + "seed": SEED, + "stream": False, + "enable_thinking": False, + "response_format": {"type": "json_object"}, + }, timeout = 240) + assert status == 200, f"json status {status}: {data}" + # Verify the response envelope shape -- this is what we + # actually want to exercise on Mac. The model output quality + # downstream of this is a Mac-Metal-quant artefact. + assert ( + isinstance(data.get("choices"), list) + and data["choices"] + and "message" in data["choices"][0] + ), f"json response envelope malformed: {data}" + content = (data["choices"][0]["message"].get("content") or "").strip() + print(f"[json] raw json_object content: {content!r}") + # Some chat templates wrap JSON in ```json fences even in JSON + # mode -- strip those before parsing. + if content.startswith("```"): + content = content.split("```", 2)[1] + if content.startswith("json"): + content = content[4:] + content = content.strip("`\n ") + if content: + try: + parsed = json.loads(content) + if "paris" in str(parsed.get("city", "")).lower(): + print(f"[json] PASS json_object -> {parsed}") + else: + print(f"[json] WARN json_object decoded but city!=Paris: {parsed}") + except json.JSONDecodeError as exc: + print(f"[json] WARN json_object content not parseable ({exc}); content={content!r}") + else: + print("[json] WARN json_object produced empty content on this Mac quant") + # Cross-check: same prompt without response_format. We care + # that the inference path stays healthy (status 200 + envelope + # shape OK); model output quality is a separate concern. + status2, data2 = post("/v1/chat/completions", { + "model": "default", + "messages": [{"role": "user", "content": "What is the capital of France? Answer with one word."}], + "temperature": TEMP, + # 1-word answer doesn't need 400 tokens; trim so a + # degenerate streaming model doesn't burn through the + # job's wallclock budget. + "max_tokens": 150, + "seed": SEED, + "stream": False, + "enable_thinking": False, + }, timeout = 240) + assert status2 == 200, f"plain status {status2}: {data2}" + plain = (data2["choices"][0]["message"].get("content") or "").lower() + print(f"[json] plain capital-of-france reply: {plain!r}") + if "paris" in plain: + print("[json] PASS plain inference path (paris mentioned)") + else: + print( + f"[json] WARN plain inference returned no 'paris' -- Mac quant " + f"degeneracy. HTTP path validated separately above." + ) + + # ── 2. OpenAI image_url (data URI base64) ─────────────────── + # 64x64 solid-red PNG. stb_image (used by Studio's image + # normaliser at routes/inference.py:3410) rejects 4x4 or + # smaller PNGs as truncated, so we go up to 64x64 -- still + # tiny in token cost. The assertion is loose: any non-empty + # response from the vision path proves multimodal end-to-end + # wiring; small VL quants are weak at colour identification. + PNG_64X64_RED_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAYklEQVR4nO3PMQ0AIADAMEAI/k" + "UhBhEcDcmqYJtn7/GzpQNeNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA" + "1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaBdCJ0BmMJ25zMAAAAASUVORK5CYII=" + ) + data_uri = f"data:image/png;base64,{PNG_64X64_RED_B64}" + + # The Mac prebuilt llama.cpp server has a known crash when + # processing image inputs alongside the gemma-4-E2B mmproj + # (server disconnects mid-completion). This is upstream + # llama.cpp behaviour, not Studio. Wrap both SDK calls in + # try/except so an upstream crash registers as a WARN rather + # than failing the whole job. Studio's contract (OpenAI/ + # Anthropic image fields are accepted and forwarded) is + # validated by the request body Studio constructs, not by + # whether llama.cpp can decode it on Mac Metal. + client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY) + try: + openai_resp = client.chat.completions.create( + model = "default", + temperature = TEMP, + max_tokens = 80, + seed = SEED, + messages = [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_uri}}, + {"type": "text", "text": "What colour dominates this image? Reply in one word."}, + ], + }], + ) + openai_text = (openai_resp.choices[0].message.content or "").lower() + print(f"[image/openai] reply: {openai_text!r}") + if openai_text: + print("[image/openai] PASS image_url accepted, non-empty response") + else: + print("[image/openai] WARN image_url accepted but empty content -- Mac quant drift") + except Exception as exc: + print( + f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: " + f"{exc}. Likely upstream llama.cpp Mac+vision crash, NOT a Studio " + f"regression. Studio successfully forwarded the request." + ) + + # ── 3. Anthropic source/base64 image ──────────────────────── + # Two SDK quirks vs. Studio: base_url must NOT include /v1 + # (the SDK appends it itself; otherwise /v1/v1/messages -> 405), + # and Studio's auth is HTTPBearer-only so the SDK's default + # x-api-key header is ignored -- send Authorization: Bearer + # via default_headers. + anthropic = Anthropic( + base_url = BASE, + api_key = "unused", + default_headers = {"Authorization": f"Bearer {KEY}"}, + ) + try: + a_msg = anthropic.messages.create( + model = "default", + max_tokens = 80, + temperature = TEMP, + extra_body = {"seed": SEED}, + messages = [{ + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": PNG_64X64_RED_B64, + }, + }, + {"type": "text", "text": "Describe this image briefly."}, + ], + }], + ) + a_text = "".join(b.text for b in a_msg.content if getattr(b, "type", None) == "text") + print(f"[image/anthropic] reply: {a_text!r}") + if a_text: + print("[image/anthropic] PASS source/base64 accepted, non-empty response") + else: + print("[image/anthropic] WARN source/base64 accepted but empty content -- Mac quant drift") + except Exception as exc: + print( + f"[image/anthropic] WARN anthropic image SDK call raised: " + f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp Mac+vision " + f"crash, NOT a Studio regression." + ) + PY + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + ss -tln | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + # Always upload so green runs are still reviewable. + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: json-images-log + path: | + logs/studio.log + logs/install.log + retention-days: 7 diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml new file mode 100644 index 0000000..362305c --- /dev/null +++ b/.github/workflows/studio-mac-install-matrix.yml @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Proves Studio's llama.cpp install loads on every supported macOS. The heavy +# app smokes stay single-OS; this matrix covers the OS-version dimension cheaply +# (install.sh + binary-load assert). Regression guard for the macOS-version +# selection in studio/install_llama_prebuilt.py. + +name: Mac Studio Install Matrix CI + +on: + pull_request: + paths: + - 'studio/install_llama_prebuilt.py' + - 'studio/setup.sh' + - 'install.sh' + - '.github/scripts/assert-llama-loads.sh' + - '.github/workflows/studio-mac-install-matrix.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + install-load: + name: Install + load (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + timeout-minutes: 25 + continue-on-error: ${{ matrix.experimental }} + strategy: + fail-fast: false + matrix: + include: + - os: macos-14 # Apple Silicon, macOS 14 Sonoma + experimental: false + - os: macos-15 # Apple Silicon, macOS 15 Sequoia + experimental: false + - os: macos-26 # Apple Silicon, macOS 26 Tahoe + experimental: false + - os: macos-15-intel # Intel x86_64, macOS 15 (informational) + experimental: true + - os: macos-26-intel # Intel x86_64, macOS 26 (last Intel macOS) + experimental: true + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Upload install log + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: mac-install-matrix-${{ matrix.os }}-log + path: logs/install.log + retention-days: 7 diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml new file mode 100644 index 0000000..20ca247 --- /dev/null +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Mac counterpart to studio-ui-smoke.yml. Same Playwright + Chromium +# end-to-end chat UI flow, but on macos-14 (M1) so we catch +# Mac-specific frontend / backend wiring regressions that the Linux +# job would miss (e.g. the Mac Tauri shell loading the same React +# bundle, or the Mac llama.cpp prebuilt's HTTP layer behaving +# differently from the Linux build). + +name: Mac Studio UI CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.sh' + - 'pyproject.toml' + - 'tests/studio/**' + - '.github/workflows/studio-mac-ui-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + ui-smoke: + name: Chat UI Tests + runs-on: macos-14 + timeout-minutes: 35 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18896' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Install Playwright + Chromium + # No --with-deps on Mac: that flag installs Linux apt packages. + # GitHub-hosted macos-14 ships the system frameworks Chromium + # needs already. + # Pinned <1.58 because all 1.55-1.58 drivers ship Node 24 on + # macos-14 and intermittently hit 'SyntaxError: Unexpected end + # of JSON input' in pipeTransport.js. Run 25491698868 showed + # the crash hitting 100% of three retry attempts -- not a + # rare race but a hard reproduction. Belt-and-suspenders fix: + # the test scripts pass --single-process to Chromium (see + # tests/studio/playwright_chat_ui.py) AND we patch + # pipeTransport.js below to swallow JSON parse errors instead + # of crashing the driver Node process. Both together let the + # in-script retry recover from any residual flakes. + run: | + pip install 'playwright>=1.55,<1.58' + python -m playwright install chromium + + - name: Patch Playwright pipeTransport.js to tolerate malformed JSON + # In Playwright 1.55-1.58, pipeTransport.js does + # `JSON.parse(message)` with no try/catch; when Chromium dies + # mid-write the partial buffer crashes the driver Node + # process and the test script exits with 'Connection closed + # while reading from the driver'. Newer Playwright versions + # added a try/catch upstream. Backport that here. + run: | + python - <<'PY' + import os, re, sys + import playwright + driver_dir = os.path.join(os.path.dirname(playwright.__file__), "driver", "package", "lib", "server") + path = os.path.join(driver_dir, "pipeTransport.js") + src = open(path).read() + # Wrap both `this.onmessage.call(null, JSON.parse(...))` sites in try/catch. + patched = re.sub( + r"this\.onmessage\.call\(null, JSON\.parse\((message2?)\)\);", + r"try { this.onmessage.call(null, JSON.parse(\1)); } " + r"catch (e) { /* swallow malformed JSON from a crashing browser */ }", + src, + ) + if patched == src: + # Already patched, or upstream changed -- either way, don't fail the build. + print(f"pipeTransport.js: no JSON.parse calls matched at {path}; skipping.") + else: + open(path, "w").write(patched) + print(f"pipeTransport.js: patched JSON.parse calls in {path}") + PY + + - name: Reset auth + boot Studio + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + + - name: Pass bootstrap password to the Playwright step + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + NEW2="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "::add-mask::$NEW2" + echo "STUDIO_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" + echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" + + - name: Drive the chat UI with Playwright + env: + BASE_URL: http://127.0.0.1:18896 + PW_ART_DIR: logs/playwright + STUDIO_UI_STRICT: '1' + # macos-14 free runner is 3 vCPU / 7 GB / no Metal-accel + # available to llama.cpp from CI; gemma-3-270m turn latency + # has been observed to crowd the 180s default. Triple it. + STUDIO_UI_TURN_TIMEOUT_MS: '540000' + # Retry up to 3 times to absorb known macos-14 free-runner + # flakes: (1) Playwright Node 24 pipeTransport.js 'Unexpected + # end of JSON input' crash when the Chromium browser process + # dies mid-test, (2) Chromium net::ERR_NO_BUFFER_SPACE when the + # runner's kernel briefly runs out of socket buffers, and (3) a + # goto 'interrupted by another navigation' when the SPA auth + # guard redirects mid-navigation. The retry FULLY resets Studio + # (kill, reset-password, reboot, wait /api/health, re-export + # bootstrap pw) before re-running the script. A real test failure + # (assertion / timeout) does NOT match any pattern so it bypasses + # retry and surfaces immediately. + run: | + mkdir -p logs/playwright + attempt=1 + max_attempts=3 + while : ; do + set +e + python tests/studio/playwright_chat_ui.py 2>&1 | tee logs/playwright_attempt_${attempt}.log + rc=${PIPESTATUS[0]} + set -e + if [ "$rc" -eq 0 ]; then + break + fi + if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \ + || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log \ + || grep -q "interrupted by another navigation" logs/playwright_attempt_${attempt}.log; } \ + && [ "$attempt" -lt "$max_attempts" ]; then + echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + unsloth studio reset-password + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > "logs/studio_retry_${attempt}.log" 2>&1 & + STUDIO_PID=$! + echo "STUDIO_PID=$STUDIO_PID" >> "$GITHUB_ENV" + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json \ + && jq -e '.status == "healthy"' /tmp/health.json >/dev/null; then + break + fi + sleep 1 + done + STUDIO_OLD_PW=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + STUDIO_NEW_PW="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + STUDIO_NEW2_PW="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$STUDIO_OLD_PW" + echo "::add-mask::$STUDIO_NEW_PW" + echo "::add-mask::$STUDIO_NEW2_PW" + export STUDIO_OLD_PW STUDIO_NEW_PW STUDIO_NEW2_PW + attempt=$((attempt + 1)) + sleep 3 + continue + fi + exit "$rc" + done + + - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders) + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + + - name: Reset auth + boot Studio for extra UI tests (port 18897) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \ + > logs/studio_extra.log 2>&1 & + echo "STUDIO_EXTRA_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18897 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18897/api/health" > /tmp/health2.json; then + jq -e '.status == "healthy"' /tmp/health2.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health2.json + + - name: Pass bootstrap pw for extra UI test + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIUiExtra-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright + env: + BASE_URL: http://127.0.0.1:18897 + STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }} + STUDIO_NEW_PW: ${{ env.STUDIO_EXTRA_NEW_PW }} + PW_ART_DIR: logs/playwright_extra + STUDIO_UI_STRICT: '1' + # See "Drive the chat UI" step. + STUDIO_UI_TURN_TIMEOUT_MS: '540000' + GGUF_REPO: ${{ env.GGUF_REPO }} + GGUF_VARIANT: ${{ env.GGUF_VARIANT }} + # Same flake-retry shape as "Drive the chat UI with Playwright" -- catches + # pipeTransport JSON crash, ERR_NO_BUFFER_SPACE, and nav interrupts. + run: | + mkdir -p logs/playwright_extra + attempt=1 + max_attempts=3 + while : ; do + set +e + python tests/studio/playwright_extra_ui.py 2>&1 | tee logs/playwright_extra_attempt_${attempt}.log + rc=${PIPESTATUS[0]} + set -e + if [ "$rc" -eq 0 ]; then + break + fi + if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \ + || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log \ + || grep -q "interrupted by another navigation" logs/playwright_extra_attempt_${attempt}.log; } \ + && [ "$attempt" -lt "$max_attempts" ]; then + echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." + kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true + sleep 2 + unsloth studio reset-password + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \ + > "logs/studio_extra_retry_${attempt}.log" 2>&1 & + STUDIO_EXTRA_PID=$! + echo "STUDIO_EXTRA_PID=$STUDIO_EXTRA_PID" >> "$GITHUB_ENV" + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18897/api/health" > /tmp/health2.json \ + && jq -e '.status == "healthy"' /tmp/health2.json >/dev/null; then + break + fi + sleep 1 + done + STUDIO_OLD_PW=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + STUDIO_NEW_PW="CIUiExtra-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$STUDIO_OLD_PW" + echo "::add-mask::$STUDIO_NEW_PW" + export STUDIO_OLD_PW STUDIO_NEW_PW + attempt=$((attempt + 1)) + sleep 3 + continue + fi + exit "$rc" + done + + - name: Stop second Studio + if: always() + run: | + kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true + sleep 2 + + - name: Upload Playwright artifacts + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: mac-studio-ui-smoke-artifacts + path: | + logs/studio.log + logs/studio_extra.log + logs/install.log + logs/playwright + logs/playwright_extra + retention-days: 7 diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml new file mode 100644 index 0000000..d104306 --- /dev/null +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Mac counterpart to studio-update-smoke.yml. Verifies that on a real +# Apple Silicon (macos-14, M1) runner: +# +# 1. install.sh --local --no-torch installs Studio AND auto-fetches +# the prebuilt llama.cpp Mac binary (llama-bNNNN-bin-macos-arm64 +# from ggml-org/llama.cpp). Hitting the source-build fallback is +# treated as an Unsloth bug -- Studio must always pick the +# prebuilt on Mac. +# 2. unsloth studio update --local is idempotent. Two consecutive +# runs both report "prebuilt up to date and validated", no +# source-build fallback. +# 3. The installed Studio still boots and /api/health returns +# healthy after the update path. + +name: Mac Studio Update CI + +on: + pull_request: + paths: + - 'install.sh' + - 'scripts/uninstall.sh' + - 'studio/setup.sh' + - 'studio/install_python_stack.py' + - 'studio/install_llama_prebuilt.py' + - 'studio/backend/requirements/**' + - 'unsloth_cli/commands/studio.py' + - 'pyproject.toml' + - '.github/workflows/studio-mac-update-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + update-idempotency: + name: Studio Updating Tests + runs-on: macos-14 + timeout-minutes: 30 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: First update should be a no-op (prebuilt already validated) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update.log + if grep -q "falling back to source build" logs/update.log; then + echo "::error::studio update fell back to source-build llama.cpp on Mac." + grep -E "llama-prebuilt|llama.cpp" logs/update.log | tail -60 + exit 1 + fi + if ! grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update.log; then + echo "::error::no prebuilt up-to-date marker in update.log." + grep -E "llama-prebuilt|llama.cpp" logs/update.log | tail -60 + exit 1 + fi + echo "update path took the prebuilt fast path" + + - name: Second update must also be a no-op + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update2.log + grep -q "falling back to source build" logs/update2.log && { + echo "::error::second update fell back to source build on Mac" + tail -60 logs/update2.log; exit 1; } || true + grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log + echo "second update was clean" + + - name: Boot Studio briefly to confirm the install is still usable + run: | + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \ + > logs/studio.log 2>&1 & + PID=$! + HEALTHY="" + for i in $(seq 1 60); do + if curl -fs http://127.0.0.1:18891/api/health > /tmp/health.json; then + if python3 -c "import json,sys; d=json.load(open('/tmp/health.json')); sys.exit(0 if d.get('status')=='healthy' else 1)"; then + HEALTHY=1 + break + fi + fi + sleep 1 + done + if [ -z "$HEALTHY" ]; then + echo "Studio failed to come up after \`update\`" + tail -200 logs/studio.log + kill "$PID" 2>/dev/null || true + exit 1 + fi + kill "$PID" 2>/dev/null || true + echo "post-update Studio /api/health OK" + + - name: Uninstall and verify clean + # Round-trip through scripts/uninstall.sh on real macOS. As a side + # effect this exercises the macOS-only .app bundle + Launch Services + # removal path (~/Applications/Unsloth Studio.app, lsregister -u) + # which is not testable from a Linux runner. Skips gracefully if + # scripts/uninstall.sh has not landed yet (lets this workflow merge + # before #5497). + run: | + set -o pipefail + if [ ! -f scripts/uninstall.sh ]; then + echo "scripts/uninstall.sh not present in this tree; skipping round-trip" + : > logs/uninstall.log + exit 0 + fi + sh scripts/uninstall.sh 2>&1 | tee logs/uninstall.log + leak=0 + for p in \ + "$HOME/.unsloth/studio" \ + "$HOME/.local/share/unsloth" \ + "$HOME/Applications/Unsloth Studio.app" \ + "$HOME/Desktop/Unsloth Studio.app" \ + "$HOME/.local/bin/unsloth"; do + if [ -e "$p" ] || [ -L "$p" ]; then + echo "::error::leak: $p" + leak=$((leak + 1)) + fi + done + [ "$leak" -eq 0 ] || exit 1 + sh scripts/uninstall.sh 2>&1 | tail -5 + sh scripts/uninstall.sh 2>&1 | tail -5 + echo "PASS: mac install -> update -> uninstall round-trip clean" + + - name: Upload update logs + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: mac-studio-update-log + path: | + logs/install.log + logs/update.log + logs/update2.log + logs/studio.log + logs/uninstall.log + retention-days: 7 diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml new file mode 100644 index 0000000..018857d --- /dev/null +++ b/.github/workflows/studio-tauri-smoke.yml @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# PR-time smoke for the Tauri desktop wrapper. Builds the frontend and the +# Tauri Linux debug binary, with no codesigning. Catches: +# - tauri.conf.json drift +# - src-tauri Cargo.toml or rust source breakage +# - Tauri CLI version drift (we pin 2.10.1, matching release-desktop.yml) +# - frontend output not picked up by Tauri's distDir +# +# Linux-only on a free `ubuntu-latest` runner. Mac and Windows desktop builds +# stay in release-desktop.yml (manual `workflow_dispatch`) because they need +# code-signing secrets and ~30 min of runner time each. + +name: Studio Tauri CI + +on: + pull_request: + paths: + - 'studio/frontend/**' + - 'studio/src-tauri/**' + # CLI rename / signature change can break Tauri's spawned + # `unsloth studio` -- include unsloth_cli in the trigger set. + - 'unsloth_cli/**' + - '.github/workflows/studio-tauri-smoke.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + linux-debug-build: + name: Tauri Linux debug build (no codesign) + runs-on: ubuntu-22.04 + timeout-minutes: 25 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux native deps for Tauri / WebKit2GTK + run: | + sudo apt-get update + sudo apt-get install -y \ + libwebkit2gtk-4.1-dev libappindicator3-dev \ + librsvg2-dev libxdo-dev libssl-dev patchelf + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '24' + + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 + + - uses: swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2.9.1 + with: + workspaces: studio/src-tauri -> target + + - name: Install pinned Tauri CLI (matches release-desktop.yml) + # Lifecycle scripts (esbuild native-binary postinstall, etc.) are + # required for `vite build`. The pre-install lockfile structural + # audit (lockfile_supply_chain_audit.py) is the practical defence + # against the npm postinstall-dropper class -- it fires BEFORE any + # tarball runs, on the injection pattern itself rather than an + # advisory-DB lookup. + run: npm install --save-dev --prefix studio @tauri-apps/cli@2.10.1 --no-fund --no-audit + + - name: Verify pinned Tauri CLI version + run: | + out="$(npx --prefix studio tauri --version)" + echo "$out" + [ "$out" = "tauri-cli 2.10.1" ] || { echo "::error::expected tauri-cli 2.10.1, got $out"; exit 1; } + + - name: Lockfile supply-chain audit (pre-install scan) + run: python3 scripts/lockfile_supply_chain_audit.py + + - name: Frontend build (npm ci, vite) + working-directory: studio/frontend + # Lifecycle scripts (esbuild native-binary postinstall, etc.) are + # required for `vite build`. The pre-install lockfile structural + # audit (lockfile_supply_chain_audit.py) is the practical defence + # against the npm postinstall-dropper class -- it fires BEFORE any + # tarball runs, on the injection pattern itself rather than an + # advisory-DB lookup. + run: | + npm ci --no-fund --no-audit + npm run build + test -f dist/index.html + + - name: Tauri debug build (Linux, no bundle, no codesign) + # `--debug` + `--no-bundle` keeps this lean: compiles the Rust crate, + # confirms the frontend dist is wired into Tauri, but skips the AppImage + # / .deb production. Code signing is irrelevant because we never produce + # a distributable artifact. + env: + TAURI_SIGNING_PRIVATE_KEY: '' + TAURI_SIGNING_PRIVATE_KEY_PASSWORD: '' + run: npx --prefix studio tauri build --debug --no-bundle + + - name: Inspect produced binary + run: | + BIN=$(find studio/src-tauri/target/debug -maxdepth 1 -type f -executable 2>/dev/null \ + | grep -Ev '\.(d|so|dylib|dll)$' \ + | grep -Ev '/(deps|build|examples)$' \ + | head -1) + echo "binary: $BIN" + if [ -z "$BIN" ]; then + echo "::error::Tauri debug binary not produced" + ls -la studio/src-tauri/target/debug/ || true + exit 1 + fi + file "$BIN" + du -h "$BIN" + + - name: Upload Tauri debug build + # Always upload so a green run leaves the binary inspectable too. + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: tauri-debug-build + path: | + studio/src-tauri/target/debug + studio/frontend/dist + retention-days: 3 diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml new file mode 100644 index 0000000..297a585 --- /dev/null +++ b/.github/workflows/studio-ui-smoke.yml @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# End-to-end Studio chat UI smoke via Playwright + Chromium against a +# headless Linux runner. Boots Studio with the smallest GGUF +# (gemma-3-270m-it UD-Q4_K_XL, ~254 MiB), drives the actual frontend +# bundle, and asserts the full bootstrap-password / change-password / +# send-message / persist-on-reload journey works end to end. +# +# This is the only workflow that catches regressions in the wiring +# between the React frontend and the FastAPI backend, e.g. assistant-ui +# version drift, /api/auth response shape changes, runtime-provider +# regressions, or chat-history persistence breaking. Backend-only and +# frontend-only CI happily pass while the actual user-visible UI is +# broken (cf. the 2026.5.1 chat-history release). + +name: Studio UI CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.sh' + - 'pyproject.toml' + # The Playwright test files themselves -- a PR that ONLY edits + # the test must still trigger UI CI. + - 'tests/studio/**' + - '.github/workflows/studio-ui-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + ui-smoke: + name: Chat UI Tests + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18892' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Install Playwright + Chromium + run: | + pip install 'playwright>=1.45' + # --with-deps installs the OS-level runtime libs Chromium + # needs (libnss3, libxkbcommon, etc.). About 30 s on a + # warm runner. + python -m playwright install --with-deps chromium + + - name: Reset auth + boot Studio + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + # 180 s -- a cold runner with venv warm-up + lazy imports has + # been seen to exceed 60 s. Failing the wait is more expensive + # than waiting an extra two minutes. + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + + - name: Pass bootstrap password to the Playwright step + # The Playwright test does its OWN /change-password through the + # UI (Setup your account / Choose a new password), then loads + # the model via page.evaluate against /api/inference/load with + # the JWT it got from change-password. So the only thing we + # have to hand it is the bootstrap password (so it can verify + # post-rotation that the OLD bootstrap pw now returns 401). + # + # NEW + NEW2 are generated freshly per CI run via secrets.token_urlsafe + # rather than hardcoded. If a workflow gets compromised, the + # attacker can't replay a known-good rotated password against + # any future / parallel Studio install -- the rotated value + # only ever exists for the lifetime of this single job, masked + # in the log via ::add-mask::. + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + NEW2="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "::add-mask::$NEW2" + echo "STUDIO_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" + echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" + + - name: Drive the chat UI with Playwright + env: + BASE_URL: http://127.0.0.1:18892 + # The test file lives in the repo so it can be run locally + # against a freshly-installed Studio (BASE_URL=...; STUDIO_OLD_PW= + # $(cat ~/.unsloth/studio/auth/.bootstrap_password); python ...). + PW_ART_DIR: logs/playwright + # Strict mode: in CI a missing button / nav / dialog must + # FAIL the test. Locally the test still runs against partial + # Studio installs without STUDIO_UI_STRICT. + STUDIO_UI_STRICT: '1' + run: | + mkdir -p logs/playwright + python tests/studio/playwright_chat_ui.py + + - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders) + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + + # The chat UI test ends by clicking the Shutdown menuitem, which + # leaves the server dead. The extra UI test (Compare / Recipes / + # Export / Studio / Settings) needs a fresh Studio, so we boot a + # second one on a different port. Boot is fast (~3-5s on the + # warm install we already did) so this adds little wall time. + - name: Reset auth + boot Studio for extra UI tests (port 18894) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18894 \ + > logs/studio_extra.log 2>&1 & + echo "STUDIO_EXTRA_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18894 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18894/api/health" > /tmp/health2.json; then + jq -e '.status == "healthy"' /tmp/health2.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health2.json + + - name: Pass bootstrap pw for extra UI test + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIUiExtra-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright + env: + BASE_URL: http://127.0.0.1:18894 + STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }} + STUDIO_NEW_PW: ${{ env.STUDIO_EXTRA_NEW_PW }} + PW_ART_DIR: logs/playwright_extra + STUDIO_UI_STRICT: '1' + GGUF_REPO: ${{ env.GGUF_REPO }} + GGUF_VARIANT: ${{ env.GGUF_VARIANT }} + run: | + mkdir -p logs/playwright_extra + python tests/studio/playwright_extra_ui.py + + - name: Stop second Studio + if: always() + run: | + kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true + sleep 2 + + # IME + multilingual paste regression (issue #5318 / PR #5327). + # Third Studio on its own port so a hang here cannot poison the + # earlier UI tests. No GGUF -- the bug surface is the composer. + - name: Reset auth + boot Studio for IME / i18n tests (port 18896) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18896 \ + > logs/studio_ime.log 2>&1 & + echo "STUDIO_IME_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18896 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18896/api/health" > /tmp/health3.json; then + jq -e '.status == "healthy"' /tmp/health3.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health3.json + + - name: Pass bootstrap pw for IME / i18n test + # IME smoke does the change-password against the bootstrap that + # Studio's frontend injects into the page, so it only needs the + # NEW password. + run: | + NEW="CIIme-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$NEW" + echo "STUDIO_IME_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive IME + multilingual paste regression with Playwright + env: + BASE_URL: http://127.0.0.1:18896 + STUDIO_NEW_PW: ${{ env.STUDIO_IME_NEW_PW }} + PW_ART_DIR: logs/playwright_ime + STUDIO_UI_STRICT: '1' + run: | + mkdir -p logs/playwright_ime + python tests/studio/playwright_chat_ime_i18n.py + + - name: Stop third Studio + if: always() + run: | + kill "${STUDIO_IME_PID}" 2>/dev/null || true + sleep 2 + # Capture backend + llama-server logs (all three Studios share this + # dir) so a stray 500 has a server-side traceback. + mkdir -p logs/server-logs + cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true + + - name: Upload Playwright artifacts + # Always upload so a green run's screenshots stay reviewable -- + # catches "passed but the UI is silently broken" regressions. + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: studio-ui-smoke-artifacts + path: | + logs/studio.log + logs/studio_extra.log + logs/studio_ime.log + logs/install.log + logs/server-logs/ + logs/playwright + logs/playwright_extra + logs/playwright_ime + retention-days: 7 diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml new file mode 100644 index 0000000..08a79af --- /dev/null +++ b/.github/workflows/studio-update-smoke.yml @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Verifies that `unsloth studio update --local` is idempotent: a fresh +# install via install.sh, followed by `unsloth studio update --local`, +# succeeds and is a no-op for the llama.cpp prebuilt (it should report +# "prebuilt up to date and validated", not re-run the source build). +# +# This catches regressions in setup.sh's update path that the existing +# GGUF / wheel jobs would miss because they only invoke install.sh once. + +name: Studio Update CI + +on: + pull_request: + paths: + - 'install.sh' + - 'scripts/uninstall.sh' + - 'studio/setup.sh' + - 'studio/install_python_stack.py' + - 'studio/install_llama_prebuilt.py' + - 'studio/backend/requirements/**' + - 'unsloth_cli/commands/studio.py' + - 'pyproject.toml' + - '.github/workflows/studio-update-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + update-idempotency: + name: Studio Updating Tests + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + # Don't cache pip: this job runs `bash install.sh` and + # `unsloth studio update --local` which both go through + # `uv` and never populate ~/.cache/pip. setup-python's + # post-step then fatal-errors with "Cache folder path is + # retrieved for pip but doesn't exist on disk". + + - name: Install Studio (--local, --no-torch) + # Pass the workflow token so the llama.cpp prebuilt installer's + # GitHub-API call to list releases isn't rate-limited (60/hr + # unauthenticated). Without this, three consecutive install + + # update + update calls in this job exceed the limit and the + # prebuilt path falls back to source build. + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: First update should be a no-op (prebuilt already validated) + # `unsloth studio update --local` runs studio/setup.sh against + # the local repo. Right after install.sh the llama.cpp prebuilt + # has just been installed and validated, so the second run must + # take the "prebuilt up to date and validated" code path. Any + # source-build fallback or re-download here means setup.sh's + # idempotency regressed. + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update.log + if grep -q "falling back to source build" logs/update.log; then + echo "::error::studio update fell back to source-build llama.cpp on a fresh install. setup.sh idempotency regressed." + grep -E "llama-prebuilt|llama.cpp" logs/update.log | tail -60 + exit 1 + fi + if ! grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update.log; then + echo "::error::no prebuilt up-to-date marker in update.log. Did setup.sh skip the prebuilt path on update?" + grep -E "llama-prebuilt|llama.cpp" logs/update.log | tail -60 + exit 1 + fi + echo "update path took the prebuilt fast path" + + - name: Second update must also be a no-op + # Two consecutive `update`s back-to-back is the usual desktop + # flow (auto-update, then user-triggered update). Asserting the + # second run is also clean rules out hidden state changes from + # the first one. + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update2.log + grep -q "falling back to source build" logs/update2.log && { + echo "::error::second update fell back to source build" + tail -60 logs/update2.log; exit 1; } || true + grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log + echo "second update was clean" + + - name: Boot Studio briefly to confirm the install is still usable + # If `update --local` accidentally broke the venv or wiped the + # llama-server binary, the server would fail to start here. + run: | + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \ + > logs/studio.log 2>&1 & + PID=$! + for i in $(seq 1 60); do + if curl -fs http://127.0.0.1:18891/api/health > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json + break + fi + sleep 1 + done + if ! jq -e '.status == "healthy"' /tmp/health.json 2>/dev/null; then + echo "Studio failed to come up after `update`" + tail -200 logs/studio.log + kill "$PID" 2>/dev/null || true + exit 1 + fi + kill "$PID" 2>/dev/null || true + echo "post-update Studio /api/health OK" + + - name: Uninstall and verify clean + # Round-trip the installer through scripts/uninstall.sh: confirms the + # uninstaller actually finds and removes everything install.sh + + # update wrote. Safety-guard scenarios (refuse-$HOME etc.) belong + # in a separate fast smoke job; this is the happy-path cleanup + # assertion that catches regressions where install.sh starts + # writing to a new location and scripts/uninstall.sh hasn't caught up. + # Skips gracefully if scripts/uninstall.sh has not landed yet (lets + # this workflow merge before #5497). + run: | + set -o pipefail + if [ ! -f scripts/uninstall.sh ]; then + echo "scripts/uninstall.sh not present in this tree; skipping round-trip" + : > logs/uninstall.log + exit 0 + fi + sh scripts/uninstall.sh 2>&1 | tee logs/uninstall.log + leak=0 + for p in \ + "$HOME/.unsloth/studio" \ + "$HOME/.local/share/unsloth" \ + "$HOME/Desktop/Unsloth Studio.desktop" \ + "$HOME/.local/bin/unsloth"; do + if [ -e "$p" ] || [ -L "$p" ]; then + echo "::error::leak: $p" + ls -la "$p" 2>&1 | head -3 + leak=$((leak + 1)) + fi + done + [ "$leak" -eq 0 ] || exit 1 + # Idempotent: re-runs exit 0 on an empty $HOME. + sh scripts/uninstall.sh 2>&1 | tail -5 + sh scripts/uninstall.sh 2>&1 | tail -5 + echo "PASS: install -> update -> uninstall round-trip clean" + + - name: Upload update logs + # Always upload so a green run still leaves the install + two + # update logs + uninstall log reviewable. + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: studio-update-log + path: | + logs/install.log + logs/update.log + logs/update2.log + logs/studio.log + logs/uninstall.log + retention-days: 7 diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml new file mode 100644 index 0000000..e9abd2d --- /dev/null +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Windows counterpart to studio-api-smoke.yml / studio-mac-api-smoke.yml. +# Same tests/studio/studio_api_smoke.py exercise (CORS hardening, auth +# state machine, JWT expiry, API key lifecycle, /v1/models / +# /v1/embeddings / /v1/responses, endpoint-by-endpoint auth audit) but +# on the FREE windows-latest runner. The file-mode hardening section +# (Section 6) is Linux-only and short-circuits on non-POSIX; the rest +# is platform-portable. + +name: Windows Studio API CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.ps1' + - 'pyproject.toml' + - 'tests/studio/**' + - '.github/workflows/studio-windows-api-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + api-smoke: + name: Studio API & Auth Tests + runs-on: windows-latest + timeout-minutes: 30 + defaults: + run: + shell: bash + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18895' + HF_HOME: ${{ github.workspace }}/hf-cache + # Force UTF-8 for stdio (Windows defaults to cp1252; hf + # download prints a "✓" checkmark and crashes otherwise). + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + # See studio-windows-update-smoke.yml for the full rationale. + # tl;dr: setup.ps1 needs npm >=11 to skip a 35 s winget Node + # reinstall, and Defender's real-time scan dominates the + # frontend / uv-pip-extract steps. + run: | + $ProgressPreference = 'SilentlyContinue' + Write-Host "npm version before upgrade: $(npm -v)" + npm install -g 'npm@^11' 2>&1 | Out-Host + Write-Host "npm version after upgrade: $(npm -v)" + # NOTE: do NOT pre-create these directories. See + # studio-windows-update-smoke.yml for the full rationale -- + # creating an empty studio/frontend/dist trips setup.ps1's + # mtime-based staleness check into "frontend up to date, skip + # rebuild" and Studio boots with an empty dist directory. + # Add-MpPreference accepts paths that do not yet exist. + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { + Add-MpPreference -ExclusionPath $p -ErrorAction Stop + Write-Host "Defender exclusion added: $p" + } catch { + Write-Host "Defender exclusion skipped ($($_.Exception.Message)): $p" + } + } + + - name: Install Studio (--local, --no-torch) + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + # *>&1 captures Write-Host (Information stream) output; + # plain 2>&1 does not. setup.ps1 emits "prebuilt installed + # and validated" via Write-Host, and we grep for that. + $ProgressPreference = 'SilentlyContinue' + & ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log + + - name: Assert install.ps1 used the Windows llama.cpp prebuilt + run: | + # Filesystem-based check (setup.ps1's stream output isn't + # captured back through this parent step's pipeline; see + # studio-windows-ui-smoke.yml for full explanation). + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp on Windows." + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + if [ ! -f "$INFO" ]; then + echo "::error::no UNSLOTH_PREBUILT_INFO.json at $INFO." + ls -la "$LLAMA_DIR" || true + exit 1 + fi + if [ ! -f "$BIN" ]; then + echo "::error::no llama-server.exe at $BIN." + ls -la "$LLAMA_DIR/build/bin" || true + exit 1 + fi + echo "install.ps1 installed the Windows prebuilt llama.cpp:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + # install.ps1's User-PATH update doesn't propagate to a + # running Git Bash session; export the shim dir so the + # next `unsloth ...` invocation finds it. + run: | + SHIM_DIR=~/.unsloth/studio/bin + if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then + echo "::error::unsloth.exe shim not found at $SHIM_DIR" + ls -la ~/.unsloth/studio/ || true + exit 1 + fi + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: Install pyjwt for the JWT-expiry forge test + run: python -m pip install 'pyjwt>=2.6' + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + + - name: Pass bootstrap password + rotated targets to the test + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="ApiSmoke-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + NEW2="ApiSmoke-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "::add-mask::$NEW2" + echo "STUDIO_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" + echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" + + - name: Run Studio API & Auth tests + # Do NOT pin STUDIO_AUTH_DIR here. The Mac/Linux mirrors + # hardcode runner-specific paths (/Users/runner/..., + # /home/runner/...), but on Windows the path is + # C:\Users\runneradmin\.unsloth\studio\auth and varies by + # runner image. studio_api_smoke.py defaults to + # Path.home()/".unsloth"/"studio"/"auth" when the env is + # unset, which is correct on every OS. + env: + BASE_URL: http://127.0.0.1:18895 + run: python tests/studio/studio_api_smoke.py + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + + - name: Upload API smoke logs + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-studio-api-smoke-log + path: | + logs/install.log + logs/studio.log + retention-days: 7 diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml new file mode 100644 index 0000000..0453c92 --- /dev/null +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -0,0 +1,1919 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Three end-to-end smoke jobs that boot a freshly-installed Studio and +# exercise the surfaces real users hit through the OpenAI / Anthropic +# SDKs and curl, on the FREE windows-latest runner. Each job picks the +# smallest model that exercises the behaviour under test, primes +# HF_HOME via actions/cache, and shares the install.ps1 --local +# --no-torch bootstrap. +# +# 1. OpenAI, Anthropic API tests +# gemma-3-270m-it UD-Q4_K_XL (~254 MiB). +# 2. Tool calling Tests +# Qwen3.5-2B UD-Q4_K_XL (~890 MiB). +# 3. JSON, images +# Qwen3-VL-2B-Instruct UD-IQ2_XXS + mmproj-F16 (~1.4 GiB total). +# Within the 14 GB windows-latest SSD budget. + +name: Windows Studio GGUF CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.ps1' + - 'pyproject.toml' + - 'tests/studio_setup_ps1/**' + - '.github/workflows/studio-windows-inference-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + # ───────────────────────────────────────────────────────────────────── + # Job 1: OpenAI, Anthropic API tests + # ───────────────────────────────────────────────────────────────────── + openai-anthropic: + name: OpenAI, Anthropic API tests + runs-on: windows-latest + timeout-minutes: 30 + defaults: + run: + shell: bash + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18888' + HF_HOME: ${{ github.workspace }}/hf-cache + # Force UTF-8 for stdio (Windows defaults to cp1252; hf + # download / Studio CLI print "✓" checkmarks and crash + # otherwise). + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + # Fast GPU-free gate: parse install.ps1 + setup.ps1 and run the PowerShell + # unit tests (CUDA-toolkit + torch-flavor helpers) before the heavy GGUF smoke. + - name: PowerShell installer unit tests + shell: pwsh + run: | + foreach ($f in @('install.ps1', 'studio/setup.ps1')) { + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path $f).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "$f parsed with no errors" + } + pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 + pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1 + pwsh -NoProfile -File tests/studio/test_node_decision.ps1 + pwsh -NoProfile -File tests/studio/test_node_probe_guard.ps1 + + # uninstall.ps1: native uninstall must keep the shared unsloth.ico while a + # WSL shortcut still references it (dual install), else that shortcut blanks. + - name: uninstall.ps1 unit test (dual-install icon preserve) + shell: pwsh + run: | + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path scripts/uninstall.ps1).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "uninstall.ps1 parsed with no errors" + pwsh -NoProfile -File tests/studio/test_uninstall_dual_install_icon.ps1 + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + # Split restore + save (rather than the one-step actions/cache) so a + # transient restore-side failure does not kill the whole job. v5 has a + # known flake where it logs "Cache hit for: " and then exits + # non-zero without actually extracting the archive (see + # actions/cache#1621 and github community discussion #163260). + # continue-on-error on restore masks that failure so the Prime step + # below can re-download from HF and the job keeps running. Save then + # populates the cache key on a real miss only; cache keys are + # immutable, so a corrupted cached entry persists until the -v1 + # suffix below is bumped. + - name: Restore HF_HOME cache for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + # Run on a real cache miss AND on the silent-restore-failure mode + # described above (outcome != success). + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} + # Only write a fresh cache entry when we actually rebuilt the + # directory (Prime ran and succeeded). Skipping when Prime is + # skipped avoids "already exists" save warnings on the happy path. + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + # See studio-windows-update-smoke.yml for the full rationale. + # tl;dr: setup.ps1 needs npm >=11 to skip a 35 s winget Node + # reinstall, and Defender's real-time scan dominates the + # frontend / uv-pip-extract steps. + run: | + $ProgressPreference = 'SilentlyContinue' + Write-Host "npm version before upgrade: $(npm -v)" + npm install -g 'npm@^11' 2>&1 | Out-Host + Write-Host "npm version after upgrade: $(npm -v)" + # NOTE: do NOT pre-create these directories. See + # studio-windows-update-smoke.yml for the full rationale -- + # creating an empty studio/frontend/dist trips setup.ps1's + # mtime-based staleness check into "frontend up to date, skip + # rebuild" and Studio boots with an empty dist directory. + # Add-MpPreference accepts paths that do not yet exist. + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { + Add-MpPreference -ExclusionPath $p -ErrorAction Stop + Write-Host "Defender exclusion added: $p" + } catch { + Write-Host "Defender exclusion skipped ($($_.Exception.Message)): $p" + } + } + + - name: Install Studio (--local, --no-torch) + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + # *>&1 captures Write-Host (Information stream) output; + # plain 2>&1 does not. setup.ps1 emits "prebuilt installed + # and validated" via Write-Host, and we grep for that. + $ProgressPreference = 'SilentlyContinue' + & ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log + + - name: Assert install.ps1 used the Windows llama.cpp prebuilt + run: | + # Filesystem check; setup.ps1's stream output isn't captured. + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp on Windows." + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + if [ ! -f "$INFO" ]; then + echo "::error::no UNSLOTH_PREBUILT_INFO.json at $INFO." + ls -la "$LLAMA_DIR" || true + exit 1 + fi + if [ ! -f "$BIN" ]; then + echo "::error::no llama-server.exe at $BIN." + ls -la "$LLAMA_DIR/build/bin" || true + exit 1 + fi + echo "install.ps1 installed the Windows prebuilt llama.cpp:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + run: | + SHIM_DIR=~/.unsloth/studio/bin + if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then + echo "::error::unsloth.exe shim not found at $SHIM_DIR" + ls -la ~/.unsloth/studio/ || true + exit 1 + fi + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: Install OpenAI + Anthropic Python SDKs + run: python -m pip install 'openai>=1.50' 'anthropic>=0.40' + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json + exit 0 + fi + sleep 1 + done + echo "Studio did not become healthy in 180s" + tail -200 logs/studio.log + exit 1 + + - name: Password rotation (old must fail, new must work) + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIRotated-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + [ -n "$OLD_TOKEN" ] && [ "$OLD_TOKEN" != "null" ] || { echo "bootstrap login failed"; exit 1; } + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + OLD_STATUS=$(curl -s -o /dev/null -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}") + if [ "$OLD_STATUS" != "401" ]; then + echo "::error::Login with old password returned $OLD_STATUS, expected 401" + exit 1 + fi + NEW_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + [ -n "$NEW_TOKEN" ] && [ "$NEW_TOKEN" != "null" ] || { echo "new login failed"; exit 1; } + echo "TOKEN=$NEW_TOKEN" >> "$GITHUB_ENV" + echo "password rotation OK (old=401, new=200)" + + - name: Load the GGUF (HF repo + variant, served from HF_HOME cache) + run: | + # Retry the load step a few times so a transient TCP RST during + # llama-server warm-up (Windows runner image churn, + # windows-latest -> windows-2025-vs2026 rollout) doesn't fail + # the whole job. The Studio backend's _wait_for_health now + # catches httpx.ReadError too; this retry layer covers the + # cases the backend can't recover from on its own. + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP; response:" + cat /tmp/load.json || true + sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_gguf, context_length}' /tmp/load.json + + - name: Multi-turn determinism via OpenAI + Anthropic SDKs + env: + BASE_URL: http://127.0.0.1:18888 + run: | + python - <<'PY' + import json + import os + from openai import OpenAI + from anthropic import Anthropic + + BASE = os.environ["BASE_URL"] + KEY = os.environ["TOKEN"] + SEED = 3407 + + PROMPTS = [ + "What is 1+1?", + "What did I ask before?", + "What is the capital of France?", + "Repeat the city name", + ] + + def run_openai(): + client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY) + history, replies = [], [] + for prompt in PROMPTS: + history.append({"role": "user", "content": prompt}) + resp = client.chat.completions.create( + model = "default", + messages = history, + temperature = 0.0, + max_tokens = 80, + seed = SEED, + extra_body = {"enable_thinking": False}, + ) + text = resp.choices[0].message.content or "" + replies.append(text) + history.append({"role": "assistant", "content": text}) + return replies + + def run_anthropic(): + client = Anthropic( + base_url = BASE, + api_key = "unused", + default_headers = {"Authorization": f"Bearer {KEY}"}, + ) + history, replies = [], [] + for prompt in PROMPTS: + history.append({"role": "user", "content": prompt}) + msg = client.messages.create( + model = "default", + max_tokens = 80, + messages = history, + temperature = 0.0, + extra_body = {"seed": SEED, "enable_thinking": False}, + ) + text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text") + replies.append(text) + history.append({"role": "assistant", "content": text}) + return replies + + for label, runner in (("openai", run_openai), ("anthropic", run_anthropic)): + first = runner() + second = runner() + for i, (a, b) in enumerate(zip(first, second), start = 1): + print(f"[{label} turn {i}] {a!r}") + assert a, f"{label}: empty turn {i} response" + # Compare on stripped content: llama-server can vary + # trailing whitespace (specifically a final '\n') between + # otherwise-identical greedy runs depending on the + # batch-flush boundary at which the stream is closed. The + # generated tokens are identical; only the trailing + # whitespace differs. Keep the raw repr in the failure + # message so a real divergence is still legible. + assert a.strip() == b.strip(), ( + f"{label} non-deterministic at turn {i} with temperature=0.0:\n" + f" run1: {a!r}\n run2: {b!r}" + ) + joined = " ".join(first).lower() + assert "1" in first[0], f"{label}: turn-1 answer should contain '1', got {first[0]!r}" + assert "paris" in joined, f"{label}: expected 'paris' somewhere in the four-turn transcript: {first}" + print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded") + PY + + - name: Stop Studio + if: always() + # Run as cmd so we are not running through the Git Bash shell; + # Git Bash on windows-latest has been observed to exit 143 + # (SIGTERM) from any inline kill/sleep block, masking a green + # test run. The runner reclaims the Studio child process at + # job end either way, so just emit a marker and exit 0. + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + + - name: Collect llama-server logs + if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true + shell: bash + # Copy llama-server's own stdout/stderr (teed by Studio under + # ~/.unsloth/studio/logs/llama-server/) into the workspace so + # upload-artifact can pick it up. Crucial for diagnosing a + # subprocess crash where Studio's traceback only shows the + # symptom (httpx ReadError) but not the cause. + run: | + mkdir -p logs/llama-server + cp -v ~/.unsloth/studio/logs/llama-server/*.log logs/llama-server/ 2>/dev/null || \ + echo "no llama-server logs to collect" + + - name: Upload logs + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-openai-anthropic-log + path: | + logs/studio.log + logs/install.log + logs/llama-server/*.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job 2: Tool calling Tests + # ───────────────────────────────────────────────────────────────────── + tool-calling: + name: Tool calling Tests + runs-on: windows-latest + timeout-minutes: 30 + defaults: + run: + shell: bash + env: + # Tool calling is the highest-volume GGUF in this workflow + # (Qwen3.5-2B at Q4_K_XL = ~1.28 GiB). The previous HF_HOME + # cache stored xet chunks + blobs + snapshots = ~4.7 GiB -- + # 3.7x file-size inflation, dominating the post-step upload + # (211 s on first run; subsequent runs hit the cache, but the + # one-time cost recurs every time the cache key bumps). Use + # main's `--local-dir gguf-cache` pattern: cache the flat .gguf + # only, pass an absolute path to Studio's /api/inference/load. + # The OpenAI/Anth and JSON+images jobs still cover the + # gguf_variant resolution path. + GGUF_REPO: unsloth/Qwen3.5-2B-GGUF + GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf + STUDIO_PORT: '18898' + # Force UTF-8 for stdio (Windows defaults to cp1252; hf + # download / Studio CLI print "✓" checkmarks and crash + # otherwise). + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + # Split restore + save so a transient restore-side failure does not + # kill the whole job. See the matching block in the tool-calling job + # above for the full rationale (actions/cache#1621). + - name: Restore GGUF model cache + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model cache + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + # See studio-windows-update-smoke.yml for the full rationale. + # tl;dr: setup.ps1 needs npm >=11 to skip a 35 s winget Node + # reinstall, and Defender's real-time scan dominates the + # frontend / uv-pip-extract steps. + run: | + $ProgressPreference = 'SilentlyContinue' + Write-Host "npm version before upgrade: $(npm -v)" + npm install -g 'npm@^11' 2>&1 | Out-Host + Write-Host "npm version after upgrade: $(npm -v)" + # NOTE: do NOT pre-create these directories. See + # studio-windows-update-smoke.yml for the full rationale -- + # creating an empty studio/frontend/dist trips setup.ps1's + # mtime-based staleness check into "frontend up to date, skip + # rebuild" and Studio boots with an empty dist directory. + # Add-MpPreference accepts paths that do not yet exist. + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { + Add-MpPreference -ExclusionPath $p -ErrorAction Stop + Write-Host "Defender exclusion added: $p" + } catch { + Write-Host "Defender exclusion skipped ($($_.Exception.Message)): $p" + } + } + + - name: Install Studio (--local, --no-torch) + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + # *>&1 captures Write-Host (Information stream) output; + # plain 2>&1 does not. setup.ps1 emits "prebuilt installed + # and validated" via Write-Host, and we grep for that. + $ProgressPreference = 'SilentlyContinue' + & ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log + + - name: Assert install.ps1 used the Windows llama.cpp prebuilt + run: | + # Filesystem check; setup.ps1's stream output isn't captured. + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp on Windows." + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + if [ ! -f "$INFO" ]; then + echo "::error::no UNSLOTH_PREBUILT_INFO.json at $INFO." + ls -la "$LLAMA_DIR" || true + exit 1 + fi + if [ ! -f "$BIN" ]; then + echo "::error::no llama-server.exe at $BIN." + ls -la "$LLAMA_DIR/build/bin" || true + exit 1 + fi + echo "install.ps1 installed the Windows prebuilt llama.cpp:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + run: | + SHIM_DIR=~/.unsloth/studio/bin + if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then + echo "::error::unsloth.exe shim not found at $SHIM_DIR" + ls -la ~/.unsloth/studio/ || true + exit 1 + fi + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: Reset auth + boot Studio (API-only, default tool policy) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, change password, load model + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CITool-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + # GITHUB_WORKSPACE on windows-latest is a Windows path with + # backslashes ("D:\a\unsloth\unsloth"). Bash handles it as a + # raw string, but we cannot embed `\a` etc. in JSON without + # JSON-string-escaping every backslash. Replace `\` with `/` + # via bash parameter expansion -- pathlib.Path on Windows + # accepts forward slashes natively, so Studio's loader sees + # a normal path. + GGUF_PATH="${GITHUB_WORKSPACE//\\//}/gguf-cache/${GGUF_FILE}" + ls -lh "$GGUF_PATH" + # Retry: same rationale as the OpenAI/Anthropic job. + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_PATH\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP; response:" + cat /tmp/load.json || true + sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name}' /tmp/load.json + + - name: Tool calling, server-side tools, thinking on/off + env: + BASE_URL: http://127.0.0.1:18898 + run: | + python - <<'PY' + import json + import os + import time + import urllib.error + import urllib.request + + BASE = os.environ["BASE_URL"] + KEY = os.environ["API_KEY"] + SEED = 3407 + # Same temperature shim as the Mac job. Small Qwen3.5-2B + # quants can degenerate at temperature=0; a small non-zero + # temperature with a fixed seed keeps the test deterministic + # while escaping the trap. + TEMP = 0.2 + + def post(path, body, *, timeout = 240): + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{BASE}{path}", + data = data, + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) + + def post_sse(path, body, *, timeout = 600): + body = {**body, "stream": True} + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{BASE}{path}", + data = data, + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + parts = [] + with urllib.request.urlopen(req, timeout = timeout) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data: "): + continue + payload = line[6:] + if payload == "[DONE]": + break + try: + chunk = json.loads(payload) + except json.JSONDecodeError: + continue + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + if delta.get("content"): + parts.append(delta["content"]) + return "".join(parts) + + # ── 1. Standard OpenAI function calling ────────────────────── + weather_tool = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + + status, data = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "What is the weather in Paris?"}], + "tools": [weather_tool], + "tool_choice": "required", + "stream": False, + "temperature": TEMP, + "seed": SEED, + "max_tokens": 600, + }) + assert status == 200, f"tool call status {status}: {data}" + choice = data["choices"][0] + tool_calls = (choice.get("message") or {}).get("tool_calls") or [] + if tool_calls: + tc = tool_calls[0] + assert tc["function"]["name"] == "get_weather", ( + f"unexpected tool name: {tc['function']['name']!r}" + ) + args = json.loads(tc["function"]["arguments"]) + assert args.get("city"), f"missing city arg: {args}" + print(f"[tools] PASS function calling -> {tc['function']['name']}({args}) finish={choice.get('finish_reason')!r}") + else: + print( + f"[tools] WARN function calling: no tool_calls (finish_reason=" + f"{choice.get('finish_reason')!r}); HTTP path OK, model output drift." + ) + + # ── 2. Server-side python tool ─────────────────────────────── + content = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], + "enable_tools": True, + "enabled_tools": ["python"], + "session_id": "ci-tool-calling-py", + "temperature": TEMP, + "seed": SEED, + "max_tokens": 600, + }) + if "56088" in content or "56,088" in content: + print(f"[tools] PASS python tool ({len(content)} chars, found 56088)") + else: + assert content, "python tool: SSE stream empty" + print( + f"[tools] WARN python tool: SSE OK ({len(content)} chars) but " + f"model didn't return 56088 -- model output drift" + ) + + # ── 3. Server-side bash (terminal) tool ────────────────────── + # On Windows the terminal tool resolves to the system shell + # (cmd.exe wrapper) and `echo hello-bash-tool` works the same + # way it does on POSIX. The model still has to choose to + # invoke the tool; assert non-empty SSE if it doesn't. + content = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Use the terminal tool to run `echo hello-bash-tool` and tell me the exact output."}], + "enable_tools": True, + "enabled_tools": ["terminal"], + "session_id": "ci-tool-calling-bash", + "temperature": TEMP, + "seed": SEED, + "max_tokens": 600, + }) + if "hello-bash-tool" in content: + print(f"[tools] PASS terminal tool ({len(content)} chars)") + else: + assert content, "terminal tool: SSE stream empty" + print( + f"[tools] WARN terminal tool: SSE OK ({len(content)} chars) but " + f"model didn't echo 'hello-bash-tool' -- model output drift" + ) + + # ── 4. Server-side web_search tool ─────────────────────────── + # DuckDuckGo can be flaky from CI runners; only assert that + # the SSE stream opens and yields any data. + try: + content = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], + "enable_tools": True, + "enabled_tools": ["web_search"], + "session_id": "ci-tool-calling-web", + "temperature": TEMP, + "seed": SEED, + "max_tokens": 400, + }) + print(f"[tools] PASS web_search stream ({len(content)} chars)") + except Exception as exc: + print(f"[tools] WARN web_search probe failed (non-blocking): {exc}") + + # ── 5. Thinking on / off ───────────────────────────────────── + def thinking_call(enable): + status, data = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Briefly: is 17 prime?"}], + "stream": False, + "enable_thinking": enable, + "temperature": TEMP, + "seed": SEED, + "max_tokens": 300, + }) + assert status == 200 + msg = data["choices"][0]["message"] + raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "") + return raw + + on_text = thinking_call(True) + off_text = thinking_call(False) + had_think_on = ("" in on_text) or len(on_text) > 80 + if not had_think_on: + print( + f"[tools] WARN enable_thinking=True produced no thinking signal: " + f"{on_text[:200]!r}" + ) + assert "" not in off_text, ( + f"enable_thinking=False but still present: {off_text!r}" + ) + print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)") + PY + + - name: Stop Studio + if: always() + # Run as cmd so we are not running through the Git Bash shell; + # Git Bash on windows-latest has been observed to exit 143 + # (SIGTERM) from any inline kill/sleep block, masking a green + # test run. The runner reclaims the Studio child process at + # job end either way, so just emit a marker and exit 0. + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + + - name: Collect llama-server logs + if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true + shell: bash + # Copy llama-server's own stdout/stderr (teed by Studio under + # ~/.unsloth/studio/logs/llama-server/) into the workspace so + # upload-artifact can pick it up. Crucial for diagnosing a + # subprocess crash where Studio's traceback only shows the + # symptom (httpx ReadError) but not the cause. + run: | + mkdir -p logs/llama-server + cp -v ~/.unsloth/studio/logs/llama-server/*.log logs/llama-server/ 2>/dev/null || \ + echo "no llama-server logs to collect" + + - name: Upload logs + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-tool-calling-log + path: | + logs/studio.log + logs/install.log + logs/llama-server/*.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job 3: JSON, images + # ───────────────────────────────────────────────────────────────────── + json-images: + name: JSON, images + runs-on: windows-latest + timeout-minutes: 35 + defaults: + run: + shell: bash + env: + GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF + GGUF_VARIANT: UD-IQ2_XXS + GGUF_FILE: Qwen3-VL-2B-Instruct-UD-IQ2_XXS.gguf + MMPROJ_FILE: mmproj-F16.gguf + STUDIO_PORT: '18899' + HF_HOME: ${{ github.workspace }}/hf-cache + # Force UTF-8 for stdio (Windows defaults to cp1252; hf + # download / Studio CLI print "✓" checkmarks and crash + # otherwise). + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + # Split restore + save so a transient restore-side failure does not + # kill the whole job. See the matching block in the tool-calling job + # for the full rationale (actions/cache#1621). This is the block that + # actually broke in run 25713577488: "Cache hit for: " was + # logged, the step exited non-zero in ~0.3 s without extracting the + # 3.4 GiB archive, and steps 6-15 were skipped. + - name: Restore HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj) + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 + + - name: Prime HF_HOME with the GGUF + mmproj + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj) + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + # See studio-windows-update-smoke.yml for the full rationale. + # tl;dr: setup.ps1 needs npm >=11 to skip a 35 s winget Node + # reinstall, and Defender's real-time scan dominates the + # frontend / uv-pip-extract steps. + run: | + $ProgressPreference = 'SilentlyContinue' + Write-Host "npm version before upgrade: $(npm -v)" + npm install -g 'npm@^11' 2>&1 | Out-Host + Write-Host "npm version after upgrade: $(npm -v)" + # NOTE: do NOT pre-create these directories. See + # studio-windows-update-smoke.yml for the full rationale -- + # creating an empty studio/frontend/dist trips setup.ps1's + # mtime-based staleness check into "frontend up to date, skip + # rebuild" and Studio boots with an empty dist directory. + # Add-MpPreference accepts paths that do not yet exist. + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { + Add-MpPreference -ExclusionPath $p -ErrorAction Stop + Write-Host "Defender exclusion added: $p" + } catch { + Write-Host "Defender exclusion skipped ($($_.Exception.Message)): $p" + } + } + + - name: Install Studio (--local, --no-torch) + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + # *>&1 captures Write-Host (Information stream) output; + # plain 2>&1 does not. setup.ps1 emits "prebuilt installed + # and validated" via Write-Host, and we grep for that. + $ProgressPreference = 'SilentlyContinue' + & ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log + + - name: Assert install.ps1 used the Windows llama.cpp prebuilt + run: | + # Filesystem check; setup.ps1's stream output isn't captured. + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp on Windows." + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + if [ ! -f "$INFO" ]; then + echo "::error::no UNSLOTH_PREBUILT_INFO.json at $INFO." + ls -la "$LLAMA_DIR" || true + exit 1 + fi + if [ ! -f "$BIN" ]; then + echo "::error::no llama-server.exe at $BIN." + ls -la "$LLAMA_DIR/build/bin" || true + exit 1 + fi + echo "install.ps1 installed the Windows prebuilt llama.cpp:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + run: | + SHIM_DIR=~/.unsloth/studio/bin + if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then + echo "::error::unsloth.exe shim not found at $SHIM_DIR" + ls -la ~/.unsloth/studio/ || true + exit 1 + fi + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: Install OpenAI + Anthropic Python SDKs + run: python -m pip install 'openai>=1.50' 'anthropic>=0.40' + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, change password, load model + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIJson-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + # Retry: same rationale as the OpenAI/Anthropic and Tool calling jobs. + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 900 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP; response:" + cat /tmp/load.json || true + sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_vision}' /tmp/load.json + + - name: JSON schema decoding + image input + env: + BASE_URL: http://127.0.0.1:18899 + run: | + python - <<'PY' + import base64 + import json + import os + import time + import urllib.error + import urllib.request + from openai import OpenAI + from anthropic import Anthropic + + BASE = os.environ["BASE_URL"] + KEY = os.environ["API_KEY"] + SEED = 3407 + TEMP = 0.2 + + def post(path, body, *, timeout = 240): + req = urllib.request.Request( + f"{BASE}{path}", + data = json.dumps(body).encode(), + method = "POST", + headers = { + "Authorization": f"Bearer {KEY}", + "Content-Type": "application/json", + }, + ) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) + + # ── 1. response_format = json_object (JSON mode) ───────────── + status, data = post("/v1/chat/completions", { + "model": "default", + "messages": [ + {"role": "system", "content": 'Reply with a single JSON object of the form {"city": "...", "country": "..."}. Output ONLY the JSON, nothing else.'}, + {"role": "user", "content": "What is the capital of France?"}, + ], + "temperature": TEMP, + "max_tokens": 600, + "seed": SEED, + "stream": False, + "enable_thinking": False, + "response_format": {"type": "json_object"}, + }, timeout = 600) + assert status == 200, f"json status {status}: {data}" + assert ( + isinstance(data.get("choices"), list) + and data["choices"] + and "message" in data["choices"][0] + ), f"json response envelope malformed: {data}" + content = (data["choices"][0]["message"].get("content") or "").strip() + print(f"[json] raw json_object content: {content!r}") + if content.startswith("```"): + content = content.split("```", 2)[1] + if content.startswith("json"): + content = content[4:] + content = content.strip("`\n ") + if content: + try: + parsed = json.loads(content) + if "paris" in str(parsed.get("city", "")).lower(): + print(f"[json] PASS json_object -> {parsed}") + else: + print(f"[json] WARN json_object decoded but city!=Paris: {parsed}") + except json.JSONDecodeError as exc: + print(f"[json] WARN json_object content not parseable ({exc}); content={content!r}") + else: + print("[json] WARN json_object produced empty content") + + status2, data2 = post("/v1/chat/completions", { + "model": "default", + "messages": [{"role": "user", "content": "What is the capital of France? Answer with one word."}], + "temperature": TEMP, + "max_tokens": 400, + "seed": SEED, + "stream": False, + "enable_thinking": False, + }, timeout = 600) + assert status2 == 200, f"plain status {status2}: {data2}" + plain = (data2["choices"][0]["message"].get("content") or "").lower() + print(f"[json] plain capital-of-france reply: {plain!r}") + if "paris" in plain: + print("[json] PASS plain inference path (paris mentioned)") + else: + print( + f"[json] WARN plain inference returned no 'paris' -- " + f"model output drift. HTTP path validated separately above." + ) + + # ── 2. OpenAI image_url (data URI base64) ─────────────────── + PNG_64X64_RED_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAYklEQVR4nO3PMQ0AIADAMEAI/k" + "UhBhEcDcmqYJtn7/GzpQNeNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA" + "1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaA1oDWgNaBdCJ0BmMJ25zMAAAAASUVORK5CYII=" + ) + data_uri = f"data:image/png;base64,{PNG_64X64_RED_B64}" + + # On Windows + the Qwen3-VL mmproj, llama.cpp's vision + # path runs on CPU (no Metal involvement). The wrapper is + # kept for resilience but the vision path is expected to + # work on Windows; an exception here is a real regression. + client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY) + try: + openai_resp = client.chat.completions.create( + model = "default", + temperature = TEMP, + max_tokens = 80, + seed = SEED, + messages = [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_uri}}, + {"type": "text", "text": "What colour dominates this image? Reply in one word."}, + ], + }], + ) + openai_text = (openai_resp.choices[0].message.content or "").lower() + print(f"[image/openai] reply: {openai_text!r}") + if openai_text: + print("[image/openai] PASS image_url accepted, non-empty response") + else: + print("[image/openai] WARN image_url accepted but empty content") + except Exception as exc: + print( + f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: " + f"{exc}. Studio successfully forwarded the request; failure here is " + f"upstream llama.cpp vision behaviour." + ) + + # ── 3. Anthropic source/base64 image ──────────────────────── + anthropic = Anthropic( + base_url = BASE, + api_key = "unused", + default_headers = {"Authorization": f"Bearer {KEY}"}, + ) + try: + a_msg = anthropic.messages.create( + model = "default", + max_tokens = 80, + temperature = TEMP, + extra_body = {"seed": SEED}, + messages = [{ + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": PNG_64X64_RED_B64, + }, + }, + {"type": "text", "text": "Describe this image briefly."}, + ], + }], + ) + a_text = "".join(b.text for b in a_msg.content if getattr(b, "type", None) == "text") + print(f"[image/anthropic] reply: {a_text!r}") + if a_text: + print("[image/anthropic] PASS source/base64 accepted, non-empty response") + else: + print("[image/anthropic] WARN source/base64 accepted but empty content") + except Exception as exc: + print( + f"[image/anthropic] WARN anthropic image SDK call raised: " + f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp vision " + f"behaviour, NOT a Studio regression." + ) + PY + + - name: Stop Studio + if: always() + # Run as cmd so we are not running through the Git Bash shell; + # Git Bash on windows-latest has been observed to exit 143 + # (SIGTERM) from any inline kill/sleep block, masking a green + # test run. The runner reclaims the Studio child process at + # job end either way, so just emit a marker and exit 0. + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + + - name: Collect llama-server logs + if: always() + # A transient Windows DLL-init crash (0xC0000142) in this diagnostic + # copy must not fail an otherwise-green job. + continue-on-error: true + shell: bash + # Copy llama-server's own stdout/stderr (teed by Studio under + # ~/.unsloth/studio/logs/llama-server/) into the workspace so + # upload-artifact can pick it up. Crucial for diagnosing a + # subprocess crash where Studio's traceback only shows the + # symptom (httpx ReadError) but not the cause. + run: | + mkdir -p logs/llama-server + cp -v ~/.unsloth/studio/logs/llama-server/*.log logs/llama-server/ 2>/dev/null || \ + echo "no llama-server logs to collect" + + - name: Upload logs + if: always() + # Diagnostic only: a transient artifact-service drop must not fail a green job. + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-json-images-log + path: | + logs/studio.log + logs/install.log + logs/llama-server/*.log + retention-days: 7 + + # ── folded from studio-windows-no-vs-smoke.yml: install + run with no Visual Studio ── + no-vs-cpu: + name: Studio install + inference without Visual Studio + runs-on: windows-latest + timeout-minutes: 35 + defaults: + run: + shell: bash + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18820' + HF_HOME: ${{ github.workspace }}/hf-cache + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + run: | + $ProgressPreference = 'SilentlyContinue' + npm install -g 'npm@^11' 2>&1 | Out-Host + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { } + } + + - name: Prepare no-build-tools simulation + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { + [void] $blocked.Add( + [Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\')) + } + } + } + } + # Normalized comparison so registry spellings (trailing slash, + # unexpanded %VAR%) still match. + function Test-Blocked([string]$p) { + $n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\') + return $blocked.Contains($n) + } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not (Test-Blocked $_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + # install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment + # rebuild the session Path from these scopes mid-install, so filter + # them too. Originals are saved for the cleanup step. + foreach ($scope in @('Machine', 'User')) { + $orig = [Environment]::GetEnvironmentVariable('Path', $scope) + if (-not $orig) { continue } + Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline + $kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';' + [Environment]::SetEnvironmentVariable('Path', $kept, $scope) + Write-Host "Filtered $scope Path scope." + } + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH<&1 | Tee-Object -FilePath logs/install.log + + - name: Assert prebuilt used AND no build tools were installed + run: | + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + fail=0 + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp without VS."; fail=1 + fi + # The deferred build-tool installs must NOT run on the prebuilt path. + for pat in "Kitware.CMake" "Microsoft.VisualStudio.2022.BuildTools" "installing via winget"; do + if grep -qi "$pat" logs/install.log; then + echo "::error::unexpected build-tool install on the prebuilt path: '$pat'"; fail=1 + fi + done + [ -f "$INFO" ] || { echo "::error::no UNSLOTH_PREBUILT_INFO.json"; ls -la "$LLAMA_DIR" || true; fail=1; } + [ -f "$BIN" ] || { echo "::error::no llama-server.exe"; ls -la "$LLAMA_DIR/build/bin" || true; fail=1; } + if [ "$fail" != "0" ]; then grep -iE "cmake|visual studio|prebuilt|source build" logs/install.log | tail -60; exit 1; fi + echo "Prebuilt installed with no build tools:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + run: | + SHIM_DIR=~/.unsloth/studio/bin + [ -f "$SHIM_DIR/unsloth.exe" ] || { echo "::error::unsloth.exe shim not found"; ls -la ~/.unsloth/studio/ || true; exit 1; } + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, load the GGUF + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json || { tail -200 logs/studio.log; exit 1; } + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CINoVS-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP"; cat /tmp/load.json || true; sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_gguf}' /tmp/load.json + + - name: Inference works via the prebuilt llama.cpp (no VS) + run: | + RESP=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/v1/chat/completions" \ + -H "Authorization: Bearer $API_KEY" -H 'content-type: application/json' \ + --max-time 240 \ + -d '{"model":"default","messages":[{"role":"user","content":"What is 1+1? Answer briefly."}],"temperature":0,"max_tokens":32,"stream":false}') + echo "$RESP" | jq '.choices[0].message' || { echo "$RESP"; exit 1; } + CONTENT=$(echo "$RESP" | jq -r '.choices[0].message.content') + [ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; } + echo "Inference OK without Visual Studio: $CONTENT" + + - name: Clean no-build-tools simulation + if: always() + shell: pwsh + run: | + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + foreach ($scope in @('Machine', 'User')) { + $saved = Join-Path $root "orig-path-$scope.txt" + if (Test-Path -LiteralPath $saved) { + [Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope) + Write-Host "Restored $scope Path scope." + } + } + Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue + + - name: Stop Studio + if: always() + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + + - name: Collect llama-server logs + if: always() + continue-on-error: true + run: | + mkdir -p logs/llama-server + cp -v ~/.unsloth/studio/logs/llama-server/*.log logs/llama-server/ 2>/dev/null || echo "no llama-server logs" + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-no-vs-cpu-log + path: | + logs/install.log + logs/studio.log + logs/llama-server/*.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job B: the GPU (CUDA) prebuilt path is also VS-free (resolve/availability) + # ───────────────────────────────────────────────────────────────────── + no-vs-gpu-resolve: + name: GPU prebuilt resolves without Visual Studio + runs-on: windows-latest + timeout-minutes: 15 + defaults: + run: + shell: bash + env: + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Prepare no-build-tools simulation + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { [void] $blocked.Add($dir) } + } + } + } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not $blocked.Contains($_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH< /tmp/rel.json + echo "release: $(jq -r .tag_name /tmp/rel.json)" + ASSETS=$(jq -r '.assets[].name' /tmp/rel.json) + echo "$ASSETS" | grep -iE 'windows-x64-cuda[0-9]' || { + echo "::error::no Windows x64 CUDA prebuilt asset found in unslothai/llama.cpp latest release" + echo "$ASSETS"; exit 1; } + # AMD parity: hosted runners have no AMD GPU, so the resolver step below + # can't exercise the ROCm path (it resolves to CPU). Pin the per-gfx + # Windows ROCm bundles here so a release that drops them fails loudly -- + # the AMD no-VS guarantee otherwise rides only on shared resolver code. + echo "$ASSETS" | grep -iE 'windows-x64-rocm-gfx' || { + echo "::error::no Windows x64 ROCm (per-gfx) prebuilt asset found in unslothai/llama.cpp latest release" + echo "$ASSETS"; exit 1; } + echo "Windows CUDA and ROCm prebuilts are available -- GPU users get them without compiling." + + - name: The prebuilt resolver runs without Visual Studio + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + $ErrorActionPreference = 'Stop' + # pwsh: bash cannot export `ProgramFiles(x86)`; set in-script so the + # python child inherits the overrides. + $env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES + ${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86 + $env:Path = $env:NO_BUILD_TOOLS_PATH + # Resolver-only (no GPU on hosted runners, so the host resolves to the + # CPU bundle). The point is that resolution needs no compiler/VS. + python -m pip install --upgrade huggingface_hub + if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 } + python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json + if ($LASTEXITCODE -ne 0) { + Write-Host "::error::resolver exited non-zero" + if (Test-Path resolve.json) { Get-Content resolve.json } + exit 1 + } + Get-Content resolve.json + Write-Host "Prebuilt resolver ran with no Visual Studio present." + + - name: Clean no-build-tools simulation + if: always() + shell: pwsh + run: | + Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue + + # ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ── + pester: + name: setup.ps1 unit tests (VS 2026 / CMake guard) + runs-on: windows-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Install Pester v5 + shell: pwsh + run: | + # PSGallery is intermittently absent from the repository list on GitHub's Windows + # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the + # name 'PSGallery' was found." Re-register the default gallery first so the policy + # change and module install below always have a repository to target. + if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) { + Register-PSRepository -Default -ErrorAction SilentlyContinue + } + Set-PSRepository PSGallery -InstallationPolicy Trusted + Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser + Import-Module Pester -MinimumVersion 5.5.0 + Get-Module Pester | Select-Object Name, Version | Format-Table + + - name: Run Pester suite + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $testDir = Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1' + if (-not (Test-Path $testDir)) { + Write-Error "Test directory not found: $testDir" + exit 1 + } + $cfg = New-PesterConfiguration + $cfg.Run.Path = $testDir + $cfg.Run.Exit = $true # non-zero exit => job fails + $cfg.Run.Throw = $true # also throw on test failure / 0 tests + $cfg.TestResult.Enabled = $true + $cfg.TestResult.OutputFormat = 'NUnitXml' + $cfg.TestResult.OutputPath = Join-Path $env:GITHUB_WORKSPACE 'pester-results.xml' + $cfg.Output.Verbosity = 'Detailed' + Invoke-Pester -Configuration $cfg + + - name: Upload Pester results + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: pester-results-setup-ps1 + path: pester-results.xml + if-no-files-found: warn + + vs-integration: + # Real detection against the VS installed on the runner image (no mocks). + name: real-VS detection (${{ matrix.label }}) + strategy: + fail-fast: false + matrix: + include: + - { os: windows-2022, label: 'VS 2022', expectGen: 'Visual Studio 17 2022', expectToolset: 'v170' } + - { os: windows-2025-vs2026, label: 'VS 2026', expectGen: 'Visual Studio 18 2026', expectToolset: 'v180' } + runs-on: ${{ matrix.os }} + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Detect the real Visual Studio with setup.ps1 functions + shell: pwsh + env: + EXPECT_GEN: ${{ matrix.expectGen }} + EXPECT_TOOLSET: ${{ matrix.expectToolset }} + run: | + $ErrorActionPreference = 'Stop' + . (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1') + $setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1' + foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Get-VcBuildCustomizationsDir', 'Find-VsBuildTools')) { + . ([scriptblock]::Create((Get-FunctionSource -Path $setup -Name $fn))) + } + + # Ground truth from the real vswhere (independent of our code), for visibility. + $vsw = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" + if (Test-Path $vsw) { + $year = (& $vsw -latest -property catalog_productLineVersion 2>$null | Select-Object -First 1) + $path = (& $vsw -latest -property installationPath 2>$null | Select-Object -First 1) + Write-Host "Real vswhere: productLineVersion='$year' installPath='$path'" + } else { + Write-Host "vswhere not present at $vsw (relying on filesystem fallback)" + } + + # Our detection must find the real VS and report the expected generator. + $r = Find-VsBuildTools + if (-not $r) { throw "Find-VsBuildTools returned null on a host with real $env:EXPECT_GEN" } + Write-Host "Find-VsBuildTools -> Generator='$($r.Generator)' Source='$($r.Source)' InstallPath='$($r.InstallPath)'" + if ($r.Generator -ne $env:EXPECT_GEN) { + throw "Detection mismatch: got '$($r.Generator)', expected '$env:EXPECT_GEN'" + } + if (-not (Test-Path $r.InstallPath)) { throw "Detected InstallPath does not exist: $($r.InstallPath)" } + + # Toolset path derivation must match the expected v-number... + $bc = Get-VcBuildCustomizationsDir -VsInstallPath $r.InstallPath -Generator $r.Generator + $derived = Split-Path (Split-Path $bc -Parent) -Leaf # e.g. v170 / v180 + Write-Host "Get-VcBuildCustomizationsDir -> '$bc' (toolset='$derived')" + if ($derived -ne $env:EXPECT_TOOLSET) { + throw "Toolset mismatch: derived '$derived', expected '$env:EXPECT_TOOLSET'" + } + + # ...and that v-number is a real folder on the VS install (where CUDA's + # BuildCustomizations would land). + $vcRoot = Join-Path $r.InstallPath 'MSBuild\Microsoft\VC' + if (Test-Path $vcRoot) { + $realToolsets = @((Get-ChildItem -Path $vcRoot -Directory -ErrorAction SilentlyContinue).Name) + Write-Host "Real VC toolset dirs: $($realToolsets -join ', ')" + if ($realToolsets -notcontains $derived) { + throw "Derived toolset '$derived' is not present on the real $env:EXPECT_GEN install (have: $($realToolsets -join ', '))" + } + Write-Host "OK: toolset '$derived' exists on the real VS install." + } else { + Write-Warning "VC MSBuild root absent ($vcRoot) - C++ workload not installed; skipping on-disk toolset check." + } + + Write-Host "PASS: real $env:EXPECT_GEN detected correctly with toolset '$derived'." + + vcredist-clean-box: + # Validate Test-VCRedistInstalled + Ensure-VCRedist on a throwaway runner: + # present on the stock image, fires on a clean box (signals removed restorably), + # then a literal uninstall/reinstall round trip. Always restored before the end. + name: VC++ runtime detect + install round-trip (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [windows-latest, windows-2025-vs2026] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Detect present, fire on a clean box, and round-trip the install + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + . (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1') + $setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1' + # Dot-source the guard + the logging closure it reaches + # (step/substep -> Write-StudioStdoutMirror / Get-StudioAnsi). + $script:StudioVtOk = $false + $script:UnslothVerbose = $false + foreach ($fn in @('Get-StudioAnsi', 'Write-StudioStdoutMirror', 'step', 'substep', + 'Invoke-SetupCommand', 'Refresh-Environment', + 'Test-VCRedistInstalled', 'Ensure-VCRedist')) { + $src = Get-FunctionSource -Path $setup -Name $fn + if (-not $src) { throw "Function '$fn' not found in setup.ps1" } + . ([scriptblock]::Create($src)) + } + + $regKeys = @( + 'HKLM\SOFTWARE\Microsoft\VisualStudio\14.0\VC\Runtimes\x64', + 'HKLM\SOFTWARE\WOW6432Node\Microsoft\VisualStudio\14.0\VC\Runtimes\x64' + ) + function Show-GroundTruth { + $dll = Join-Path $env:SystemRoot 'System32\vcruntime140_1.dll' + Write-Host (" System32\vcruntime140_1.dll present: {0}" -f (Test-Path $dll)) + foreach ($k in $regKeys) { + $r = Get-ItemProperty -Path "HKLM:\$($k.Substring(5))" -ErrorAction SilentlyContinue + if ($r) { Write-Host (" {0}: Installed={1} {2}.{3}" -f $k, $r.Installed, $r.Major, $r.Minor) } + else { Write-Host (" {0}: (absent)" -f $k) } + } + } + + Write-Host '== A. Detection on the stock runner (expect present) ==' + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'Test-VCRedistInstalled reported ABSENT on a stock runner that ships the VC++ runtime (detection regression).' } + Write-Host ' Test-VCRedistInstalled -> present OK' + + Write-Host '== B. Genuinely clean box (restorable): detection must FIRE ==' + $scratch = Join-Path $env:RUNNER_TEMP 'cleanwin' + New-Item -ItemType Directory -Force -Path (Join-Path $scratch 'System32') | Out-Null + $backup = Join-Path $env:RUNNER_TEMP 'vcreg_backup' + New-Item -ItemType Directory -Force -Path $backup | Out-Null + $origSysRoot = $env:SystemRoot + try { + for ($i = 0; $i -lt $regKeys.Count; $i++) { + reg query $regKeys[$i] *> $null + if ($LASTEXITCODE -eq 0) { + reg export $regKeys[$i] (Join-Path $backup "$i.reg") /y *> $null + reg delete $regKeys[$i] /f *> $null + } + } + $env:SystemRoot = $scratch + if (Test-VCRedistInstalled) { throw 'Detection still PRESENT after both signals were removed (it would never trigger an install on a clean box).' } + Write-Host ' Test-VCRedistInstalled -> absent OK (detection fires on a clean box)' + } finally { + $env:SystemRoot = $origSysRoot + for ($i = 0; $i -lt $regKeys.Count; $i++) { + $f = Join-Path $backup "$i.reg" + if (Test-Path $f) { reg import $f *> $null } + } + } + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'Detection did not recover after restoring the registry (test restore bug).' } + + Write-Host '== C. Literal uninstall on this throwaway VM (official installer), observe detection ==' + $exe = Join-Path $env:RUNNER_TEMP 'vc_redist.x64.exe' + Invoke-WebRequest -Uri 'https://aka.ms/vs/17/release/vc_redist.x64.exe' -OutFile $exe + Start-Process -FilePath $exe -ArgumentList '/uninstall', '/quiet', '/norestart' -Wait + Show-GroundTruth + Write-Host (" Test-VCRedistInstalled after uninstall -> {0}" -f (Test-VCRedistInstalled)) + if (Test-VCRedistInstalled) { + Write-Host ' Note: the Visual Studio on this image ref-counts the runtime, so the package' + Write-Host ' uninstall is a no-op here; section B already proved detection on a clean box.' + } + + Write-Host '== D. Restore via Ensure-VCRedist (winget product path), installer fallback if needed ==' + Ensure-VCRedist + if (-not (Test-VCRedistInstalled)) { + Write-Host ' winget path did not restore it; using the official installer to close the round trip.' + Start-Process -FilePath $exe -ArgumentList '/install', '/quiet', '/norestart' -Wait + } + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'VC++ runtime could not be restored after the uninstall round-trip.' } + Write-Host ' Test-VCRedistInstalled -> present OK' + Write-Host 'PASS: detection is correct on a real install, fires on a clean box, and the install round-trip restores the runtime.' diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml new file mode 100644 index 0000000..4053099 --- /dev/null +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Windows counterpart to studio-ui-smoke.yml / studio-mac-ui-smoke.yml. +# Same Playwright + Chromium end-to-end chat UI flow + extra UI flow, +# but on the FREE windows-latest runner so we catch Windows-specific +# regressions in the install path (install.ps1), the Studio CLI's +# Windows process-management branches, and the llama.cpp prebuilt's +# Windows HTTP layer. + +name: Windows Studio UI CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.ps1' + - 'pyproject.toml' + - 'tests/studio/**' + - '.github/workflows/studio-windows-ui-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + ui-smoke: + name: Chat UI Tests + runs-on: windows-latest + timeout-minutes: 45 + # Default every step's shell to Git Bash. windows-latest's default + # shell is pwsh; without this each curl / heredoc / `kill $PID` + # step would need its own `shell: bash`. Steps that genuinely + # need PowerShell (install.ps1 invocation) override per-step. + defaults: + run: + shell: bash + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18896' + HF_HOME: ${{ github.workspace }}/hf-cache + # Force UTF-8 for stdio so Python tools (hf download, Studio + # CLI, etc.) can print Unicode characters like the success + # checkmark "✓". Windows defaults to cp1252 / charmap and + # any tool that prints "OK ✓" hits a UnicodeEncodeError. + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + # No `cache: 'npm'`. setup-node's npm cache restore silently + # aborts the entire job on Windows runners when the npm cache + # path (`C:\npm\cache` per `npm config get cache`) doesn't yet + # exist on a fresh runner -- the step exits without an error + # message and every following step gets skipped. See + # npm/cli#7308. The frontend `npm ci` is fast enough without + # the cache that the reliability gain is worth the ~30s. + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + # No `cache: 'pip'`. install.ps1 / setup.ps1 use uv and + # never populate ~/.cache/pip; setup-python's post-step + # then fatal-errors with "Cache folder path is retrieved + # for pip but doesn't exist on disk". + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + # See studio-windows-update-smoke.yml for the full rationale. + # tl;dr: setup.ps1 needs npm >=11 to skip a 35 s winget Node + # reinstall, and Defender's real-time scan dominates the + # frontend / uv-pip-extract steps. + run: | + $ProgressPreference = 'SilentlyContinue' + Write-Host "npm version before upgrade: $(npm -v)" + npm install -g 'npm@^11' 2>&1 | Out-Host + Write-Host "npm version after upgrade: $(npm -v)" + # NOTE: do NOT pre-create these directories. See + # studio-windows-update-smoke.yml for the full rationale -- + # creating an empty studio/frontend/dist trips setup.ps1's + # mtime-based staleness check into "frontend up to date, skip + # rebuild" and Studio boots with an empty dist directory. + # Add-MpPreference accepts paths that do not yet exist. + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { + Add-MpPreference -ExclusionPath $p -ErrorAction Stop + Write-Host "Defender exclusion added: $p" + } catch { + Write-Host "Defender exclusion skipped ($($_.Exception.Message)): $p" + } + } + + - name: Seed a legacy launch-studio.vbs (upgrade-cleanup check) + # Simulate a pre-hardening install so the post-install assertion below + # proves the installer DELETES an existing launch-studio.vbs (the exact + # Kaspersky-flagged file), not merely stops generating it. + shell: pwsh + run: | + $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio' + New-Item -ItemType Directory -Force -Path $appDir | Out-Null + Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode + Write-Host "seeded legacy launch-studio.vbs at $appDir" + + - name: Install Studio (--local, --no-torch) + # install.ps1 is the supported Windows installer. install.sh + # has no Windows branch (apt-get / brew calls). The PS1 + # script's `Install-UnslothStudio @args` line at the bottom + # forwards `--local --no-torch` correctly. + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + # *>&1 redirects ALL PowerShell streams (stdout, stderr, + # warning, verbose, debug, information) into the success + # stream so Tee-Object captures everything. install.ps1 + # and setup.ps1 emit step/substep markers via Write-Host + # which lands on the Information stream (PS 5+); without + # the wildcard redirect, those markers (including + # "prebuilt installed and validated") never reach + # logs/install.log and the post-step grep asserter fails. + $ProgressPreference = 'SilentlyContinue' + & ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log + + - name: Assert install.ps1 used the Windows llama.cpp prebuilt + run: | + # install.ps1's setup.ps1 child writes "prebuilt installed + # and validated" to its own console host -- that output + # does NOT come back through this parent step's stdout + # pipeline (no matter how aggressively we redirect: *>&1, + # tee, etc.). Verify the install via the filesystem + # instead. setup.ps1 writes UNSLOTH_PREBUILT_INFO.json + # next to the install dir on success, and lays the + # binaries under build/bin/Release/ on Windows. + STUDIO_HOME=~/.unsloth/studio + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + # Source-build fallback grep stays as a fast bail-out. + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp on Windows." + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + if [ ! -f "$INFO" ]; then + echo "::error::no UNSLOTH_PREBUILT_INFO.json at $INFO; setup.ps1 didn't install the prebuilt." + ls -la "$LLAMA_DIR" || true + exit 1 + fi + if [ ! -f "$BIN" ]; then + echo "::error::no llama-server.exe at $BIN; prebuilt extraction incomplete." + ls -la "$LLAMA_DIR/build/bin" || true + ls -la "$LLAMA_DIR/build/bin/Release" || true + exit 1 + fi + echo "install.ps1 installed the Windows prebuilt llama.cpp:" + cat "$INFO" + + - name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut) + # The shortcut launch path is otherwise untested here (the steps below + # boot `unsloth studio` directly). Guard against re-introducing the VBS + # that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk + # pointing anywhere other than hidden PowerShell over launch-studio.ps1. + shell: pwsh + run: | + $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio' + if (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.vbs')) { + throw "regression: launch-studio.vbs exists (the Kaspersky VBS-FP shape)" + } + if (-not (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.ps1'))) { + throw "missing launch-studio.ps1 in $appDir" + } + $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk' + if (-not (Test-Path -LiteralPath $lnk)) { + $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk' + } + if (-not (Test-Path -LiteralPath $lnk)) { throw "no Unsloth Studio.lnk on Desktop or Start Menu" } + $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk) + Write-Host "shortcut target: $($sc.TargetPath)" + Write-Host "shortcut args: $($sc.Arguments)" + if ($sc.TargetPath -match 'wscript\.exe$') { throw "shortcut still targets wscript.exe (VBS host)" } + if ($sc.TargetPath -notmatch 'powershell\.exe$') { throw "unexpected shortcut target: $($sc.TargetPath)" } + if ($sc.Arguments -notmatch '-WindowStyle Hidden') { + throw "shortcut must launch windowless (-WindowStyle Hidden)" + } + Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)" + + - name: Launch Studio via the shortcut and assert health + # Run the exact command the .lnk stores (hidden PowerShell over + # launch-studio.ps1) and confirm it brings the backend up. This is the + # only step that proves the shortcut launch is not silently broken. + # Default port range is 8888-8908; the later UI tests use 18896/18897, so + # there is no conflict, and we tear this server down before they boot. + shell: pwsh + run: | + $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk' + if (-not (Test-Path -LiteralPath $lnk)) { + $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk' + } + $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk) + Write-Host "launching: $($sc.TargetPath) $($sc.Arguments)" + Start-Process -FilePath $sc.TargetPath -ArgumentList $sc.Arguments -WorkingDirectory $sc.WorkingDirectory + $foundPort = 0 + foreach ($i in 1..180) { + foreach ($port in 8888..8908) { + try { + $r = Invoke-RestMethod -Uri "http://127.0.0.1:$port/api/health" -TimeoutSec 1 + if ($r.status -eq 'healthy' -and $r.service -eq 'Unsloth UI Backend') { $foundPort = $port; break } + } catch {} + } + if ($foundPort) { break } + Start-Sleep -Seconds 1 + } + # Tear down the shortcut-launched server before the main UI tests boot. + try { + $owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess + if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null } + } catch {} + if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" } + Write-Host "Studio healthy on port $foundPort (launched via the shortcut)" + + - name: Add Studio shim to GITHUB_PATH + # install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe + # and adds that dir to the User PATH via the Windows registry. + # Registry-level PATH updates don't propagate to a running + # Git Bash session, so the next step's `unsloth ...` invocation + # would hit "command not found". Re-export the shim dir to + # GITHUB_PATH so every subsequent step in this job sees it. + run: | + SHIM_DIR=~/.unsloth/studio/bin + if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then + echo "::error::unsloth.exe shim not found at $SHIM_DIR" + ls -la ~/.unsloth/studio/ || true + exit 1 + fi + # GITHUB_PATH wants Windows-style paths; convert via cygpath. + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + echo "Added Studio shim dir to PATH: $(cygpath -w "$SHIM_DIR")" + + - name: Install Playwright + Chromium + # No --with-deps on Windows: that flag installs Linux apt + # packages. windows-latest ships the system frameworks + # Chromium needs (Edge / WebView2) already. + run: | + python -m pip install 'playwright>=1.45' + python -m playwright install chromium + + - name: Reset auth + boot Studio + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json + + - name: Pass bootstrap password to the Playwright step + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + NEW2="CIUi-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "::add-mask::$NEW2" + echo "STUDIO_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" + echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" + + - name: Drive the chat UI with Playwright + env: + BASE_URL: http://127.0.0.1:18896 + PW_ART_DIR: logs/playwright + STUDIO_UI_STRICT: '1' + # windows-latest free runner is 4 vCPU / 16 GB; gemma-3- + # 270m turn latency under llama-server's CPU backend can + # crowd the 180s default (slower than ubuntu-latest on + # the same model). Keep the same generous budget the Mac + # job uses. + STUDIO_UI_TURN_TIMEOUT_MS: '540000' + run: | + mkdir -p logs/playwright + python tests/studio/playwright_chat_ui.py + + - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders) + if: always() + run: | + kill "${STUDIO_PID}" 2>/dev/null || true + sleep 2 + + - name: Reset auth + boot Studio for extra UI tests (port 18897) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \ + > logs/studio_extra.log 2>&1 & + echo "STUDIO_EXTRA_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18897 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18897/api/health" > /tmp/health2.json; then + jq -e '.status == "healthy"' /tmp/health2.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health2.json + + - name: Pass bootstrap pw for extra UI test + run: | + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIUiExtra-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV" + echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright + env: + BASE_URL: http://127.0.0.1:18897 + STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }} + STUDIO_NEW_PW: ${{ env.STUDIO_EXTRA_NEW_PW }} + PW_ART_DIR: logs/playwright_extra + STUDIO_UI_STRICT: '1' + STUDIO_UI_TURN_TIMEOUT_MS: '540000' + GGUF_REPO: ${{ env.GGUF_REPO }} + GGUF_VARIANT: ${{ env.GGUF_VARIANT }} + run: | + mkdir -p logs/playwright_extra + python tests/studio/playwright_extra_ui.py + + - name: Stop second Studio + if: always() + run: | + kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true + sleep 2 + + - name: Upload Playwright artifacts + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-studio-ui-smoke-artifacts + path: | + logs/studio.log + logs/studio_extra.log + logs/install.log + logs/playwright + logs/playwright_extra + retention-days: 7 diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml new file mode 100644 index 0000000..5b92f1a --- /dev/null +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Windows counterpart to studio-update-smoke.yml / +# studio-mac-update-smoke.yml. Verifies that on the FREE +# windows-latest runner: +# +# 1. install.ps1 --local --no-torch installs Studio AND auto-fetches +# the prebuilt llama.cpp Windows binary (app--windows-x64-cpu +# from unslothai/llama.cpp). Hitting the source-build fallback is +# treated as an Unsloth bug -- Studio must always pick the +# prebuilt on Windows. +# 2. unsloth studio update --local is idempotent. Two consecutive +# runs both report "prebuilt up to date and validated", no +# source-build fallback. The CLI's _find_setup_script picks +# setup.ps1 on Windows automatically. +# 3. The installed Studio still boots and /api/health returns +# healthy after the update path. + +name: Windows Studio Update CI + +on: + pull_request: + paths: + - 'install.ps1' + - 'scripts/uninstall.ps1' + - 'studio/setup.ps1' + - 'studio/setup.bat' + - 'studio/install_python_stack.py' + - 'studio/install_llama_prebuilt.py' + - 'studio/backend/requirements/**' + - 'unsloth_cli/commands/studio.py' + - 'pyproject.toml' + - '.github/workflows/studio-windows-update-smoke.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + update-idempotency: + name: Studio Updating Tests + runs-on: windows-latest + timeout-minutes: 30 + defaults: + run: + shell: bash + env: + # Force UTF-8 for stdio (Windows defaults to cp1252; hf + # download / Studio CLI print "✓" checkmarks and crash + # otherwise). + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + # Don't cache pip: install.ps1 + setup.ps1 go through uv + # and never populate ~/.cache/pip; setup-python's post-step + # then fatal-errors with "Cache folder path is retrieved + # for pip but doesn't exist on disk". + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + # Two surgical fixes against measured Windows-only install + # waste (vs Mac/Linux on the same SHA): + # + # (1) npm. setup.ps1's Get-NodeDecision requires Node 22.12+ + # (or 20.19+ / 23+) AND npm >=11 because Vite 8 needs both. + # actions/setup-node@v4 with `node-version: '22'` lands + # Node 22.22.2 + the npm 10.9.7 it bundles, so the decision + # is "bundled" and setup.ps1 downloads an isolated Node (~30 + # MB) we don't need on a runner that already has a fine Node. + # `npm install -g npm@^11` updates the runner's npm in-place + # in ~5 s, flipping the decision to "system" so setup.ps1 + # reuses the existing Node with no download. + # + # (2) Defender. windows-latest's real-time scan opens / hashes + # every file Studio writes during install (Vite output = + # thousands of small chunks, uv pip = wheel-extraction = + # thousands of small files). The latency dominates the + # 200 s frontend build and the 90 s deps install. Adding + # ExclusionPath entries for the directories the install + # writes to drops per-file open latency from ~ms to ~us. + # Add-MpPreference needs admin; the runneradmin user has + # it, but wrap in try/catch so a permission flake leaves + # the install otherwise unaffected. + run: | + $ProgressPreference = 'SilentlyContinue' + Write-Host "npm version before upgrade: $(npm -v)" + npm install -g 'npm@^11' 2>&1 | Out-Host + Write-Host "npm version after upgrade: $(npm -v)" + # NOTE: do NOT pre-create these directories before adding the + # exclusion -- creating an empty studio/frontend/dist trips + # setup.ps1 line 1281-1296's mtime-based "is the frontend + # stale?" check into "up to date, skip rebuild", because the + # newly-created dist's mtime is younger than every source + # file. Studio then boots with an empty dist and 500s on + # GET / with FileNotFoundError: dist\index.html. See run + # 25546676715 / job 74984469728. + # Add-MpPreference accepts paths that do not yet exist; the + # exclusion is registered and applies when the path + # materialises. + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { + Add-MpPreference -ExclusionPath $p -ErrorAction Stop + Write-Host "Defender exclusion added: $p" + } catch { + Write-Host "Defender exclusion skipped ($($_.Exception.Message)): $p" + } + } + + - name: Install Studio (--local, --no-torch) + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + # *>&1 captures Write-Host (Information stream) output; + # plain 2>&1 does not. setup.ps1 emits "prebuilt installed + # and validated" via Write-Host, and we grep for that. + $ProgressPreference = 'SilentlyContinue' + & ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log + + - name: Assert install.ps1 used the Windows llama.cpp prebuilt + run: | + # Filesystem-based check (setup.ps1's stream output isn't + # captured back through the parent pipeline). + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp on Windows." + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + if [ ! -f "$INFO" ]; then + echo "::error::no UNSLOTH_PREBUILT_INFO.json at $INFO." + ls -la "$LLAMA_DIR" || true + exit 1 + fi + if [ ! -f "$BIN" ]; then + echo "::error::no llama-server.exe at $BIN." + ls -la "$LLAMA_DIR/build/bin" || true + exit 1 + fi + echo "install.ps1 installed the Windows prebuilt llama.cpp:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + run: | + SHIM_DIR=~/.unsloth/studio/bin + if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then + echo "::error::unsloth.exe shim not found at $SHIM_DIR" + ls -la ~/.unsloth/studio/ || true + exit 1 + fi + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: First update should be a no-op (prebuilt already validated) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update.log + if grep -q "falling back to source build" logs/update.log; then + echo "::error::studio update fell back to source-build llama.cpp on Windows." + grep -E "llama-prebuilt|llama.cpp" logs/update.log | tail -60 + exit 1 + fi + if ! grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update.log; then + echo "::error::no prebuilt up-to-date marker in update.log." + grep -E "llama-prebuilt|llama.cpp" logs/update.log | tail -60 + exit 1 + fi + echo "update path took the prebuilt fast path" + + - name: Second update must also be a no-op + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update2.log + grep -q "falling back to source build" logs/update2.log && { + echo "::error::second update fell back to source build on Windows" + tail -60 logs/update2.log; exit 1; } || true + grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log + echo "second update was clean" + + - name: Boot Studio briefly to confirm the install is still usable + run: | + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \ + > logs/studio.log 2>&1 & + PID=$! + HEALTHY="" + # Use jq (a Git Bash builtin) instead of `python -c + # open('/tmp/health.json')` to read the saved health + # response. Bash on windows-latest is MSYS Git Bash, which + # resolves `/tmp/...` against the MSYS root, while the + # python interpreter is Windows-native and resolves it + # against the current drive's root. The two paths don't + # agree, so python never finds the file curl just wrote. + # jq reads through MSYS, so the path matches. Mirrors what + # studio-windows-api-smoke.yml and the other Windows smoke + # workflows already do. + for i in $(seq 1 60); do + if curl -fs http://127.0.0.1:18891/api/health > /tmp/health.json; then + if jq -e '.status == "healthy"' /tmp/health.json >/dev/null; then + HEALTHY=1 + break + fi + fi + sleep 1 + done + if [ -z "$HEALTHY" ]; then + echo "Studio failed to come up after \`update\`" + tail -200 logs/studio.log + kill "$PID" 2>/dev/null || true + exit 1 + fi + kill "$PID" 2>/dev/null || true + echo "post-update Studio /api/health OK" + + - name: Uninstall and verify clean + # Round-trip through scripts/uninstall.ps1 against the default + # install tree at %USERPROFILE%\.unsloth\studio. Catches + # regressions where install.ps1 starts writing under a new key + # (registry, Start Menu, %APPDATA%) and scripts/uninstall.ps1 has + # not been updated to match. Skips gracefully if + # scripts/uninstall.ps1 has not landed yet (lets this workflow + # merge before #5513). + shell: pwsh + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + if (-not (Test-Path "$PWD\scripts\uninstall.ps1")) { + Write-Host "scripts/uninstall.ps1 not present in this tree; skipping round-trip" + "" | Set-Content logs/uninstall.log + exit 0 + } + pwsh -NoProfile -File "$PWD\scripts\uninstall.ps1" *>&1 | Tee-Object -FilePath logs/uninstall.log + $leak = 0 + foreach ($p in @( + "$env:USERPROFILE\.unsloth\studio", + "$env:USERPROFILE\.unsloth\studio\unsloth_studio", + "$env:USERPROFILE\.unsloth\studio\bin\unsloth.exe" + )) { + if (Test-Path -LiteralPath $p) { + Write-Host "::error::leak: $p" + $leak++ + } + } + if ($leak -gt 0) { exit 1 } + # Idempotency. + pwsh -NoProfile -File "$PWD\scripts\uninstall.ps1" *>&1 | Select-Object -Last 5 + pwsh -NoProfile -File "$PWD\scripts\uninstall.ps1" *>&1 | Select-Object -Last 5 + Write-Host "PASS: windows install -> update -> uninstall round-trip clean" + + - name: Upload update logs + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-studio-update-log + path: | + logs/install.log + logs/update.log + logs/update2.log + logs/studio.log + logs/uninstall.log + retention-days: 7 diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml new file mode 100644 index 0000000..6becccc --- /dev/null +++ b/.github/workflows/version-compat-ci.yml @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Cross-version compat canary for the four upstream packages whose +# release cadence regularly breaks unsloth + unsloth-zoo: +# +# 1. vLLM (LoRA worker manager, BnB loader, cumem allocator) +# 2. TRL / GRPO (trainer source rewriters in unsloth.models.rl*) +# 3. PEFT (LoraConfig, get_peft_model, LoraLayer, bnb integration) +# 4. sentence-transformers (Transformer/Pooling/Normalize, Trainer) +# 5. bitsandbytes (Linear4bit, dequantize_4bit) +# +# Strategy: GitHub raw-fetch + symbol grep against every tracked +# version (no pip install, CPU-only). When upstream renames a symbol +# we depend on, the matching test fails BEFORE a user hits it. The +# `main` branch entries give us a few-day lead on PyPI releases. +# +# Cross-references: +# tests/vllm_compat/test_vllm_pinned_symbols.py (vLLM symbols) +# tests/version_compat/test_trl_grpo_pinned_symbols.py +# tests/version_compat/test_peft_pinned_symbols.py +# tests/version_compat/test_sentence_transformers_pinned_symbols.py +# tests/version_compat/test_bitsandbytes_pinned_symbols.py + +name: Version Compat CI + +on: + pull_request: + # Trigger on any unsloth source change, not just the three previously + # named files. The symbol-existence tests verify that EVERY pinned + # upstream reference in unsloth still resolves; a new + # `from peft.foo import Bar` added in unsloth/kernels/whatever.py + # is just as much a compat regression risk as one added in + # unsloth/models/rl.py. + paths: + - 'unsloth/**' + - 'tests/vllm_compat/**' + - 'tests/version_compat/**' + - 'pyproject.toml' + - '.github/workflows/version-compat-ci.yml' + schedule: + # Daily 06:43 UTC. Catches upstream PyPI releases roughly within + # 24 h. Off the :00 / :30 fleet-collision spots. + - cron: '43 6 * * *' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + vllm-pinned-symbols: + name: vLLM pinned-symbol matrix (≥ 0.9.0 + main) + runs-on: ubuntu-latest + timeout-minutes: 12 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install pytest only + # The test fetches from raw.githubusercontent.com and greps + # source. No pip install of vllm / torch / transformers is + # needed — that's the whole point of this canary. + run: | + python -m pip install --upgrade pip + pip install 'pytest>=8' + - name: Run vllm-compat suite + env: + # Authenticated requests get a 5000-req/h quota on raw + # fetches; unauthenticated is 60/h and trips on the matrix. + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + python -m pytest tests/vllm_compat/test_vllm_pinned_symbols.py -v --tb=short + + trl-grpo-pinned-symbols: + name: TRL / GRPO pinned-symbol matrix + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install pytest only + run: | + python -m pip install --upgrade pip + pip install 'pytest>=8' + - name: Run trl-compat suite + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + # PYTHONPATH=. so `from tests.version_compat._fetch import …` + # works without an editable install of unsloth itself. + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_pinned_symbols.py \ + -v --tb=short + + peft-pinned-symbols: + name: PEFT pinned-symbol matrix (pyproject window + main) + runs-on: ubuntu-latest + timeout-minutes: 8 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install pytest only + run: | + python -m pip install --upgrade pip + pip install 'pytest>=8' + - name: Run peft-compat suite + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_peft_pinned_symbols.py \ + tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py \ + -v --tb=short + + st-pinned-symbols: + name: sentence-transformers pinned-symbol matrix + runs-on: ubuntu-latest + timeout-minutes: 8 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install pytest only + run: | + python -m pip install --upgrade pip + pip install 'pytest>=8' + - name: Run sentence-transformers compat suite + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_sentence_transformers_pinned_symbols.py \ + -v --tb=short + + bitsandbytes-pinned-symbols: + name: bitsandbytes pinned-symbol matrix + runs-on: ubuntu-latest + timeout-minutes: 8 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install pytest only + run: | + python -m pip install --upgrade pip + pip install 'pytest>=8' + - name: Run bitsandbytes compat suite + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_bitsandbytes_pinned_symbols.py \ + -v --tb=short + + transformers-pinned-symbols: + name: transformers pinned-symbol matrix (4.57.6 + 5.x + main) + runs-on: ubuntu-latest + timeout-minutes: 12 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install pytest only + run: | + python -m pip install --upgrade pip + pip install 'pytest>=8' + - name: Run transformers compat suite + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_transformers_pinned_symbols.py \ + -v --tb=short + + # Optional second layer: actually `pip install` ONE representative + # version of each package and verify unsloth + unsloth-zoo modules + # import on it under the existing CUDA spoof. CPU-only, runs on + # ubuntu-latest. Catches the small set of breakages that the static + # symbol check misses (e.g. import-time side effects). + zoo-imports-under-spoof: + name: unsloth_zoo vllm/grpo/peft/st modules import under CUDA spoof + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - name: Clone unsloth-zoo @ main + run: | + # github.com occasionally 500s on the git fetch; retry so a + # single upstream blip does not fail CI. + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install CPU torch + supported pkg pins + run: | + python -m pip install --upgrade pip + # CPU torch (vllm/peft/st all depend on it). + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' + # torchcodec is a hard requirement on transformers 5.x: + # transformers/audio_utils.py:55 does + # `importlib.metadata.version("torchcodec")` UNCONDITIONALLY, + # which raises PackageNotFoundError on a CPU runner that + # otherwise has no audio path -- and that error trickles up + # through every `import unsloth_zoo.` because + # unsloth-zoo's vision_utils transitively pulls + # transformers.processing_utils (-> audio_utils). The 0.10 + # cap mirrors the torch 2.10 / torchvision 0.26 ABI window + # we already pin above. + # Ladder of supported floor versions per pyproject.toml. + pip install \ + 'transformers>=4.56,<5.6' 'trl>=0.22,<0.26' \ + 'peft>=0.18.0' 'sentence-transformers>=5.0' \ + 'accelerate>=1.0' 'datasets>=3.4,<5' \ + 'bitsandbytes>=0.45.5' \ + sentencepiece protobuf safetensors numpy 'pytest>=8' \ + 'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow + # Editable-install both repos so the test imports the + # checkouts (not whatever stale PyPI version pip resolved). + pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo" + pip install --no-deps -e ./unsloth + - name: Run vllm_compat zoo-imports tests under spoof + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + cd unsloth + # tests/vllm_compat/test_unsloth_zoo_imports.py: narrow vllm/grpo + # import gates (5 tests). + # tests/vllm_compat/test_extended_module_imports.py: full sweep + # of unsloth_zoo + unsloth.models.* modules + RL dispatch + # table population + FastModel API surface under spoof + # (~30 tests). Catches transformers / peft / bnb symbol pin + # drift at module-top BEFORE any runtime call. + PYTHONPATH=. python -m pytest \ + tests/vllm_compat/test_unsloth_zoo_imports.py \ + tests/vllm_compat/test_extended_module_imports.py \ + -v --tb=short + + # Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike + # the static symbol/source greps above, this drives unsloth's actual + # source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only + # runner under the tests/conftest.py spoof harness -- no GPU, no training. + # Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple + # per-token-logps return, restructured PEFT ref-adapter block) by asserting + # the generated Unsloth trainer still satisfies the transform contracts. + grpo-fake-run: + name: GRPO fake-run (latest + main TRL, CPU spoof) + runs-on: ubuntu-latest + timeout-minutes: 18 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - name: Clone unsloth-zoo @ main + run: | + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install CPU torch + ecosystem + TRL latest + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' + # Ecosystem floors unsloth needs; TRL itself is installed last so it + # can pull the transformers/peft it requires. + pip install \ + 'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \ + 'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \ + 'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow + pip install --upgrade trl + pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo" + pip install --no-deps -e ./unsloth + - name: Fake-run vs TRL latest + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + # Disable dynamo/inductor at the process level, before conftest.py's early + # `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner + # (defense in depth; the CPU fake-train also flips this at runtime). + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge + # TRL break does not red every PR. github.event_name is valid in a step if. + - name: Fake-run vs TRL main (scheduled / dispatch only) + if: ${{ github.event_name != 'pull_request' }} + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + pip install --upgrade "git+https://github.com/huggingface/trl" + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + + # Daily-only: same suites but with --strict on importable upstream + # tags. Schedule-only so PR jobs stay fast; cron tolerates a flake. + daily-fresh-fetch: + name: daily fresh-fetch sweep (cron only) + if: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }} + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install pytest + run: pip install 'pytest>=8' + - name: Run all version-compat suites in one process (no cache) + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + PYTHONPATH=. python -m pytest \ + tests/vllm_compat/test_vllm_pinned_symbols.py \ + tests/version_compat/ \ + -v --tb=short diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml new file mode 100644 index 0000000..3de3c33 --- /dev/null +++ b/.github/workflows/wheel-smoke.yml @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Builds the PyPI wheel from the PR branch, then verifies the built wheel +# actually contains what we expect to ship and does NOT contain the broken +# Studio bundle that 2026.5.1 published. This is the single workflow that +# would have blocked the 2026.5.1 release before twine upload. +# +# Verified locally end-to-end against this branch: +# - python -m build produces unsloth--py3-none-any.whl in 13s +# - wheel content sanity passes: +# lockfile shipped, frontend dist shipped, +# no node_modules in wheel, no bun.lock in wheel, +# main bundle has unstable_Provider hits=1 (assistant-ui internals only). +# - Studio backend imports cleanly from the installed wheel with the +# lightweight dep set below. + +name: Wheel CI + +on: + pull_request: + paths: + - 'pyproject.toml' + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - '.github/workflows/wheel-smoke.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + wheel: + name: Wheel build + content sanity + import smoke + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Lockfile supply-chain audit (pre-install scan) + run: python3 scripts/lockfile_supply_chain_audit.py + + - name: Build frontend + # Lifecycle scripts (esbuild native-binary postinstall, etc.) are + # required for `vite build`. The pre-install lockfile structural + # audit (lockfile_supply_chain_audit.py) is the practical defence + # against the npm postinstall-dropper class -- it fires BEFORE any + # tarball runs, on the injection pattern itself rather than an + # advisory-DB lookup. + run: | + cd studio/frontend + npm ci --no-fund --no-audit + npm run build + + - name: Build wheel + sdist + run: | + python -m pip install --upgrade pip build + rm -rf dist build ./*.egg-info + python -m build + + - name: Wheel content sanity + run: | + python - <<'PY' + import zipfile, glob, sys + w = glob.glob("dist/unsloth-*.whl") + if not w: + print("FAIL: no wheel produced"); sys.exit(2) + w = w[0] + print(f"wheel: {w}") + with zipfile.ZipFile(w) as z: + n = z.namelist() + checks = { + "lockfile shipped": any(s.endswith("studio/frontend/package-lock.json") for s in n), + "frontend dist shipped": any(s.endswith("studio/frontend/dist/index.html") for s in n), + "no node_modules": not any("studio/frontend/node_modules/" in s for s in n), + "no bun.lock": not any(s.endswith("studio/frontend/bun.lock") for s in n), + } + js = [s for s in n + if "studio/frontend/dist/assets/" in s + and s.endswith(".js") + and "/index-" in s] + if not js: + print("FAIL: no main bundle index-*.js in wheel"); sys.exit(2) + data = z.read(js[0]).decode("utf-8", "replace") + hits = data.count("unstable_Provider:") + print(f"main bundle: {js[0]}") + print(f"unstable_Provider hits: {hits} (>=4 indicates 2026.5.1 regression)") + checks["bundle has no Studio unstable_Provider call site"] = (hits < 4) + + print() + for k, v in checks.items(): + print(f" [{'PASS' if v else 'FAIL'}] {k}") + sys.exit(0 if all(checks.values()) else 1) + PY + + - name: Studio backend import smoke + # Imports `studio.backend.main:app` from the freshly-installed wheel in + # a clean venv. This catches the class of bug that 2026.5.1 shipped with: + # frontend dist missing, package-lock.json missing, or the wheel's Python + # source tree broken in a way that surfaces only at app construction time. + run: | + python -m venv /tmp/v + /tmp/v/bin/pip install --upgrade pip + /tmp/v/bin/pip install -r studio/backend/requirements/studio.txt + /tmp/v/bin/pip install \ + python-multipart aiofiles sqlalchemy cryptography \ + pyyaml jinja2 mammoth unpdf requests \ + 'numpy<3' + /tmp/v/bin/pip install --no-deps dist/unsloth-*.whl + # Run from /tmp so Python imports the installed package, not the source tree. + cd /tmp + /tmp/v/bin/python -c "from studio.backend.main import app; print('Studio backend OK:', app.title)" + + - name: Upload wheel on failure + if: failure() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: unsloth-wheel + path: dist/ + retention-days: 7 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..39ca222 --- /dev/null +++ b/.gitignore @@ -0,0 +1,241 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*.class +unsloth_compiled_cache/ +# Notebook-validator runtime PyPI metadata cache (CI repopulates). +scripts/data/pypi_cache/ +# ML artifacts (large files) +feature/ +outputs/ +exports/ +/datasets/ +studio/backend/assets/datasets/ +# Generated async worker / reviewer transcripts (never part of the product). +studio/backend/async_task_outputs/ +unsloth_training_checkpoints/ +*.gguf +*.safetensors + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +/lib/ +/lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +.venv_overlay/ +.venv_t5/ +environment.yaml + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# Ruff stuff: +.ruff_cache/ +.pre-commit-cache/ + +# PyPI configuration file and IDE/Editors +.pypirc +.vscode +.idea/ +.claude/ +*.swp +*.swo + +# oh-my-codex +.omx/ + +# Firebase +firebase-debug.log + +# Other +resources/ +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 +**/AGENT.md +**/agent.md +docs/canvas-lab-architecture.md +log_rtx.txt +log.txt +setup_leo.sh +server.pid +*.log +# Ignore stray lockfiles; real npm projects opt back in below (npm ci needs them). +package-lock.json +!studio/frontend/package-lock.json +!studio/backend/core/data_recipe/oxc-validator/package-lock.json +!studio/package-lock.json +llama.cpp/ +# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo. +/~/ diff --git a/.pre-commit-ci.yaml b/.pre-commit-ci.yaml new file mode 100644 index 0000000..dbcf58a --- /dev/null +++ b/.pre-commit-ci.yaml @@ -0,0 +1,6 @@ +ci: + autofix_prs: true + autofix_prs_limit: 5 + autoupdate_schedule: monthly + autoupdate_commit_msg: "chore: pre-commit autoupdate" + skip: [] diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..8dcb913 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,33 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.18 + hooks: + - id: ruff + args: + - --fix + - --exit-non-zero-on-fix + exclude: '\.ipynb$' + - repo: local + hooks: + - id: ruff-format-with-kwargs + name: Ruff format with kwarg spacing + entry: scripts/run_ruff_format.py + language: python + types: [python] + # Mirror ruff's [tool.ruff] extend-exclude so this hook does not + # half-process files ruff itself skips (which produced churn). + exclude: '(chat_templates|ollama_template_mappers|_auto_install|mapper)\.py$' + additional_dependencies: + - ruff==0.6.9 + # Re-pins allowScripts entries after dependency bumps. pre-commit.ci + # pushes the fix to PR branches, Dependabot's included, so stale pins + # heal without a human in the loop. + - id: sync-allow-scripts-pins + name: Sync allowScripts pins with the frontend lockfile + # `python + + diff --git a/studio/backend/auth/.gitkeep b/studio/backend/auth/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/studio/backend/auth/__init__.py b/studio/backend/auth/__init__.py new file mode 100644 index 0000000..a63d530 --- /dev/null +++ b/studio/backend/auth/__init__.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Authentication module for JWT-based auth with SQLite storage.""" + +from .authentication import ( + create_access_token, + create_refresh_token, + refresh_access_token, + get_current_subject, + get_current_subject_allow_password_change, + reload_secret, +) +from .storage import ( + DEFAULT_ADMIN_USERNAME, + clear_bootstrap_password, + generate_bootstrap_password, + get_bootstrap_password, + is_initialized, + create_initial_user, + ensure_default_admin, + get_jwt_secret, + get_user_and_secret, + load_jwt_secret, + requires_password_change, + save_refresh_token, + update_password, + verify_refresh_token, + revoke_user_refresh_tokens, +) +from .hashing import hash_password, verify_password + +__all__ = [ + "create_access_token", + "create_refresh_token", + "refresh_access_token", + "get_current_subject", + "get_current_subject_allow_password_change", + "reload_secret", + "DEFAULT_ADMIN_USERNAME", + "clear_bootstrap_password", + "generate_bootstrap_password", + "get_bootstrap_password", + "is_initialized", + "create_initial_user", + "ensure_default_admin", + "get_jwt_secret", + "get_user_and_secret", + "load_jwt_secret", + "requires_password_change", + "save_refresh_token", + "update_password", + "verify_refresh_token", + "revoke_user_refresh_tokens", + "hash_password", + "verify_password", +] diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py new file mode 100644 index 0000000..b13cd1c --- /dev/null +++ b/studio/backend/auth/authentication.py @@ -0,0 +1,217 @@ +# 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 secrets +from datetime import datetime, timedelta, timezone +from typing import Optional, Tuple + +from fastapi import Depends, HTTPException, status +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, +) + +ALGORITHM = "HS256" +ACCESS_TOKEN_EXPIRE_MINUTES = 60 +REFRESH_TOKEN_EXPIRE_DAYS = 7 + +security = HTTPBearer() # Reads Authorization: Bearer + + +def _get_secret_for_subject(subject: str) -> str: + secret = get_jwt_secret(subject) + if secret is None: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid or expired token", + ) + return secret + + +def _decode_subject_without_verification(token: str) -> Optional[str]: + try: + payload = jwt.decode( + token, + options = {"verify_signature": False, "verify_exp": False}, + ) + except jwt.InvalidTokenError: + return None + + subject = payload.get("sub") + return subject if isinstance(subject, str) else None + + +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). + + Valid across restarts: 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) + ) + to_encode.update({"exp": expire}) + return jwt.encode( + to_encode, + _get_secret_for_subject(subject), + algorithm = ALGORITHM, + ) + + +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. + + Refresh tokens are opaque (not JWTs); expire after REFRESH_TOKEN_EXPIRE_DAYS. + """ + token = secrets.token_urlsafe(48) + expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS) + 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], bool]: + """ + Validate a refresh token and issue a new access token. + + The refresh token is NOT consumed; it stays valid until expiry. + Returns a new access_token, or None if the refresh token is invalid/expired. + """ + verified = verify_refresh_token(refresh_token) + if verified is None: + return None, None, False + username, is_desktop = verified + return ( + create_access_token(subject = username, desktop = is_desktop), + username, + is_desktop, + ) + + +def reload_secret() -> None: + """ + Legacy API compat for callers expecting auth storage init. + + Auth now resolves the current signing secret directly from SQLite. + """ + load_jwt_secret() + + +async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str: + """Validate JWT and require the password-change flow to be completed.""" + return await _get_current_subject( + credentials, + allow_password_change = False, + ) + + +async def authenticated_via_api_key( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> bool: + """True when the caller used an sk-unsloth API key, not a UI session JWT. + + Lets routes treat programmatic API callers differently from the Studio UI + (e.g. refuse a teardown the UI would allow). + """ + return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX)) + + +async def get_current_subject_allow_password_change( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> str: + """Validate JWT but allow access to the password-change endpoint.""" + return await _get_current_subject( + credentials, + allow_password_change = True, + ) + + +async def _get_current_subject( + credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool +) -> str: + """FastAPI dependency: validate the JWT and return the subject. Use on protected routes.""" + 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( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid token payload", + ) + + record = get_user_and_secret(subject) + if record is None: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid or expired token", + ) + + _salt, _pwd_hash, jwt_secret, must_change_password = record + try: + payload = jwt.decode(token, jwt_secret, algorithms = [ALGORITHM]) + if payload.get("sub") != subject: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid token payload", + ) + 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", + ) + return subject + except jwt.InvalidTokenError: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid or expired token", + ) diff --git a/studio/backend/auth/bootstrap_timeout.py b/studio/backend/auth/bootstrap_timeout.py new file mode 100644 index 0000000..728433d --- /dev/null +++ b/studio/backend/auth/bootstrap_timeout.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged. + +On a fresh install the seeded bootstrap admin password stays a valid login +credential until first login changes it. When the web UI is put on the network +(``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within +a deadline, tear Studio down so a fresh, unconfigured instance does not stay +publicly reachable indefinitely. If the password was changed, Studio keeps +running. + +Scope: web UI launches only (never ``--api-only``, which authenticates by API +key rather than the admin password, and never Colab). Configurable via +``UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT`` (seconds; default 3600; ``0`` disables). +""" + +import os +import sys +import threading + +BOOTSTRAP_TIMEOUT_ENV_VAR = "UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT" +DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS = 3600 + + +def bootstrap_timeout_seconds(env = None) -> int: + """Resolve the deadline in seconds. ``0`` (or invalid/negative) disables it. + + A malformed value falls back to the default rather than disabling, so a typo + cannot silently remove the protection. + """ + env = os.environ if env is None else env + raw = env.get(BOOTSTRAP_TIMEOUT_ENV_VAR) + if raw is None or raw.strip() == "": + return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + try: + value = int(raw) + except ValueError: + return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + return value if value > 0 else 0 + + +def _is_exposed_bind(host: str, secure: bool) -> bool: + """True when this launch puts the web UI on the network (tunnel or non-loopback).""" + if secure: + return True + if host in ("0.0.0.0", "::"): + return True + try: + from utils.host_policy import is_external_host + except Exception: + return False + return bool(is_external_host(host)) + + +def should_arm_bootstrap_timeout( + *, + host: str, + secure: bool, + api_only: bool, + frontend_served: bool, + is_colab: bool, + requires_change: bool, + timeout_seconds: int, +) -> bool: + """Whether to arm the deadline: only for an exposed web UI whose seeded admin + password is still unchanged. Pure decision (no I/O) for cheap unit testing.""" + if timeout_seconds <= 0: + return False + if api_only or not frontend_served or is_colab: + return False + if not requires_change: + return False + return _is_exposed_bind(host, secure) + + +def _format_duration(seconds: int) -> str: + """Human-friendly duration for the shutdown message (seconds under a minute).""" + + def _plural(n: int, unit: str) -> str: + return f"{n} {unit}{'' if n == 1 else 's'}" + + if seconds < 60: + return _plural(seconds, "second") + minutes, rem = divmod(seconds, 60) + label = _plural(minutes, "minute") + if rem: + label += f" {_plural(rem, 'second')}" + return label + + +def enforce_bootstrap_password_deadline( + storage, + trigger_shutdown, + *, + timeout_seconds: int, + logger = None, +) -> bool: + """Deadline handler: shut down iff the seeded admin password is still unchanged. + + Returns True if it shut Studio down, False if it left it running (the + password was changed in time). + """ + try: + still_default = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME) + except Exception: + return False + if not still_default: + return False # password changed in time -> leave Studio running + + message = ( + "\nUnsloth Studio was exposed on the network but its default admin " + f"password was not changed within {_format_duration(timeout_seconds)}. " + "Shutting down to avoid leaving an unsecured public instance running.\n" + "Next time, sign in and change the password on first login, or set " + f"{BOOTSTRAP_TIMEOUT_ENV_VAR}=0 to disable this timeout." + ) + if logger is not None: + logger.warning(message) + print(message, file = sys.stderr, flush = True) + try: + trigger_shutdown() + except Exception as e: # shutdown is best-effort; never raise from the timer + if logger is not None: + logger.warning("Bootstrap-timeout shutdown failed: %s", e) + return True + + +def arm_bootstrap_timeout( + storage, + trigger_shutdown, + *, + timeout_seconds: int, + logger = None, +) -> "threading.Timer": + """Start a daemon timer that enforces the deadline. Returns the Timer.""" + timer = threading.Timer( + timeout_seconds, + enforce_bootstrap_password_deadline, + args = (storage, trigger_shutdown), + kwargs = {"timeout_seconds": timeout_seconds, "logger": logger}, + ) + timer.daemon = True + timer.start() + return timer diff --git a/studio/backend/auth/hashing.py b/studio/backend/auth/hashing.py new file mode 100644 index 0000000..0a1b0c1 --- /dev/null +++ b/studio/backend/auth/hashing.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Password hashing utilities using PBKDF2. +""" + +import hashlib +import hmac +import secrets +from typing import Tuple + + +def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]: + """ + Hash a password using PBKDF2-HMAC-SHA256. + + Returns (salt, hex_hash) tuple. + """ + if salt is None: + salt = secrets.token_hex(16) + dk = hashlib.pbkdf2_hmac( + "sha256", + password.encode("utf-8"), + salt.encode("utf-8"), + 100_000, + ) + return salt, dk.hex() + + +def verify_password(password: str, salt: str, hashed: str) -> bool: + """ + Verify a password against a stored salt and hash. + + Uses constant-time comparison to prevent timing attacks. + """ + dk = hashlib.pbkdf2_hmac( + "sha256", + password.encode("utf-8"), + salt.encode("utf-8"), + 100_000, + ) + return hmac.compare_digest(dk.hex(), hashed) diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py new file mode 100644 index 0000000..a0da2b2 --- /dev/null +++ b/studio/backend/auth/storage.py @@ -0,0 +1,891 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""SQLite storage for auth data (user credentials + JWT secret).""" + +import hashlib +import hmac +import ipaddress +import os +import secrets +import sqlite3 +import threading +from datetime import datetime, timezone +from typing import Optional, Tuple + +from utils.paths import auth_db_path, ensure_dir + +DB_PATH = auth_db_path() +DEFAULT_ADMIN_USERNAME = "unsloth" + +# Plaintext bootstrap password file beside auth.db, deleted on first password +# change so the credential never lingers on disk. +_BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password" + +# In-process cache to avoid re-reading the file on every HTML serve. +_bootstrap_password: Optional[str] = None + + +def generate_bootstrap_password() -> str: + """Generate a 4-word diceware passphrase and persist it to disk. + + Persisted (the DB stores only the hash) so it survives restarts; later + calls return the persisted value. + """ + global _bootstrap_password + + # Cached in this process? + if _bootstrap_password is not None: + return _bootstrap_password + + # Persisted from a previous run? + if _BOOTSTRAP_PW_PATH.is_file(): + _bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip() + if _bootstrap_password: + return _bootstrap_password + + # First startup: generate a fresh passphrase. + import diceware + + _bootstrap_password = diceware.get_passphrase( + options = diceware.handle_options(args = ["-n", "4", "-d", "", "-c"]) + ) + + # Persist so the same passphrase survives restarts until password change. + 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 + + +def get_bootstrap_password() -> Optional[str]: + """Return the cached bootstrap password, or None if not yet generated.""" + 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 + _bootstrap_password = None + if _BOOTSTRAP_PW_PATH.is_file(): + _BOOTSTRAP_PW_PATH.unlink(missing_ok = True) + + +def _hash_token(token: str) -> str: + """SHA-256 hash helper for refresh token storage. + + Plain SHA-256 is intentional: refresh tokens are 384-bit random strings, so + a slow KDF adds no security while costing per-refresh latency. API keys use + the separate ``_pbkdf2_api_key`` helper, only to satisfy CodeQL's + ``py/weak-sensitive-data-hashing`` query, not for crypto reasons. + """ + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def get_connection() -> sqlite3.Connection: + """Get a connection to the auth database, creating tables if needed.""" + ensure_dir(DB_PATH.parent) + conn = sqlite3.connect(DB_PATH) + # Keep the auth dir + DB private (they hold the JWT/identity secrets and + # password hashes); sqlite3.connect would otherwise create the DB 0644 under + # a 022 umask, letting another OS user read the identity secret and forge proofs. + for _path, _mode in ((DB_PATH.parent, 0o700), (DB_PATH, 0o600)): + try: + os.chmod(_path, _mode) + except OSError: + pass + conn.row_factory = sqlite3.Row + # WAL lets token reads run concurrently with refresh-token writes; + # busy_timeout bounds lock waits. Matches the other Studio SQLite stores. + # Set busy_timeout first: switching journal_mode needs a lock, so if a + # refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY; + # with busy_timeout already in effect it waits instead of failing and leaving + # this connection on SQLite's default zero lock wait. + try: + conn.execute("PRAGMA busy_timeout=5000") + conn.execute("PRAGMA journal_mode=WAL") + except sqlite3.Error: + pass + conn.execute( + """ + CREATE TABLE IF NOT EXISTS 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( + """ + CREATE TABLE IF NOT EXISTS refresh_tokens ( + id INTEGER PRIMARY KEY, + token_hash TEXT NOT NULL, + username 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 + ); + """ + ) + columns = {row["name"] for row in conn.execute("PRAGMA table_info(auth_user)")} + if "must_change_password" not in columns: + 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 +# via ``_get_or_create_api_key_pbkdf2_salt``. No lock needed: (a) ``INSERT OR +# IGNORE`` is atomic at the SQLite layer and (b) concurrent populations +# converge on the same value, so the worst case is a harmless duplicate read +# on startup. +_api_key_pbkdf2_salt_cache: Optional[bytes] = None + + +def _get_or_create_api_key_pbkdf2_salt() -> bytes: + """Return the persistent API-key PBKDF2 salt, generating it once if missing. + + Hex-encoded 32-byte random value in ``app_secrets``. Regenerated only when + the row is missing (fresh install, or operator deleted it). + """ + global _api_key_pbkdf2_salt_cache + if _api_key_pbkdf2_salt_cache is not None: + 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 + + +# Secret answering the /api/auth/identity challenge (HMAC(secret, nonce)). Lives +# in this same-user DB so a port squatter or remote/fake server can't forge a +# proof. Separate from the per-user JWT secret. +_IDENTITY_SECRET_DB_KEY = "studio_identity_secret" +_identity_secret_cache: Optional[bytes] = None + + +def get_or_create_identity_secret() -> bytes: + """Return the identity secret (hex 32-byte row in app_secrets), creating it once.""" + global _identity_secret_cache + if _identity_secret_cache is not None: + return _identity_secret_cache + + conn = get_connection() + try: + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + if row is None: + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + (_IDENTITY_SECRET_DB_KEY, secrets.token_hex(32)), + ) + conn.commit() + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + secret = bytes.fromhex(row["value"]) + finally: + conn.close() + + _identity_secret_cache = secret + return secret + + +def compute_identity_proof(nonce: bytes, host: str, port: int) -> str: + """HMAC-SHA256 proof that the caller holds this install's identity secret, + bound to the loopback address and port the connection landed on. A proof + relayed from a Studio on a different address/port (a squatter proxying to the + real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was + computed for that other endpoint and won't match the one the client dialed.""" + try: + host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms + except ValueError: + host = (host or "").lower() + msg = b"|".join([nonce, host.encode(), str(int(port)).encode()]) + return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest() + + +# Capability secret for public ``/p`` preview share links. HMAC(secret, ref) +# turns the deterministic preview ref into an unguessable bearer capability, so a +# guessed run/checkpoint name can't reach inference. Dedicated (not the per-user +# JWT secret) so rotating it revokes every shared link without touching logins. +_PREVIEW_LINK_SECRET_DB_KEY = "preview_link_secret" +_preview_link_secret_cache: Optional[bytes] = None + + +def get_or_create_preview_link_secret() -> bytes: + """Return the preview-link signing secret (hex 32-byte row in app_secrets), creating it once.""" + global _preview_link_secret_cache + if _preview_link_secret_cache is not None: + return _preview_link_secret_cache + + conn = get_connection() + try: + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_PREVIEW_LINK_SECRET_DB_KEY,), + ).fetchone() + if row is None: + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + (_PREVIEW_LINK_SECRET_DB_KEY, secrets.token_hex(32)), + ) + conn.commit() + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_PREVIEW_LINK_SECRET_DB_KEY,), + ).fetchone() + secret = bytes.fromhex(row["value"]) + finally: + conn.close() + + _preview_link_secret_cache = secret + return secret + + +def rotate_preview_link_secret() -> bytes: + """Rotate the preview-link secret, immediately revoking every outstanding ``/p`` share link.""" + global _preview_link_secret_cache + new_secret_hex = secrets.token_hex(32) + conn = get_connection() + try: + conn.execute( + "INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)", + (_PREVIEW_LINK_SECRET_DB_KEY, new_secret_hex), + ) + conn.commit() + finally: + conn.close() + + secret = bytes.fromhex(new_secret_hex) + _preview_link_secret_cache = secret + return secret + + +_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. + + For API-key storage ONLY, not refresh tokens. The slow KDF is only to + appease CodeQL's ``py/weak-sensitive-data-hashing`` query, not a crypto + requirement (API keys are random 128-bit tokens). The salt lives in + ``app_secrets`` so dumping ``api_keys`` alone can't derive hashes. + """ + salt = _get_or_create_api_key_pbkdf2_salt() + dk = hashlib.pbkdf2_hmac( + "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) + + +# Memoize the deterministic raw-key -> PBKDF2-hash derivation so the 100k-round +# KDF runs once per key instead of on every authenticated request. Keyed by a +# salted HMAC of the key (not the key itself); revocation/expiry are still +# enforced by the SQLite read on every call, so a cache hit only skips the KDF. +# Only keys present in the DB are cached, so unknown-key spam can't grow it. +_api_key_hash_cache: dict[str, str] = {} +_API_KEY_HASH_CACHE_MAX = 4096 +_api_key_hash_cache_lock = threading.Lock() + + +def _api_key_cache_id(raw_key: str) -> str: + """Cache id for a raw key: salted HMAC-SHA256 (not the key itself).""" + return hmac.new( + _get_or_create_api_key_pbkdf2_salt(), raw_key.encode("utf-8"), hashlib.sha256 + ).hexdigest() + + +def _reset_api_key_hash_cache() -> None: + """Drop memoized derivations (tests / salt change).""" + with _api_key_hash_cache_lock: + _api_key_hash_cache.clear() + + +def is_initialized() -> bool: + """Check if auth is ready for login (at least one user exists in DB).""" + conn = get_connection() + cur = conn.execute("SELECT COUNT(*) AS c FROM auth_user") + row = cur.fetchone() + conn.close() + return bool(row["c"]) + + +def create_initial_user( + username: str, + password: str, + jwt_secret: str, + *, + must_change_password: bool = False, +) -> None: + """ + Create the initial admin user in the database. + + Raises sqlite3.IntegrityError if username already exists. + """ + from .hashing import hash_password + + salt, pwd_hash = hash_password(password) + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO auth_user ( + username, + password_salt, + password_hash, + jwt_secret, + must_change_password + ) + VALUES (?, ?, ?, ?, ?) + """, + (username, salt, pwd_hash, jwt_secret, int(must_change_password)), + ) + conn.commit() + finally: + conn.close() + + +def delete_user(username: str) -> None: + """ + Delete a user from the database. + + Used for rollback when user creation fails partway through bootstrap. + """ + conn = get_connection() + try: + conn.execute("DELETE FROM auth_user WHERE username = ?", (username,)) + conn.commit() + finally: + conn.close() + + +def get_user_and_secret(username: str) -> Optional[Tuple[str, str, str, bool]]: + """ + Get user's password salt, hash, and JWT secret. + + Returns (password_salt, password_hash, jwt_secret, must_change_password) + or None if user not found. + """ + conn = get_connection() + try: + cur = conn.execute( + """ + SELECT password_salt, password_hash, jwt_secret, must_change_password + FROM auth_user + WHERE username = ? + """, + (username,), + ) + row = cur.fetchone() + if not row: + return None + return ( + row["password_salt"], + row["password_hash"], + row["jwt_secret"], + bool(row["must_change_password"]), + ) + finally: + conn.close() + + +def get_jwt_secret(username: str) -> Optional[str]: + """Return the current JWT signing secret for a user.""" + conn = get_connection() + try: + cur = conn.execute( + "SELECT jwt_secret FROM auth_user WHERE username = ?", + (username,), + ) + row = cur.fetchone() + return row["jwt_secret"] if row else None + finally: + conn.close() + + +def requires_password_change(username: str) -> bool: + """Return whether the user must change the seeded default password.""" + conn = get_connection() + try: + cur = conn.execute( + "SELECT must_change_password FROM auth_user WHERE username = ?", + (username,), + ) + row = cur.fetchone() + return bool(row and row["must_change_password"]) + finally: + conn.close() + + +def load_jwt_secret() -> str: + """ + Load the JWT secret from the database. + + Raises RuntimeError if no auth user has been created yet. + """ + conn = get_connection() + try: + cur = conn.execute("SELECT jwt_secret FROM auth_user LIMIT 1") + row = cur.fetchone() + if not row: + raise RuntimeError( + "Auth is not initialized. Wait for the seeded admin bootstrap to complete." + ) + return row["jwt_secret"] + finally: + conn.close() + + +def ensure_default_admin() -> bool: + """Seed the default admin account on first startup. + + 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( + username = DEFAULT_ADMIN_USERNAME, + password = bootstrap_pw, + jwt_secret = secrets.token_urlsafe(64), + must_change_password = True, + ) + return True + except sqlite3.IntegrityError: + return False + + +def update_password(username: str, new_password: str) -> bool: + """Update password, clear first-login requirement, rotate JWT secret.""" + from .hashing import hash_password + + salt, pwd_hash = hash_password(new_password) + jwt_secret = secrets.token_urlsafe(64) + conn = get_connection() + try: + cursor = conn.execute( + """ + UPDATE auth_user + SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0 + WHERE username = ? + """, + (salt, pwd_hash, jwt_secret, username), + ) + 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, + *, + is_desktop: bool = False, +) -> None: + """ + Store a hashed refresh token with its associated username and expiry. + """ + token_hash = _hash_token(token) + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop) + VALUES (?, ?, ?, ?) + """, + (token_hash, username, expires_at, int(is_desktop)), + ) + conn.commit() + finally: + conn.close() + + +def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]: + """Atomically validate-and-delete a refresh token for single-use rotation. + + DELETE RETURNING fuses validate and delete into one statement so two + concurrent refresh requests cannot both consume the same token. + """ + token_hash = _hash_token(token) + now = datetime.now(timezone.utc).isoformat() + conn = get_connection() + try: + conn.execute( + "DELETE FROM refresh_tokens WHERE expires_at < ?", + (now,), + ) + cur = conn.execute( + """ + DELETE FROM refresh_tokens + WHERE token_hash = ? AND expires_at >= ? + RETURNING username, is_desktop + """, + (token_hash, now), + ) + row = cur.fetchone() + conn.commit() + if row is None: + return None + return row["username"], bool(row["is_desktop"]) + finally: + conn.close() + + +def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]: + """ + Verify a refresh token and return the username plus desktop marker. + + 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) + conn = get_connection() + try: + # Opportunistically clean up expired tokens + conn.execute( + "DELETE FROM refresh_tokens WHERE expires_at < ?", + (datetime.now(timezone.utc).isoformat(),), + ) + conn.commit() + + cur = conn.execute( + """ + SELECT id, username, expires_at, is_desktop FROM refresh_tokens + WHERE token_hash = ? + """, + (token_hash,), + ) + row = cur.fetchone() + if row is None: + return None + + # Check expiry + expires_at = datetime.fromisoformat(row["expires_at"]) + if datetime.now(timezone.utc) > expires_at: + conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],)) + conn.commit() + return None + + return row["username"], bool(row["is_desktop"]) + finally: + conn.close() + + +def revoke_user_refresh_tokens(username: str) -> None: + """Revoke all refresh tokens for a user (e.g. on logout).""" + conn = get_connection() + try: + conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,)) + 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. + """ + cache_id = _api_key_cache_id(raw_key) + cached_hash = _api_key_hash_cache.get(cache_id) + key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key) + conn = get_connection() + try: + cur = conn.execute( + "SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?", + (key_hash,), + ) + row = cur.fetchone() + if row is None: + return None + # Real key: memoize so later requests skip the KDF. Bounded; clear on overflow. + if cached_hash is None: + with _api_key_hash_cache_lock: + if len(_api_key_hash_cache) >= _API_KEY_HASH_CACHE_MAX: + _api_key_hash_cache.clear() + _api_key_hash_cache[cache_id] = key_hash + if not row["is_active"]: + return None + if row["expires_at"] is not None: + 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() diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py new file mode 100644 index 0000000..ef7bacb --- /dev/null +++ b/studio/backend/cloudflare_tunnel.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Free Cloudflare quick tunnel for Studio's 0.0.0.0 launches. + +The raw http://: is often unreachable (https-vs-http, blocked ports, +closed security groups); a cloudflared quick tunnel gives a free +https://*.trycloudflare.com URL that works anywhere, with no account or domain. + +Best-effort throughout: any failure collapses to "no URL" and Studio keeps +running. Stdlib only (back-end imports are lazy) so it is safe to import early. +""" + +from __future__ import annotations + +import os +import platform +import re +import shutil +import subprocess +import sys +import threading +from pathlib import Path +from typing import Optional, Tuple + +# cloudflared logs the quick-tunnel URL; match only the URL so we do not depend +# on the surrounding wording, which Cloudflare may change. The negative lookahead +# drops cloudflared's own API host, which appears in failure lines such as +# failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel" +# and must never be mistaken for a usable tunnel URL. +_URL_RE = re.compile(r"https://(?!api\.)[A-Za-z0-9-]+\.trycloudflare\.com") + +# cloudflared logs this once per edge connection it establishes. Until at least +# one appears the quick-tunnel URL returns Cloudflare error 1033 (HTTP 530), so +# we wait for it before advertising the URL. +_REGISTERED_MARKER = "Registered tunnel connection" + +_RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download" + +_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection +_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download + + +def _windows_hidden_kwargs() -> dict: + """Suppress a child console window on Windows; no-op elsewhere.""" + if sys.platform != "win32": + return {} + flags = getattr(subprocess, "CREATE_NO_WINDOW", 0) + return {"creationflags": flags} if flags else {} + + +def _lifetime_kwargs() -> dict: + """Bind cloudflared to the parent's lifetime (Linux PDEATHSIG). Lazy + + best-effort so this module still loads standalone (storage_roots-style).""" + try: + from utils.process_lifetime import child_popen_kwargs + return child_popen_kwargs() + except Exception: + return {} + + +def _asset_name() -> Optional[Tuple[str, bool]]: + """(release asset filename, is_tgz) for this OS/arch, or None if unsupported.""" + system = platform.system().lower() + machine = platform.machine().lower() + is_x64 = machine in ("x86_64", "amd64", "x64") + is_arm64 = machine in ("aarch64", "arm64") + is_x86 = machine in ("i386", "i686", "x86") + if system == "linux": + if is_x64: + return ("cloudflared-linux-amd64", False) + if is_arm64: + return ("cloudflared-linux-arm64", False) + elif system == "darwin": + if is_arm64: + return ("cloudflared-darwin-arm64.tgz", True) + if is_x64: + return ("cloudflared-darwin-amd64.tgz", True) + elif system == "windows": + if is_x64: + return ("cloudflared-windows-amd64.exe", False) + if is_x86: + return ("cloudflared-windows-386.exe", False) + return None + + +def _cache_path() -> Optional[Path]: + """studio_bin_root()/cloudflared(.exe), or None if the studio home is unresolvable.""" + try: + from utils.paths.storage_roots import studio_bin_root # lazy: backend-only import + except Exception: + return None + name = "cloudflared.exe" if sys.platform == "win32" else "cloudflared" + return studio_bin_root() / name + + +def find_cloudflared() -> Optional[str]: + """Locate an existing cloudflared: PATH first, then the Studio bin cache.""" + on_path = shutil.which("cloudflared") + if on_path: + return on_path + cached = _cache_path() + if cached is not None and cached.is_file() and os.access(cached, os.X_OK): + return str(cached) + return None + + +def _download(url: str, dest: Path) -> bool: + """Download url to dest via urllib (temp file + atomic rename). Best-effort -> bool.""" + import tempfile + import urllib.request + + tmp_path: Optional[Path] = None + try: + dest.parent.mkdir(parents = True, exist_ok = True) + with tempfile.NamedTemporaryFile( + prefix = dest.name + ".tmp-", dir = dest.parent, delete = False + ) as handle: + tmp_path = Path(handle.name) + # GitHub's CDN 403s the default Python-urllib User-Agent. + req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = _DOWNLOAD_TIMEOUT) as response: + shutil.copyfileobj(response, handle) + if tmp_path.stat().st_size == 0: + raise RuntimeError("empty download") + os.replace(tmp_path, dest) + return True + except Exception: + if tmp_path is not None: + try: + tmp_path.unlink(missing_ok = True) + except Exception: + pass + return False + + +def _extract_tgz_member(tgz_path: Path, dest: Path) -> bool: + """Extract just the `cloudflared` member from a darwin .tgz to dest. + + Rejects absolute paths and `..` traversal so a hostile archive cannot write + outside dest. Best-effort -> bool. + """ + import tarfile + try: + with tarfile.open(tgz_path, "r:gz") as tar: + member = None + for m in tar.getmembers(): + if not m.isfile() or os.path.basename(m.name) != "cloudflared": + continue + if m.name.startswith("/") or ".." in Path(m.name).parts: + continue + member = m + break + if member is None: + return False + src = tar.extractfile(member) + if src is None: + return False + with src, open(dest, "wb") as out: + shutil.copyfileobj(src, out) + return True + except Exception: + return False + + +def ensure_cloudflared() -> Optional[str]: + """Return a cloudflared path, downloading + caching the binary once if missing.""" + existing = find_cloudflared() + if existing: + return existing + asset = _asset_name() + cached = _cache_path() + if asset is None or cached is None: + return None + name, is_tgz = asset + url = f"{_RELEASE_BASE}/{name}" + try: + cached.parent.mkdir(parents = True, exist_ok = True) + if is_tgz: + tgz = cached.with_suffix(".tgz") + if not _download(url, tgz) or not _extract_tgz_member(tgz, cached): + tgz.unlink(missing_ok = True) + return None + tgz.unlink(missing_ok = True) + elif not _download(url, cached): + return None + if sys.platform != "win32": + os.chmod(cached, 0o755) + return str(cached) + except Exception: + return None + + +class CloudflareTunnel: + """A cloudflared quick tunnel to http://localhost:. Best-effort throughout. + + Use localhost (not the wildcard bind) as the tunnel origin so cloudflared's + upstream stays local-only. + """ + + def __init__( + self, + port: int, + binary: str, + protocol: Optional[str] = None, + ): + self.port = port + self.binary = binary + # None lets cloudflared pick its default (quic, with its own http2 + # fallback); set to "http2" to force it when quic is blocked. + self.protocol = protocol + self._proc: Optional[subprocess.Popen] = None + self._lock = threading.Lock() + self._stopped = False + self._url_event = threading.Event() + self._ready_event = threading.Event() + self.url: Optional[str] = None + self.ready = False + self.error: Optional[str] = None + + def start(self) -> None: + cmd = [ + self.binary, + "tunnel", + "--url", + f"http://localhost:{self.port}", + "--no-autoupdate", + ] + if self.protocol: + cmd += ["--protocol", self.protocol] + with self._lock: + # A stop() that landed before us (e.g. a shutdown in the caller's + # register->start window) marks the tunnel stopped; spawning now would + # orphan a process nobody owns, so refuse. + if self._stopped: + return + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + stdin = subprocess.DEVNULL, + text = True, + errors = "replace", + bufsize = 1, + **_windows_hidden_kwargs(), + **_lifetime_kwargs(), + ) + self._proc = proc + threading.Thread( + target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True + ).start() + + def _reader(self, proc: subprocess.Popen) -> None: + # Drain cloudflared's output: capture the first trycloudflare URL and the + # first edge-connection registration, and keep draining so it never + # blocks on a full pipe. + try: + if proc.stdout is not None: + for line in proc.stdout: + if self.url is None: + match = _URL_RE.search(line) + if match: + self.url = match.group(0) + self._url_event.set() + if not self.ready and _REGISTERED_MARKER in line: + self.ready = True + self._ready_event.set() + except Exception: + pass + finally: + # stdout closed -> cloudflared has exited. Record why, and unblock any + # waiters at once instead of letting them wait out the full timeout. + if self.url is None: + self.error = "cloudflared exited before emitting a tunnel URL" + elif not self.ready: + self.error = "cloudflared exited before the tunnel connection registered" + self._url_event.set() + self._ready_event.set() + + def wait_for_ready(self, timeout: float = _READY_TIMEOUT) -> Optional[str]: + """Block until the tunnel is actually serving -- the URL has been minted + *and* at least one edge connection has registered -- or until timeout. + + Returns the URL only when ready, so callers never advertise a URL that + would return Cloudflare error 1033 (HTTP 530).""" + self._ready_event.wait(timeout) + return self.url if self.ready else None + + def stop(self) -> None: + """Terminate the tunnel. Idempotent and safe to call from a signal handler.""" + with self._lock: + # Mark stopped so a start() racing behind us refuses to spawn. + self._stopped = True + proc, self._proc = self._proc, None + if proc is None: + return + try: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout = 5) + except subprocess.TimeoutExpired: + proc.kill() + try: + proc.wait(timeout = 5) + except Exception: + pass + except Exception: + pass + + +# Single serving process per Studio launch, so one module-level tunnel handle is +# enough; the lock guards the start/stop/shutdown races. +_active_tunnel: Optional[CloudflareTunnel] = None +_active_lock = threading.Lock() +# Latched by stop_studio_tunnel so a shutdown landing *between* a start's retry +# attempts aborts the loop instead of starting a tunnel nobody will ever stop. +_shutdown_requested = False + + +def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[str]: + """Start a quick tunnel and return its public URL once it is actually + serving, or None (best-effort). + + Waits for cloudflared to both mint the URL and register an edge connection + before returning, so the caller never advertises a URL that yields Cloudflare + error 1033 (HTTP 530). If a URL is minted but no connection registers within + the window (e.g. quic is blocked on this network), retries once forcing the + http2 protocol. On any failure the tunnel is stopped and None is returned. + """ + global _active_tunnel, _shutdown_requested + binary = ensure_cloudflared() + if not binary: + return None + with _active_lock: + _shutdown_requested = False # fresh session + # Default protocol first (quic, with cloudflared's own http2 fallback); if a + # URL appears but no connection registers, quic is likely blocked -> retry + # once forcing http2. + for protocol in (None, "http2"): + # Create + register under the lock, and bail if a stop already landed + # (e.g. between this and the previous attempt) so we never start a tunnel + # after shutdown has run. + with _active_lock: + if _shutdown_requested: + _active_tunnel = None + return None + tunnel = CloudflareTunnel(port, binary, protocol = protocol) + prior, _active_tunnel = _active_tunnel, tunnel + if prior is not None: + prior.stop() + try: + tunnel.start() + url = tunnel.wait_for_ready(timeout) + except Exception: + url = None + if url: + return url + saw_url = tunnel.url is not None + # Not ready: drop it, but only if we are still the active tunnel. + with _active_lock: + was_active = _active_tunnel is tunnel + if was_active: + _active_tunnel = None + tunnel.stop() + # A concurrent shutdown or start took over while we waited; retrying would + # spawn a tunnel nobody owns (orphaned after shutdown), so bail instead. + if not was_active: + return None + # No URL at all is an API/network failure, not a protocol one; forcing + # http2 will not help, so do not burn another window on it. + if not saw_url: + return None + return None + + +def stop_studio_tunnel() -> None: + """Terminate the active tunnel, if any. Idempotent.""" + global _active_tunnel, _shutdown_requested + with _active_lock: + # Latch so an in-flight start_studio_tunnel won't start a fresh tunnel + # (e.g. its http2 retry) after we have already torn down. + _shutdown_requested = True + tunnel, _active_tunnel = _active_tunnel, None + if tunnel is not None: + tunnel.stop() diff --git a/studio/backend/colab.py b/studio/backend/colab.py new file mode 100644 index 0000000..dd27439 --- /dev/null +++ b/studio/backend/colab.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Colab helpers for Unsloth Studio. Uses Colab's built-in proxy. +""" + +from pathlib import Path +import sys + +# Seed platform._sys_version_cache before attrs->rich->structlog->platform crash on conda Python. +# See: https://github.com/python/cpython/issues/102396 +_backend_dir = str(Path(__file__).parent) +if _backend_dir not in sys.path: + sys.path.insert(0, _backend_dir) +import _platform_compat # noqa: F401 + + +from loggers import get_logger + +logger = get_logger(__name__) + + +def get_colab_url(port: int = 8888) -> str: + """ + Get the Colab proxy URL for a port. + + Retries up to 3 times, validating the result is a real HTTPS Colab URL. + Falls back to http://localhost:{port} only when all attempts fail. + """ + import time as _time + + fallback = f"http://localhost:{port}" + + try: + from google.colab.output import eval_js + except ImportError: + return fallback + + for attempt in range(3): + try: + url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10) + # Valid proxy URL is https:// and embeds the port. + if url and isinstance(url, str) and url.startswith("https://") and str(port) in url: + return url.rstrip("/") + except Exception as e: + logger.info(f"Note: Could not get Colab URL (attempt {attempt + 1}/3: {e})") + if attempt < 2: + _time.sleep(1) + + logger.warning( + f"Could not get a valid Colab proxy URL after 3 attempts — using localhost fallback. " + f"The link/iframe may not work from outside the runtime." + ) + return fallback + + +def show_link(port: int = 8888, *, _url: "str | None" = None): + """Display a styled clickable link to the UI. + + *_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip. + """ + from IPython.display import display, HTML + + url = _url if _url is not None else get_colab_url(port) + + # Truncated display URL; try/except so an odd URL shape still renders the link. + try: + port_prefix = f"{port}-" + idx = url.index(port_prefix) + next_dash = url.index("-", idx + len(port_prefix)) + short_url = url[: next_dash + 1] + "..." + except (ValueError, IndexError): + short_url = url + + # Plain-text line so the URL shows even if HTML display fails. + logger.info(f"🌐 Unsloth Studio URL: {url}") + + html = f""" +

+

+ + Unsloth Studio is Ready! +

+ + + Open Unsloth Studio + +

+ If the link doesn't work, you can scroll down to view the UI generated directly in Colab. +

+

+ {short_url} +

+
+ """ + display(HTML(html)) + + +def _bootstrap_password_pending() -> bool: + """True while the default admin still owes a bootstrap-password change. + + While pending, main.py injects that password into same-origin GETs, and a public + tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin + access. Fails safe to pending if the state cannot be read. + """ + try: + from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME + return bool(requires_password_change(DEFAULT_ADMIN_USERNAME)) + except Exception as e: + logger.info(f"Could not check admin password state ({e}); refusing tunnel to be safe.") + return True + + +def start_cloudflare_tunnel(port: int) -> "str | None": + """Open a shareable Cloudflare quick tunnel to localhost:*port*, or None. + + run_server suppresses the tunnel on Colab by design, so we start it directly. + Refused while the bootstrap password is pending; any failure collapses to None + and the Colab proxy still works. + """ + if _bootstrap_password_pending(): + logger.warning( + "Cloudflare link not started: the admin account still has its temporary " + "bootstrap password, which is exposed to anyone who can load the page. " + "Open Studio in this tab, log in and change the admin password, then re-run " + "start(cloudflare=True) to get the shareable link." + ) + return None + try: + from cloudflare_tunnel import start_studio_tunnel + except Exception as e: + logger.info(f"Cloudflare tunnel unavailable ({e}); using Colab proxy only.") + return None + try: + url = start_studio_tunnel(port) + except Exception as e: + logger.info(f"Cloudflare tunnel failed to start ({e}); using Colab proxy only.") + return None + # Success is logged by _show_and_embed; note only misses here. + if not url: + logger.info("Cloudflare tunnel did not produce a URL; using Colab proxy only.") + return url + + +def _publish_cloudflare_url(cloudflare_url: "str | None") -> None: + """Publish a directly-started tunnel URL onto app.state so /api/health advertises it. + + run_server only sets this when it opens the tunnel itself, which it skips on Colab, + so we set it here. Otherwise the frontend's API examples fall back to an + unreachable server_url. Best-effort. + """ + if not cloudflare_url: + return + try: + from main import app as _studio_app + _studio_app.state.cloudflare_url = cloudflare_url + except Exception as e: + logger.info(f"Could not publish Cloudflare URL to /api/health ({e}).") + + +def _stop_cloudflare_tunnel() -> None: + """Best-effort teardown of the Cloudflare tunnel started by start_cloudflare_tunnel.""" + try: + from cloudflare_tunnel import stop_studio_tunnel + stop_studio_tunnel() + except Exception: + pass + # Stop /api/health advertising a dead tunnel. + try: + from main import app as _studio_app + _studio_app.state.cloudflare_url = None + except Exception: + pass + + +def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: + """True only if Unsloth Studio (not some other app) answers /api/health on *port*. + + The service-marker check stops the reuse path reusing or tunneling a foreign + process that merely serves /api/health. + """ + import json, urllib.request + try: + with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout) as r: + return json.loads(r.read()).get("service") == "Unsloth UI Backend" + except Exception: + return False + + +def _shareable_link_html(cloudflare_url: str) -> str: + """Branded card for the shareable Cloudflare link, styled like the show_link banner.""" + return f""" +
+

+ + Shareable Studio Link is Ready! +

+ + + Open Unsloth Studio + +

+ This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab. +

+

+ 🔗 {cloudflare_url} +

+
+ """ + + +def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None): + """Render the Studio header + iframe for *port*, with a shareable-link card above + when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe.""" + url = get_colab_url(port) + logger.info(f"🌐 Unsloth Studio URL: {url}") + if cloudflare_url: + logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") + + try: + from IPython.display import HTML, display + + iframe_id = f"unsloth-studio-{port}" + + # Truncated header URL — best-effort, falls back to full URL. + try: + port_prefix = f"{port}-" + idx = url.index(port_prefix) + next_dash = url.index("-", idx + len(port_prefix)) + short_url = url[: next_dash + 1] + "..." + except (ValueError, IndexError): + short_url = url + + if cloudflare_url: + display(HTML(_shareable_link_html(cloudflare_url))) + + display( + HTML(f""" +
+
+ + Unsloth Studio + {short_url} +
+ +
+""") + ) + except Exception: + # Fallback: Colab's built-in helper. + try: + from google.colab import output as colab_output + colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") + except ImportError: + pass + + +def start(port: int = 8888, *, cloudflare: bool = False): + """Start Unsloth Studio in Colab and display the URL. + + Args: + port: Port to bind/serve on. + cloudflare: Opt in to a shareable Cloudflare HTTPS link reachable from any + device (default OFF). It exposes Studio's login page beyond Colab, so it + stays an explicit opt-in; the default shows only the in-tab proxy iframe. + + Usage: + start() # Colab-proxy iframe only (default) + start(cloudflare=True) # also open a shareable Cloudflare link + """ + import time + + logger.info("🦥 Starting Unsloth Studio...") + + # Fast path: Studio already running (cell re-run). Re-launching would collide on + # the port, so just re-show the link and iframe. + if _is_studio_healthy(port): + logger.info(f" Studio is already running on port {port} — reusing existing server.") + # try/finally: tear the tunnel down even if interrupted mid-start/render. + try: + cf_url = start_cloudflare_tunnel(port) if cloudflare else None + _publish_cloudflare_url(cf_url) + _show_and_embed(port, cloudflare_url = cf_url) + for _ in range(10000): + time.sleep(300) + print("=", end = "", flush = True) + except KeyboardInterrupt: + logger.info("\nUnsloth Studio keepalive stopped.") + finally: + _stop_cloudflare_tunnel() + return + + logger.info(" Loading backend...") + from run import run_server + + # Auto-detect frontend path + repo_root = Path(__file__).parent.parent + frontend_path = repo_root / "frontend" / "dist" + + if not (frontend_path / "index.html").exists(): + logger.info("❌ Frontend not built! Please run the setup cell first.") + return + + logger.info(" Starting server...") + try: + # cloudflare=False: this helper owns the tunnel. run_server's default True + # would tunnel this 0.0.0.0 bind if Colab detection fails, breaking the opt-out. + app = run_server( + host = "0.0.0.0", + port = port, + frontend_path = frontend_path, + silent = True, + cloudflare = False, + ) + except SystemExit as exc: + logger.error(f"❌ Unsloth Studio failed to start: {exc}") + return + except Exception as exc: + logger.error(f"❌ Unsloth Studio failed to start: {exc}") + return + + # run_server auto-increments the port if in use; read back the bound port so the + # proxy URL and iframe point at the right place. + actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port + + logger.info(f" Server started on port {actual_port}!") + + # Poll health endpoint before showing the link — avoids the race where ready_event + # fires but the process hasn't finished binding. + import urllib.request + + server_ready = False + for _ in range(40): + try: + with urllib.request.urlopen(f"http://localhost:{actual_port}/api/health", timeout = 1): + server_ready = True + break + except Exception: + time.sleep(0.5) + + if not server_ready: + logger.error( + f"❌ Unsloth Studio did not become healthy on port {actual_port}. " + "Check for errors above." + ) + return + + # Open the tunnel now the server is healthy, publish its URL for /api/health, and + # tear it down on interrupt (try/finally) rather than orphan the process. + try: + cf_url = start_cloudflare_tunnel(actual_port) if cloudflare else None + _publish_cloudflare_url(cf_url) + _show_and_embed(actual_port, cloudflare_url = cf_url) + + # Keep kernel alive so the daemon server thread runs. + for _ in range(10000): + time.sleep(300) + print("=", end = "", flush = True) + except KeyboardInterrupt: + logger.info("\nUnsloth Studio keepalive stopped.") + finally: + _stop_cloudflare_tunnel() + + +if __name__ == "__main__": + start() diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py new file mode 100644 index 0000000..30ecb27 --- /dev/null +++ b/studio/backend/core/__init__.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unified core module for Unsloth backend + +Imports are LAZY (via __getattr__) so training subprocesses can import +core.training.worker without pulling in heavy ML deps (unsloth, transformers, +torch) before the version-activation code runs. +""" + +import sys +from pathlib import Path + +# Add backend dir to sys.path so bare "from utils.*" imports work when core +# is imported as a package. +_backend_dir = str(Path(__file__).resolve().parent.parent) +if _backend_dir not in sys.path: + sys.path.insert(0, _backend_dir) + +__all__ = [ + # Inference + "InferenceBackend", + "get_inference_backend", + # Training + "get_training_backend", + "TrainingBackend", + "TrainingProgress", + # Config + "ModelConfig", + "is_vision_model", + "scan_trained_models", + "scan_trained_loras", + "load_model_defaults", + "get_base_model_from_lora", + # Utils + "format_and_template_dataset", + "normalize_path", + "is_local_path", + "is_model_cached", + "without_hf_auth", + "format_error_message", + "get_gpu_memory_info", + "log_gpu_memory", + "get_device", + "is_apple_silicon", + "clear_gpu_cache", + "DeviceType", +] + + +def __getattr__(name): + # Inference + if name in ("InferenceBackend", "get_inference_backend"): + from .inference import InferenceBackend, get_inference_backend + + globals()["InferenceBackend"] = InferenceBackend + globals()["get_inference_backend"] = get_inference_backend + return globals()[name] + + # Training + if name in ("TrainingBackend", "get_training_backend", "TrainingProgress"): + from .training import TrainingBackend, get_training_backend, TrainingProgress + + globals()["TrainingBackend"] = TrainingBackend + globals()["get_training_backend"] = get_training_backend + globals()["TrainingProgress"] = TrainingProgress + return globals()[name] + + # Config (utils.models) + if name in ( + "is_vision_model", + "ModelConfig", + "scan_trained_models", + "scan_trained_loras", + "load_model_defaults", + "get_base_model_from_lora", + ): + from utils.models import ( + is_vision_model, + ModelConfig, + scan_trained_models, + load_model_defaults, + get_base_model_from_lora, + ) + + globals()["is_vision_model"] = is_vision_model + globals()["ModelConfig"] = ModelConfig + 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] + + # Paths + if name in ("normalize_path", "is_local_path", "is_model_cached"): + from utils.paths import normalize_path, is_local_path, is_model_cached + + globals()["normalize_path"] = normalize_path + globals()["is_local_path"] = is_local_path + globals()["is_model_cached"] = is_model_cached + return globals()[name] + + # Utils + if name in ("without_hf_auth", "format_error_message"): + from utils.utils import without_hf_auth, format_error_message + + globals()["without_hf_auth"] = without_hf_auth + globals()["format_error_message"] = format_error_message + return globals()[name] + + # Hardware + if name in ( + "get_device", + "is_apple_silicon", + "clear_gpu_cache", + "get_gpu_memory_info", + "log_gpu_memory", + "DeviceType", + ): + from utils.hardware import ( + get_device, + is_apple_silicon, + clear_gpu_cache, + get_gpu_memory_info, + log_gpu_memory, + DeviceType, + ) + + globals()["get_device"] = get_device + globals()["is_apple_silicon"] = is_apple_silicon + globals()["clear_gpu_cache"] = clear_gpu_cache + globals()["get_gpu_memory_info"] = get_gpu_memory_info + globals()["log_gpu_memory"] = log_gpu_memory + globals()["DeviceType"] = DeviceType + return globals()[name] + + # Datasets + if name == "format_and_template_dataset": + from utils.datasets import format_and_template_dataset + globals()["format_and_template_dataset"] = format_and_template_dataset + return format_and_template_dataset + + raise AttributeError(f"module 'core' has no attribute {name!r}") diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py new file mode 100644 index 0000000..6336954 --- /dev/null +++ b/studio/backend/core/_torchao_stub.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared torchao Windows-ROCm import stub. + +torchao (pulled in by transformers.quantizers) imports distributed_c10d.py +unconditionally, which crashes on Windows ROCm because the RCCL backend +(torch._C._distributed_c10d) is absent. Stubbing torchao short-circuits its +import chain; _StubSubpackageFinder handles any depth of torchao.xxx.yyy. +Worker subprocesses call install_torchao_windows_rocm_stub() before importing +transformers / unsloth_zoo. +""" + +from __future__ import annotations + +import sys +import types +import importlib.abc +import importlib.machinery + +_STUB_SENTINEL = object() + + +# Metaclass for stub types so isinstance(x, StubClass) returns False instead of +# raising TypeError -- peft's lora/torchao.py does isinstance() against torchao +# types, which fails if those names resolve to stub modules rather than types. +class _StubTypeMeta(type): + def __instancecheck__(cls, instance): + return False + + def __subclasscheck__(cls, subclass): + return False + + def __getattr__(cls, attr): + if attr.startswith("__"): + raise AttributeError(attr) + child = _StubTypeMeta(attr, (), {}) + setattr(cls, attr, child) + return child + + def __call__(cls, *args, **kwargs): + return None + + +def _make_stub_type(name): + """Stub class: accepted by isinstance() (always False), supports attr access.""" + return _StubTypeMeta(name, (), {}) + + +def _make_mod_stub(mod_name): + m = types.ModuleType(mod_name) + m.__path__ = [] + m.__package__ = mod_name + m._unsloth_stub = _STUB_SENTINEL + m.__spec__ = importlib.machinery.ModuleSpec(mod_name, loader = None, is_package = True) + + def _ga( + attr, + _m = m, + _n = mod_name, + ): + if attr.startswith("__"): + raise AttributeError(attr) + # Return a stub CLASS (not module) so isinstance() returns False, not TypeError. + child = _make_stub_type(f"{_n}.{attr}") + setattr(_m, attr, child) + return child + + m.__getattr__ = _ga + return m + + +class _StubSubpackageLoader(importlib.abc.Loader): + def __init__(self, mod_name): + self._mod_name = mod_name + + def create_module(self, spec): + return _make_mod_stub(self._mod_name) + + def exec_module(self, module): + pass + + +class _StubSubpackageFinder(importlib.abc.MetaPathFinder): + def find_spec( + self, + fullname, + path, + target = None, + ): + if "." not in fullname: + return None + parent = sys.modules.get(fullname.rsplit(".", 1)[0]) + if parent is None: + return None + if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL: + return None + return importlib.machinery.ModuleSpec( + fullname, _StubSubpackageLoader(fullname), is_package = True + ) + + +def install_torchao_windows_rocm_stub() -> None: + """Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash. + + No-op elsewhere (incl. Windows CUDA, where torchao is real). Must run before + importing transformers / unsloth_zoo. Safe to call once per worker. + """ + # Gate on the active torch runtime, not env-var presence -- HIP_PATH/ROCM_PATH + # persist after reverting to a CUDA wheel. Some ROCm wheels lack + # torch.version.hip but still encode "rocm" in __version__, so accept either. + _is_win32_rocm = False + if sys.platform == "win32": + try: + import torch as _torch_probe + _is_win32_rocm = bool( + getattr(getattr(_torch_probe, "version", None), "hip", None) + or "rocm" in getattr(_torch_probe, "__version__", "").lower() + ) + del _torch_probe + except Exception: + pass + if _is_win32_rocm: + # Register the finder only on Windows ROCm. + sys.meta_path.append(_StubSubpackageFinder()) + # Seed torchao top-level + key submodules; the finder handles the rest. + for _tao_name in ( + "torchao", + "torchao.quantization", + "torchao.dtypes", + "torchao.float8", + "torchao.utils", + ): + if _tao_name not in sys.modules: + sys.modules[_tao_name] = _make_mod_stub(_tao_name) diff --git a/studio/backend/core/data_recipe/__init__.py b/studio/backend/core/data_recipe/__init__.py new file mode 100644 index 0000000..ebc2798 --- /dev/null +++ b/studio/backend/core/data_recipe/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Data Recipe core (DataDesigner wrapper + job runner).""" + +from .jobs import JobManager, get_job_manager + +__all__ = ["JobManager", "get_job_manager"] diff --git a/studio/backend/core/data_recipe/huggingface.py b/studio/backend/core/data_recipe/huggingface.py new file mode 100644 index 0000000..7a1219b --- /dev/null +++ b/studio/backend/core/data_recipe/huggingface.py @@ -0,0 +1,121 @@ +# 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 json +from pathlib import Path + +from utils.paths import recipe_datasets_root, resolve_dataset_path + +_DATA_DESIGNER_FOOTER = ( + 'Made with ❤️ using 🎨 ' + 'NeMo Data Designer' +) +_UNSLOTH_STUDIO_FOOTER = ( + 'Made with ❤️ using 🦥 ' "Unsloth Studio" +) + + +class RecipeDatasetPublishError(ValueError): + """Raised when a recipe dataset cannot be published to Hugging Face.""" + + +def _resolve_recipe_artifact_path(artifact_path: str) -> Path: + root = recipe_datasets_root().expanduser().resolve() + candidate = resolve_dataset_path(artifact_path).expanduser() + resolved = candidate.resolve(strict = False) + + try: + resolved.relative_to(root) + except ValueError as exc: + raise RecipeDatasetPublishError( + "This execution artifact is outside the Recipe Studio dataset storage." + ) from exc + + if not resolved.exists(): + raise RecipeDatasetPublishError("Execution artifacts are no longer available.") + if not resolved.is_dir(): + raise RecipeDatasetPublishError("Execution artifact path is not a dataset folder.") + + return resolved + + +def publish_recipe_dataset( + *, + artifact_path: str, + repo_id: str, + description: str, + hf_token: str | None = None, + private: bool = False, +) -> str: + dataset_path = _resolve_recipe_artifact_path(artifact_path) + + try: + from data_designer.engine.storage.artifact_storage import ( + FINAL_DATASET_FOLDER_NAME, + METADATA_FILENAME, + PROCESSORS_OUTPUTS_FOLDER_NAME, + SDG_CONFIG_FILENAME, + ) + from data_designer.integrations.huggingface.client import ( + HuggingFaceHubClient, + HuggingFaceHubClientUploadError, + ) + from data_designer.integrations.huggingface.dataset_card import ( + DataDesignerDatasetCard, + ) + except ImportError as exc: + raise RecipeDatasetPublishError( + "NeMo Data Designer Hugging Face integration is not installed." + ) from exc + + try: + client = HuggingFaceHubClient(token = hf_token) + client._validate_repo_id(repo_id = repo_id) + client._validate_dataset_path(base_dataset_path = dataset_path) + client._create_or_get_repo(repo_id = repo_id, private = private) + + metadata_path = dataset_path / METADATA_FILENAME + builder_config_path = dataset_path / SDG_CONFIG_FILENAME + + with metadata_path.open(encoding = "utf-8") as fh: + metadata = json.load(fh) + + builder_config = None + if builder_config_path.exists(): + with builder_config_path.open(encoding = "utf-8") as fh: + builder_config = json.load(fh) + + card = DataDesignerDatasetCard.from_metadata( + metadata = metadata, + builder_config = builder_config, + repo_id = repo_id, + description = description, + tags = None, + ) + card.text = card.text.replace(_DATA_DESIGNER_FOOTER, _UNSLOTH_STUDIO_FOOTER) + # Data Designer drops the explicit token, so push the card ourselves to keep auth request-local. + card.push_to_hub(repo_id, token = hf_token, repo_type = "dataset") + + client._upload_main_dataset_files( + repo_id = repo_id, + parquet_folder = dataset_path / FINAL_DATASET_FOLDER_NAME, + ) + client._upload_images_folder( + repo_id = repo_id, + images_folder = dataset_path / "images", + ) + client._upload_processor_files( + repo_id = repo_id, + processors_folder = dataset_path / PROCESSORS_OUTPUTS_FOLDER_NAME, + ) + client._upload_config_files( + repo_id = repo_id, + metadata_path = metadata_path, + builder_config_path = builder_config_path, + ) + + return f"https://huggingface.co/datasets/{repo_id}" + except HuggingFaceHubClientUploadError as exc: + raise RecipeDatasetPublishError(str(exc)) from exc diff --git a/studio/backend/core/data_recipe/jobs/__init__.py b/studio/backend/core/data_recipe/jobs/__init__.py new file mode 100644 index 0000000..cf03d62 --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from .manager import JobManager, get_job_manager + +__all__ = ["JobManager", "get_job_manager"] diff --git a/studio/backend/core/data_recipe/jobs/constants.py b/studio/backend/core/data_recipe/jobs/constants.py new file mode 100644 index 0000000..0045276 --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/constants.py @@ -0,0 +1,34 @@ +# 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 + +# stages parsed from data-designer logs +STAGE_CREATE = "create" +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" +STAGE_PROFILING = "profiling" + +USAGE_RESET_STAGES = { + STAGE_CREATE, + STAGE_PREVIEW, + STAGE_DAG, + STAGE_HEALTHCHECK, + STAGE_SAMPLING, + STAGE_GENERATING, + STAGE_PROFILING, +} + +# job event types emitted by worker/manager +EVENT_JOB_ENQUEUED = "job.enqueued" +EVENT_JOB_STARTED = "job.started" +EVENT_JOB_CANCELLING = "job.cancelling" +EVENT_JOB_CANCELLED = "job.cancelled" +EVENT_JOB_COMPLETED = "job.completed" +EVENT_JOB_ERROR = "job.error" diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py new file mode 100644 index 0000000..0e00447 --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -0,0 +1,599 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import asyncio +import json +import queue +import threading +import time +import uuid +from pathlib import Path +from collections import deque +from dataclasses import dataclass +from typing import Any + +import multiprocessing as mp + +from ..jsonable import to_preview_jsonable +from .constants import ( + EVENT_JOB_CANCELLING, + EVENT_JOB_CANCELLED, + EVENT_JOB_COMPLETED, + EVENT_JOB_ENQUEUED, + EVENT_JOB_ERROR, + EVENT_JOB_STARTED, +) +from .parse import apply_update, coerce_event, parse_log_message +from .types import Job +from .worker import run_job_process +from loggers import get_logger + +logger = get_logger(__name__) + + +_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] + _q: queue.Queue + _next_id: int = 0 + + async def next_event(self, *, timeout_sec: float) -> dict | None: + """Wait for next event (SSE), w/ timeout so we can check disconnects.""" + try: + return await asyncio.to_thread(self._q.get, True, timeout_sec) + except queue.Empty: + return None + + def format_sse(self, event: dict) -> bytes: + """Turn event dict into SSE bytes (id/event/data).""" + event_id = event.get("seq") + if event_id is None: + self._next_id += 1 + event_id = self._next_id + body = json.dumps(event, separators = (",", ":"), ensure_ascii = False) + event_type = event.get("type") or "message" + return (f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n").encode("utf-8") + + +class JobManager: + def __init__(self) -> None: + """Single-job runner (in-mem). Simple on purpose, not a whole platform.""" + self._lock = threading.Lock() + self._job: Job | None = None + self._proc: mp.Process | None = None + self._mp_q: Any | None = None + self._events: deque[dict] = deque(maxlen = 5000) + self._subs: list[queue.Queue] = [] + self._pump_thread: threading.Thread | None = None + self._seq: int = 0 + + 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 a workflow-scoped sk-unsloth-* key row id + minted by the route layer; revoked on terminal state so the key's + live window is no longer than the run. + """ + llm_columns = recipe.get("columns") or [] + llm_column_count = 0 + if isinstance(llm_columns, list): + for column in llm_columns: + if not isinstance(column, dict): + continue + column_type = str(column.get("column_type") or "").strip().lower() + if column_type.startswith("llm"): + llm_column_count += 1 + if llm_column_count <= 0: + llm_column_count = 1 + + with self._lock: + if self._proc is not None and self._proc.is_alive(): + raise RuntimeError("job already running") + + 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 + from utils.native_path_leases import ( + native_path_secret_removed_for_child_start, + run_without_native_path_secret, + ) + + 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() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) + + self._mp_q = mp_q + self._proc = proc + self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True) + self._pump_thread.start() + + self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id}) + return job_id + + def cancel(self, job_id: str) -> bool: + """Hard stop. We terminate the subprocess. Quick + reliable.""" + with self._lock: + if self._job is None or self._job.job_id != job_id: + return False + if self._proc is None or not self._proc.is_alive(): + return True + self._job.status = "cancelling" + self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id}) + try: + self._proc.terminate() + except (AttributeError, OSError): + pass + return True + + def get_status(self, job_id: str) -> dict | None: + """UI-friendly structured snapshot; an alternative to SSE.""" + with self._lock: + if self._job is None or self._job.job_id != job_id: + return None + job = self._job + return { + "job_id": job.job_id, + "status": job.status, + "stage": job.stage, + "current_column": job.current_column, + "completed_columns": list(job.completed_columns), + "batch": {"idx": job.batch.idx, "total": job.batch.total}, + "progress": { + "done": job.progress.done, + "total": job.progress.total, + "percent": job.progress.percent, + "eta_sec": job.progress.eta_sec, + "rate": job.progress.rate, + "ok": job.progress.ok, + "failed": job.progress.failed, + }, + "column_progress": { + "done": job.column_progress.done, + "total": job.column_progress.total, + "percent": job.column_progress.percent, + "eta_sec": job.column_progress.eta_sec, + "rate": job.column_progress.rate, + "ok": job.column_progress.ok, + "failed": job.column_progress.failed, + }, + "source_progress": _source_progress_status(job), + "model_usage": { + name: { + "model": usage.model, + "tokens": { + "input": usage.input_tokens, + "output": usage.output_tokens, + "total": usage.total_tokens, + "tps": usage.tps, + }, + "requests": { + "success": usage.requests_success, + "failed": usage.requests_failed, + "total": usage.requests_total, + "rpm": usage.rpm, + }, + } + for name, usage in job.model_usage.items() + }, + "rows": job.rows, + "cols": job.cols, + "error": job.error, + "has_analysis": job.analysis is not None, + "dataset_rows": None if job.dataset is None else len(job.dataset), + "artifact_path": job.artifact_path, + "execution_type": job.execution_type, + "started_at": job.started_at, + "finished_at": job.finished_at, + } + + def get_current_status(self) -> dict | None: + """Single-job convenience (last/current).""" + job_id = self.get_current_job_id() + if job_id is None: + return None + return self.get_status(job_id) + + def get_current_job_id(self) -> str | None: + """Return current job_id (or None).""" + with self._lock: + return None if self._job is None else self._job.job_id + + def get_analysis(self, job_id: str) -> dict | None: + """Final profiling output (only after job completes).""" + with self._lock: + if self._job is None or self._job.job_id != job_id: + return None + return self._job.analysis + + def get_dataset( + self, + job_id: str, + *, + limit: int, + offset: int = 0, + ) -> dict[str, Any] | None: + """Load dataset page (offset + limit) and include total rows.""" + with self._lock: + if self._job is None or self._job.job_id != job_id: + return None + in_memory_dataset = self._job.dataset + artifact_path = self._job.artifact_path + job_status = self._job.status + + if in_memory_dataset is not None: + total = len(in_memory_dataset) + rows = in_memory_dataset[offset : offset + limit] + return {"dataset": rows, "total": total} + if not artifact_path: + if job_status in {"completed", "error", "cancelled"}: + return {"error": "artifact path missing"} + return None + + try: + base_dataset_path = Path(artifact_path) + parquet_dir = base_dataset_path / "parquet-files" + if not parquet_dir.exists(): + return {"error": f"dataset path missing: {parquet_dir}"} + + return self._load_dataset_page(parquet_dir = parquet_dir, limit = limit, offset = offset) + except Exception as exc: + return {"error": f"dataset load failed: {exc}"} + + @staticmethod + def _load_dataset_page(*, parquet_dir: Path, limit: int, offset: int) -> dict[str, Any]: + dataset_page = JobManager._load_dataset_page_with_duckdb( + parquet_dir = parquet_dir, + limit = limit, + offset = offset, + ) + if dataset_page is not None: + return dataset_page + return JobManager._load_dataset_page_with_data_designer( + parquet_dir = parquet_dir, + limit = limit, + offset = offset, + ) + + @staticmethod + def _load_dataset_page_with_duckdb( + *, parquet_dir: Path, limit: int, offset: int + ) -> dict[str, Any] | None: + parquet_glob = str((parquet_dir / "*.parquet").resolve()) + try: + import duckdb # type: ignore + except Exception: + return None + + try: + conn = duckdb.connect(":memory:") + try: + total_row = conn.execute( + "SELECT COUNT(*) FROM read_parquet(?)", + [parquet_glob], + ).fetchone() + total = int(total_row[0] if total_row else 0) + dataframe = conn.execute( + ( + "SELECT *, row_number() OVER (PARTITION BY filename) AS __row_num__ " + "FROM read_parquet(?, filename=true) " + "ORDER BY filename, __row_num__ " + "LIMIT ? OFFSET ?" + ), + [parquet_glob, int(limit), int(offset)], + ).fetchdf() + finally: + conn.close() + except (RuntimeError, ValueError, duckdb.Error): + return None + + for helper_col in ("filename", "__row_num__"): + if helper_col in dataframe.columns: + dataframe = dataframe.drop(columns = [helper_col]) + + rows = dataframe.to_dict(orient = "records") + return {"dataset": to_preview_jsonable(rows), "total": total} + + @staticmethod + def _load_dataset_page_with_data_designer( + *, parquet_dir: Path, limit: int, offset: int + ) -> dict[str, Any]: + from data_designer.config.utils.io_helpers import read_parquet_dataset + + dataframe = read_parquet_dataset(parquet_dir) + total = int(len(dataframe.index)) + rows = dataframe.iloc[offset : offset + limit].to_dict(orient = "records") + return {"dataset": to_preview_jsonable(rows), "total": total} + + def subscribe( + self, + job_id: str, + *, + after_seq: int | None = None, + ) -> Subscription | None: + """SSE subscribe: get replay buffer + live events stream.""" + with self._lock: + if self._job is None or self._job.job_id != job_id: + return None + q: queue.Queue = queue.Queue(maxsize = 2000) + self._subs.append(q) + if after_seq is None: + replay = list(self._events) + else: + replay = [e for e in self._events if int(e.get("seq") or 0) > after_seq] + return Subscription(replay = replay, _q = q) + + def unsubscribe(self, sub: Subscription) -> None: + """Drop SSE subscriber (client disconnected).""" + with self._lock: + self._subs = [q for q in self._subs if q is not sub._q] + + def _emit(self, event: dict) -> None: + """Broadcast event to replay buffer + all subscribers.""" + self._seq += 1 + event["seq"] = self._seq + self._events.append(event) + stale: list[queue.Queue] = [] + for q in self._subs: + try: + q.put_nowait(event) + except queue.Full: + stale.append(q) + if stale: + self._subs = [q for q in self._subs if q not in stale] + + def _snapshot(self) -> tuple[Job, mp.Process, Any] | None: + """Grab pointers for the pump loop (avoid holding lock too long).""" + with self._lock: + if self._job is None or self._proc is None or self._mp_q is None: + return None + return self._job, self._proc, self._mp_q + + @staticmethod + def _read_queue_with_timeout(q: Any, *, timeout_sec: float) -> dict | None: + """Try read 1 event from mp queue. Timeout = pump stays responsive.""" + try: + return coerce_event(q.get(timeout = timeout_sec)) + except queue.Empty: + return None + except (EOFError, OSError, ValueError): + return None + + @staticmethod + def _drain_queue(q: Any) -> list[dict]: + """Drain mp queue fast (used on process exit).""" + events: list[dict] = [] + while True: + try: + events.append(coerce_event(q.get_nowait())) + except queue.Empty: + return events + except Exception: + # Return what we have so the run still finalizes rather than wedging "active". + logger.exception( + "Data-recipe job pump: queue drain failed; finalizing with drained events" + ) + return events + + def _safe_handle_event(self, job: Job, event: dict) -> None: + """Apply one event, swallowing any handler error so the pump can't die.""" + try: + self._handle_event(job, event) + except Exception: + etype = event.get("type") if isinstance(event, dict) else type(event).__name__ + logger.exception("Data-recipe job pump: failed to handle %s event; skipping", etype) + + def _pump_loop(self) -> None: + """Background thread: consume worker events and update the job snapshot. + + Guarded so no single event can end the loop; it is the sole writer of the + snapshot the UI polls, so its death would freeze status/SSE. + """ + while True: + snap = self._snapshot() + if snap is None: + return + job, proc, mp_q = snap + + try: + event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25) + except Exception: + # If a read keeps raising after the worker died, finalize instead + # of spinning forever; only retry while the worker is still alive. + logger.exception("Data-recipe job pump: queue read failed; continuing") + if proc.is_alive(): + time.sleep(0.1) + continue + event = None + + if event is not None: + self._safe_handle_event(job, event) + continue + + if proc.is_alive(): + continue + + # Worker exited: drain + finalize, guarded so an error can't strand the run "active". + try: + for e in self._drain_queue(mp_q): + self._safe_handle_event(job, e) + + retired_job: Job | None = None + with self._lock: + if self._job and self._job.status in { + "pending", + "active", + "cancelling", + }: + if self._job.status == "cancelling": + self._job.status = "cancelled" + else: + self._job.status = "error" + self._job.error = self._job.error or "process exited" + self._job.finished_at = time.time() + event_type = ( + EVENT_JOB_CANCELLED + if self._job.status == "cancelled" + else EVENT_JOB_ERROR + ) + self._emit( + { + "type": event_type, + "ts": time.time(), + "job_id": self._job.job_id, + } + ) + retired_job = self._job + if retired_job is not None: + self._retire_workflow_key(retired_job) + except Exception: + logger.exception("Data-recipe job pump: finalization after worker exit failed") + return + + def _handle_event(self, job: Job, event: dict) -> None: + """Apply event -> job state + forward to SSE.""" + 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 + if et == EVENT_JOB_STARTED: + self._job.status = "active" + if et == EVENT_JOB_COMPLETED: + self._job.status = "completed" + self._job.finished_at = time.time() + self._job.analysis = event.get("analysis") + self._job.artifact_path = event.get("artifact_path") + self._job.execution_type = event.get("execution_type") + self._job.dataset = event.get("dataset") + self._job.processor_artifacts = event.get("processor_artifacts") + 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: failures are swallowed. The key expires after 24h, so a + missed revoke is a latency, not correctness, concern. + """ + key_id = getattr(job, "internal_api_key_id", None) + if not key_id: + return + try: + from auth import storage # deferred: avoid circular import + storage.revoke_internal_api_key(int(key_id)) + except Exception: + pass + job.internal_api_key_id = None + + +_JOB_MANAGER: JobManager | None = None + + +def get_job_manager() -> JobManager: + """Singleton JobManager (we only run 1 job anyway).""" + global _JOB_MANAGER + if _JOB_MANAGER is None: + _JOB_MANAGER = JobManager() + return _JOB_MANAGER diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py new file mode 100644 index 0000000..3be830d --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -0,0 +1,468 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import re +import time +from dataclasses import dataclass +from typing import Any + +from .constants import ( + STAGE_BATCH, + STAGE_COLUMN_CONFIG, + STAGE_CREATE, + STAGE_DAG, + STAGE_GENERATING, + STAGE_HEALTHCHECK, + STAGE_PREVIEW, + STAGE_PROFILING, + STAGE_SAMPLING, + STAGE_SOURCE, + USAGE_RESET_STAGES, +) +from .types import Job, ModelUsage, Progress, SourceProgress + + +@dataclass(frozen = True) +class ParsedUpdate: + stage: str | None = None + current_column: str | None = None + progress: Progress | None = None + rows: int | None = None + cols: int | None = None + batch_idx: int | None = None + batch_total: int | None = None + usage_model: str | None = None + usage_input_tokens: int | None = None + usage_output_tokens: int | None = None + usage_total_tokens: int | None = None + usage_tps: float | None = None + usage_requests_success: int | None = None + usage_requests_failed: int | None = None + usage_requests_total: int | None = None + usage_rpm: float | None = None + usage_section_start: bool | None = None + source_progress: SourceProgress | None = None + + +# Best-effort parser from data-designer logs -> structured status for UI. +_RE_SAMPLERS = re.compile( + r"Preparing samplers to generate (?P\d+) records across (?P\d+) columns" +) +_RE_COLCFG = re.compile(r"model config for column '(?P[^']+)'") +_RE_PROCESSING_COL = re.compile(r"Processing .* column '(?P[^']+)'") +_RE_PROGRESS = re.compile( + r"progress: (?P\d+)/(?P\d+) \((?P\d+)%\) complete, " + r"(?P\d+) ok, (?P\d+) failed, (?P[0-9.]+) rec/s, eta (?P[0-9.]+)s" +) +_RE_BATCH = re.compile(r"Processing batch (?P\d+) of (?P\d+)") +_RE_USAGE_MODEL = re.compile(r"model:\s*(?P.+)$") +_RE_USAGE_TOKENS = re.compile( + r"tokens:\s*input=(?P\d+),\s*output=(?P\d+),\s*total=(?P\d+),\s*tps=(?P[0-9.]+)" +) +_RE_USAGE_REQUESTS = re.compile( + r"requests:\s*success=(?P\d+),\s*failed=(?P\d+),\s*total=(?P\d+),\s*rpm=(?P[0-9.]+)" +) +_RE_GITHUB_PAGE = re.compile( + r"^\[(?P[^\]\s]+/[^\]\s]+)\]\s+" + r"(?Pissues|PRs|commits)\s+page\s+(?P\d+)\s+" + r"\(\+(?P\d+)\).*?\bremaining=(?P\d+)", + re.IGNORECASE, +) +_RE_GITHUB_RATE_LIMIT = re.compile( + r"Rate limit hit\. Sleeping (?P\d+)s until reset\.", + re.IGNORECASE, +) +_RE_GITHUB_SECONDARY_RATE_LIMIT = re.compile( + r"Secondary rate limit(?: on REST)?\. Sleep (?P\d+)s\.", + re.IGNORECASE, +) +_RE_GITHUB_REST_RATE_LIMIT = re.compile( + r"REST 403/429, sleep (?P\d+)", + re.IGNORECASE, +) +_RE_GITHUB_TRANSIENT = re.compile( + r"^(?PGraphQL|REST) (?P\d{3}) transient, retrying", + re.IGNORECASE, +) +_RE_GITHUB_NETWORK_RETRY = re.compile( + r"^(?PGraphQL|REST) network error: .* Retry\.", + re.IGNORECASE, +) +_RE_GITHUB_TRIAL_LIMIT = re.compile( + r"Trial limit reached for (?Pissues|PRs|commits) \((?P\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( + stage = STAGE_SAMPLING, + rows = int(m.group("rows")), + cols = int(m.group("cols")), + ) + + if "Sorting column configs into a Directed Acyclic Graph" in msg: + return ParsedUpdate(stage = STAGE_DAG) + if "Running health checks for models" in msg: + return ParsedUpdate(stage = STAGE_HEALTHCHECK) + if "Preview generation in progress" in msg: + return ParsedUpdate(stage = STAGE_PREVIEW) + if "Creating Data Designer dataset" in msg: + return ParsedUpdate(stage = STAGE_CREATE) + if "Measuring dataset column statistics" in msg: + return ParsedUpdate(stage = STAGE_PROFILING) + + m = _RE_COLCFG.search(msg) + if m: + col = m.group("col") + return ParsedUpdate(stage = STAGE_COLUMN_CONFIG, current_column = col) + + m = _RE_PROCESSING_COL.search(msg) + if m: + col = m.group("col") + return ParsedUpdate(stage = STAGE_GENERATING, current_column = col) + + m = _RE_PROGRESS.search(msg) + if m: + p = Progress( + done = int(m.group("done")), + total = int(m.group("total")), + percent = float(m.group("pct")), + ok = int(m.group("ok")), + failed = int(m.group("failed")), + rate = float(m.group("rate")), + eta_sec = float(m.group("eta")), + ) + return ParsedUpdate(stage = STAGE_GENERATING, progress = p) + + m = _RE_BATCH.search(msg) + if m: + return ParsedUpdate( + stage = STAGE_BATCH, + batch_idx = int(m.group("idx")), + batch_total = int(m.group("total")), + ) + + if "Model usage summary" in msg: + return ParsedUpdate(usage_section_start = True) + + m = _RE_USAGE_MODEL.search(msg) + if m and "|-- model:" in msg: + return ParsedUpdate(usage_model = str(m.group("model")).strip()) + + m = _RE_USAGE_TOKENS.search(msg) + if m: + return ParsedUpdate( + usage_input_tokens = int(m.group("input")), + usage_output_tokens = int(m.group("output")), + usage_total_tokens = int(m.group("total")), + usage_tps = float(m.group("tps")), + ) + + m = _RE_USAGE_REQUESTS.search(msg) + if m: + return ParsedUpdate( + usage_requests_success = int(m.group("success")), + usage_requests_failed = int(m.group("failed")), + usage_requests_total = int(m.group("total")), + usage_rpm = float(m.group("rpm")), + ) + + return None + + +def apply_update(job: Job, update: ParsedUpdate) -> None: + if update.stage is not None: + job.stage = update.stage + if update.current_column is not None: + job.current_column = update.current_column + if ( + update.stage == STAGE_GENERATING + and update.current_column not in job._seen_generation_columns + ): + job._seen_generation_columns.append(update.current_column) + if update.rows is not None: + job.rows = update.rows + if update.cols is not None: + job.cols = update.cols + if update.progress is not None: + job.column_progress = update.progress + if ( + job.current_column + and update.progress.done is not None + and update.progress.total is not None + and update.progress.total > 0 + and update.progress.done >= update.progress.total + and job.current_column not in job.completed_columns + ): + job.completed_columns.append(job.current_column) + job.progress = _compute_overall_progress(job, update.progress) + if update.batch_idx is not 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; reset on the next stage. + job._in_usage_summary = False + + if update.usage_section_start is not None: + job._in_usage_summary = update.usage_section_start + if update.usage_section_start: + job._current_usage_model = None + + if not job._in_usage_summary: + return + + if update.usage_model is not None: + name = update.usage_model.strip().strip("'").strip('"') + job._current_usage_model = name + if name not in job.model_usage: + job.model_usage[name] = ModelUsage(model = name) + + if job._current_usage_model is None: + return + + usage = job.model_usage.get(job._current_usage_model) + if usage is None: + return + + if update.usage_input_tokens is not None: + usage.input_tokens = update.usage_input_tokens + if update.usage_output_tokens is not None: + usage.output_tokens = update.usage_output_tokens + if update.usage_total_tokens is not None: + usage.total_tokens = update.usage_total_tokens + if update.usage_tps is not None: + usage.tps = update.usage_tps + if update.usage_requests_success is not None: + usage.requests_success = update.usage_requests_success + if update.usage_requests_failed is not None: + usage.requests_failed = update.usage_requests_failed + if update.usage_requests_total is not None: + usage.requests_total = update.usage_requests_total + if update.usage_rpm is not 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 + + total_rows = max(1, int(job.rows)) + current_done = 0 if column_progress.done is None else int(column_progress.done) + current_done = max(0, min(current_done, total_rows)) + total_columns = max(1, int(job.progress_columns_total or 1)) + + if job.current_column: + job._column_done[job.current_column] = current_done + + if len(job._column_done) == 0: + done = current_done + else: + sum_done = sum(max(0, min(value, total_rows)) for value in job._column_done.values()) + done = int(sum_done / total_columns) + + prev_done = int(job.progress.done or 0) + if done < prev_done: + done = prev_done + if done > total_rows: + done = total_rows + percent = (done / total_rows) * 100 if total_rows > 0 else 100.0 + prev_percent = float(job.progress.percent or 0.0) + if percent < prev_percent: + percent = prev_percent + + return Progress( + done = done, + total = total_rows, + percent = percent, + eta_sec = column_progress.eta_sec, + rate = column_progress.rate, + ok = column_progress.ok, + failed = column_progress.failed, + ) + + +def coerce_event(obj: Any) -> dict: + """Normalize worker payload into event dict.""" + return obj if isinstance(obj, dict) else {"type": "log", "message": str(obj)} diff --git a/studio/backend/core/data_recipe/jobs/types.py b/studio/backend/core/data_recipe/jobs/types.py new file mode 100644 index 0000000..c30606b --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/types.py @@ -0,0 +1,101 @@ +# 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 dataclasses import dataclass, field +from typing import Any, Literal + + +JobStatus = Literal[ + "created", + "pending", + "active", + "cancelling", + "cancelled", + "error", + "completed", +] + + +@dataclass +class Progress: + done: int | None = None + total: int | None = None + percent: float | None = None + eta_sec: float | None = None + rate: float | None = None + ok: int | None = None + failed: int | None = None + + +@dataclass +class BatchProgress: + idx: int | None = None + 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 + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + tps: float | None = None + requests_success: int | None = None + requests_failed: int | None = None + requests_total: int | None = None + rpm: float | None = None + + +@dataclass +class Job: + job_id: str + status: JobStatus = "created" + stage: str | None = None + current_column: str | None = None + 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 + started_at: float | None = None + finished_at: float | None = None + + analysis: dict[str, Any] | None = None + artifact_path: str | None = None + execution_type: str | None = None + dataset: list[dict[str, Any]] | None = None + 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 ends so the key's window matches the run, not its 24h TTL. + internal_api_key_id: int | None = None + _current_usage_model: str | None = None + _in_usage_summary: bool = False + _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) diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py new file mode 100644 index 0000000..4073288 --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -0,0 +1,238 @@ +# 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 json +import structlog +import loggers +import logging +import re +import shutil +import time +import traceback +import unicodedata +from pathlib import Path +from typing import Any + +from ..jsonable import to_jsonable, to_preview_jsonable +from .constants import EVENT_JOB_COMPLETED, EVENT_JOB_ERROR, EVENT_JOB_STARTED +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=", message) + return _RE_SECRET_TOKEN.sub("", message) + + +class _QueueLogHandler(logging.Handler): + def __init__(self, event_queue): + super().__init__() + self._q = event_queue + + def emit(self, record: logging.LogRecord) -> None: + try: + event = { + "type": "log", + "ts": record.created, + "level": record.levelname, + "logger": record.name, + "message": _sanitize_log_message(record.getMessage()), + } + self._q.put(event) + except (OSError, RuntimeError, ValueError): + pass + + +def _slugify_run_name(value: str) -> str: + normalized = unicodedata.normalize("NFKD", value) + ascii_only = normalized.encode("ascii", "ignore").decode("ascii") + slug = re.sub(r"[^a-zA-Z0-9]+", "-", ascii_only).strip("-").lower() + if not slug: + return "" + return slug[:80].strip("-") + + +def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Path) -> str: + fallback = f"recipe_{job_id}" + slug = _slugify_run_name(run_name or "") + base_name = f"recipe_{slug}" if slug else fallback + candidate = base_name + suffix = 2 + while (artifact_root / candidate).exists(): + candidate = f"{base_name}_{suffix}" + suffix += 1 + return candidate + + +def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) -> None: + """Subprocess entrypoint. Sends events to `event_queue`.""" + import os + + os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports + + import warnings + from loggers.config import LogConfig + + if os.getenv("ENVIRONMENT_TYPE", "production") == "production": + warnings.filterwarnings("ignore") + + LogConfig.setup_logging( + service_name = "unsloth-studio-data-worker", + env = os.getenv("ENVIRONMENT_TYPE", "production"), + ) + + event_queue.put({"type": EVENT_JOB_STARTED, "ts": time.time()}) + + try: + from data_designer.config.run_config import RunConfig + + rows = int(run.get("rows") or 1000) + job_id = str(run.get("_job_id") or "").strip() + if not job_id: + job_id = f"{int(time.time())}" + run_name_raw = run.get("run_name") + run_name = run_name_raw if isinstance(run_name_raw, str) else None + dataset_name = _build_dataset_name( + run_name = run_name, + job_id = job_id, + artifact_root = _ARTIFACT_ROOT, + ) + merge_batches = bool(run.get("merge_batches")) + ensure_dir(_ARTIFACT_ROOT) + run_config_raw = run.get("run_config") or {} + + builder = build_config_builder(recipe) + designer = create_data_designer(recipe, artifact_path = str(_ARTIFACT_ROOT)) + + # DataDesigner resets root logging in __init__; attach the queue handler + # to the named loggers directly so parser events survive. + handler = _QueueLogHandler(event_queue) + handler.setLevel(logging.INFO) + for logger_name in ( + "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)) + + execution_type = str(run.get("execution_type") or "full").strip().lower() + if execution_type == "preview": + results = designer.preview(builder, num_records = rows) + analysis = ( + None + if results.analysis is None + else to_jsonable(results.analysis.model_dump(mode = "json")) + ) + dataset = ( + [] + if results.dataset is None + else to_preview_jsonable(results.dataset.to_dict(orient = "records")) + ) + processor_artifacts = ( + None + if results.processor_artifacts is None + else to_jsonable(results.processor_artifacts) + ) + event_queue.put( + { + "type": EVENT_JOB_COMPLETED, + "ts": time.time(), + "analysis": analysis, + "dataset": dataset, + "processor_artifacts": processor_artifacts, + "artifact_path": None, + "execution_type": execution_type, + } + ) + else: + results = designer.create(builder, num_records = rows, dataset_name = dataset_name) + analysis = to_jsonable(results.load_analysis().model_dump(mode = "json")) + if merge_batches: + _merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path) + artifact_path = str(results.artifact_storage.base_dataset_path) + event_queue.put( + { + "type": EVENT_JOB_COMPLETED, + "ts": time.time(), + "analysis": analysis, + "artifact_path": artifact_path, + "execution_type": execution_type, + } + ) + except Exception as exc: + event_queue.put( + { + "type": EVENT_JOB_ERROR, + "ts": time.time(), + "error": _sanitize_log_message(str(exc)), + "stack": _sanitize_log_message(traceback.format_exc(limit = 20)), + } + ) + + +def _merge_batches_to_single_parquet(base_dataset_path: Path) -> None: + parquet_dir = base_dataset_path / "parquet-files" + parquet_files = sorted(parquet_dir.glob("*.parquet")) + if len(parquet_files) <= 1: + return + + try: + from data_designer.config.utils.io_helpers import read_parquet_dataset + except ImportError: + return + + dataframe = read_parquet_dataset(parquet_dir) + shutil.rmtree(parquet_dir) + parquet_dir.mkdir(parents = True, exist_ok = True) + merged_file = parquet_dir / "batch_00000.parquet" + dataframe.to_parquet(merged_file, index = False) + _rewrite_merged_metadata( + base_dataset_path = base_dataset_path, + parquet_file = merged_file, + ) + + +def _rewrite_merged_metadata(*, base_dataset_path: Path, parquet_file: Path) -> None: + metadata_path = base_dataset_path / "metadata.json" + if not metadata_path.exists(): + return + + try: + metadata = json.loads(metadata_path.read_text(encoding = "utf-8")) + except (OSError, TypeError, ValueError): + return + + if not isinstance(metadata, dict): + return + + relative_parquet_path = str(parquet_file.relative_to(base_dataset_path)) + file_paths = metadata.get("file_paths") + if not isinstance(file_paths, dict): + file_paths = {} + file_paths["parquet-files"] = [relative_parquet_path] + metadata["file_paths"] = file_paths + metadata["total_num_batches"] = 1 + metadata["num_completed_batches"] = 1 + + try: + metadata_path.write_text( + json.dumps(metadata, indent = 2, sort_keys = True), + encoding = "utf-8", + ) + except OSError: + return diff --git a/studio/backend/core/data_recipe/jsonable.py b/studio/backend/core/data_recipe/jsonable.py new file mode 100644 index 0000000..5a4fcd3 --- /dev/null +++ b/studio/backend/core/data_recipe/jsonable.py @@ -0,0 +1,119 @@ +# 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 base64 +import io +from pathlib import Path +from typing import Any + + +def _pil_to_preview_payload(image: Any) -> dict[str, Any]: + buffer = io.BytesIO() + image.convert("RGB").save(buffer, format = "JPEG", quality = 85) + return { + "type": "image", + "mime": "image/jpeg", + "width": image.width, + "height": image.height, + "data": base64.b64encode(buffer.getvalue()).decode("ascii"), + } + + +def _open_pil_image_from_bytes(raw_bytes: bytes): + from PIL import Image # type: ignore + with Image.open(io.BytesIO(raw_bytes)) as image: + return image.copy() + + +def _to_pil_from_hf_image_dict(value: Any) -> Any | None: + if not isinstance(value, dict): + return None + + raw_bytes = value.get("bytes") + if isinstance(raw_bytes, (bytes, bytearray)) and len(raw_bytes) > 0: + try: + return _open_pil_image_from_bytes(bytes(raw_bytes)) + except (OSError, ValueError): + pass + if ( + isinstance(raw_bytes, list) + and len(raw_bytes) > 0 + and all(isinstance(item, int) and 0 <= item <= 255 for item in raw_bytes) + ): + try: + return _open_pil_image_from_bytes(bytes(raw_bytes)) + except (OSError, ValueError): + pass + + path_value = value.get("path") + if isinstance(path_value, str) and path_value.strip(): + try: + from PIL import Image # type: ignore + with Image.open(Path(path_value)) as image: + return image.copy() + except (OSError, ValueError, TypeError): + return None + + return None + + +def to_jsonable(value: Any) -> Any: + """Convert numpy/pandas-ish values into plain JSON-safe values.""" + try: + import numpy as np # type: ignore + except ImportError: # pragma: no cover + np = None # type: ignore + + if np is not None: + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + + if isinstance(value, dict): + return {str(k): to_jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [to_jsonable(v) for v in value] + + if hasattr(value, "isoformat") and callable(value.isoformat): + try: + return value.isoformat() + except (TypeError, ValueError): + return value + + return value + + +def _to_preview_image_payload(value: Any) -> dict[str, Any] | None: + try: + from PIL.Image import Image as PILImage # type: ignore + except ImportError: # pragma: no cover + return None + + if not isinstance(value, PILImage): + hf_image = _to_pil_from_hf_image_dict(value) + if hf_image is None: + return None + value = hf_image + + return _pil_to_preview_payload(value) + + +def to_preview_jsonable(value: Any) -> Any: + """Convert values into JSON-safe preview values, including PIL images.""" + image_payload = _to_preview_image_payload(value) + if image_payload is not None: + return image_payload + + converted = to_jsonable(value) + if converted is None or isinstance(converted, (str, int, float, bool)): + return converted + if isinstance(converted, dict): + return {str(k): to_preview_jsonable(v) for k, v in converted.items()} + if isinstance(converted, (list, tuple, set)): + return [to_preview_jsonable(v) for v in converted] + if isinstance(converted, (bytes, bytearray)): + return base64.b64encode(bytes(converted)).decode("ascii") + return str(converted) diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py new file mode 100644 index 0000000..ebb1d39 --- /dev/null +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -0,0 +1,337 @@ +# 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 json +import os +import structlog +import subprocess +from copy import deepcopy +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Any + +from loggers import get_logger +from utils.node_runtime import resolve_node_executable +from utils.paths import ensure_dir, oxc_validator_tmp_root + +logger = get_logger(__name__) + +OXC_VALIDATION_FN_MARKER = "unsloth_oxc_validator" + +_OXC_LANG_TO_NODE_LANG = { + "javascript": "js", + "typescript": "ts", + "jsx": "jsx", + "tsx": "tsx", +} +_OXC_VALIDATION_MODES = {"syntax", "lint", "syntax+lint"} +_OXC_CODE_SHAPES = {"auto", "module", "snippet"} + +_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 + drop: bool + target_columns: list[str] + batch_size: int + code_lang: str + validation_mode: str + code_shape: str + + +def split_oxc_local_callable_validators( + recipe_core: dict[str, Any], +) -> tuple[dict[str, Any], list[OxcLocalCallableValidatorSpec]]: + columns = recipe_core.get("columns") + if not isinstance(columns, list): + return recipe_core, [] + + sanitized = deepcopy(recipe_core) + sanitized_columns = sanitized.get("columns") + if not isinstance(sanitized_columns, list): + return sanitized, [] + + kept_columns: list[Any] = [] + oxc_specs: list[OxcLocalCallableValidatorSpec] = [] + + for column in sanitized_columns: + if not isinstance(column, dict): + kept_columns.append(column) + continue + + maybe_spec = _parse_oxc_spec(column = column) + if maybe_spec is None: + kept_columns.append(column) + continue + oxc_specs.append(maybe_spec) + + sanitized["columns"] = kept_columns + return sanitized, oxc_specs + + +def register_oxc_local_callable_validators( + *, builder, specs: list[OxcLocalCallableValidatorSpec] +) -> None: + if not specs: + return + + from data_designer.config.column_configs import ValidationColumnConfig + from data_designer.config.validator_params import ( + LocalCallableValidatorParams, + ValidatorType, + ) + + for spec in specs: + validation_function = _build_oxc_validation_function( + spec.code_lang, + spec.validation_mode, + spec.code_shape, + ) + builder.add_column( + ValidationColumnConfig( + name = spec.name, + drop = spec.drop, + target_columns = spec.target_columns, + validator_type = ValidatorType.LOCAL_CALLABLE, + validator_params = LocalCallableValidatorParams( + validation_function = validation_function, + ), + batch_size = spec.batch_size, + ) + ) + + +def _parse_oxc_spec(*, column: dict[str, Any]) -> OxcLocalCallableValidatorSpec | None: + if str(column.get("column_type") or "").strip() != "validation": + return None + if str(column.get("validator_type") or "").strip() != "local_callable": + return None + + params = column.get("validator_params") + if not isinstance(params, dict): + return None + + fn_raw = params.get("validation_function") + fn_name = fn_raw.strip() if isinstance(fn_raw, str) else "" + if not fn_name.startswith(OXC_VALIDATION_FN_MARKER): + return None + + name = str(column.get("name") or "").strip() + if not name: + return None + + target_columns_raw = column.get("target_columns") + target_columns = ( + [value.strip() for value in target_columns_raw if isinstance(value, str) and value.strip()] + if isinstance(target_columns_raw, list) + else [] + ) + if not target_columns: + return None + + code_lang, validation_mode, code_shape = _parse_oxc_validation_marker(fn_name) + batch_size = _parse_batch_size(column.get("batch_size")) + drop = bool(column.get("drop") is True) + + return OxcLocalCallableValidatorSpec( + name = name, + drop = drop, + target_columns = target_columns, + batch_size = batch_size, + code_lang = code_lang, + validation_mode = validation_mode, + code_shape = code_shape, + ) + + +def _parse_batch_size(value: Any) -> int: + try: + parsed = int(value) + except (TypeError, ValueError): + return 10 + return parsed if parsed >= 1 else 10 + + +def _parse_oxc_validation_marker(fn_name: str) -> tuple[str, str, str]: + marker = f"{OXC_VALIDATION_FN_MARKER}:" + if not fn_name.startswith(marker): + return "javascript", "syntax", "auto" + suffix = fn_name[len(marker) :] + parts = [part.strip() for part in suffix.split(":") if part.strip()] + if len(parts) < 2: + return "javascript", "syntax", "auto" + code_lang = parts[0] if parts[0] in _OXC_LANG_TO_NODE_LANG else "javascript" + mode = parts[1] if parts[1] in _OXC_VALIDATION_MODES else "syntax" + code_shape = parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto" + return code_lang, mode, code_shape + + +@lru_cache(maxsize = 8) +def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape: str): + node_lang = _OXC_LANG_TO_NODE_LANG.get(lang, "js") + mode = validation_mode if validation_mode in _OXC_VALIDATION_MODES else "syntax" + normalized_code_shape = code_shape if code_shape in _OXC_CODE_SHAPES else "auto" + + def _validator(df): + import pandas as pd # lazy import for local callable runtime + + row_count = int(len(df.index)) + if row_count == 0: + return pd.DataFrame({"is_valid": []}) + + code_column = str(df.columns[0]) if len(df.columns) > 0 else "" + code_values = ( + ["" for _ in range(row_count)] + if not code_column + else ["" if value is None else str(value) for value in df[code_column].tolist()] + ) + + results = _run_oxc_batch( + node_lang = node_lang, + validation_mode = mode, + code_shape = normalized_code_shape, + code_values = code_values, + ) + if len(results) != row_count: + results = _fallback_results( + row_count, + "OXC validator returned mismatched result size.", + ) + return pd.DataFrame(results) + + _validator.__name__ = ( + f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}" + ) + return _validator + + +def _run_oxc_batch( + *, node_lang: str, validation_mode: str, code_shape: str, code_values: list[str] +) -> list[dict[str, Any]]: + if not _OXC_RUNNER_PATH.exists(): + return _fallback_results( + len(code_values), + f"OXC runner missing at {_OXC_RUNNER_PATH}", + ) + + payload = { + "lang": node_lang, + "mode": validation_mode, + "code_shape": code_shape, + "codes": code_values, + } + # Resolve a usable Node (system or the isolated install, which is not on the + # user's PATH); a bare "node" would fail for isolated-Node users. + node_executable = resolve_node_executable() + if not node_executable: + return _fallback_results( + len(code_values), + "Node.js not found (install Node >= 20.19, or re-run Studio setup to provision it).", + ) + try: + tmp_dir = ensure_dir(oxc_validator_tmp_root()) + env = child_env_without_native_path_secret() + tmp_dir_str = str(tmp_dir) + env["TMPDIR"] = tmp_dir_str + env["TMP"] = tmp_dir_str + env["TEMP"] = tmp_dir_str + # Resolved node's dir first on the child PATH so it finds its own npm/npx. + node_bin_dir = os.path.dirname(node_executable) + if node_bin_dir: + env["PATH"] = node_bin_dir + os.pathsep + env.get("PATH", "") + env.pop("NODE_PATH", None) + proc = subprocess.run( + [node_executable, str(_OXC_RUNNER_PATH)], + cwd = str(_OXC_TOOL_DIR), + input = json.dumps(payload), + text = True, + capture_output = True, + check = False, + env = env, + **_windows_hidden_subprocess_kwargs(), + ) + except (OSError, ValueError) as exc: + logger.warning("OXC subprocess launch failed: %s", exc) + return _fallback_results(len(code_values), f"OXC launch failed: {exc}") + + if proc.returncode != 0: + message = (proc.stderr or proc.stdout or "unknown error").strip() + if len(message) > 300: + message = f"{message[:300]}..." + return _fallback_results(len(code_values), f"OXC failed: {message}") + + try: + raw = json.loads(proc.stdout) + except json.JSONDecodeError: + return _fallback_results(len(code_values), "OXC output parse failed.") + + if not isinstance(raw, list): + return _fallback_results(len(code_values), "OXC output must be an array.") + + out: list[dict[str, Any]] = [] + for item in raw: + if not isinstance(item, dict): + out.append( + { + "is_valid": False, + "error_count": 1, + "error_message": "Invalid OXC result entry.", + "severity": None, + "code": None, + "labels": [], + "codeframe": None, + "warning_count": 0, + } + ) + continue + is_valid_raw = item.get("is_valid") + error_count_raw = item.get("error_count") + message_raw = item.get("error_message") + severity_raw = item.get("severity") + code_raw = item.get("code") + labels_raw = item.get("labels") + codeframe_raw = item.get("codeframe") + warning_count_raw = item.get("warning_count") + out.append( + { + "is_valid": bool(is_valid_raw) if isinstance(is_valid_raw, bool) else False, + "error_count": int(error_count_raw) if isinstance(error_count_raw, int) else 0, + "error_message": str(message_raw or ""), + "severity": str(severity_raw) if isinstance(severity_raw, str) else None, + "code": str(code_raw) if isinstance(code_raw, str) else None, + "labels": labels_raw if isinstance(labels_raw, list) else [], + "codeframe": str(codeframe_raw) if isinstance(codeframe_raw, str) else None, + "warning_count": int(warning_count_raw) + if isinstance(warning_count_raw, int) + else 0, + } + ) + return out + + +def _fallback_results(row_count: int, message: str) -> list[dict[str, Any]]: + return [ + { + "is_valid": False, + "error_count": 1, + "error_message": message, + "severity": None, + "code": None, + "labels": [], + "codeframe": None, + "warning_count": 0, + } + for _ in range(row_count) + ] diff --git a/studio/backend/core/data_recipe/oxc-validator/package-lock.json b/studio/backend/core/data_recipe/oxc-validator/package-lock.json new file mode 100644 index 0000000..1e630cf --- /dev/null +++ b/studio/backend/core/data_recipe/oxc-validator/package-lock.json @@ -0,0 +1,798 @@ +{ + "name": "unsloth-oxc-validator-runtime", + "version": "0.0.1", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "unsloth-oxc-validator-runtime", + "version": "0.0.1", + "dependencies": { + "oxc-parser": "^0.131.0", + "oxlint": "^1.65.0" + } + }, + "node_modules/@emnapi/core": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.10.0.tgz", + "integrity": "sha512-yq6OkJ4p82CAfPl0u9mQebQHKPJkY7WrIuk205cTYnYe+k2Z8YBh11FrbRG/H6ihirqcacOgl2BIO8oyMQLeXw==", + "license": "MIT", + "optional": true, + "dependencies": { + "@emnapi/wasi-threads": "1.2.1", + "tslib": "^2.4.0" + } + }, + "node_modules/@emnapi/runtime": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.10.0.tgz", + "integrity": "sha512-ewvYlk86xUoGI0zQRNq/mC+16R1QeDlKQy21Ki3oSYXNgLb45GV1P6A0M+/s6nyCuNDqe5VpaY84BzXGwVbwFA==", + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@emnapi/wasi-threads": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz", + "integrity": "sha512-uTII7OYF+/Mes/MrcIOYp5yOtSMLBWSIoLPpcgwipoiKbli6k322tcoFsxoIIxPDqW01SQGAgko4EzZi2BNv2w==", + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@napi-rs/wasm-runtime": { + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.4.tgz", + "integrity": "sha512-3NQNNgA1YSlJb/kMH1ildASP9HW7/7kYnRI2szWJaofaS1hWmbGI4H+d3+22aGzXXN9IJ+n+GiFVcGipJP18ow==", + "license": "MIT", + "optional": true, + "dependencies": { + "@tybys/wasm-util": "^0.10.1" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + }, + "peerDependencies": { + "@emnapi/core": "^1.7.1", + "@emnapi/runtime": "^1.7.1" + } + }, + "node_modules/@oxc-parser/binding-android-arm-eabi": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-android-arm-eabi/-/binding-android-arm-eabi-0.131.0.tgz", + "integrity": "sha512-t2xicr9pfzkSRYx5aPqZqlLaayIwJTqgQ81Jor31Xep2nGyL2Aq3d0K5wOfeR7VevaSdxaS9dzSQP9xDwn8fDg==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-android-arm64": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-android-arm64/-/binding-android-arm64-0.131.0.tgz", + "integrity": "sha512-nlGIod6gw75x1aEDgLS+srj+JRGY0HHm9MI9YgzE/B64l6d6+H3MSP9NOgp0+HTg8tp4vV9rVfgQGgd+TfVZcA==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-darwin-arm64": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-darwin-arm64/-/binding-darwin-arm64-0.131.0.tgz", + "integrity": "sha512-jukuV6xe5RbQKFo7QD34NDCLDZp4PSOm8rmckhNdH/60ymG5zXbDzGBEyc+nTkuLQNama2aSGCt+CPfpjNTqyw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-darwin-x64": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-darwin-x64/-/binding-darwin-x64-0.131.0.tgz", + "integrity": "sha512-g3JOo4khe9rslHm5WYaVDWb0HS/M1MLR3I9S8560MkKIcC96VQY00QjOlsuRyfSj/JDXj8i9T7ryPO2RidiXVg==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-freebsd-x64": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-freebsd-x64/-/binding-freebsd-x64-0.131.0.tgz", + "integrity": "sha512-1hziITDTxjMePnX+dR9ocVT+EuZkQ8wm4FPAbmbEiKG+Phbo73J1ZnPAA6Y/aGsWF3McOFnQuZIktAFwalkfJQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm-gnueabihf": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-0.131.0.tgz", + "integrity": "sha512-9uRxfXwyKG9+MwmGQBo2ncPNwZH5HTmCETFM2WiuDBNDCW4NC5ttSQkwCAMrTAWgwMzVBH1CP8pM0v7nebCWXQ==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm-musleabihf": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-0.131.0.tgz", + "integrity": "sha512-mgbLvzRShXOLBdWGInf08Af4q+pfj1xD8hSgLClDZ9of/BXkB6+LIhTH7fihiDUipqB3yoSkKBWaZ3Ejlf5Yag==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm64-gnu": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-0.131.0.tgz", + "integrity": "sha512-OPT8++4aN6j2GJ8+3IZHS/byXoZP4aSBn+FoG6rgBJ2fKwPKXWF3MqrFMNW7NKHM28FLY579xYLxJSfgobEqPA==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-arm64-musl": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-arm64-musl/-/binding-linux-arm64-musl-0.131.0.tgz", + "integrity": "sha512-vtPiwmfVTAXzaxDKsOXG+LwgRAA7WEnaeHzhS5z0GE89gAK18KSXnly7Z6saXXq6L3dVMyK44uoTI03zKxrpmw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-ppc64-gnu": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-0.131.0.tgz", + "integrity": "sha512-8AW8L7w5cGHSdZPcyZX2yR0+GUODsT15rbRjfdD54rv6DMbtuEB19ysLOpKJlRGfH6UNYNpCHaU1uJWgTWf1/w==", + "cpu": [ + "ppc64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-riscv64-gnu": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-0.131.0.tgz", + "integrity": "sha512-vvpjkjEOUsPcsYf8evE4MO3aGx9+3wodXEBOicGNnOwTuAik8eBONNkgSdhkGsAblQmfVHJyanRnpxglddTXIA==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-riscv64-musl": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-0.131.0.tgz", + "integrity": "sha512-AqmcNC3fClXX+fxQ6VGEN1667xVFiRBkY0CZmDMSiaeFUsv1+UkBPYYi48IUKcA9/ivvoKNRzQl2I4//kT9F/w==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-s390x-gnu": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-0.131.0.tgz", + "integrity": "sha512-7d3jOMKy7RSQCcDLIci+ySll2FgsOMl/GiRux4q2JNv0zg4EdhFISa9idvrdN/HEUIQQJNg6dmveUeJl2YErGA==", + "cpu": [ + "s390x" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-x64-gnu": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-x64-gnu/-/binding-linux-x64-gnu-0.131.0.tgz", + "integrity": "sha512-JHK/h95qVqVQ+ITER837kcTdwBDFpFaNnOTYGCP0zdUSX/mLKC7tXOoyrTb6vG7iRPwGlcgBil3v2IjYw1FqJA==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-linux-x64-musl": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-linux-x64-musl/-/binding-linux-x64-musl-0.131.0.tgz", + "integrity": "sha512-b2BO82O8azXAyf7EUgOPKu145nWypbNyk07HbU09fkzhm9lEA5oPvaN/M8Nlo7tOErVTa2WOgS4QbOnxAPXdDQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-openharmony-arm64": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-openharmony-arm64/-/binding-openharmony-arm64-0.131.0.tgz", + "integrity": "sha512-GHO9glZaX7LkX/OGfluEPf1yjg+ehiFbUdowbX6uNWOQhmwKWU4m4+nZ9FJkrHNKuxyI1KKertMdGjVKCApKWA==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "openharmony" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-wasm32-wasi": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-wasm32-wasi/-/binding-wasm32-wasi-0.131.0.tgz", + "integrity": "sha512-3SkikPaEFoih1N83qLVEDLRLeY4nYsf6JT9SnWiMCQ5lGQdKup6bEuKCqkRiG9dD1IIaFeYz9RjlciPmYoFIWA==", + "cpu": [ + "wasm32" + ], + "license": "MIT", + "optional": true, + "dependencies": { + "@emnapi/core": "1.10.0", + "@emnapi/runtime": "1.10.0", + "@napi-rs/wasm-runtime": "^1.1.4" + }, + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-win32-arm64-msvc": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-0.131.0.tgz", + "integrity": "sha512-Os5bEhryeA2jkH+ZrnZyAC1EP5gs+X4YB1Fjqml7UPD5kU7ecsK1MPEVMfCrdt/GDNpDbavYXiOXOdyJ5b3OPw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-win32-ia32-msvc": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-ia32-msvc/-/binding-win32-ia32-msvc-0.131.0.tgz", + "integrity": "sha512-m+jNz9EuF0NXoiptc6B9h5yompZQVW/a5MJeOu5zojfH5yWk82tvF2ccrHkfhgtrS9h9DD5l1Qv8dWlfY7Nz8g==", + "cpu": [ + "ia32" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-parser/binding-win32-x64-msvc": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-parser/binding-win32-x64-msvc/-/binding-win32-x64-msvc-0.131.0.tgz", + "integrity": "sha512-o14Hk8dAyiEUMFEWEgmAwFZvBt1RzAYLM3xeQ+5315JXgVYhoemivgYcbYVRbsFkS71ShMGlAFE0kPnr460rww==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxc-project/types": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.131.0.tgz", + "integrity": "sha512-PgnWDfV0h+b16XNKbXU7Daib/BFSt/J2mEzfYIBu6JB/wNdlU+kVYXCkGA1A9fWkTbOgbjh4e6NhPeQOYvFhEA==", + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/Boshen" + } + }, + "node_modules/@oxlint/binding-android-arm-eabi": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-android-arm-eabi/-/binding-android-arm-eabi-1.65.0.tgz", + "integrity": "sha512-jDVaGNURT5pEA9qcabh6WusIoBNybOMMDPCx+EFt+gxo6rVvoUf0+73Xy5x81+ZrxU+ewk5uRBYifjy5pgkcnA==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-android-arm64": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-android-arm64/-/binding-android-arm64-1.65.0.tgz", + "integrity": "sha512-v0z80IWNA7c9RhUydq9YprBxCVZrQ6Ixls2tdxUC1F/1FFqSfa7xTX+EJf0mj6+BKRg2zWXqWfcbJUnETlLlIw==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "android" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-darwin-arm64": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-darwin-arm64/-/binding-darwin-arm64-1.65.0.tgz", + "integrity": "sha512-pL/mG/5gMzBwp1gdc5+Cwi87F9j3XRnPxHGyVj5Zd+dCEV5YkKt0L70PB3EGmEEHxgn4H+jnMS3xLuXs6mZW/Q==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-darwin-x64": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-darwin-x64/-/binding-darwin-x64-1.65.0.tgz", + "integrity": "sha512-jVTneaeuHtqTrKYnhrdH1buhnSorinvpy1sv43ayclfWx/e/DfdRWv+h1fopJcHQbYr5WMcZMmDvnfEBkPZ+1A==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-freebsd-x64": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-freebsd-x64/-/binding-freebsd-x64-1.65.0.tgz", + "integrity": "sha512-8lJQ7B6RloYDUhwVdbSpwT2eKsCN5KP1Scn18ly1tytCuhXhbs0nkfKHT4jWWZBJqmynWuzd+78bF7wILrj6pw==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "freebsd" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm-gnueabihf": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.65.0.tgz", + "integrity": "sha512-EgmZY+DeWhLLEnNl70/49j3ltA8I6X9kxMfexupWi2Vwfp6RonGsBaHtGoedLolaU37ne7eDUgoxa3CFB95GZA==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm-musleabihf": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm-musleabihf/-/binding-linux-arm-musleabihf-1.65.0.tgz", + "integrity": "sha512-OJMWmAYRVBCPPxnYr3j5sXRwHPh1bAuMlTStGco1Z8q3HkvSH4h+A10E9MiRNYmLhUuli5a2P5wmfj8cagiF5Q==", + "cpu": [ + "arm" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm64-gnu": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.65.0.tgz", + "integrity": "sha512-D8uNi50LsYKgS0vGARZDRx05TBZeSxAVdLGddSEqQLSU7xsiqdImHPEw55xq8sKA5rCc/4au/5uS7FQALWdLCg==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-arm64-musl": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.65.0.tgz", + "integrity": "sha512-IpbA8QGbwFehQhO+YaHwmoI81f93xvywpspf8HrdPCWOIeKwYfM1dhVhO4YKfZewTRRQEPY/JFjTOXTgkwhKrA==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-ppc64-gnu": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.65.0.tgz", + "integrity": "sha512-ZSe8HgaZdgyHSv2+/pTG68z10+OarB18CkFKQOhRs3lmmP/p2vuigedK2e9d0ztoG2DU/duJzhxXBSjy/492HQ==", + "cpu": [ + "ppc64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-riscv64-gnu": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-riscv64-gnu/-/binding-linux-riscv64-gnu-1.65.0.tgz", + "integrity": "sha512-DcTERf++v6HyPHukKAr0JFTRqB+YeDEvqzRgNDMaz7jITPf+tlJIwRxodlAqoXMYhNVEZhXdQM5RAAYH8/oPuw==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-riscv64-musl": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-riscv64-musl/-/binding-linux-riscv64-musl-1.65.0.tgz", + "integrity": "sha512-xjhMwuFJwRh40NOBzol4gM5gqAa0xPCJU+GQLM6BydV8TbfkIA7JeyCFNhyfbE9Q/5EWcKYTx62R0cRcjP7DAA==", + "cpu": [ + "riscv64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-s390x-gnu": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.65.0.tgz", + "integrity": "sha512-lrWSXb8JzboPWYBG6Kunt/eemvjo2oCFXktShsm3yMToY7HjzKLjxh7CljSvGnnZH9oohNFHOKc9xYpGKCPm6w==", + "cpu": [ + "s390x" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-x64-gnu": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.65.0.tgz", + "integrity": "sha512-A7xfghw250m4a1sPV+q44Mow2G5bhiC9FBvhAuIhJS6QovWnqzuL5AFQPEuwOB+PM4DhABkqxVa3Iwe3Y/nFlQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-linux-x64-musl": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-linux-x64-musl/-/binding-linux-x64-musl-1.65.0.tgz", + "integrity": "sha512-reqOun1+pWO3fW6cv7bsa8hHG0TN3t/82qPdaoJo90FwugXiMjKhZMChmH5Z01cFNRHmxN4+543Fy8478cM/iA==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "linux" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-openharmony-arm64": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-openharmony-arm64/-/binding-openharmony-arm64-1.65.0.tgz", + "integrity": "sha512-KQpqOb/juDBO0xyloDkVDhOVxDUgAfZ2OAAVq99TJScJDzT319xry1QzB9LQohV9QGnA7p6m/XATZkMXc84lwA==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "openharmony" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-win32-arm64-msvc": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.65.0.tgz", + "integrity": "sha512-xfqcOc3nJFeAd1kDY4T9d3XeJIhr00twaaW0kOAzGPyUHkruXtNJv6zz1Ra9fRtSek5VpW2Yoj5AcwPIlT0ZiQ==", + "cpu": [ + "arm64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-win32-ia32-msvc": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-win32-ia32-msvc/-/binding-win32-ia32-msvc-1.65.0.tgz", + "integrity": "sha512-JV+pXm45p8sdgs3c7LOPAohW23optCNZETFOXUcjn6cS4PYZhEU/RI54Z5dHdMudab3nw7T48PZILthM+Q0COQ==", + "cpu": [ + "ia32" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@oxlint/binding-win32-x64-msvc": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/@oxlint/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.65.0.tgz", + "integrity": "sha512-D7L/oBbskLss21bYrRbFuIs81AiSQV+wRzwck54dOkHIlq2qu1xjLz8u6jCqGH8Fltk8bB5DLBpVhE7v/fA8XQ==", + "cpu": [ + "x64" + ], + "license": "MIT", + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": "^20.19.0 || >=22.12.0" + } + }, + "node_modules/@tybys/wasm-util": { + "version": "0.10.2", + "resolved": "https://registry.npmjs.org/@tybys/wasm-util/-/wasm-util-0.10.2.tgz", + "integrity": "sha512-RoBvJ2X0wuKlWFIjrwffGw1IqZHKQqzIchKaadZZfnNpsAYp2mM0h36JtPCjNDAHGgYez/15uMBpfGwchhiMgg==", + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/oxc-parser": { + "version": "0.131.0", + "resolved": "https://registry.npmjs.org/oxc-parser/-/oxc-parser-0.131.0.tgz", + "integrity": "sha512-SJ3/7ZPbgie8dr5Z9BI/M51zZbpXba+hRSG0MDzVwMW5CRQg2fjYE0jHGlLX4eeiibGgC/mzoDFKSDHwVZEHRQ==", + "license": "MIT", + "dependencies": { + "@oxc-project/types": "^0.131.0" + }, + "engines": { + "node": "^20.19.0 || >=22.12.0" + }, + "funding": { + "url": "https://github.com/sponsors/Boshen" + }, + "optionalDependencies": { + "@oxc-parser/binding-android-arm-eabi": "0.131.0", + "@oxc-parser/binding-android-arm64": "0.131.0", + "@oxc-parser/binding-darwin-arm64": "0.131.0", + "@oxc-parser/binding-darwin-x64": "0.131.0", + "@oxc-parser/binding-freebsd-x64": "0.131.0", + "@oxc-parser/binding-linux-arm-gnueabihf": "0.131.0", + "@oxc-parser/binding-linux-arm-musleabihf": "0.131.0", + "@oxc-parser/binding-linux-arm64-gnu": "0.131.0", + "@oxc-parser/binding-linux-arm64-musl": "0.131.0", + "@oxc-parser/binding-linux-ppc64-gnu": "0.131.0", + "@oxc-parser/binding-linux-riscv64-gnu": "0.131.0", + "@oxc-parser/binding-linux-riscv64-musl": "0.131.0", + "@oxc-parser/binding-linux-s390x-gnu": "0.131.0", + "@oxc-parser/binding-linux-x64-gnu": "0.131.0", + "@oxc-parser/binding-linux-x64-musl": "0.131.0", + "@oxc-parser/binding-openharmony-arm64": "0.131.0", + "@oxc-parser/binding-wasm32-wasi": "0.131.0", + "@oxc-parser/binding-win32-arm64-msvc": "0.131.0", + "@oxc-parser/binding-win32-ia32-msvc": "0.131.0", + "@oxc-parser/binding-win32-x64-msvc": "0.131.0" + } + }, + "node_modules/oxlint": { + "version": "1.65.0", + "resolved": "https://registry.npmjs.org/oxlint/-/oxlint-1.65.0.tgz", + "integrity": "sha512-ChUuE3Q7XnAbscvT4XLMsH7HFJmLgLVv9lu+RRgFL5wSXnDqUOzTp5IS8qWDBGd/ZDSzQ2tbX8fjAmijlGLC7A==", + "license": "MIT", + "bin": { + "oxlint": "bin/oxlint" + }, + "engines": { + "node": "^20.19.0 || >=22.12.0" + }, + "funding": { + "url": "https://github.com/sponsors/Boshen" + }, + "optionalDependencies": { + "@oxlint/binding-android-arm-eabi": "1.65.0", + "@oxlint/binding-android-arm64": "1.65.0", + "@oxlint/binding-darwin-arm64": "1.65.0", + "@oxlint/binding-darwin-x64": "1.65.0", + "@oxlint/binding-freebsd-x64": "1.65.0", + "@oxlint/binding-linux-arm-gnueabihf": "1.65.0", + "@oxlint/binding-linux-arm-musleabihf": "1.65.0", + "@oxlint/binding-linux-arm64-gnu": "1.65.0", + "@oxlint/binding-linux-arm64-musl": "1.65.0", + "@oxlint/binding-linux-ppc64-gnu": "1.65.0", + "@oxlint/binding-linux-riscv64-gnu": "1.65.0", + "@oxlint/binding-linux-riscv64-musl": "1.65.0", + "@oxlint/binding-linux-s390x-gnu": "1.65.0", + "@oxlint/binding-linux-x64-gnu": "1.65.0", + "@oxlint/binding-linux-x64-musl": "1.65.0", + "@oxlint/binding-openharmony-arm64": "1.65.0", + "@oxlint/binding-win32-arm64-msvc": "1.65.0", + "@oxlint/binding-win32-ia32-msvc": "1.65.0", + "@oxlint/binding-win32-x64-msvc": "1.65.0" + }, + "peerDependencies": { + "oxlint-tsgolint": ">=0.22.1" + }, + "peerDependenciesMeta": { + "oxlint-tsgolint": { + "optional": true + } + } + }, + "node_modules/tslib": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", + "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", + "license": "0BSD", + "optional": true + } + } +} diff --git a/studio/backend/core/data_recipe/oxc-validator/package.json b/studio/backend/core/data_recipe/oxc-validator/package.json new file mode 100644 index 0000000..2817b6c --- /dev/null +++ b/studio/backend/core/data_recipe/oxc-validator/package.json @@ -0,0 +1,10 @@ +{ + "name": "unsloth-oxc-validator-runtime", + "private": true, + "version": "0.0.1", + "type": "module", + "dependencies": { + "oxc-parser": "^0.131.0", + "oxlint": "^1.65.0" + } +} diff --git a/studio/backend/core/data_recipe/oxc-validator/validate.mjs b/studio/backend/core/data_recipe/oxc-validator/validate.mjs new file mode 100644 index 0000000..ad61fb5 --- /dev/null +++ b/studio/backend/core/data_recipe/oxc-validator/validate.mjs @@ -0,0 +1,576 @@ +// 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 { spawnSync } from "node:child_process"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { basename, dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { parseSync } from "oxc-parser"; + +const LANG_TO_EXT = { + js: "js", + jsx: "jsx", + ts: "ts", + tsx: "tsx", +}; + +const VALIDATION_MODES = new Set(["syntax", "lint", "syntax+lint"]); +const CODE_SHAPES = new Set(["auto", "module", "snippet"]); +const SNIPPET_PREFIX = "(() => {\n"; +const SNIPPET_SUFFIX = "\n})();\nexport {};\n"; +const OXLINT_SUPPRESSED_RULES = ["no-unused-vars", "no-new-array"]; +const TOOL_DIR = dirname(fileURLToPath(import.meta.url)); + +function mapLang(value) { + const normalized = String(value || "").trim().toLowerCase(); + if (normalized === "javascript" || normalized === "js") { + return "js"; + } + if (normalized === "typescript" || normalized === "ts") { + return "ts"; + } + if (normalized === "jsx") { + return "jsx"; + } + if (normalized === "tsx") { + return "tsx"; + } + return "js"; +} + +function mapMode(value) { + const normalized = String(value || "").trim().toLowerCase(); + if (VALIDATION_MODES.has(normalized)) { + return normalized; + } + return "syntax"; +} + +function mapCodeShape(value) { + const normalized = String(value || "").trim().toLowerCase(); + if (CODE_SHAPES.has(normalized)) { + return normalized; + } + return "auto"; +} + +function parseFileIndex(filePath) { + if (typeof filePath !== "string") { + return null; + } + const match = basename(filePath).match(/^snippet_(\d+)\./); + if (!match) { + return null; + } + const parsed = Number.parseInt(match[1], 10); + return Number.isFinite(parsed) ? parsed : null; +} + +function toCodeString(code) { + return typeof code === "string" ? code : String(code ?? ""); +} + +function makeValidationEntry({ code, index, lang, codeShape }) { + const source = toCodeString(code); + if (codeShape === "snippet") { + return { + index, + lang, + code: `${SNIPPET_PREFIX}${source}${SNIPPET_SUFFIX}`, + offset: SNIPPET_PREFIX.length, + }; + } + return { + index, + lang, + code: source, + offset: 0, + }; +} + +function shiftOffset(value, offset) { + if (!Number.isInteger(value)) { + return null; + } + const shifted = value - offset; + return shifted >= 0 ? shifted : null; +} + +function remapDiagnosticOffsets(diagnostic, offset) { + if (!diagnostic || typeof diagnostic !== "object" || offset <= 0) { + return diagnostic; + } + return { + ...diagnostic, + labels: Array.isArray(diagnostic.labels) + ? diagnostic.labels.map((label) => ({ + ...label, + start: shiftOffset(label.start, offset), + end: shiftOffset(label.end, offset), + })) + : [], + }; +} + +function normalizeParserError(error) { + if (typeof error === "string") { + return { + code: null, + message: error.trim() || "Unknown parser error", + severity: null, + labels: [], + codeframe: null, + }; + } + if (!error || typeof error !== "object") { + return { + code: null, + message: "Unknown parser error", + severity: null, + labels: [], + codeframe: null, + }; + } + const code = typeof error.code === "string" ? error.code : null; + const message = String(error.message || error.reason || "").trim() || "Unknown parser error"; + const severity = typeof error.severity === "string" ? error.severity : null; + const labels = Array.isArray(error.labels) + ? error.labels.map((label) => ({ + message: + label && typeof label === "object" && typeof label.message === "string" + ? label.message + : null, + start: + label && typeof label === "object" && Number.isInteger(label.start) + ? label.start + : null, + end: + label && typeof label === "object" && Number.isInteger(label.end) + ? label.end + : null, + })) + : []; + const codeframe = typeof error.codeframe === "string" ? error.codeframe : null; + return { + code, + message, + severity, + labels, + codeframe, + }; +} + +function normalizeLintDiagnostic(diagnostic) { + if (!diagnostic || typeof diagnostic !== "object") { + return null; + } + + const readString = (value) => + typeof value === "string" ? value : null; + const readInt = (value) => + Number.isInteger(value) ? value : null; + const asObject = (value) => + value && typeof value === "object" ? value : null; + + const message = String(diagnostic.message || "").trim(); + if (!message) { + return null; + } + + const severityRaw = String(diagnostic.severity || "").trim().toLowerCase(); + const severity = severityRaw === "error" ? "error" : "warning"; + + const labels = []; + if (Array.isArray(diagnostic.labels)) { + for (const label of diagnostic.labels) { + const labelObj = asObject(label); + const span = asObject(labelObj?.span); + const start = readInt(span?.offset); + const length = readInt(span?.length); + labels.push({ + message: readString(labelObj?.label), + start, + end: start !== null && length !== null ? start + length : null, + }); + } + } + + const code = typeof diagnostic.code === "string" ? diagnostic.code : null; + return { + code, + message: code ? `${code}: ${message}` : message, + severity, + labels, + codeframe: null, + }; +} + +function makeResult({ + isValid, + errorCount, + warningCount = 0, + message = "", + severity = null, + code = null, + labels = [], + codeframe = null, +}) { + return { + is_valid: Boolean(isValid), + error_count: Number.isInteger(errorCount) ? errorCount : 0, + warning_count: Number.isInteger(warningCount) ? warningCount : 0, + error_message: String(message || ""), + severity: typeof severity === "string" ? severity : null, + code: typeof code === "string" ? code : null, + labels: Array.isArray(labels) ? labels : [], + codeframe: typeof codeframe === "string" ? codeframe : null, + }; +} + +function syntaxResultFromErrors(errors) { + const first = errors[0] ?? null; + return makeResult({ + isValid: errors.length === 0, + errorCount: errors.length, + warningCount: 0, + message: errors.slice(0, 3).map((error) => error.message).join(" | "), + severity: first ? first.severity : null, + code: first ? first.code : null, + labels: first ? first.labels : [], + codeframe: first ? first.codeframe : null, + }); +} + +function runSyntaxParse(entry) { + const ext = LANG_TO_EXT[entry.lang] ?? "js"; + const filename = `snippet_${entry.index}.${ext}`; + try { + const parsed = parseSync(filename, entry.code, { + lang: entry.lang, + sourceType: "module", + showSemanticErrors: true, + }); + const errors = Array.isArray(parsed?.errors) + ? parsed.errors + .map(normalizeParserError) + .filter(Boolean) + .map((error) => remapDiagnosticOffsets(error, entry.offset)) + : []; + return errors; + } catch (error) { + return [ + remapDiagnosticOffsets( + normalizeParserError(error), + entry.offset, + ), + ]; + } +} + +function pickPreferredErrorList(firstErrors, secondErrors) { + if (secondErrors.length < firstErrors.length) { + return secondErrors; + } + return firstErrors; +} + +function validateSyntaxOne({ code, lang, index, codeShape }) { + if (codeShape !== "auto") { + const lintEntry = makeValidationEntry({ + code, + index, + lang, + codeShape, + }); + const errors = runSyntaxParse(lintEntry); + return { + result: syntaxResultFromErrors(errors), + lintEntry, + }; + } + + const moduleEntry = makeValidationEntry({ + code, + index, + lang, + codeShape: "module", + }); + const moduleErrors = runSyntaxParse(moduleEntry); + if (moduleErrors.length === 0) { + return { + result: syntaxResultFromErrors(moduleErrors), + lintEntry: moduleEntry, + }; + } + + const snippetEntry = makeValidationEntry({ + code, + index, + lang, + codeShape: "snippet", + }); + const snippetErrors = runSyntaxParse(snippetEntry); + if (snippetErrors.length === 0) { + return { + result: syntaxResultFromErrors(snippetErrors), + lintEntry: snippetEntry, + }; + } + + const chosenErrors = pickPreferredErrorList(moduleErrors, snippetErrors); + const lintEntry = chosenErrors === snippetErrors ? snippetEntry : moduleEntry; + return { + result: syntaxResultFromErrors(chosenErrors), + lintEntry, + }; +} + +function resolveLintEntry({ code, lang, index, codeShape }) { + if (codeShape !== "auto") { + return makeValidationEntry({ + code, + index, + lang, + codeShape, + }); + } + + const moduleEntry = makeValidationEntry({ + code, + index, + lang, + codeShape: "module", + }); + if (runSyntaxParse(moduleEntry).length === 0) { + return moduleEntry; + } + + const snippetEntry = makeValidationEntry({ + code, + index, + lang, + codeShape: "snippet", + }); + if (runSyntaxParse(snippetEntry).length === 0) { + return snippetEntry; + } + + return moduleEntry; +} + +function fallbackLintResults(entries, message) { + return new Map( + entries.map((entry) => [ + entry.index, + makeResult({ + isValid: false, + errorCount: 1, + warningCount: 0, + message, + severity: "error", + }), + ]), + ); +} + +function runLintBatch(entries) { + if (entries.length === 0) { + return new Map(); + } + + const entryByIndex = new Map(entries.map((entry) => [entry.index, entry])); + const tempDir = mkdtempSync(join(tmpdir(), "oxlint-")); + try { + for (const entry of entries) { + const ext = LANG_TO_EXT[entry.lang] ?? "js"; + const filePath = join(tempDir, `snippet_${entry.index}.${ext}`); + writeFileSync(filePath, entry.code, "utf8"); + } + + const oxlintBin = join(TOOL_DIR, "node_modules", ".bin", "oxlint"); + const oxlintArgs = [ + ...OXLINT_SUPPRESSED_RULES.flatMap((rule) => ["-A", rule]), + "--format", + "json", + tempDir, + ]; + const exec = spawnSync(oxlintBin, oxlintArgs, { + encoding: "utf8", + cwd: TOOL_DIR, + }); + if (exec.error) { + return fallbackLintResults( + entries, + `oxlint execution failed: ${exec.error.message}`, + ); + } + const stdout = String(exec.stdout || "").trim(); + if (!stdout) { + const stderr = String(exec.stderr || "").trim(); + return fallbackLintResults( + entries, + stderr || "oxlint returned empty output", + ); + } + + let parsed; + try { + parsed = JSON.parse(stdout); + } catch { + return fallbackLintResults(entries, "oxlint JSON parse failed"); + } + + const rawDiagnostics = Array.isArray(parsed?.diagnostics) + ? parsed.diagnostics + : []; + const byIndex = new Map(); + + for (const diag of rawDiagnostics) { + const filenameRaw = + typeof diag?.filename === "string" ? diag.filename : ""; + const filename = filenameRaw.startsWith("file://") + ? filenameRaw.replace("file://", "") + : filenameRaw; + const index = parseFileIndex(filename); + if (index === null) { + continue; + } + const normalized = normalizeLintDiagnostic(diag); + if (!normalized) { + continue; + } + const entry = entryByIndex.get(index); + const remapped = remapDiagnosticOffsets(normalized, entry?.offset ?? 0); + const list = byIndex.get(index) ?? []; + list.push(remapped); + byIndex.set(index, list); + } + + const results = new Map(); + for (const entry of entries) { + const diagnostics = byIndex.get(entry.index) ?? []; + const errorDiagnostics = diagnostics.filter( + (diag) => diag.severity === "error", + ); + const warningDiagnostics = diagnostics.filter( + (diag) => diag.severity !== "error", + ); + const top = errorDiagnostics[0] ?? warningDiagnostics[0] ?? null; + const messageSource = + errorDiagnostics.length > 0 ? errorDiagnostics : warningDiagnostics; + results.set( + entry.index, + makeResult({ + isValid: errorDiagnostics.length === 0, + errorCount: errorDiagnostics.length, + warningCount: warningDiagnostics.length, + message: messageSource + .slice(0, 3) + .map((diag) => diag.message) + .join(" | "), + severity: top ? top.severity : null, + code: top ? top.code : null, + labels: top ? top.labels : [], + codeframe: top ? top.codeframe : null, + }), + ); + } + return results; + } catch (error) { + return fallbackLintResults(entries, `oxlint execution failed: ${error}`); + } finally { + rmSync(tempDir, { recursive: true, force: true }); + } +} + +function readStdin() { + return new Promise((resolve, reject) => { + let data = ""; + process.stdin.setEncoding("utf8"); + process.stdin.on("data", (chunk) => { + data += chunk; + }); + process.stdin.on("end", () => resolve(data)); + process.stdin.on("error", (error) => reject(error)); + }); +} + +function runValidation({ codes, lang, mode, codeShape }) { + if (mode === "syntax") { + return codes.map((code, index) => + validateSyntaxOne({ code, lang, index, codeShape }).result, + ); + } + + if (mode === "lint") { + const entries = codes.map((code, index) => + resolveLintEntry({ code, lang, index, codeShape }), + ); + const lintMap = runLintBatch(entries); + return entries.map( + (entry) => + lintMap.get(entry.index) ?? + makeResult({ + isValid: true, + errorCount: 0, + warningCount: 0, + }), + ); + } + + const syntaxRuns = codes.map((code, index) => + validateSyntaxOne({ code, lang, index, codeShape }), + ); + const lintTargets = syntaxRuns + .filter((run) => run.result.is_valid === true) + .map((run) => run.lintEntry); + const lintMap = runLintBatch(lintTargets); + + return syntaxRuns.map((run) => { + if (run.result.is_valid !== true) { + return run.result; + } + return ( + lintMap.get(run.lintEntry.index) ?? + makeResult({ + isValid: true, + errorCount: 0, + warningCount: 0, + }) + ); + }); +} + +async function main() { + const raw = await readStdin(); + let payload; + try { + payload = JSON.parse(raw || "{}"); + } catch { + process.stdout.write( + JSON.stringify([ + makeResult({ + isValid: false, + errorCount: 1, + warningCount: 0, + message: "Invalid JSON payload", + severity: "error", + }), + ]), + ); + return; + } + + const lang = mapLang(payload?.lang); + const mode = mapMode(payload?.mode); + const codeShape = mapCodeShape(payload?.code_shape); + const codes = Array.isArray(payload?.codes) ? payload.codes : []; + const out = runValidation({ codes, lang, mode, codeShape }); + process.stdout.write(JSON.stringify(out)); +} + +main().catch((error) => { + process.stderr.write(String(error?.stack || error)); + process.exit(1); +}); diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py new file mode 100644 index 0000000..9d8ca5c --- /dev/null +++ b/studio/backend/core/data_recipe/service.py @@ -0,0 +1,329 @@ +# 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 base64 +import io +import os +from pathlib import Path +from typing import Any + +from .jsonable import to_jsonable +from .local_callable_validators import ( + register_oxc_local_callable_validators, + split_oxc_local_callable_validators, +) + +_IMAGE_CONTEXT_PATCHED = False + + +def _encode_bytes_to_base64(value: bytes | bytearray) -> str: + return base64.b64encode(bytes(value)).decode("utf-8") + + +def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None: + try: + path = Path(path_value) + candidates: list[Path] = [] + if path.is_absolute(): + candidates.append(path) + else: + if base_path: + candidates.append(Path(base_path) / path) + candidates.append(Path.cwd() / path) + + for candidate in candidates: + if not candidate.exists() or not candidate.is_file(): + continue + with candidate.open("rb") as f: + return _encode_bytes_to_base64(f.read()) + except (OSError, TypeError, ValueError): + return None + return None + + +def _pil_image_to_base64(value: Any) -> str | None: + try: + from PIL.Image import Image as PILImage # type: ignore + except ImportError: + return None + if not isinstance(value, PILImage): + return None + buffer = io.BytesIO() + image_format = str(getattr(value, "format", "") or "").upper() + if image_format not in {"PNG", "JPEG", "JPG", "WEBP", "GIF"}: + image_format = "PNG" + value.save(buffer, format = image_format) + return _encode_bytes_to_base64(buffer.getvalue()) + + +def _normalize_image_context_value(value: Any, *, base_path: str | None = None) -> Any: + if isinstance(value, str): + return value + + if isinstance(value, (bytes, bytearray)): + return _encode_bytes_to_base64(value) + + pil_base64 = _pil_image_to_base64(value) + if pil_base64 is not None: + return pil_base64 + + if isinstance(value, dict): + url = value.get("url") + if isinstance(url, str): + return url + + image_url = value.get("image_url") + if isinstance(image_url, str): + return image_url + if isinstance(image_url, dict): + nested_url = image_url.get("url") + if isinstance(nested_url, str): + return nested_url + + inline_data = value.get("data") + if isinstance(inline_data, str): + return inline_data + + raw_bytes = value.get("bytes") + if isinstance(raw_bytes, (bytes, bytearray)): + return _encode_bytes_to_base64(raw_bytes) + if isinstance(raw_bytes, str) and raw_bytes.strip(): + return raw_bytes + + path_value = value.get("path") + if isinstance(path_value, str) and path_value.strip(): + if as_base64 := _load_image_file_to_base64(path_value, base_path = base_path): + return as_base64 + return path_value + + return value + + +def _apply_data_designer_image_context_patch() -> None: + global _IMAGE_CONTEXT_PATCHED + if _IMAGE_CONTEXT_PATCHED: + return + + try: + from data_designer.config.models import ImageContext # pyright: ignore[reportMissingImports] + except ImportError: + return + + if getattr(ImageContext, "_unsloth_image_context_patch_applied", False): + _IMAGE_CONTEXT_PATCHED = True + return + + original_auto_resolve = ImageContext._auto_resolve_context_value + + def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any: + normalized = _normalize_image_context_value(context_value, base_path = base_path) + return original_auto_resolve(self, normalized, base_path) + + ImageContext._auto_resolve_context_value = _patched_auto_resolve + setattr(ImageContext, "_unsloth_image_context_patch_applied", True) + _IMAGE_CONTEXT_PATCHED = True + + +def build_model_providers(recipe: dict[str, Any]): + from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] + + providers: list[ModelProvider] = [] + for provider in recipe.get("model_providers", []): + api_key = provider.get("api_key") + api_key_env = provider.get("api_key_env") + if not api_key and api_key_env: + api_key = os.getenv(api_key_env) + providers.append( + ModelProvider( + name = provider["name"], + endpoint = provider["endpoint"], + provider_type = provider.get("provider_type", "openai"), + api_key = api_key, + extra_headers = provider.get("extra_headers"), + extra_body = provider.get("extra_body"), + ) + ) + + return providers + + +def _recipe_has_llm_columns(recipe: dict[str, Any]) -> bool: + for column in recipe.get("columns", []): + if not isinstance(column, dict): + continue + column_type = column.get("column_type") + if isinstance(column_type, str) and column_type.startswith("llm-"): + return True + return False + + +def _validate_recipe_runtime_support(recipe: dict[str, Any], model_providers: list[Any]) -> None: + if _recipe_has_llm_columns(recipe) and not model_providers: + raise ValueError("Add a Provider connection block before running this recipe.") + + +def build_mcp_providers(recipe: dict[str, Any]) -> list: + from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports] + + # Same gate as the chat MCP path: stdio providers spawn a local subprocess, + # so build them only when this host allows it (desktop / explicit opt-in). + from core.inference.mcp_client import stdio_mcp_enabled + + stdio_allowed = stdio_mcp_enabled() + + providers: list[MCPProvider | LocalStdioMCPProvider] = [] + for provider in recipe.get("mcp_providers", []): + if not isinstance(provider, dict): + continue + provider_type = provider.get("provider_type") + if provider_type == "stdio": + if not stdio_allowed: + continue + env = provider.get("env") + if not isinstance(env, dict): + env = {} + args = provider.get("args") + if not isinstance(args, list): + args = [] + providers.append( + LocalStdioMCPProvider( + name = str(provider.get("name", "")), + command = str(provider.get("command", "")), + args = [str(value) for value in args], + env = {str(key): str(value) for key, value in env.items()}, + ) + ) + continue + + if provider_type in {"sse", "streamable_http"}: + api_key = provider.get("api_key") + api_key_env = provider.get("api_key_env") + if not api_key and api_key_env: + api_key = os.getenv(str(api_key_env)) + providers.append( + MCPProvider( + name = str(provider.get("name", "")), + endpoint = str(provider.get("endpoint", "")), + provider_type = str(provider_type), + api_key = str(api_key) if api_key else None, + ) + ) + return providers + + +def _strip_frontend_model_config_metadata(recipe: dict[str, Any]) -> dict[str, Any]: + model_configs = recipe.get("model_configs") + if not isinstance(model_configs, list): + return recipe + + changed = False + next_model_configs: list[Any] = [] + for model_config in model_configs: + if isinstance(model_config, dict) and "gguf_variant" in model_config: + next_model_config = dict(model_config) + next_model_config.pop("gguf_variant", None) + next_model_configs.append(next_model_config) + changed = True + continue + next_model_configs.append(model_config) + + if not changed: + return recipe + + return { + **recipe, + "model_configs": next_model_configs, + } + + +def build_config_builder(recipe: dict[str, Any]): + _apply_data_designer_image_context_patch() + from data_designer.config import DataDesignerConfigBuilder # pyright: ignore[reportMissingImports] + from data_designer.config.processors import ProcessorType # pyright: ignore[reportMissingImports] + + recipe_core = { + key: value + for key, value in recipe.items() + if key not in {"model_providers", "mcp_providers"} + } + recipe_core = _strip_frontend_model_config_metadata(recipe_core) + recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators(recipe_core) + builder = DataDesignerConfigBuilder.from_config({"data_designer": recipe_core}) + register_oxc_local_callable_validators( + builder = builder, + specs = oxc_local_callable_specs, + ) + + # DataDesignerConfigBuilder.from_config currently skips processors. + # Re-attach so drop_columns/schema_transform survive the API payload. + for processor in recipe_core.get("processors") or []: + if not isinstance(processor, dict): + continue + processor_type_raw = processor.get("processor_type") + if not isinstance(processor_type_raw, str): + continue + kwargs = {k: v for k, v in processor.items() if k != "processor_type"} + builder.add_processor( + processor_type = ProcessorType(processor_type_raw), + **kwargs, + ) + + return builder + + +def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = None): + _apply_data_designer_image_context_patch() + from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports] + + recipe = _strip_frontend_model_config_metadata(recipe) + model_providers = build_model_providers(recipe) + _validate_recipe_runtime_support(recipe, model_providers) + + # DataDesigner requires >=1 model provider even with no LLM columns; stub + # one so sampler/expression-only recipes run without a real provider. + if not model_providers: + from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] + model_providers = [ + ModelProvider( + name = "_unused", + endpoint = "http://localhost", + provider_type = "openai", + api_key = None, + ) + ] + + return DataDesigner( + artifact_path = artifact_path, + model_providers = model_providers, + mcp_providers = build_mcp_providers(recipe), + ) + + +def validate_recipe(recipe: dict[str, Any]) -> None: + builder = build_config_builder(recipe) + designer = create_data_designer(recipe) + designer.validate(builder) + + +def preview_recipe( + recipe: dict[str, Any], num_records: int +) -> tuple[list[dict[str, Any]], dict[str, Any] | None, dict[str, Any] | None]: + builder = build_config_builder(recipe) + designer = create_data_designer(recipe) + results = designer.preview(builder, num_records = num_records) + + dataset: list[dict[str, Any]] = [] + if results.dataset is not None: + raw_rows = results.dataset.to_dict(orient = "records") + dataset = [to_jsonable(row) for row in raw_rows] + + artifacts = ( + None if results.processor_artifacts is None else to_jsonable(results.processor_artifacts) + ) + analysis = ( + None if results.analysis is None else to_jsonable(results.analysis.model_dump(mode = "json")) + ) + + return dataset, artifacts, analysis diff --git a/studio/backend/core/export/__init__.py b/studio/backend/core/export/__init__.py new file mode 100644 index 0000000..f7241cd --- /dev/null +++ b/studio/backend/core/export/__init__.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Export submodule - model export operations. + +get_export_backend() returns an ExportOrchestrator that delegates to a +subprocess. The original ExportBackend runs inside the subprocess and can be +imported directly from .export when needed. +""" + +from .orchestrator import ExportOrchestrator, get_export_backend + +# Expose ExportOrchestrator as ExportBackend for backward compat +ExportBackend = ExportOrchestrator + +__all__ = [ + "ExportBackend", + "ExportOrchestrator", + "get_export_backend", +] diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py new file mode 100644 index 0000000..c8be50b --- /dev/null +++ b/studio/backend/core/export/export.py @@ -0,0 +1,1166 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Export backend - exports models in various formats.""" + +import glob +import json +import structlog +import tempfile +from loggers import get_logger +import os +import shutil +import contextlib +from pathlib import Path +from typing import Optional, Tuple, List + +# unsloth imports torch on non-MLX hosts, so a --no-torch install raises here. Stay importable +# (null the classes) so exports return a clean "PyTorch is not installed" error, not an import crash. +try: + from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + _UNSLOTH_IMPORT_ERROR = None +except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load + FastLanguageModel = None + FastVisionModel = None + _IS_MLX = False + _UNSLOTH_IMPORT_ERROR = _unsloth_exc + +from huggingface_hub import HfApi, ModelCard +from utils.hardware import clear_gpu_cache + +from utils.models import is_vision_model, get_base_model_from_lora +from utils.models.model_config import detect_audio_type +from utils.paths import ( + ensure_dir, + outputs_root, + resolve_export_write_dir, + resolve_output_dir, +) +from core.inference import get_inference_backend + +# GPU/PyTorch-only imports, skipped on MLX and on a --no-torch install so the module stays +# importable; export then degrades to a clear "PyTorch is not installed" error. +torch = None +_TORCH_IMPORT_ERROR: Optional[BaseException] = None +if not _IS_MLX: + try: + from peft import PeftModel, PeftModelForCausalLM + from transformers.modeling_utils import PushToHubMixin + import torch + except Exception as _torch_exc: # ImportError, or a broken native torch load + _TORCH_IMPORT_ERROR = _torch_exc + +logger = get_logger(__name__) + + +def _export_runtime_available() -> bool: + """True if export can run: MLX active, or Unsloth imported (only succeeds on a GPU host).""" + return bool(_IS_MLX) or (FastLanguageModel is not None) + + +def _export_runtime_message() -> str: + """Precise reason the export runtime is unavailable, mirroring hardware.export_capability().""" + if torch is None: + return ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + return ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export on " + "CPU only.)" + ) + + +# Kept for call sites / tests referencing the PyTorch-missing text. +_PYTORCH_MISSING_MESSAGE = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." +) + +_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False + + +def _supports_kwarg(fn, name): + """True if `fn` accepts keyword `name` directly or via **kwargs.""" + import inspect + + try: + params = inspect.signature(fn).parameters + except (TypeError, ValueError): + return False + return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + + +def _compressed_export_supported(): + """True if the installed unsloth build can do FP8/NVFP4 compressed-tensors export.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_compressed_method") + except Exception: + return False + + +def _torchao_export_supported(): + """True if the installed unsloth build has the portable torchao FP8/INT8 export path.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_torchao_method") + except Exception: + return False + + +def _has_nvidia_gpu(): + """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" + try: + from utils.hardware import hardware as _hw + return _hw.DEVICE == _hw.DeviceType.CUDA and not _hw.IS_ROCM + except Exception: + try: + import torch + return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None + except Exception: + return False + + +def _hf_offline(timeout = 3): + """True if export should avoid the Hub: honors the HF offline env vars, else does one + cheap TCP reachability probe so a network-down load uses local files / the HF cache + instead of hanging on connection timeouts. Proxy-aware (probes the proxy egress when + one is configured); disable the probe with UNSLOTH_OFFLINE_PROBE=0.""" + _offline = {"1", "true", "yes", "on"} + if ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline + ): + return True + if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {"0", "false", "no", "off"}: + return False # probe disabled -> assume online; loads still pass local_files_only on env + + # Shared bounded, proxy-aware probe (also used by the export worker before version activation). + from utils.transformers_version import hf_endpoint_unreachable + + if hf_endpoint_unreachable(timeout): + logger.warning("Hugging Face endpoint unreachable; loading checkpoint in offline mode") + return True + return False + + +# Reuse Unsloth's lock-guarded forced-offline context; no-op fallback if it moves. +try: + from unsloth.models.loader_utils import _force_hf_offline +except Exception: + import contextlib as _contextlib + + @_contextlib.contextmanager + def _force_hf_offline(): + yield + + +def _offline_window_if(local_files_only): + """Forced-offline window when offline was detected, else a no-op context.""" + return _force_hf_offline() if local_files_only else contextlib.nullcontext() + + +def _is_wsl(): + """Detect if running under Windows Subsystem for Linux.""" + try: + return "microsoft" in open("/proc/version").read().lower() + except Exception: + return False + + +def _apply_wsl_sudo_patch(): + """On WSL, monkey-patch do_we_need_sudo() to return False. + + WSL lacks passwordless sudo and do_we_need_sudo()'s `sudo apt-get update` + hangs on a stdin password; setup.sh pre-installs the build deps anyway. + """ + if not _is_wsl(): + return + + try: + import unsloth_zoo.llama_cpp as llama_cpp_module + + def _wsl_do_we_need_sudo(system_type = "debian"): + logger.info("WSL detected — skipping sudo check (build deps pre-installed by setup.sh)") + return False + + llama_cpp_module.do_we_need_sudo = _wsl_do_we_need_sudo + logger.info("Applied WSL sudo patch to unsloth_zoo.llama_cpp.do_we_need_sudo") + except Exception as e: + logger.warning(f"Could not apply WSL sudo patch: {e}") + + +# Model card template +MODEL_CARD = """--- +base_model: {base_model} +tags: +- text-generation-inference +- transformers +- unsloth +- {model_type} +- {extra} +license: apache-2.0 +language: +- en +--- + +# Uploaded finetuned {method} model + +- **Developed by:** {username} +- **License:** apache-2.0 +- **Finetuned from model :** {base_model} + +This {model_type} model was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth) and Huggingface's TRL library. + +[](https://github.com/unslothai/unsloth) +""" + + +class ExportBackend: + """Handles model export operations""" + + def __init__(self): + self.inference_backend = get_inference_backend() + self.current_checkpoint = None + self.current_model = None + self.current_tokenizer = None + self.is_vision = False + self.is_peft = False + self._audio_type = None + + def cleanup_memory(self): + """Offload and delete all models from memory""" + try: + logger.info("Starting memory cleanup...") + + model_names = list(self.inference_backend.models.keys()) + for model_name in model_names: + self.inference_backend.unload_model(model_name) + + self.current_model = None + self.current_tokenizer = None + self.current_checkpoint = None + self._audio_type = None + + clear_gpu_cache() + + logger.info("Memory cleanup completed successfully") + return True + + except Exception as e: + logger.error(f"Error during memory cleanup: {e}") + return False + + def scan_checkpoints( + self, outputs_dir: str = str(outputs_root()) + ) -> List[Tuple[str, List[Tuple[str, str]]]]: + """ + Scan outputs folder for training runs and their checkpoints. + + Returns: [(model_name, [(display_name, checkpoint_path), ...]), ...] + """ + from utils.models.checkpoints import scan_checkpoints + return scan_checkpoints(outputs_dir = outputs_dir) + + def load_checkpoint( + self, + checkpoint_path: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + trust_remote_code: bool = False, + hf_token: Optional[str] = None, + ) -> Tuple[bool, str]: + """ + Load a checkpoint for export. + + ``hf_token`` authenticates the actual weight load for gated/private + checkpoints, matching the token the worker used for the security preflight + (otherwise a gated repo passes scanning then 401s at from_pretrained). + + Returns: + Tuple of (success: bool, message: str) + """ + token = hf_token if hf_token and hf_token.strip() else None + try: + logger.info(f"Loading checkpoint: {checkpoint_path}") + + self.cleanup_memory() + + checkpoint_path_obj = Path(checkpoint_path) + + # Model identity for type detection + adapter_config = checkpoint_path_obj / "adapter_config.json" + base_model = None + if adapter_config.exists(): + base_model = get_base_model_from_lora(checkpoint_path) + if not base_model: + return False, "Could not determine base model for adapter" + + model_id = base_model or checkpoint_path + + # Skip the Hub when offline so a no-internet export uses the local cache. + local_files_only = _hf_offline() + + # Run the type-detection probes in the forced-offline window (else a gated + # base 404s); it covers is_vision_model's Hub reads + the transformers-5 + # subprocess, and local_files_only makes detect_audio_type's requests.get skip. + with _offline_window_if(local_files_only): + self._audio_type = detect_audio_type( + model_id, hf_token = token, local_files_only = local_files_only + ) + self.is_vision = not self._audio_type and is_vision_model( + model_id, hf_token = token, local_files_only = local_files_only + ) + + if self._audio_type == "csm": + from unsloth import FastModel + from transformers import CsmForConditionalGeneration + + logger.info("Loading as CSM audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name = checkpoint_path, + max_seq_length = max_seq_length, + dtype = None, + auto_model = CsmForConditionalGeneration, + load_in_4bit = False, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, + ) + + elif self._audio_type == "whisper": + from unsloth import FastModel + from transformers import WhisperForConditionalGeneration + + logger.info("Loading as Whisper audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name = checkpoint_path, + dtype = None, + load_in_4bit = False, + auto_model = WhisperForConditionalGeneration, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, + ) + + elif self._audio_type == "snac": + logger.info("Loading as SNAC (Orpheus) audio model...") + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = checkpoint_path, + max_seq_length = max_seq_length, + dtype = None, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, + ) + + elif self._audio_type == "bicodec": + from unsloth import FastModel + logger.info("Loading as BiCodec (Spark-TTS) audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name = checkpoint_path, + max_seq_length = max_seq_length, + dtype = None if _IS_MLX else torch.float32, + load_in_4bit = False, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, + ) + + elif self._audio_type == "dac": + from unsloth import FastModel + logger.info("Loading as DAC (OuteTTS) audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name = checkpoint_path, + max_seq_length = max_seq_length, + load_in_4bit = False, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, + ) + + elif self.is_vision: + logger.info("Loading as vision model...") + model, processor = FastVisionModel.from_pretrained( + model_name = checkpoint_path, + max_seq_length = max_seq_length, + dtype = None, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, + ) + tokenizer = processor # vision: processor acts as tokenizer + + else: + logger.info("Loading as text model...") + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = checkpoint_path, + max_seq_length = max_seq_length, + dtype = None, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, + ) + + if _IS_MLX: + # MLX doesn't use PeftModel — detect LoRA via adapter_config.json + self.is_peft = adapter_config.exists() + else: + self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM)) + + self.current_model = model + self.current_tokenizer = tokenizer + self.current_checkpoint = checkpoint_path + + if self._audio_type: + model_type = f"Audio ({self._audio_type})" + elif self.is_vision: + model_type = "Vision" + else: + model_type = "Text" + peft_info = " (PEFT Adapter)" if self.is_peft else " (Merged Model)" + + logger.info(f"Successfully loaded {model_type} model{peft_info}") + return True, f"Loaded {model_type} model{peft_info} successfully" + + except Exception as e: + logger.error(f"Error loading checkpoint: {e}") + import traceback + + logger.error(traceback.format_exc()) + return False, f"Failed to load checkpoint: {str(e)}" + + def _write_export_metadata(self, save_directory: str): + """Write export_metadata.json with base model info for Chat page discovery.""" + try: + base_model = ( + get_base_model_from_lora(self.current_checkpoint) + if self.current_checkpoint + else None + ) + metadata = {"base_model": base_model} + metadata_path = os.path.join(save_directory, "export_metadata.json") + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent = 2) + logger.info(f"Wrote export metadata to {metadata_path}") + except Exception as e: + logger.warning(f"Could not write export metadata: {e}") + + def export_merged_model( + self, + save_directory: str, + format_type: str = "16-bit (FP16)", + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + compressed_method: Optional[str] = None, + ) -> Tuple[bool, str, Optional[str]]: + """ + Export merged model (for PEFT models). + + Args: + save_directory: Local directory to save model + format_type: "16-bit (FP16)", "4-bit (FP4)", or a compressed-tensors label + compressed_method: Optional compressed-tensors scheme alias (e.g. "fp8", + "fp8_static", "w8a8", "w4a16", "mxfp4", "mxfp8", "nvfp4"). Overrides + format_type and is resolved against unsloth.save COMPRESSED_EXPORT_SCHEMES. + push_to_hub: Whether to push to Hugging Face Hub + repo_id: Hub repository ID (username/model-name) + hf_token: Hugging Face token + private: Whether to make the repo private + + Returns: + Tuple of (success: bool, message: str, output_path: Optional[str]) + """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None + if not self.current_model or not self.current_tokenizer: + return False, "No model loaded. Please select a checkpoint first.", None + + # Merged export works for PEFT adapters and non-PEFT Local/HF base models alike + # (save_pretrained_merged is a no-op merge that just saves the base). + + output_path: Optional[str] = None + # Quantized formats save to a sibling "-". Two backends: compressed-tensors + # (llm-compressor, NVIDIA-only) and portable torchao FP8/INT8 (device-agnostic). The alias + # comes from `compressed_method` (the "all formats" dropdown) or the `format_type` label. + _LABEL_TO_ALIAS = { + "FP8 (compressed-tensors)": "fp8", + "NVFP4 (compressed-tensors)": "nvfp4", + } + compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type) + compressed_suffix: Optional[str] = None + # Classify the alias: torchao-portable vs compressed-tensors. + torchao_info = None + if compressed_alias and _torchao_export_supported(): + try: + import unsloth.save as _us_t + torchao_info = _us_t._normalize_torchao_method(compressed_alias) + except Exception: + torchao_info = None + is_torchao = torchao_info is not None + is_compressed = compressed_alias is not None and not is_torchao + try: + if _IS_MLX and (is_compressed or is_torchao): + return ( + False, + "Quantized (FP8/FP4/INT) export is not supported on macOS/MLX. " + "Use 16-bit or GGUF.", + None, + ) + + if is_torchao: + # Portable torchao: no NVIDIA GPU, no calibration. + compressed_suffix = torchao_info[1] + + if is_compressed: + # compressed-tensors needs CUDA; enforce in the backend even if the UI gate is bypassed. + if not _has_nvidia_gpu(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an NVIDIA GPU. On other " + "hardware use the portable FP8/INT8 (torchao) formats or 16-bit.", + None, + ) + if not _compressed_export_supported(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an Unsloth build with " + "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", + None, + ) + import unsloth.save as _us + + # Prefer the llm-compressor-main shadow (transformers 5.x): it quantizes newer models + # (Qwen3.5, Gemma-4, ...) the shipped 0.10.x cannot. Route all compressed exports + # through it when available; else fall back to the workspace 0.10.x path below. + _shadow_pp = None + try: + from utils.transformers_version import llmcompressor_shadow_pythonpath + _shadow_pp = llmcompressor_shadow_pythonpath() + except Exception as e: + logger.warning(f"llm-compressor-main shadow unavailable: {e}") + if _shadow_pp: + os.environ[_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV] = _shadow_pp + else: + # No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its + # transformers ceiling, so fail fast for sidecar models; default-tier still works. + os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None) + _exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + return ( + False, + "FP8/FP4 compressed-tensors export is not available for this model: it " + f"runs under transformers {_tf_ver}, but the installed llm-compressor " + f"supports transformers <= {_us._LLM_COMPRESSOR_MAX_TRANSFORMERS} and the " + "llm-compressor-main runtime could not be provisioned (offline or " + "UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN). Export to GGUF or 16-bit instead.", + None, + ) + + try: + info = _us._normalize_compressed_method(compressed_alias) + except Exception as e: + return False, f"Unsupported compressed export '{compressed_alias}': {e}", None + if info is None: + return ( + False, + f"'{compressed_alias}' is not a recognized compressed-tensors export.", + None, + ) + compressed_suffix = info[2] + + if _IS_MLX: + mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + elif is_compressed or is_torchao: + save_method = compressed_alias + elif format_type == "4-bit (FP4)": + save_method = "merged_4bit_forced" + elif self._audio_type == "whisper": + save_method = None + else: + save_method = "merged_16bit" + + if save_directory: + save_directory = str(resolve_export_write_dir(save_directory)) + logger.info(f"Saving merged model locally to: {save_directory}") + ensure_dir(Path(save_directory)) + + if _IS_MLX: + self.current_model.save_pretrained_merged( + save_directory, + self.current_tokenizer, + save_method = mlx_save_method, + ) + else: + self.current_model.save_pretrained_merged( + save_directory, self.current_tokenizer, save_method = save_method + ) + + # Compressed / torchao writes to the "-" sibling; report that as output. + final_dir = ( + f"{save_directory}-{compressed_suffix}" + if (is_compressed or is_torchao) + else save_directory + ) + self._write_export_metadata(final_dir) + logger.info(f"Model saved successfully to {final_dir}") + output_path = str(Path(final_dir).resolve()) + + if push_to_hub: + if not repo_id or not hf_token: + return ( + False, + "Repository ID and Hugging Face token required for Hub upload", + None, + ) + + logger.info(f"Pushing merged model to Hub: {repo_id}") + + if _IS_MLX: + if save_directory: + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = save_directory, + token = hf_token, + private = private, + ) + else: + with tempfile.TemporaryDirectory() as tmp_dir: + self.current_model.save_pretrained_merged( + tmp_dir, + self.current_tokenizer, + save_method = mlx_save_method, + ) + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = tmp_dir, + token = hf_token, + private = private, + ) + elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir(): + # Already built in output_path; upload it directly instead of re-running the + # expensive quantization that push_to_hub_merged(save_method=...) would redo. + hf_api = HfApi(token = hf_token) + repo_id = PushToHubMixin._create_repo( + PushToHubMixin, + repo_id = repo_id, + private = private, + token = hf_token, + ) + content = MODEL_CARD.format( + username = repo_id.split("/")[0], + base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), + model_type = getattr(self.current_model.config, "model_type", "llm"), + method = compressed_alias or format_type, + extra = "unsloth", + ) + ModelCard(content).push_to_hub( + repo_id, token = hf_token, commit_message = "Unsloth Model Card" + ) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) + else: + hub_save_method = save_method if save_method is not None else "merged_16bit" + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_method = hub_save_method, + token = hf_token, + private = private, + ) + logger.info(f"Model pushed successfully to {repo_id}") + + 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)}", None + + def export_base_model( + self, + save_directory: str, + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + base_model_id: Optional[str] = None, + ) -> Tuple[bool, str, Optional[str]]: + """ + Export base model (for non-PEFT models). + + Returns: + Tuple of (success: bool, message: str, output_path: Optional[str]) + """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None + if not self.current_model or not self.current_tokenizer: + 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: + if save_directory: + save_directory = str(resolve_export_write_dir(save_directory)) + logger.info(f"Saving base model locally to: {save_directory}") + ensure_dir(Path(save_directory)) + + if _IS_MLX: + # MLX: save_pretrained_merged handles non-LoRA models too + # (fuse() is a no-op without LoRA layers) + self.current_model.save_pretrained_merged( + save_directory, + self.current_tokenizer, + save_method = "merged_16bit", + ) + else: + self.current_model.save_pretrained(save_directory) + self.current_tokenizer.save_pretrained(save_directory) + + # 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()) + + if push_to_hub: + if not repo_id or not hf_token: + return ( + False, + "Repository ID and Hugging Face token required for Hub upload", + None, + ) + + logger.info(f"Pushing base model to Hub: {repo_id}") + + if _IS_MLX: + if save_directory: + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = save_directory, + token = hf_token, + private = private, + ) + else: + with tempfile.TemporaryDirectory() as tmp_dir: + self.current_model.save_pretrained_merged( + tmp_dir, + self.current_tokenizer, + save_method = "merged_16bit", + ) + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = tmp_dir, + token = hf_token, + private = private, + ) + else: + # Base model name from request or model config + base_model = ( + base_model_id or self.current_model.config._name_or_path or "unknown" + ) + + hf_api = HfApi(token = hf_token) + repo_id = PushToHubMixin._create_repo( + PushToHubMixin, + repo_id = repo_id, + private = private, + token = hf_token, + ) + username = repo_id.split("/")[0] + + content = MODEL_CARD.format( + username = username, + base_model = base_model, + model_type = self.current_model.config.model_type, + method = "", + extra = "unsloth", + ) + card = ModelCard(content) + card.push_to_hub(repo_id, token = hf_token, commit_message = "Unsloth Model Card") + + if save_directory: + hf_api.upload_folder( + folder_path = save_directory, + repo_id = repo_id, + repo_type = "model", + ) + logger.info(f"Model pushed successfully to {repo_id}") + else: + return ( + False, + "Local save directory required for Hub upload", + None, + ) + + 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)}", None + + def export_gguf( + self, + save_directory: str, + quantization_method = "Q4_K_M", + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + imatrix_file = None, + ) -> Tuple[bool, str, Optional[str]]: + """ + Export model in GGUF format. + + Args: + save_directory: Local directory to save model + quantization_method: A single GGUF quant method (e.g., "Q4_K_M") or a list of them + (e.g., ["Q4_K_M", "Q8_0"]). A list produces one GGUF per quant from a single + model load (unsloth save_to_gguf loops internally). + push_to_hub: Whether to push to Hugging Face Hub + repo_id: Hub repository ID + hf_token: Hugging Face token + + Returns: + Tuple of (success: bool, message: str, output_path: Optional[str]) + """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None + if not self.current_model or not self.current_tokenizer: + return False, "No model loaded. Please select a checkpoint first.", None + + # Only forward imatrix_file to an unsloth build that accepts it, else older builds raise + # an unexpected-keyword error even for a plain no-imatrix export. + if imatrix_file is not None and not _supports_kwarg( + self.current_model.save_pretrained_gguf, "imatrix_file" + ): + return ( + False, + "This Unsloth build does not support GGUF imatrix export. " + "Upgrade unsloth and unsloth_zoo, or disable the imatrix option.", + None, + ) + imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {} + + output_path: Optional[str] = None + model_tmp_to_cleanup: Optional[str] = None + try: + # Normalize to a lowercased list so multiple quants come from one model load. + if isinstance(quantization_method, (list, tuple)): + quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()] + else: + quant_methods = [str(quantization_method).lower()] + if not quant_methods: + quant_methods = ["q4_k_m"] + quant_method = quant_methods if len(quant_methods) > 1 else quant_methods[0] + + # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it + # can't drift past the pinned llama-quantize binary's gguf API. + global _LLAMA_CPP_SCRIPTS_WARNING_EMITTED + try: + from unsloth_zoo.llama_cpp import ( + LLAMA_CPP_DEFAULT_DIR, + _resolve_local_convert_script, # noqa: F401 + ) + os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR) + 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 + + if save_directory: + save_directory = str(resolve_export_write_dir(save_directory)) + # Keep unsloth relative-path internals anchored to the repo cwd. + abs_save_dir = os.path.abspath(save_directory) + logger.info(f"Saving GGUF model locally to: {abs_save_dir}") + + ensure_dir(Path(abs_save_dir)) + + # On WSL, patch out sudo check before llama.cpp build + _apply_wsl_sudo_patch() + + # convert_to_gguf writes output relative to cwd (repo root); + # snapshot existing .gguf so we can diff and relocate afterwards. + cwd = os.getcwd() + pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) + + pre_existing_subs = {d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()} + + # Avoid clobbering an existing user-owned model/ directory. + import uuid + + _model_tmp = os.path.join(abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}") + model_tmp_to_cleanup = _model_tmp + self.current_model.save_pretrained_gguf( + _model_tmp, + self.current_tokenizer, + quantization_method = quant_method, + **imatrix_kw, + ) + + # Relocate the .gguf that convert_to_gguf wrote to cwd (repo root). + new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs + for src in sorted(new_ggufs): + dest = os.path.join(abs_save_dir, os.path.basename(src)) + shutil.move(src, dest) + logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/") + + # Flatten GGUF files from subdirs created during this export. + for sub in list(Path(abs_save_dir).iterdir()): + if not sub.is_dir(): + continue + if sub.name in pre_existing_subs: + continue + for src in sub.glob("*.gguf"): + dest = os.path.join(abs_save_dir, src.name) + shutil.move(str(src), dest) + logger.info(f"Relocated GGUF: {src.name} → {abs_save_dir}/") + shutil.rmtree(str(sub), ignore_errors = True) + logger.info(f"Cleaned up subdirectory: {sub.name}") + + # For non-PEFT models, save_pretrained_gguf leaves a *_gguf dir at + # the checkpoint path; relocate its GGUFs and clean it up. + if self.current_checkpoint: + ckpt = Path(self.current_checkpoint) + gguf_dir = ckpt.parent / f"{ckpt.name}_gguf" + if gguf_dir.is_dir() and gguf_dir.resolve() != Path(abs_save_dir).resolve(): + for src in gguf_dir.glob("*.gguf"): + dest = os.path.join(abs_save_dir, src.name) + shutil.move(str(src), dest) + logger.info(f"Relocated GGUF: {src.name} → {abs_save_dir}/") + # Also relocate Ollama Modelfile if present + modelfile = gguf_dir / "Modelfile" + if modelfile.is_file(): + shutil.move(str(modelfile), os.path.join(abs_save_dir, "Modelfile")) + logger.info(f"Relocated Modelfile → {abs_save_dir}/") + shutil.rmtree(str(gguf_dir), ignore_errors = True) + logger.info(f"Cleaned up intermediate GGUF dir: {gguf_dir}") + + # Write export metadata so the Chat page can identify the base model + self._write_export_metadata(abs_save_dir) + + final_ggufs = sorted(glob.glob(os.path.join(abs_save_dir, "*.gguf"))) + logger.info( + "GGUF export complete. Final files in %s:\n %s", + abs_save_dir, + "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + ) + output_path = str(Path(abs_save_dir).resolve()) + + if push_to_hub: + if not repo_id or not hf_token: + return ( + False, + "Repository ID and Hugging Face token required for Hub upload", + None, + ) + + logger.info(f"Pushing GGUF model to Hub: {repo_id}") + + self.current_model.push_to_hub_gguf( + repo_id, + self.current_tokenizer, + quantization_method = quant_method, + token = hf_token, + **imatrix_kw, + ) + logger.info(f"GGUF model pushed successfully to {repo_id}") + + return ( + True, + f"GGUF model exported successfully ({', '.join(quant_methods)})", + output_path, + ) + + except Exception as e: + if model_tmp_to_cleanup: + shutil.rmtree(model_tmp_to_cleanup, ignore_errors = True) + logger.error(f"Error exporting GGUF model: {e}") + import traceback + + logger.error(traceback.format_exc()) + return False, f"GGUF export failed: {str(e)}", None + + def export_lora_adapter( + self, + save_directory: str, + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", + ) -> Tuple[bool, str, Optional[str]]: + """ + Export LoRA adapter only (not merged). + + Args: + gguf: If True, also convert the adapter to a GGUF LoRA file (llama.cpp + convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`. + gguf_outtype: GGUF LoRA output float type; one of q8_0/f16/bf16/f32. + + Returns: + Tuple of (success: bool, message: str, output_path: Optional[str]) + """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None + if not self.current_model or not self.current_tokenizer: + 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.", None + + _GGUF_LORA_OUTTYPES = ("q8_0", "f16", "bf16", "f32") + if gguf: + if _IS_MLX: + return ( + False, + "GGUF LoRA adapter export is not supported on macOS/MLX. " + "Use the safetensors adapter instead.", + None, + ) + outtype = str(gguf_outtype).lower() + if outtype not in _GGUF_LORA_OUTTYPES: + return ( + False, + f"Invalid GGUF LoRA outtype '{gguf_outtype}'. " + f"Choose one of {', '.join(_GGUF_LORA_OUTTYPES)}.", + None, + ) + # getattr so an older build without save_pretrained_gguf returns a clean message + # instead of an AttributeError (a generic 500). + _save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None) + if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"): + return ( + False, + "This Unsloth build does not support GGUF LoRA adapter export. " + "Upgrade unsloth and unsloth_zoo, or export the safetensors adapter.", + None, + ) + + output_path: Optional[str] = None + try: + if save_directory: + save_directory = str(resolve_export_write_dir(save_directory)) + logger.info(f"Saving LoRA adapter locally to: {save_directory}") + ensure_dir(Path(save_directory)) + + if gguf: + # Writes the adapter files plus "-lora-.gguf". + _apply_wsl_sudo_patch() + self.current_model.save_pretrained_gguf( + save_directory, + self.current_tokenizer, + save_method = "lora", + quantization_method = outtype, + # Forward the token so convert_lora_to_gguf.py can fetch a gated base's config. + token = hf_token or None, + ) + final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf"))) + logger.info( + "LoRA GGUF export complete. Files in %s:\n %s", + save_directory, + "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + ) + elif _IS_MLX: + # MLX: save adapters.safetensors + tokenizer files + self.current_model.save_lora_adapters(save_directory) + self.current_tokenizer.save_pretrained(save_directory) + else: + 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()) + + if push_to_hub: + if not repo_id or not hf_token: + return ( + False, + "Repository ID and Hugging Face token required for Hub upload", + None, + ) + + logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") + + if gguf: + # Upload the locally-built GGUF folder; needs a local save_directory so the + # conversion is not re-run. + if not (output_path and Path(output_path).is_dir()): + return ( + False, + "GGUF LoRA Hub upload requires a local save directory; set one and " + "retry.", + None, + ) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) + elif _IS_MLX: + with tempfile.TemporaryDirectory() as tmp_dir: + self.current_model.save_lora_adapters(tmp_dir) + self.current_tokenizer.save_pretrained(tmp_dir) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = tmp_dir, + repo_id = repo_id, + repo_type = "model", + ) + else: + self.current_model.push_to_hub(repo_id, token = hf_token, private = private) + self.current_tokenizer.push_to_hub(repo_id, token = hf_token, private = private) + logger.info(f"Adapter pushed successfully to {repo_id}") + + 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)}", None + + +# Global export backend instance +_export_backend = None + + +def get_export_backend() -> ExportBackend: + """Get or create the global export backend instance""" + global _export_backend + if _export_backend is None: + _export_backend = ExportBackend() + return _export_backend diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py new file mode 100644 index 0000000..671ef36 --- /dev/null +++ b/studio/backend/core/export/orchestrator.py @@ -0,0 +1,633 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Export orchestrator — subprocess-based. + +Same API as ExportBackend, but delegates all ML work to a persistent +subprocess spawned on first checkpoint load and reused for later exports. + +When switching between checkpoints needing different transformers +versions, the old subprocess is killed and a new one spawned. + +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, Deque, Dict, List, Optional, Tuple +from utils.paths import outputs_root + +logger = get_logger(__name__) + +_CTX = mp.get_context("spawn") + +# Max log lines kept per orchestrator (live log panel scrollback); ~1 MB worst-case. +_LOG_BUFFER_MAXLEN = 4000 + + +class ExportOrchestrator: + """ + Export backend orchestrator — subprocess-based. + + Exposes the same API surface as ExportBackend so routes/export.py + needs minimal changes. All heavy ML work happens in a persistent + subprocess. + """ + + def __init__(self): + self._proc: Optional[mp.Process] = None + self._cmd_queue: Any = None + self._resp_queue: Any = None + # Serializes export ops so concurrent HTTP requests can't interleave commands. + self._lock = threading.Lock() + + # Local state mirrors (updated from subprocess responses). + self.current_checkpoint: Optional[str] = None + self.is_vision: bool = False + self.is_peft: bool = False + + # Thread-safe ring buffer of worker log lines; powers the export logs SSE endpoint. + self._log_buffer: Deque[Dict[str, Any]] = deque(maxlen = _LOG_BUFFER_MAXLEN) + self._log_lock = threading.Lock() + # Monotonic seq, never reset, so SSE clients have a stable cursor across clear_logs(). + self._log_seq: int = 0 + # _log_seq snapshot at the current run's start; SSE defaults its cursor here so a + # late-connecting client still sees the full run. Current run has seq > this. + self._run_start_seq: int = 0 + # True while an export op runs; SSE ends the stream 1s after this flips False. + self._export_active: bool = False + # Set by cancel_export(); reset when a new load/export run starts. Lets the + # caller distinguish a user cancel from a genuine subprocess crash. + self._cancel_requested: bool = False + + # Last finished operation, so a client whose blocking POST was cut off by a + # Cloudflare tunnel timeout (524 at ~100s, while the op runs for minutes) can + # poll /api/export/status and still learn the real outcome. Guarded by + # _op_lock. `_op_seq` is a monotonic counter the client uses as a baseline to + # tell "my op finished" (seq grew) from a stale previous result. + self._op_lock = threading.Lock() + self._op_seq: int = 0 + self._active_op_kind: Optional[str] = None + self._last_op: Optional[Dict[str, Any]] = None + + 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 worker log line to the buffer, stamped with a monotonic seq.""" + 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 buffered log lines from a previous op so the UI shows only this run. + + The seq counter is NOT reset (clients keep a stable cursor); the current seq + is snapshotted into ``_run_start_seq`` to anchor the SSE default cursor. + """ + with self._log_lock: + self._log_buffer.clear() + 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 captured at the current run's start (SSE default cursor).""" + 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 + + def was_cancelled(self) -> bool: + """True if the in-flight (or most recent) run was cancelled by the user.""" + return self._cancel_requested + + def _record_op_finished(self, success: bool, message: str, output_path: Optional[str]) -> None: + """Snapshot the just-finished op so status pollers can recover its outcome. + + Called from each op's ``finally`` (with ``_active_op_kind`` still set) BEFORE + ``_export_active`` is cleared, so a status read that observes the op as + inactive is guaranteed to also see this matching result. + """ + with self._op_lock: + self._op_seq += 1 + status = "cancelled" if self._cancel_requested else ("success" if success else "error") + self._last_op = { + "seq": self._op_seq, + "kind": self._active_op_kind, + "status": status, + "output_path": output_path if success else None, + "error": None if success else (message or None), + } + + def get_last_op(self) -> Optional[Dict[str, Any]]: + """Return the last finished op record (or None), for status recovery.""" + with self._op_lock: + return dict(self._last_op) if self._last_op is not None else None + + def get_active_op_kind(self) -> Optional[str]: + """Return the kind of the currently running op (or None when idle).""" + return self._active_op_kind + + def cancel_export(self) -> bool: + """Terminate the in-flight export subprocess immediately. + + An export op holds ``self._lock`` for its whole duration (blocked in + ``_wait_response``), so we deliberately do NOT take the lock here -- we + kill the worker process directly, which unblocks that wait and makes the + in-flight op return a failure the caller surfaces as "cancelled". + + Only the export subprocess is touched; training and inference run in + their own subprocesses and are left untouched. + + Returns True if a live subprocess was terminated, False if none ran. + """ + self._cancel_requested = True + proc = self._proc + if proc is None or not proc.is_alive(): + return False + logger.info( + "Export cancel requested: terminating export subprocess (pid=%s)", + proc.pid, + ) + try: + proc.terminate() + proc.join(timeout = 5) + except Exception: + pass + if proc.is_alive(): + logger.warning("Export subprocess survived terminate, killing") + try: + proc.kill() + proc.join(timeout = 3) + except Exception: + pass + return True + + # ------------------------------------------------------------------ + # 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 + + with native_path_secret_removed_for_child_start(): + self._cmd_queue = _CTX.Queue() + self._resp_queue = _CTX.Queue() + + 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() + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) + logger.info("Export subprocess started (pid=%s)", self._proc.pid) + + def _shutdown_subprocess(self, timeout: float = 10.0) -> None: + """Gracefully shut down the export subprocess.""" + if self._proc is None or not self._proc.is_alive(): + self._proc = None + return + + self._drain_queue() + + try: + self._cmd_queue.put({"type": "shutdown"}) + except (OSError, ValueError): + pass + + try: + self._proc.join(timeout = timeout) + except Exception: + pass + + # Force kill if still alive. + if self._proc is not None and self._proc.is_alive(): + logger.warning("Export subprocess did not exit gracefully, terminating") + try: + self._proc.terminate() + self._proc.join(timeout = 5) + except Exception: + pass + if self._proc is not None and self._proc.is_alive(): + logger.warning("Subprocess still alive after terminate, killing") + try: + self._proc.kill() + self._proc.join(timeout = 3) + except Exception: + pass + + self._proc = None + self._cmd_queue = None + self._resp_queue = None + logger.info("Export subprocess shut down") + + def _cleanup(self): + """atexit handler.""" + self._shutdown_subprocess(timeout = 5.0) + + def _ensure_subprocess_alive(self) -> bool: + """Check if subprocess is alive.""" + return self._proc is not None and self._proc.is_alive() + + # ------------------------------------------------------------------ + # Queue helpers + # ------------------------------------------------------------------ + + def _send_cmd(self, cmd: dict) -> None: + """Send a command to the subprocess.""" + if self._cmd_queue is None: + raise RuntimeError("No export subprocess running") + try: + self._cmd_queue.put(cmd) + except (OSError, ValueError) as exc: + raise RuntimeError(f"Failed to send command to subprocess: {exc}") + + def _read_resp(self, timeout: float = 1.0) -> Optional[dict]: + """Read a response from the subprocess (non-blocking with timeout).""" + if self._resp_queue is None: + return None + try: + return self._resp_queue.get(timeout = timeout) + except queue.Empty: + return None + except (EOFError, OSError, ValueError): + return None + + def _wait_response( + self, + expected_type: str, + timeout: float = 3600.0, + ) -> dict: + """Block until a response of the expected type arrives. + + Export ops can take a long time — GGUF conversion for large + models (30B+) easily takes 20-30 minutes. Default timeout 1 hour. + """ + deadline = time.monotonic() + timeout + + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 2.0)) + + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError("Export subprocess crashed during wait") + continue + + rtype = resp.get("type", "") + + if rtype == expected_type: + return resp + + if rtype == "error": + error_msg = resp.get("error", "Unknown error") + raise RuntimeError(f"Subprocess error: {error_msg}") + + if rtype == "log": + # Forwarded stdout/stderr line from the worker. + self._append_log(resp) + continue + + if rtype == "status": + message = resp.get("message", "") + logger.info("Export subprocess status: %s", message) + # Surface status in the live log panel for high-level progress. + if message: + self._append_log( + { + "stream": "status", + "line": message, + "ts": resp.get("ts", time.time()), + } + ) + continue + + # Other response types during wait — skip. + logger.debug( + "Skipping response type '%s' while waiting for '%s'", + rtype, + expected_type, + ) + + raise RuntimeError(f"Timeout waiting for '{expected_type}' response after {timeout}s") + + def _drain_queue(self) -> list: + """Drain all pending responses.""" + events = [] + if self._resp_queue is None: + return events + while True: + try: + events.append(self._resp_queue.get_nowait()) + except queue.Empty: + return events + except (EOFError, OSError, ValueError): + return events + + # ------------------------------------------------------------------ + # Public API — same interface as ExportBackend + # ------------------------------------------------------------------ + + def load_checkpoint( + self, + checkpoint_path: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, + hf_token: Optional[str] = None, + subject: Optional[str] = None, + ) -> Tuple[bool, str]: + """Load a checkpoint for export. + + Always spawns a fresh subprocess to ensure a clean Python interpreter. + """ + sub_config = { + "checkpoint_path": checkpoint_path, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, + "subject": subject, + "hf_token": hf_token, + } + + with self._lock: + # Fresh log buffer so the UI sees only this run's output. + self.clear_logs() + self._cancel_requested = False + self._active_op_kind = "load_checkpoint" + self._export_active = True + op_success, op_message = False, "" + try: + # Always kill any existing subprocess and spawn fresh. + if self._ensure_subprocess_alive(): + self._shutdown_subprocess() + elif self._proc is not None: + self._shutdown_subprocess(timeout = 2) + + logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path) + self._spawn_subprocess(sub_config) + + 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 + op_success, op_message = False, str(exc) + 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) + op_success, op_message = True, resp.get("message", "Loaded successfully") + return True, op_message + 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 + op_success, op_message = False, error + return False, error + finally: + self._record_op_finished(op_success, op_message, None) + self._active_op_kind = None + self._export_active = False + + def export_merged_model( + self, + save_directory: str, + format_type: str = "16-bit (FP16)", + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + compressed_method: Optional[str] = None, + ) -> Tuple[bool, str, Optional[str]]: + """Export merged PEFT model.""" + return self._run_export( + "merged", + { + "save_directory": save_directory, + "format_type": format_type, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + "private": private, + "compressed_method": compressed_method, + }, + ) + + def export_base_model( + self, + save_directory: str, + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + base_model_id: Optional[str] = None, + ) -> Tuple[bool, str, Optional[str]]: + """Export base model (non-PEFT).""" + return self._run_export( + "base", + { + "save_directory": save_directory, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + "private": private, + "base_model_id": base_model_id, + }, + ) + + def export_gguf( + self, + save_directory: str, + quantization_method = "Q4_K_M", + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + imatrix_file = None, + ) -> Tuple[bool, str, Optional[str]]: + """Export model in GGUF format. `quantization_method` may be a single method or a list.""" + return self._run_export( + "gguf", + { + "save_directory": save_directory, + "quantization_method": quantization_method, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + "imatrix_file": imatrix_file, + }, + ) + + def export_lora_adapter( + self, + save_directory: str, + push_to_hub: bool = False, + repo_id: Optional[str] = None, + hf_token: Optional[str] = None, + private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", + ) -> Tuple[bool, str, Optional[str]]: + """Export LoRA adapter only (optionally also as a GGUF LoRA file).""" + return self._run_export( + "lora", + { + "save_directory": save_directory, + "push_to_hub": push_to_hub, + "repo_id": repo_id, + "hf_token": hf_token, + "private": private, + "gguf": gguf, + "gguf_outtype": gguf_outtype, + }, + ) + + def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str, Optional[str]]: + """Send an export command and wait for the result. + + Returns ``(success, message, output_path)``. ``output_path`` is the on-disk + dir the worker wrote to (None if it only pushed to Hub or failed pre-write). + """ + with self._lock: + if not self._ensure_subprocess_alive(): + return ( + False, + "No export subprocess running. Load a checkpoint first.", + None, + ) + + self.clear_logs() + self._cancel_requested = False + self._active_op_kind = f"export_{export_type}" + self._export_active = True + op_success, op_message, op_output_path = False, "", None + try: + cmd = {"type": "export", "export_type": export_type, **params} + try: + self._send_cmd(cmd) + # GGUF for 30B+ models can take 30+ min per quant; a multi-quant list runs them + # all in one op off a single merge, so scale the timeout by the quant count. + _qm = params.get("quantization_method") + _n = len(_qm) if isinstance(_qm, (list, tuple)) and _qm else 1 + resp = self._wait_response( + f"export_{export_type}_done", + timeout = 3600 * max(1, _n), + ) + op_success = resp.get("success", False) + op_message = resp.get("message", "") + op_output_path = resp.get("output_path") + return op_success, op_message, op_output_path + except RuntimeError as exc: + op_success, op_message = False, str(exc) + return False, str(exc), None + finally: + self._record_op_finished(op_success, op_message, op_output_path) + self._active_op_kind = None + self._export_active = False + + def cleanup_memory(self) -> bool: + """Cleanup export-related models from memory.""" + with self._lock: + if not self._ensure_subprocess_alive(): + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return True + + self._active_op_kind = "cleanup" + self._export_active = True + success = False + 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() + + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return success + finally: + self._record_op_finished(success, "", None) + self._active_op_kind = None + self._export_active = False + + def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]: + """Scan for checkpoints — runs locally, no ML imports.""" + from utils.models.checkpoints import scan_checkpoints + return scan_checkpoints(outputs_dir = outputs_dir) + + +# ========== GLOBAL INSTANCE ========== +_export_backend = None + + +def get_export_backend() -> ExportOrchestrator: + """Get global export backend instance (orchestrator).""" + global _export_backend + if _export_backend is None: + _export_backend = ExportOrchestrator() + return _export_backend diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py new file mode 100644 index 0000000..7828116 --- /dev/null +++ b/studio/backend/core/export/worker.py @@ -0,0 +1,691 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Export subprocess entry point. + +Each export session runs in a persistent subprocess (mp spawn), giving a clean +interpreter with no stale module state, which solves transformers version +switching. The subprocess stays alive while a model is loaded, accepting commands +(load, export_*, cleanup, shutdown) via mp.Queue. + +Pattern follows core/inference/worker.py and core/training/worker.py. +""" + +from __future__ import annotations + +import contextlib +import errno +import structlog +from loggers import get_logger +import os +import sys +import threading +import time +import traceback +from pathlib import Path +from typing import Any + +logger = get_logger(__name__) + + +# Gate controlling whether captured stdout/stderr lines are forwarded to the +# parent's resp_queue (and on to the export-dialog SSE stream). Closed by default +# so the noisy bootstrap phase (imports, model resolution, loading bars) is +# suppressed in the UI; _handle_export() opens it when export work starts. The +# orchestrator spawns a fresh subprocess per checkpoint load, resetting this. +# Dropped lines are still echoed to the saved fds so the server log keeps them. +_log_forward_gate = threading.Event() + + +def _setup_log_capture(resp_queue: Any) -> None: + """Redirect fds 1 and 2 through pipes so every line printed by this worker + and any child it spawns is forwarded to the parent via resp_queue as + {"type": "log", ...} messages. + + Must run BEFORE LogConfig.setup_logging and any ML imports, else library + handlers may capture the original stderr reference and bypass the pipe. + Lines are also echoed back to the original fds so the server console keeps + the full output even while ``_log_forward_gate`` is closed. + """ + + try: + saved_out_fd = os.dup(1) + saved_err_fd = os.dup(2) + except OSError: + # dup failed; give up quietly (export still works, no live 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). + os.close(w_out) + os.close(w_err) + + # Replace sys.stdout/sys.stderr with line-buffered writers on fds 1 and 2. + try: + sys.stdout = os.fdopen(1, "w", buffering = 1, encoding = "utf-8", errors = "replace") + sys.stderr = os.fdopen(2, "w", buffering = 1, encoding = "utf-8", errors = "replace") + 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 keeps 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): already echoed above; drop the + # line so the export dialog skips import noise. + continue + try: + resp_queue.put_nowait( + { + "type": "log", + "stream": stream_name, + "line": line, + "ts": time.time(), + } + ) + except Exception: + # Queue put failed; drop the line rather than crash the 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, hf_token: str | None = None) -> None: + """Activate the correct transformers version BEFORE any ML imports.""" + # Ensure backend is on sys.path for utils imports. + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.transformers_version import activate_transformers_for_subprocess + + activate_transformers_for_subprocess(model_name, hf_token) + + +@contextlib.contextmanager +def _offline_window_if_unreachable(step = "loading"): + """Force HF offline for a network-touching step (transformers version activation, or the + load preflights that hit the Hub) when the endpoint is unreachable, then restore the prior + env. Keeps a no-network export from hanging on Hub calls that run before load_checkpoint's + own probe, while letting this persistent worker re-decide per operation once back online. + + Post-ML-import (the load preflights), huggingface_hub has already read its in-process + offline constant and cached sessions, so env alone is too late: defer to the loader's + _force_hf_offline (env + in-process flags + session reset). Pre-import (activation), + huggingface_hub is not loaded yet, so setting the env vars suffices for its urllib probes.""" + saved: dict[str, str | None] = {} + force_ctx = None + try: + from utils.transformers_version import _env_offline, hf_endpoint_unreachable + probe_enabled = os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() not in ( + "0", + "false", + "no", + "off", + ) + if not _env_offline() and probe_enabled and hf_endpoint_unreachable(): + logger.warning("Hugging Face endpoint unreachable; %s offline", step) + if "huggingface_hub" in sys.modules: + try: + from unsloth.models.loader_utils import _force_hf_offline + force_ctx = _force_hf_offline() + force_ctx.__enter__() # sets env + in-process flags + resets sessions + except Exception: + force_ctx = None + if force_ctx is None: + for k in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"): + saved[k] = os.environ.get(k) + os.environ[k] = "1" + except Exception: + pass + try: + yield + finally: + if force_ctx is not None: + try: + force_ctx.__exit__(None, None, None) + except Exception: + pass + for k, v in saved.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + +def _send_response(resp_queue: Any, response: dict) -> None: + """Send a response to the parent process.""" + try: + resp_queue.put(response) + except (OSError, ValueError) as exc: + logger.error("Failed to send response: %s", exc) + + +def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: + """Handle a load_checkpoint command.""" + checkpoint_path = cmd["checkpoint_path"] + max_seq_length = cmd.get("max_seq_length", 2048) + 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: + from utils.security.trusted_org import is_trusted_org_repo + + _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/")) + # Genuine first-party Hub repo only (not a local/spoof name starting + # with "unsloth/"); authenticated so private repos resolve. + and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token")) + ): + trust_remote_code = True + logger.info( + "Auto-enabled trust_remote_code for Nemotron model: %s", + checkpoint_path, + ) + + # Malware gate: a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) every + # load. Local checkpoints have no Hub scan and are skipped in the helper; a + # LoRA merges its base weights, so gate that repo too. + from utils.security import evaluate_file_security, security_load_subdirs + + malware_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + _hf_token = cmd.get("hf_token") + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token) + ) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + }, + ) + return + + # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless + # pinned-approved. A LoRA merges its base model, whose code runs, so gate it too. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a local or remote adapter's base so its base repo is gated too. + base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = cmd.get("hf_token"), + trust_remote_code = True, + approved_fingerprint = cmd.get("approved_remote_code_fingerprint"), + subject = cmd.get("subject"), + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Checkpoint '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review and " + f"approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + }, + ) + return + + try: + _send_response( + resp_queue, + { + "type": "status", + "message": f"Loading checkpoint: {checkpoint_path}", + "ts": time.time(), + }, + ) + + success, message = backend.load_checkpoint( + checkpoint_path = checkpoint_path, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + hf_token = cmd.get("hf_token"), + ) + + _send_response( + resp_queue, + { + "type": "loaded", + "success": success, + "message": message, + "checkpoint": checkpoint_path if success else None, + "is_vision": backend.is_vision if success else False, + "is_peft": backend.is_peft if success else False, + "ts": time.time(), + }, + ) + + except Exception as exc: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + + +def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: + """Handle any export command (merged, base, gguf, lora).""" + 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 export progress in the live + # log panel. Stays open for the rest of this subprocess's life; the + # orchestrator spawns a fresh subprocess per checkpoint load, resetting it. + _log_forward_gate.set() + + output_path: Any = None + try: + if export_type == "merged": + 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), + repo_id = cmd.get("repo_id"), + hf_token = cmd.get("hf_token"), + private = cmd.get("private", False), + compressed_method = cmd.get("compressed_method"), + ) + elif export_type == "base": + 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"), + hf_token = cmd.get("hf_token"), + private = cmd.get("private", False), + base_model_id = cmd.get("base_model_id"), + ) + elif export_type == "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), + repo_id = cmd.get("repo_id"), + hf_token = cmd.get("hf_token"), + imatrix_file = cmd.get("imatrix_file"), + ) + elif export_type == "lora": + 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"), + hf_token = cmd.get("hf_token"), + private = cmd.get("private", False), + gguf = cmd.get("gguf", False), + gguf_outtype = cmd.get("gguf_outtype", "q8_0"), + ) + else: + success, message = False, f"Unknown export type: {export_type}" + + _send_response( + resp_queue, + { + "type": response_type, + "success": success, + "message": message, + "output_path": output_path, + "ts": time.time(), + }, + ) + + except Exception as exc: + _send_response( + resp_queue, + { + "type": response_type, + "success": False, + "message": str(exc), + "output_path": None, + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + + +def _handle_cleanup(backend, resp_queue: Any) -> None: + """Handle a cleanup command.""" + try: + success = backend.cleanup_memory() + _send_response( + resp_queue, + { + "type": "cleanup_done", + "success": success, + "ts": time.time(), + }, + ) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "cleanup_done", + "success": False, + "message": str(exc), + "ts": time.time(), + }, + ) + + +def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None: + """Subprocess entrypoint. Persistent — runs command loop until shutdown. + + Args: + cmd_queue: mp.Queue for receiving commands from parent. + resp_queue: mp.Queue for sending responses to parent. + config: Initial configuration dict with checkpoint_path. + """ + import queue as _queue + + # Install fd-level stdout/stderr capture FIRST so every subsequent print and + # every child process inherits the redirected fds (powers the live log stream). + _setup_log_capture(resp_queue) + + os.environ["TOKENIZERS_PARALLELISM"] = "false" + os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports + # Unbuffered output from child Python (e.g. GGUF converter) so prints surface live. + os.environ["PYTHONUNBUFFERED"] = "1" + # tqdm defaults to a 10s mininterval when stdout isn't a tty (we redirected + # fd 1/2 to a pipe), making multi-step bars look frozen; force frequent flushes. + os.environ.setdefault("TQDM_MININTERVAL", "0.5") + + import warnings + from loggers.config import LogConfig + + if os.getenv("ENVIRONMENT_TYPE", "production") == "production": + warnings.filterwarnings("ignore") + + LogConfig.setup_logging( + service_name = "unsloth-studio-export-worker", + env = os.getenv("ENVIRONMENT_TYPE", "production"), + ) + + checkpoint_path = config["checkpoint_path"] + + # ── 1. Activate correct transformers version BEFORE any ML imports ── + with _offline_window_if_unreachable(step = "activating transformers"): + try: + _activate_transformers_version(checkpoint_path, config.get("hf_token") or None) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "error", + "error": f"Failed to activate transformers version: {exc}", + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + return + + # ── 1b. Check Triton on Windows (must precede import torch) ── + if sys.platform == "win32": + try: + import triton # noqa: F401 + logger.info("Triton available — torch.compile enabled") + except ImportError: + os.environ["TORCHDYNAMO_DISABLE"] = "1" + logger.warning( + "Triton not found on Windows — torch.compile disabled. " + 'Install for better performance: pip install "triton-windows<3.7"' + ) + + # ── 1c. Stub torchao on Windows ROCm ── + # See core/_torchao_stub.py: torchao crashes on Windows ROCm (RCCL absent). + # No-op off Windows ROCm. Must run before importing transformers / unsloth_zoo. + from core._torchao_stub import install_torchao_windows_rocm_stub + + install_torchao_windows_rocm_stub() + + # ── 2. Import ML libraries (fresh in this clean process) ── + try: + _send_response( + resp_queue, + { + "type": "status", + "message": "Importing Unsloth...", + "ts": time.time(), + }, + ) + + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + # Recover from any namespace-package shadow before importing Unsloth. + from core.import_guards import ensure_real_packages + + ensure_real_packages("unsloth_zoo", "unsloth") + + from core.export.export import ExportBackend + + import transformers + + logger.info("Export subprocess loaded transformers %s", transformers.__version__) + + except Exception as exc: + _send_response( + resp_queue, + { + "type": "error", + "error": f"Failed to import ML libraries: {exc}", + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + return + + # ── 3. Create export backend and load initial checkpoint ── + try: + backend = ExportBackend() + + # Offline window covers the load preflights (malware/consent scans hit the Hub) + # before load_checkpoint runs its own probe; restored after so later loads re-decide. + with _offline_window_if_unreachable(): + _handle_load(backend, config, resp_queue) + + except Exception as exc: + _send_response( + resp_queue, + { + "type": "error", + "error": f"Failed to initialize export backend: {exc}", + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + return + + # ── 4. Command loop — process commands until shutdown ── + logger.info("Export subprocess ready, entering command loop") + + while True: + try: + cmd = cmd_queue.get(timeout = 1.0) + except _queue.Empty: + continue + except (EOFError, OSError): + logger.info("Command queue closed, shutting down") + return + + if cmd is None: + continue + + cmd_type = cmd.get("type", "") + logger.info("Received command: %s", cmd_type) + + try: + if cmd_type == "load": + # Load a new checkpoint, reusing this subprocess. + backend.cleanup_memory() + # Offline window also covers this load's Hub preflights (re-probed per load). + with _offline_window_if_unreachable(): + _handle_load(backend, cmd, resp_queue) + + elif cmd_type == "export": + _handle_export(backend, cmd, resp_queue) + + elif cmd_type == "cleanup": + _handle_cleanup(backend, resp_queue) + + elif cmd_type == "status": + _send_response( + resp_queue, + { + "type": "status_response", + "checkpoint": backend.current_checkpoint, + "is_vision": backend.is_vision, + "is_peft": backend.is_peft, + "ts": time.time(), + }, + ) + + elif cmd_type == "shutdown": + logger.info("Shutdown command received, cleaning up and exiting") + try: + backend.cleanup_memory() + except Exception: + pass + _send_response( + resp_queue, + { + "type": "shutdown_ack", + "ts": time.time(), + }, + ) + return + + else: + logger.warning("Unknown command type: %s", cmd_type) + _send_response( + resp_queue, + { + "type": "error", + "error": f"Unknown command type: {cmd_type}", + "ts": time.time(), + }, + ) + + except Exception as exc: + logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info = True) + _send_response( + resp_queue, + { + "type": "error", + "error": f"Command '{cmd_type}' failed: {exc}", + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) diff --git a/studio/backend/core/import_guards.py b/studio/backend/core/import_guards.py new file mode 100644 index 0000000..5b85a96 --- /dev/null +++ b/studio/backend/core/import_guards.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Recover `unsloth`/`unsloth_zoo` from a namespace-package shadow. Stdlib-only.""" + +from __future__ import annotations + +import os +import sys + + +def ensure_real_packages(*names: str) -> None: + """Drop sys.path entries where a bare `/` dir (no __init__.py) shadows + the installed package as a namespace, import the real packages, restore + sys.path. No-op without a shadow. Pass dependency-first (e.g. "unsloth_zoo", + "unsloth"); imports run dependency-last.""" + import importlib + import importlib.util + + bad: set = set() + shadowed: list = [] + for name in names: + try: + spec = importlib.util.find_spec(name) + except (ImportError, ValueError, AttributeError): + spec = None + # real package -> spec.origin is its __init__; namespace shadow -> None/"namespace" + if spec is None or spec.origin not in (None, "namespace"): + continue + dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])} + if not dirs: + continue + shadowed.append(name) + for entry in sys.path: + pkg = os.path.join(entry or os.getcwd(), name) + if os.path.realpath(pkg) in dirs and not os.path.isfile( + os.path.join(pkg, "__init__.py") + ): + bad.add(entry) + if not bad: + return + saved = list(sys.path) + sys.path[:] = [e for e in sys.path if e not in bad] + for name in shadowed: + for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]: + del sys.modules[cached] + try: + importlib.invalidate_caches() + # import unsloth before unsloth_zoo: unsloth.__init__ runs GPU/bnb fixes zoo relies on + for name in reversed(names): + importlib.import_module(name) + finally: + sys.path[:] = saved diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py new file mode 100644 index 0000000..ad78157 --- /dev/null +++ b/studio/backend/core/inference/__init__.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Inference submodule - backend for model loading and generation. + +The default get_inference_backend() returns an InferenceOrchestrator that +delegates to a subprocess. The original InferenceBackend runs inside the +subprocess and can be imported directly from .inference when needed. + +Public names are resolved lazily (PEP 562): importing this package -- or a +dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull +the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML +backend and its Studio dependencies). Those load only when a public name is +actually accessed, so standalone helpers stay unit-testable without the full +inference stack. +""" + +from typing import TYPE_CHECKING + +__all__ = [ + "InferenceBackend", + "InferenceOrchestrator", + "get_inference_backend", + "LlamaCppBackend", +] + +# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator. +_LAZY_ATTRS = { + "InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"), + "InferenceBackend": ("orchestrator", "InferenceOrchestrator"), + "get_inference_backend": ("orchestrator", "get_inference_backend"), + "LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"), +} + + +def __getattr__(name): + try: + submodule, attr = _LAZY_ATTRS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + from importlib import import_module + + value = getattr(import_module(f"{__name__}.{submodule}"), attr) + globals()[name] = value # cache so later access skips __getattr__ + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) + + +if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names + from .llama_cpp import LlamaCppBackend + from .orchestrator import InferenceOrchestrator, get_inference_backend + InferenceBackend = InferenceOrchestrator diff --git a/studio/backend/core/inference/_html_to_md.py b/studio/backend/core/inference/_html_to_md.py new file mode 100644 index 0000000..e7de4a5 --- /dev/null +++ b/studio/backend/core/inference/_html_to_md.py @@ -0,0 +1,415 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Minimal HTML-to-Markdown converter using only the standard library. + +Replaces the external ``html2text`` (GPL-3.0) dependency with a ~250-line +``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic, +lists, tables, blockquotes, code blocks, and entity decoding. +""" + +from __future__ import annotations + +import html +import re +from html.parser import HTMLParser + +__all__ = ["html_to_markdown"] + +_SKIP_TAGS = frozenset( + { + "script", + "style", + "head", + "noscript", + "svg", + "math", + "nav", + "footer", + } +) +_BLOCK_TAGS = frozenset( + { + "p", + "div", + "section", + "article", + "main", + "aside", + "figure", + "figcaption", + "details", + "summary", + "dl", + "dt", + "dd", + } +) +_HEADING_TAGS = frozenset({"h1", "h2", "h3", "h4", "h5", "h6"}) +_INLINE_EMPHASIS = {"strong": "**", "b": "**", "em": "*", "i": "*"} + + +class _MarkdownRenderer(HTMLParser): + """HTMLParser subclass that emits Markdown tokens into a list.""" + + def __init__(self): + super().__init__(convert_charrefs = False) + self._out: list[str] = [] + self._skip_depth: int = 0 + + # Link state + self._link_href: str | None = None + self._link_text_parts: list[str] = [] + self._in_link: bool = False + + # List state + self._list_stack: list[str] = [] # "ul" or "ol" + self._ol_counter: list[int] = [] + + # Table state + self._in_table: bool = False + self._current_row: list[str] = [] + self._cell_parts: list[str] = [] + self._in_cell: bool = False + self._header_row_done: bool = False + self._row_has_th: bool = False + self._is_first_row: bool = False + + # Pre/code state + self._in_pre: bool = False + self._pre_parts: list[str] = [] + self._in_inline_code: bool = False + + # Blockquote state: stack of buffers so nested blockquotes get the right ">" depth. + self._bq_stack: list[list[str]] = [] + + # ------------------------------------------------------------------ + def _emit(self, text: str) -> None: + if self._in_link: + self._link_text_parts.append(text) + elif self._in_cell: + self._cell_parts.append(text) + elif self._in_pre: + self._pre_parts.append(text) + elif self._bq_stack: + self._bq_stack[-1].append(text) + else: + self._out.append(text) + + # ------------------------------------------------------------------ + def _prefix_blockquote(self, content: str) -> str: + """Prefix every line of *content* with ``> ``.""" + # Strip trailing whitespace, then collapse blank lines. + content = re.sub(r"[ \t]+$", "", content, flags = re.MULTILINE) + content = re.sub(r"\n{3,}", "\n\n", content).strip() + if not content: + return "" + lines = content.split("\n") + prefixed: list[str] = [] + for line in lines: + if line.strip(): + prefixed.append("> " + line) + else: + prefixed.append(">") + return "\n".join(prefixed) + + # Table helpers: flush open cells/rows so omitted / don't lose data. + def _finish_cell(self) -> None: + if not self._in_cell: + return + self._in_cell = False + cell_text = "".join(self._cell_parts).strip().replace("\n", " ") + cell_text = cell_text.replace("|", "\\|") + self._current_row.append(cell_text) + self._cell_parts = [] + + def _finish_row(self) -> None: + if not self._current_row: + return + line = "| " + " | ".join(self._current_row) + " |" + self._emit(line + "\n") + if not self._header_row_done and (self._row_has_th or self._is_first_row): + sep = "| " + " | ".join("---" for _ in self._current_row) + " |" + self._emit(sep + "\n") + self._header_row_done = True + self._is_first_row = False + self._current_row = [] + self._row_has_th = False + + # Link text helper: normalize whitespace so block content in stays single-line. + def _finish_link(self) -> None: + text = re.sub(r"\s+", " ", "".join(self._link_text_parts)).strip() + href = self._link_href or "" + self._in_link = False + if href and text: + self._emit(f"[{text}]({href})") + elif text: + self._emit(text) + + # ------------------------------------------------------------------ + # Tag handlers + # ------------------------------------------------------------------ + def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: + tag = tag.lower() + + if tag in _SKIP_TAGS: + self._skip_depth += 1 + return + if self._skip_depth: + return + + attr_dict = dict(attrs) + + if tag in _HEADING_TAGS: + level = int(tag[1]) + self._emit("\n\n" + "#" * level + " ") + + elif tag == "a": + self._link_href = attr_dict.get("href") + self._link_text_parts = [] + self._in_link = True + + elif tag in _INLINE_EMPHASIS: + self._emit(_INLINE_EMPHASIS[tag]) + + elif tag == "br": + self._emit("\n") + + elif tag in _BLOCK_TAGS: + self._emit("\n\n") + + elif tag == "hr": + self._emit("\n\n---\n\n") + + elif tag == "blockquote": + self._emit("\n\n") + self._bq_stack.append([]) + + elif tag == "ul": + self._list_stack.append("ul") + self._emit("\n") + + elif tag == "ol": + self._list_stack.append("ol") + start_attr = attr_dict.get("start") + try: + start = int(start_attr) if start_attr is not None else 1 + except (ValueError, TypeError): + start = 1 + self._ol_counter.append(start - 1) + self._emit("\n") + + elif tag == "li": + indent = " " * max(0, len(self._list_stack) - 1) + if self._list_stack and self._list_stack[-1] == "ol": + if self._ol_counter: + self._ol_counter[-1] += 1 + self._emit(f"\n{indent}{self._ol_counter[-1]}. ") + else: + self._emit(f"\n{indent}1. ") + else: + self._emit(f"\n{indent}* ") + + elif tag == "pre": + self._pre_parts = [] + self._in_pre = True + + elif tag == "code" and not self._in_pre: + self._in_inline_code = True + self._emit("`") + + elif tag == "table": + self._in_table = True + self._header_row_done = False + self._is_first_row = True + self._emit("\n\n") + + elif tag == "tr": + # Flush open cell/row from a prior row that omitted /. + self._finish_cell() + self._finish_row() + + elif tag in ("th", "td"): + self._finish_cell() # handles omitted / + self._cell_parts = [] + self._in_cell = True + if tag == "th": + self._row_has_th = True + + elif tag == "img": + # Skip images: keeps text readable, avoids data-URI amplification. + return + + def handle_endtag(self, tag: str) -> None: + tag = tag.lower() + + if tag in _SKIP_TAGS: + self._skip_depth = max(0, self._skip_depth - 1) + return + if self._skip_depth: + return + + if tag in _HEADING_TAGS: + self._emit("\n\n") + + elif tag == "a": + self._finish_link() + + elif tag in _INLINE_EMPHASIS: + self._emit(_INLINE_EMPHASIS[tag]) + + elif tag in _BLOCK_TAGS: + self._emit("\n\n") + + elif tag == "blockquote": + if self._bq_stack: + content = "".join(self._bq_stack.pop()) + prefixed = self._prefix_blockquote(content) + if prefixed: + self._emit("\n\n" + prefixed + "\n\n") + + elif tag == "ul": + if self._list_stack and self._list_stack[-1] == "ul": + self._list_stack.pop() + self._emit("\n") + + elif tag == "ol": + if self._list_stack and self._list_stack[-1] == "ol": + self._list_stack.pop() + if self._ol_counter: + self._ol_counter.pop() + self._emit("\n") + + elif tag == "pre": + raw = "".join(self._pre_parts) + self._in_pre = False + block = "```\n" + raw + "\n```" + self._emit("\n\n" + block + "\n\n") + + elif tag == "code" and not self._in_pre: + self._in_inline_code = False + self._emit("`") + + elif tag in ("th", "td"): + self._finish_cell() + + elif tag == "tr": + self._finish_cell() + self._finish_row() + + elif tag == "table": + # Flush remaining row (handles omitted ). + self._finish_cell() + self._finish_row() + self._in_table = False + self._emit("\n") + + # ------------------------------------------------------------------ + # Text / entity handlers + # ------------------------------------------------------------------ + def handle_data(self, data: str) -> None: + if self._skip_depth: + return + if self._in_pre: + self._pre_parts.append(data) + return + # Preserve literal whitespace inside inline spans. + if self._in_inline_code: + self._emit(data) + return + # Collapse all whitespace (including newlines) per HTML rules. + text = re.sub(r"\s+", " ", data) + # Suppress whitespace-only nodes between table elements (source indentation). + if self._in_table and not self._in_cell and not text.strip(): + return + self._emit(text) + + def handle_entityref(self, name: str) -> None: + if self._skip_depth: + return + self._emit(html.unescape(f"&{name};")) + + def handle_charref(self, name: str) -> None: + if self._skip_depth: + return + self._emit(html.unescape(f"&#{name};")) + + # Flush pending buffers (handles truncated HTML from capped fetches) + def flush_pending(self) -> None: + """Flush open side-buffers into ``_out`` after close(), recovering truncated HTML.""" + # Flush innermost buffers first so their content propagates outward. + if self._in_link: + self._finish_link() + + if self._in_inline_code: + self._in_inline_code = False + self._emit("`") + + self._finish_cell() + self._finish_row() + + if self._in_pre: + raw = "".join(self._pre_parts) + self._in_pre = False + block = "```\n" + raw + "\n```" + self._emit("\n\n" + block + "\n\n") + + # Flatten any open blockquote buffers (innermost first). + while self._bq_stack: + content = "".join(self._bq_stack.pop()) + prefixed = self._prefix_blockquote(content) + if not prefixed: + continue + if self._bq_stack: + self._bq_stack[-1].append("\n\n" + prefixed + "\n\n") + else: + self._out.append("\n\n" + prefixed + "\n\n") + + +# Post-processing +def _cleanup(text: str) -> str: + """Normalize whitespace and blank lines, preserving fenced code blocks verbatim.""" + lines = text.split("\n") + out: list[str] = [] + in_fence = False + blank_run = 0 + + for line in lines: + stripped = line.rstrip(" \t") + if stripped.startswith("```"): + in_fence = not in_fence + blank_run = 0 + out.append(stripped) + continue + + if in_fence: + out.append(line) + continue + + if not stripped: + blank_run += 1 + if blank_run <= 1: + out.append("") + continue + + blank_run = 0 + out.append(stripped) + + return "\n".join(out).strip() + + +# Public API +def html_to_markdown(source_html: str) -> str: + """Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities). + + ``' + html = html_bytes.decode("utf-8") + html = html.replace("", f"{tag}", 1) + return html.encode("utf-8"), nonce + + +_DEFAULT_PORTS = {"http": 80, "https": 443, "ws": 80, "wss": 443} + + +def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]]: + """Canonicalise an Origin to ``(scheme, host, port)`` for equality. + Browsers strip default ports (RFC 6454 sec 6.1) and scheme/host are + case-insensitive (RFC 3986), so a bare string compare misclassifies + same-origin requests as cross-origin. Returns ``None`` on unparseable input + so callers fall to the safer cross-origin default. + """ + scheme = (scheme or "").strip().lower() + if not scheme or not netloc: + return None + # Strip userinfo (RFC 3986); Origin never carries credentials. + if "@" in netloc: + netloc = netloc.rsplit("@", 1)[1] + # IPv6 hosts use brackets (RFC 3986 sec 3.2.2): ``[::1]:8902``. Bare + # ``partition(":")`` mis-parses these, breaking ``unsloth studio -H ::1``. + if netloc.startswith("["): + close = netloc.find("]") + if close == -1: + return None + host = netloc[1:close] + rest = netloc[close + 1 :] + if rest.startswith(":"): + port_str = rest[1:] + elif rest == "": + port_str = "" + else: + return None + else: + host, _, port_str = netloc.partition(":") + host = host.strip().lower() + if not host: + return None + if port_str: + try: + port = int(port_str) + except ValueError: + return None + else: + port = _DEFAULT_PORTS.get(scheme, 0) + return (scheme, host, port) + + +def _is_same_origin_request(request: Request) -> bool: + """True when Origin is missing or matches request's scheme://host:port. + + Missing Origin counts as same-origin (top-level GETs omit it). Both sides + are canonicalised via :func:`_canonical_origin`; callers must emit + ``Vary: Origin``. + """ + origin = request.headers.get("origin") + if origin is None: + # Missing header: top-level same-document GETs omit Origin. + return True + # Empty string is not a valid serialised origin (RFC 6454 sec 6.1). + if not origin: + return False + # "null" token (sandboxed iframes, file:// pages) is never same-origin. + if origin == "null": + return False + # ``urlparse`` raises ``ValueError`` on malformed IPv6 brackets; swallow + # so a garbage Origin doesn't 500 the SPA handler. + try: + parsed = urlparse(origin) + except ValueError: + return False + origin_canon = _canonical_origin(parsed.scheme, parsed.netloc) + if origin_canon is None: + return False + try: + self_canon = _canonical_origin(request.url.scheme, request.url.netloc) + except ValueError: + return False + if self_canon is None: + return False + return origin_canon == self_canon + + +def setup_frontend(app: FastAPI, build_path: Path): + """Mount frontend static files (optional)""" + if not build_path.exists(): + return False + + assets_dir = build_path / "assets" + if assets_dir.exists(): + app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets") + + def _build_index_response(request: Request) -> Response: + content = (build_path / "index.html").read_bytes() + content = _strip_crossorigin(content) + # Bootstrap pw is same-origin only; Vary: Origin keeps caches honest. + if _is_same_origin_request(request): + content, nonce = _inject_bootstrap(content, app) + else: + nonce = None + headers = { + "Cache-Control": "no-cache, no-store, must-revalidate", + "Vary": "Origin", + } + if nonce: + headers[_CSP_SCRIPT_NONCE_HEADER] = nonce + return Response( + content = content, + media_type = "text/html", + headers = headers, + ) + + @app.get("/") + async def serve_root(request: Request): + return _build_index_response(request) + + @app.get("/{full_path:path}") + async def serve_frontend(request: Request, full_path: str): + # Unknown API paths: raise a real 404 so the api_errors handlers can + # render the correct envelope for /v1/* (and {"detail":...} for /api/*). + # This handler only sees paths NOT matched by a real route. The full + # request path is "/" + full_path. + if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")): + raise HTTPException(status_code = 404, detail = "API endpoint not found") + + file_path = (build_path / full_path).resolve() + + # Block path traversal — resolved path must stay inside build_path + if not file_path.is_relative_to(build_path.resolve()): + return Response(status_code = 403) + + if file_path.is_file(): + return FileResponse(file_path) + + # Serve index.html as bytes — avoids Content-Length mismatch + return _build_index_response(request) + + return True diff --git a/studio/backend/models/.gitkeep b/studio/backend/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py new file mode 100644 index 0000000..dbedeec --- /dev/null +++ b/studio/backend/models/__init__.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic models for API request/response schemas.""" + +from .training import ( + TrainingStartRequest, + TrainingJobResponse, + TrainingStatus, + TrainingProgress, + TrainingRunSummary, + TrainingRunListResponse, + TrainingRunMetrics, + TrainingRunDetailResponse, + TrainingRunDeleteResponse, + TrainingRunUpdateRequest, +) +from .models import ( + CheckpointInfo, + ModelCheckpoints, + CheckpointListResponse, + ModelDetails, + LocalModelInfo, + LocalModelListResponse, + LoRAInfo, + LoRAScanResponse, + ModelListResponse, +) +from .auth import ( + AuthLoginRequest, + RefreshTokenRequest, + AuthStatusResponse, + ChangePasswordRequest, +) +from .export import ( + LoadCheckpointRequest, + ExportStatusResponse, + ExportOperationResponse, + ExportMergedModelRequest, + ExportBaseModelRequest, + ExportGGUFRequest, + ExportLoRAAdapterRequest, +) +from .users import Token +from .datasets import ( + CheckFormatRequest, + CheckFormatResponse, +) +from .inference import ( + LoadRequest, + UnloadRequest, + GenerateRequest, + LoadResponse, + UnloadResponse, + InferenceStatusResponse, +) +from .responses import ( + TrainingStopResponse, + TrainingMetricsResponse, + LoRABaseModelResponse, + VisionCheckResponse, + EmbeddingCheckResponse, +) +from .data_recipe import ( + RecipePayload, + PreviewResponse, + ValidateError, + ValidateResponse, + JobCreateResponse, +) + +__all__ = [ + # Training schemas + "TrainingStartRequest", + "TrainingJobResponse", + "TrainingStatus", + "TrainingProgress", + "TrainingRunSummary", + "TrainingRunListResponse", + "TrainingRunMetrics", + "TrainingRunDetailResponse", + "TrainingRunDeleteResponse", + "TrainingRunUpdateRequest", + # Model management schemas + "ModelDetails", + "LocalModelInfo", + "LocalModelListResponse", + "LoRAInfo", + "LoRAScanResponse", + "ModelListResponse", + # Auth schemas + "AuthLoginRequest", + "RefreshTokenRequest", + "AuthStatusResponse", + "ChangePasswordRequest", + # Export schemas + "CheckpointInfo", + "ModelCheckpoints", + "CheckpointListResponse", + "LoadCheckpointRequest", + "ExportStatusResponse", + "ExportOperationResponse", + "ExportMergedModelRequest", + "ExportBaseModelRequest", + "ExportGGUFRequest", + "ExportLoRAAdapterRequest", + "Token", + # Dataset schemas + "CheckFormatRequest", + "CheckFormatResponse", + # Inference schemas + "LoadRequest", + "UnloadRequest", + "GenerateRequest", + "LoadResponse", + "UnloadResponse", + "InferenceStatusResponse", + # Response schemas + "TrainingStopResponse", + "TrainingMetricsResponse", + "LoRABaseModelResponse", + "VisionCheckResponse", + "EmbeddingCheckResponse", + # Data recipe + "RecipePayload", + "PreviewResponse", + "ValidateError", + "ValidateResponse", + "JobCreateResponse", +] diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py new file mode 100644 index 0000000..f451f8d --- /dev/null +++ b/studio/backend/models/auth.py @@ -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 + +"""Pydantic schemas for the Authentication API.""" + +from typing import Optional + +from pydantic import BaseModel, Field + + +class AuthLoginRequest(BaseModel): + """Login payload: username/password to obtain a JWT.""" + + username: str = Field(..., description = "Username") + 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.""" + + refresh_token: str = Field(..., description = "Refresh token from a previous login or refresh") + + +class AuthStatusResponse(BaseModel): + """Indicate whether the seeded admin auth flow is ready.""" + + initialized: bool = Field(..., description = "True if the auth database contains a login user") + default_username: str = Field( + "unsloth", + description = "Default admin username for first-boot UI prefill.", + ) + requires_password_change: bool = Field( + ..., + description = "True if the seeded admin must still change the default password", + ) + + +class ChangePasswordRequest(BaseModel): + """Change the current user's password, typically on first login.""" + + current_password: str = Field( + ..., min_length = 8, description = "Existing password for the authenticated user" + ) + 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] diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py new file mode 100644 index 0000000..e6f27e6 --- /dev/null +++ b/studio/backend/models/data_recipe.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic schemas for Data Recipe (DataDesigner) API.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field, model_validator + + +class RecipePayload(BaseModel): + recipe: dict[str, Any] = Field(default_factory = dict) + run: dict[str, Any] | None = None + ui: dict[str, Any] | None = None + + +class PreviewResponse(BaseModel): + dataset: list[dict[str, Any]] = Field(default_factory = list) + processor_artifacts: dict[str, Any] | None = None + analysis: dict[str, Any] | None = None + + +class ValidateError(BaseModel): + message: str + path: str | None = None + code: str | None = None + + +class ValidateResponse(BaseModel): + valid: bool + errors: list[ValidateError] = Field(default_factory = list) + raw_detail: str | None = None + + +class JobCreateResponse(BaseModel): + job_id: str + + +class PublishDatasetRequest(BaseModel): + repo_id: str = Field(min_length = 3, description = "Hugging Face dataset repo ID") + description: str = Field( + min_length = 1, + max_length = 4000, + description = "Short dataset description for the dataset card", + ) + hf_token: str | None = Field( + default = None, + description = "Optional Hugging Face token for private or write-protected repos", + ) + private: bool = Field( + default = False, + description = "Create or update the dataset repo as private", + ) + artifact_path: str | None = Field( + default = None, + description = "Execution artifact path captured by the UI for completed runs", + ) + + +class PublishDatasetResponse(BaseModel): + success: bool = True + url: str + message: str + + +class SeedInspectRequest(BaseModel): + dataset_name: str = Field(min_length = 1) + hf_token: str | None = None + subset: str | None = None + split: str | None = "train" + preview_size: int = Field(default = 10, ge = 1, le = 50) + + +class SeedInspectUploadRequest(BaseModel): + # Legacy single-file flow (mutually exclusive with file_ids) + filename: str | None = None + content_base64: str | None = None + # Multi-file flow (mutually exclusive with content_base64) + block_id: str | None = None + file_ids: list[str] | None = None + file_names: list[str] | None = None + # Shared fields + preview_size: int = Field(default = 10, ge = 1, le = 50) + seed_source_type: str | None = None + unstructured_chunk_size: int | None = Field(default = None, ge = 1, le = 20000) + unstructured_chunk_overlap: int | None = Field(default = None, ge = 0, le = 20000) + + @model_validator(mode = "after") + def _check_mutual_exclusivity(self) -> "SeedInspectUploadRequest": + has_legacy = self.content_base64 is not None + has_multi = self.file_ids is not None + if has_legacy and has_multi: + raise ValueError("Provide either content_base64 or file_ids, not both") + if not has_legacy and not has_multi: + raise ValueError("Provide either content_base64 or file_ids") + if has_multi: + if len(self.file_ids) == 0: + raise ValueError("file_ids must not be empty") + if not self.block_id: + raise ValueError("block_id is required when using file_ids") + if self.file_names is None or len(self.file_ids) != len(self.file_names): + raise ValueError("file_names must be provided and same length as file_ids") + if has_legacy: + if not self.filename: + raise ValueError("filename is required when using content_base64") + return self + + +class SeedInspectResponse(BaseModel): + dataset_name: str + resolved_path: str + columns: list[str] = Field(default_factory = list) + preview_rows: list[dict[str, Any]] = Field(default_factory = list) + split: str | None = None + subset: str | None = None + resolved_paths: list[str] | None = None + + +class UnstructuredFileUploadResponse(BaseModel): + file_id: str + filename: str + size_bytes: int + status: str # "ok" or "error" + error: str | None = None + + +class McpToolsListRequest(BaseModel): + mcp_providers: list[dict[str, Any]] = Field(default_factory = list) + timeout_sec: float | None = Field(default = None, gt = 0) + + +class McpToolsProviderResult(BaseModel): + name: str + tools: list[str] = Field(default_factory = list) + error: str | None = None + + +class McpToolsListResponse(BaseModel): + providers: list[McpToolsProviderResult] = Field(default_factory = list) + duplicate_tools: dict[str, list[str]] = Field(default_factory = dict) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py new file mode 100644 index 0000000..b40a841 --- /dev/null +++ b/studio/backend/models/datasets.py @@ -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 + +"""Dataset Pydantic models for API requests and responses.""" + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field, model_validator + + +class CheckFormatRequest(BaseModel): + """Request for dataset format check""" + + dataset_name: str # HuggingFace dataset name or local path + is_vlm: bool = False + hf_token: Optional[str] = None + subset: Optional[str] = None + train_split: Optional[str] = "train" + + @model_validator(mode = "before") + @classmethod + def _compat_split(cls, values: Any) -> Any: + """Accept legacy 'split' field as alias for 'train_split'.""" + if isinstance(values, dict) and "split" in values: + values.setdefault("train_split", values.pop("split")) + return values + + +class CheckFormatResponse(BaseModel): + """Response for dataset format check""" + + requires_manual_mapping: bool + detected_format: str + columns: List[str] + is_image: bool = False + is_audio: bool = False + multimodal_columns: Optional[List[str]] = None + suggested_mapping: Optional[Dict[str, str]] = None + detected_image_column: Optional[str] = None + detected_audio_column: Optional[str] = None + detected_text_column: Optional[str] = None + detected_speaker_column: Optional[str] = None + chat_column: Optional[str] = None + preview_samples: Optional[List[Dict]] = None + total_rows: Optional[int] = None + warning: Optional[str] = None + + +class AiAssistMappingRequest(BaseModel): + """Request for LLM-assisted column classification (user-triggered).""" + + columns: List[str] + samples: List[Dict[str, Any]] # Preview rows already loaded in the dialog + dataset_name: Optional[str] = None # For LLM context + hf_token: Optional[str] = None # For fetching dataset card + model_name: Optional[str] = None + model_type: Optional[str] = None + + +class AiAssistMappingResponse(BaseModel): + """Response from LLM-assisted column classification and conversion advice.""" + + success: bool + suggested_mapping: Optional[Dict[str, str]] = None + warning: Optional[str] = None + # Conversion advisor fields + system_prompt: Optional[str] = None + label_mapping: Optional[Dict[str, Dict[str, str]]] = None + dataset_type: Optional[str] = None + is_conversational: Optional[bool] = None + user_notification: Optional[str] = None + + +class UploadDatasetResponse(BaseModel): + """Response with stored dataset path for training.""" + + filename: str = Field(..., description = "Original filename") + stored_path: str = Field(..., description = "Absolute path stored on backend") + + +class LocalDatasetItem(BaseModel): + class Metadata(BaseModel): + actual_num_records: Optional[int] = None + target_num_records: Optional[int] = None + total_num_batches: Optional[int] = None + num_completed_batches: Optional[int] = None + columns: Optional[List[str]] = None + + id: str + label: str + path: str + rows: Optional[int] = None + updated_at: Optional[float] = None + metadata: Optional[Metadata] = None + + +class LocalDatasetsResponse(BaseModel): + datasets: List[LocalDatasetItem] = Field(default_factory = list) diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py new file mode 100644 index 0000000..9dc4d94 --- /dev/null +++ b/studio/backend/models/export.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic schemas for Export API.""" + +from pathlib import Path, PureWindowsPath + +from pydantic import BaseModel, Field, field_validator +from typing import List, Optional, Literal, Dict, Any, Union + + +def _validate_save_directory(value: str) -> str: + """Validate save_directory — allows absolute paths (user may want a different drive).""" + if value is None: + raise ValueError("save_directory is required") + raw = str(value).strip() + if not raw: + raise ValueError("save_directory must not be empty") + if "\x00" in raw: + raise ValueError("save_directory may not contain null bytes") + if any(ch in raw for ch in ("\r", "\n")): + raise ValueError("save_directory may not contain control characters") + path = Path(raw).expanduser() + path_parts = (*path.parts, *PureWindowsPath(raw).parts, *raw.replace("\\", "/").split("/")) + if any(len(part) > 255 for part in path_parts if part not in ("", ".", "/", "\\")): + raise ValueError("save_directory path components must be <= 255 characters") + if ( + ".." in path.parts + or ".." in PureWindowsPath(raw).parts + or ".." in raw.replace("\\", "/").split("/") + ): + raise ValueError("save_directory may not contain '..' segments") + return raw + + +class LoadCheckpointRequest(BaseModel): + """Request for loading a checkpoint into the export backend.""" + + checkpoint_path: str = Field(..., description = "Path to the checkpoint directory") + max_seq_length: int = Field( + 2048, + ge = 128, + le = 32768, + description = "Maximum sequence length for loading the model", + ) + load_in_4bit: bool = Field( + True, + description = "Whether to load the model in 4-bit quantization", + ) + trust_remote_code: bool = Field( + False, + description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.", + ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) + hf_token: Optional[str] = Field( + None, + description = "Hugging Face token used to scan/load gated checkpoints and their base models.", + ) + + +class ExportStatusResponse(BaseModel): + """Current export backend status.""" + + current_checkpoint: Optional[str] = Field( + None, + description = "Path to the currently loaded checkpoint, if any", + ) + is_vision: bool = Field( + False, + description = "True if the loaded checkpoint is a vision model", + ) + is_peft: bool = Field( + False, + description = "True if the loaded checkpoint is a PEFT (LoRA) model", + ) + is_export_active: bool = Field( + False, + description = "True while a load / export / cleanup operation is running", + ) + # Recovery fields: when a blocking export POST is cut off by a Cloudflare tunnel + # timeout (524 at ~100s), the client polls this endpoint to learn the real + # outcome of the operation that kept running on the backend. + active_op_kind: Optional[str] = Field( + None, + description = "Kind of the currently running op (load_checkpoint / export_* / cleanup)", + ) + last_op_seq: int = Field( + 0, + description = "Monotonic counter of finished ops; client baseline to detect 'my op finished'", + ) + last_op_kind: Optional[str] = Field( + None, + description = "Kind of the most recently finished op", + ) + last_op_status: Optional[str] = Field( + None, + description = "Outcome of the most recently finished op: success / error / cancelled", + ) + last_op_output_path: Optional[str] = Field( + None, + description = "Output path of the most recently finished op, if it produced one", + ) + last_op_error: Optional[str] = Field( + None, + description = "Error message of the most recently finished op, if it failed", + ) + + +class ExportOperationResponse(BaseModel): + """Generic response for export operations.""" + + success: bool = Field(..., description = "True if the operation succeeded") + message: str = Field(..., description = "Human-readable status or error message") + details: Optional[Dict[str, Any]] = Field( + default = None, + description = "Optional extra details about the operation", + ) + + +class ExportCommonOptions(BaseModel): + """Common options for export operations that save locally and/or push to Hub.""" + + save_directory: str = Field( + ..., + description = "Local directory where the exported artifacts will be written", + ) + + @field_validator("save_directory", mode = "before") + @classmethod + def _check_save_directory(cls, v): + return _validate_save_directory(v) + + push_to_hub: bool = Field( + False, + description = "If True, also push the exported model to the Hugging Face Hub", + ) + repo_id: Optional[str] = Field( + None, + description = "Hugging Face Hub repository ID (username/model-name)", + ) + hf_token: Optional[str] = Field( + None, + description = "Hugging Face access token used for Hub operations", + ) + private: bool = Field( + False, + description = "If True, create a private repository on the Hub (where applicable)", + ) + base_model_id: Optional[str] = Field( + None, + description = "HuggingFace model ID of the base model (for model card metadata)", + ) + + +class ExportMergedModelRequest(ExportCommonOptions): + """Request for exporting a merged PEFT model.""" + + format_type: Literal[ + "16-bit (FP16)", + "4-bit (FP4)", + "FP8 (compressed-tensors)", + "NVFP4 (compressed-tensors)", + ] = Field( + "16-bit (FP16)", + description = "Export precision / format for the merged model. The compressed-tensors " + "options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).", + ) + compressed_method: Optional[str] = Field( + None, + description = "Optional quantized-export alias. Either a compressed-tensors scheme " + "(e.g. 'fp8', 'fp8_static', 'w8a8', 'w4a16', 'mxfp4', 'mxfp8', 'nvfp4' - NVIDIA only) " + "from unsloth.save COMPRESSED_EXPORT_SCHEMES, or a portable torchao alias " + "('torchao_fp8', 'torchao_int8') from TORCHAO_EXPORT_SCHEMES that needs no NVIDIA GPU. " + "When set, it overrides format_type. Lets the export UI expose the full set of formats " + "beyond the quick buttons.", + ) + + +class ExportBaseModelRequest(ExportCommonOptions): + """Request for exporting a non-PEFT (base) model.""" + + # Uses fields from ExportCommonOptions only + + +class ExportGGUFRequest(BaseModel): + """Request for exporting the current model to GGUF format.""" + + save_directory: str = Field( + ..., + description = "Directory where GGUF files will be saved", + ) + + @field_validator("save_directory", mode = "before") + @classmethod + def _check_save_directory(cls, v): + return _validate_save_directory(v) + + quantization_method: Union[str, List[str]] = Field( + "Q4_K_M", + description = 'GGUF quantization method(s). A single method (e.g. "Q4_K_M") or a list ' + '(e.g. ["Q4_K_M", "Q8_0"]) to produce multiple GGUFs from one model load.', + ) + push_to_hub: bool = Field( + False, + description = "If True, also push GGUF artifacts to the Hugging Face Hub", + ) + repo_id: Optional[str] = Field( + None, + description = "Hugging Face Hub repository ID for GGUF upload", + ) + hf_token: Optional[str] = Field( + None, + description = "Hugging Face token for GGUF upload", + ) + imatrix: bool = Field( + False, + description = "Use an importance matrix (auto-downloads the upstream unsloth GGUF " + "imatrix). Required for the IQ low-bit quants such as iq2_xxs / iq4_xs.", + ) + imatrix_path: Optional[str] = Field( + None, + description = "Path to a custom imatrix file; overrides the auto-download when set.", + ) + + +class ExportLoRAAdapterRequest(ExportCommonOptions): + """Request for exporting only the LoRA adapter (not merged).""" + + gguf: bool = Field( + False, + description = "If True, also convert the adapter to a GGUF LoRA file " + "(llama.cpp convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.", + ) + gguf_outtype: Literal["q8_0", "f16", "bf16", "f32"] = Field( + "q8_0", + description = "GGUF LoRA output float type (only used when gguf=True). " + "Q8_0 falls back to F16 per tensor for dims not divisible by the block size (32).", + ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py new file mode 100644 index 0000000..53b0f14 --- /dev/null +++ b/studio/backend/models/inference.py @@ -0,0 +1,1765 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic schemas for the Inference API.""" + +from __future__ import annotations + +import time +import uuid +from typing import Annotated, Any, Dict, Literal, Optional, List, Union + +from pydantic import ( + BaseModel, + Discriminator, + Field, + Tag, + field_validator, + 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") + max_seq_length: int = Field( + 0, + ge = 0, + le = 1048576, + description = "Maximum sequence length (0 = model default for GGUF)", + ) + load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") + is_lora: bool = Field(False, description = "Whether this is a LoRA adapter") + gguf_variant: Optional[str] = Field( + None, description = "GGUF quantization variant (e.g. 'Q4_K_M')" + ) + trust_remote_code: bool = Field( + False, + description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", + ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) + chat_template_override: Optional[str] = Field( + None, + description = "Custom Jinja2 chat template to use instead of the model's default", + ) + + @field_validator("chat_template_override") + @classmethod + def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]: + if value is not None and value.strip() == "": + return None + return value + + cache_type_kv: Optional[str] = Field( + None, + description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')", + ) + gpu_ids: Optional[List[int]] = Field( + 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. Canonical values: " + "'auto' (platform-aware: MTP on MTP GGUFs, ngram-mod fallback " + "for sub-3B), 'mtp' (force draft-mtp only on both GPU and CPU), " + "'ngram' (force ngram-mod only), 'mtp+ngram' (force " + "ngram-mod+draft-mtp chain on both platforms), 'off' (disabled). " + "Legacy values 'default' (-> auto), 'draft-mtp' (-> mtp), " + "'ngram-mod' (-> ngram), and 'ngram-simple' (kept as-is) are " + "still accepted. Ignored for non-GGUF models." + ), + ) + spec_draft_n_max: Optional[int] = Field( + None, + ge = 1, + le = 16, + description = ( + "Max draft tokens per step for MTP speculative decoding " + "(--spec-draft-n-max). Defaults to 2 on GPU and 3 on CPU/Mac " + "when unset (upstream-bench sweet spot for dense Qwen3.6 MTP " + "quants). Only applied when speculative_type resolves to " + "'mtp' or 'mtp+ngram'." + ), + ) + tensor_parallel: bool = Field( + False, + description = ( + "Split the model across GPUs by tensor (--split-mode tensor) " + "instead of by layer for GGUF models. Only affects multi-GPU " + "setups, where it can make generation significantly faster. " + "No effect on a single GPU. Ignored for non-GGUF 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, UI/server mode) are rejected. Ignored for non-GGUF models." + ), + ) + + +class UnloadRequest(BaseModel): + """Request to unload a model""" + + model_path: str = Field(..., description = "Model identifier to unload") + + +class ValidateModelRequest(BaseModel): + """Check whether an identifier resolves to a ModelConfig; does NOT load weights.""" + + 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") + gguf_variant: Optional[str] = Field( + None, description = "GGUF quantization variant (e.g. 'Q4_K_M')" + ) + # Intended load settings so validate's coexistence check matches the follow-up + # /load; defaults preserve old behavior for callers that omit them. + max_seq_length: int = Field(0, ge = 0, le = 1048576) + load_in_4bit: bool = Field(True) + gpu_ids: Optional[List[int]] = Field(None) + include_context_length: bool = Field( + False, + description = "Also read the native context length from the local GGUF header. " + "Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.", + ) + + +class ValidateModelResponse(BaseModel): + """Result of model validation. + + valid == True means from_identifier() succeeded and GGUF/LoRA/vision flags are available. + """ + + valid: bool = Field(..., description = "Whether the model identifier looks valid") + message: str = Field(..., description = "Human-readable validation message") + identifier: Optional[str] = Field(None, description = "Resolved model identifier") + display_name: Optional[str] = Field(None, description = "Display name derived from identifier") + 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.", + ) + requires_security_review: bool = Field( + False, + description = "Whether Hugging Face's security scan flagged unsafe files (e.g. a " + "malicious pickle), so the load is hard-blocked pending review.", + ) + context_length: Optional[int] = Field( + None, + description = "Native training context length, read from the GGUF header when the file " + "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", + ) + + +class GenerateRequest(BaseModel): + """Request for text generation (legacy /generate/stream endpoint)""" + + messages: List[dict] = Field(..., description = "Chat messages in OpenAI format") + system_prompt: str = Field("", description = "System prompt") + temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature") + top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling") + top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling") + min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p sampling") + max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate") + repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty") + presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty") + image_base64: Optional[str] = Field(None, description = "Base64 encoded image for vision models") + + +class LoadResponse(BaseModel): + """Response after loading a model""" + + status: str = Field(..., description = "Load status") + model: str = Field(..., description = "Model identifier") + display_name: str = Field(..., description = "Display name of the model") + is_vision: bool = Field(False, description = "Whether model is a vision model") + is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") + is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether model is a block-diffusion model (DiffusionGemma)" + ) + is_audio: bool = Field(False, description = "Whether model is a TTS audio model") + audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") + has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") + 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 = "Runtime context length in tokens for the loaded model" + ) + max_context_length: Optional[int] = Field( + None, description = "Maximum context length currently available on this hardware" + ) + native_context_length: Optional[int] = Field( + None, + description = "Model's native context length from GGUF metadata (not capped by VRAM)", + ) + supports_reasoning: bool = Field( + False, + description = "Whether model supports thinking/reasoning mode (enable_thinking or reasoning_effort)", + ) + reasoning_style: Literal["enable_thinking", "reasoning_effort", "enable_thinking_effort"] = ( + Field( + "enable_thinking", + description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)", + ) + ) + reasoning_effort_levels: List[str] = Field( + default_factory = list, + description = "Discrete reasoning_effort levels the template offers when reasoning_style is 'enable_thinking_effort' (e.g. ['high', 'max']); empty otherwise", + ) + reasoning_always_on: bool = Field( + False, + description = "Whether reasoning is always on (hardcoded 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.)", + ) + cache_type_kv: Optional[str] = Field( + None, + description = "KV cache data type for K and V (e.g. 'f16', 'bf16', 'q8_0')", + ) + chat_template: Optional[str] = Field( + None, + description = "Jinja2 chat template string (from GGUF metadata or tokenizer)", + ) + speculative_type: Optional[str] = Field( + None, + description = ( + "Canonical UI-facing requested speculative decoding mode " + "('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / " + "'ngram-simple'), round-tripped from the original LoadRequest " + "via _canonicalize_spec_mode. None when no model is loaded." + ), + ) + spec_draft_n_max: Optional[int] = Field( + None, + description = ( + "Active --spec-draft-n-max for MTP speculative decoding, or " + "None when the platform default is in effect." + ), + ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) + + +class UnloadResponse(BaseModel): + """Response after unloading a model""" + + status: str = Field(..., description = "Unload status") + model: str = Field(..., description = "Model identifier that was unloaded") + + +class LoadProgressResponse(BaseModel): + """Progress of the active GGUF load, sampled on demand. + + Drives a real progress bar during the post-download warmup (mmap + CUDA upload) + instead of a 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""" + + active_model: Optional[str] = Field( + None, description = "Currently active model display identifier" + ) + model_identifier: Optional[str] = Field( + None, + description = "Loadable identifier for the active model.", + ) + is_vision: bool = Field(False, description = "Whether the active model is a vision model") + is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)" + ) + gguf_variant: Optional[str] = Field(None, description = "GGUF quantization variant (e.g. Q4_K_M)") + is_audio: bool = Field(False, description = "Whether the active model is a TTS audio model") + audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") + has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") + loading: List[str] = Field(default_factory = list, description = "Models currently being loaded") + loaded: List[str] = Field(default_factory = list, description = "Models currently loaded") + 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", "enable_thinking_effort"] = ( + Field( + "enable_thinking", + description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)", + ) + ) + reasoning_effort_levels: List[str] = Field( + default_factory = list, + description = "Discrete reasoning_effort levels the template offers when reasoning_style is 'enable_thinking_effort' (e.g. ['high', 'max']); empty otherwise", + ) + 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" + ) + context_length: Optional[int] = Field(None, description = "Context length of the active model") + max_context_length: Optional[int] = Field( + None, + description = "Maximum context length currently available for the active model", + ) + native_context_length: Optional[int] = Field( + None, + description = "Model's native context length from GGUF metadata (not capped by VRAM)", + ) + cache_type_kv: Optional[str] = Field( + None, + description = "KV cache quantization dtype (e.g. 'q8_0'), or None for default", + ) + chat_template: Optional[str] = Field( + None, description = "Model's default chat template (Jinja2 source), if any" + ) + chat_template_override: Optional[str] = Field( + None, + description = "Active chat template override applied at load time, or None if model is using its default", + ) + speculative_type: Optional[str] = Field( + None, + description = ( + "Canonical UI-facing requested speculative decoding mode " + "('auto' / 'mtp' / 'ngram' / 'mtp+ngram' / 'off' / " + "'ngram-simple'), round-tripped from the original LoadRequest. " + "None when no model is loaded." + ), + ) + spec_draft_n_max: Optional[int] = Field( + None, + description = ( + "Active --spec-draft-n-max for MTP speculative decoding, or " + "None when the platform default is in effect." + ), + ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) + llama_cpp_supports_mtp: bool = Field( + True, + description = ( + "Whether llama.cpp supports MTP (--spec-type mtp/draft-mtp). " + "False -> recommend `unsloth studio update`." + ), + ) + spec_fallback_reason: Optional[str] = Field( + None, + description = ( + "Why MTP was disabled on the loaded model despite being requested " + "(auto on an MTP model, or forced mtp / mtp+ngram). " + "'binary_no_mtp' / 'binary_outdated' -> a newer prebuilt would " + "re-enable it (show the update affordance); 'runtime_error' -> the " + "current build could not run it; 'drafter_not_found' -> the model's " + "separate MTP drafter could not be resolved; 'mla_mtp_disabled' -> " + "an Auto-mode policy downgrade: the model is MLA (GLM-5.2 et al.) " + "whose llama.cpp MTP path runs slower than no speculation, so Auto " + "used ngram-mod or spec-off instead -- updating won't help; choose " + "MTP in Settings (or set UNSLOTH_MLA_MTP_ENABLED=1) to force it. " + "None when MTP engaged or was not requested." + ), + ) + llama_cpp_prebuilt_stale: bool = Field( + False, + description = ( + "Installed llama.cpp prebuilt is >=3 days behind the latest " + "release. True -> show `unsloth studio update` banner." + ), + ) + llama_cpp_installed_tag: Optional[str] = Field( + None, + description = "Installed llama.cpp tag, or None if unknown.", + ) + llama_cpp_latest_tag: Optional[str] = Field( + None, + description = "Latest published llama.cpp tag, or None if GitHub unreachable.", + ) + + +# ===================================================================== +# OpenAI-Compatible Chat Completions Models +# ===================================================================== + + +# ── Multimodal content parts (OpenAI vision format) ────────────── + + +class TextContentPart(BaseModel): + """Text content part in a multimodal message.""" + + type: Literal["text"] + text: str + + +class ImageUrl(BaseModel): + """Image URL object — supports data URIs and remote URLs.""" + + url: str = Field(..., description = "data:image/png;base64,... or https://...") + detail: Optional[Literal["auto", "low", "high", "original"]] = "auto" + + +class ImageContentPart(BaseModel): + """Image content part in a multimodal message.""" + + type: Literal["image_url"] + image_url: ImageUrl + + +class InputDocumentContentPart(BaseModel): + """Document (PDF / file) content part in a multimodal message. + + Studio-normalised shape (file_data or file_url, plus optional filename/media_type). + Mapped onto Anthropic ``document`` / OpenAI ``input_file`` for vision providers; + dropped for non-vision providers. + """ + + type: Literal["input_document"] + file_data: Optional[str] = Field( + None, + description = "data:;base64, URI for inline payloads. Either file_data or file_url must be set; otherwise the part is dropped.", + ) + file_url: Optional[str] = Field( + None, + description = "Remote URL pointing to the document (https://...).", + ) + filename: Optional[str] = Field( + None, + description = "Display filename, forwarded to providers as `title`/`filename`.", + ) + media_type: Optional[str] = Field( + None, + description = 'Override the media type sniffed from the data URI (e.g. "application/pdf").', + ) + + +class OpenAIReasoningContentPart(BaseModel): + """OpenAI Responses reasoning item paired with a tool output. + + Reasoning models may require this replayed before an ``image_generation_call`` + id. OpenAI-only; routes strip it for other providers before proxying. + """ + + type: Literal["reasoning"] + id: str = Field(..., description = "OpenAI reasoning output item id.") + summary: list[dict[str, Any]] = Field(default_factory = list) + status: Optional[Literal["in_progress", "completed", "incomplete"]] = None + + +class ImageGenerationCallContentPart(BaseModel): + """OpenAI Responses image_generation call reference. + + Prior ``image_generation_call`` items let follow-up prompts edit a generated + image without resending the payload. The frontend forwards it as a synthetic + assistant part; ``external_provider`` maps it back to a top-level input item. + """ + + type: Literal["image_generation_call"] + id: str = Field(..., description = "OpenAI image_generation_call output item id.") + response_id: Optional[str] = Field( + None, + description = "OpenAI Responses response id to use as previous_response_id for follow-up edits.", + ) + + +class CompactionContentPart(BaseModel): + """Anthropic server-side compaction state, round-tripped on the next turn. + + Anthropic returns a ``compaction`` block on the assistant message; the next + request must forward it back so Anthropic reuses the compaction state instead + of re-summarising. See ``external_provider._stream_anthropic`` and + https://platform.claude.com/docs/en/build-with-claude/compaction + """ + + type: Literal["compaction"] + content: str = Field( + ..., + description = "Anthropic-produced summary of the compacted-away conversation prefix.", + ) + + +def _content_part_discriminator(v): + if isinstance(v, dict): + return v.get("type") + return getattr(v, "type", None) + + +ContentPart = Annotated[ + Union[ + Annotated[TextContentPart, Tag("text")], + Annotated[ImageContentPart, Tag("image_url")], + Annotated[InputDocumentContentPart, Tag("input_document")], + Annotated[OpenAIReasoningContentPart, Tag("reasoning")], + Annotated[ImageGenerationCallContentPart, Tag("image_generation_call")], + Annotated[CompactionContentPart, Tag("compaction")], + ], + Discriminator(_content_part_discriminator), +] +"""Union type for multimodal content parts, discriminated by the 'type' field.""" + + +# ── Messages ───────────────────────────────────────────────────── + + +class ChatMessage(BaseModel): + """Single message in a chat conversation. + + ``content`` is a string or list of multimodal parts. Assistant messages with + only ``tool_calls`` may set ``content=None``. Missing ``tool_call_id`` on + ``role="tool"`` is resolved at the ``ChatCompletionRequest`` layer. + """ + + role: Literal["system", "user", "assistant", "tool", "developer"] = Field( + ..., description = "Message role" + ) + 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.", + ) + extra_content: Optional[dict] = Field( + None, + description = ( + "Provider-specific extra fields the translator may read. " + "Gemini reads `extra_content.google.thought_signature` " + "from assistant messages to replay text-part signatures." + ), + ) + + @model_validator(mode = "after") + def _validate_role_shape(self) -> "ChatMessage": + 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.') + + if self.role == "tool": + # tool_call_id resolution happens at ChatCompletionRequest scope. + # OpenAI accepts empty tool results (commands with no output); + # normalize to "" instead of a 400 agentic clients treat as fatal. + if self.content is None or self.content == []: + self.content = "" + elif self.role == "assistant": + # Post-Stop sentinel: collapse content="" / [] to None. + if (self.content == "" or self.content == []) and not self.tool_calls: + self.content = None + else: # "user" | "system" + if self.content is None or self.content == []: + raise ValueError(f'role="{self.role}" messages require "content".') + return self + + +class ThinkingConfig(BaseModel): + """Anthropic-compatible thinking/reasoning configuration. + Use type='disabled' to turn off thinking, or type='enabled' to turn it on. + Only type is read; extra fields (e.g. budget_tokens) are ignored, since + Studio sets provider thinking budgets itself. + """ + + type: Literal["disabled", "enabled"] = "disabled" + + +class ChatCompletionRequest(BaseModel): + """OpenAI-compatible chat completion request. + + Non-OpenAI extension fields are marked with 'x-unsloth'. + """ + + # Accept unknown fields so future OpenAI fields aren't dropped 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( + 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': ...}}" + ), + ) + max_completion_tokens: Optional[int] = Field( + None, + ge = 1, + description = "OpenAI upper bound on generated tokens (supersedes the deprecated max_tokens).", + ) + n: Optional[int] = Field( + None, + ge = 1, + le = 128, + description = "Number of chat completion choices to generate.", + ) + logprobs: Optional[bool] = Field( + None, description = "Whether to return log probabilities of the output tokens." + ) + top_logprobs: Optional[int] = Field( + None, + ge = 0, + le = 20, + description = "Number of most likely tokens (0-20) to return per position; requires logprobs=true.", + ) + parallel_tool_calls: Optional[bool] = Field( + None, description = "Whether to enable parallel function calling during tool use." + ) + seed: Optional[int] = Field(None, description = "Best-effort deterministic sampling seed.") + stream_options: Optional[dict] = Field( + None, + description = 'Streaming options, e.g. {"include_usage": true} to emit a final usage chunk.', + ) + + # ── Unsloth extensions (ignored by standard OpenAI clients) ── + top_k: int = Field(20, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling") + min_p: float = Field(0.01, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold") + repetition_penalty: float = Field( + 1.0, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty" + ) + image_base64: Optional[str] = Field( + None, description = "[x-unsloth] Base64-encoded image for vision models" + ) + audio_base64: Optional[str] = Field( + None, + description = "[x-unsloth] Base64-encoded audio (wav/mp3/ogg/flac/m4a) for audio-input models", + ) + use_adapter: Optional[Union[bool, str]] = Field( + None, + description = ( + "[x-unsloth] Adapter control for compare mode. " + "null = no change (default), " + "false = disable adapters (base model), " + "true = enable the current adapter, " + "string = enable a specific adapter by name." + ), + ) + enable_thinking: Optional[bool] = Field( + None, + description = "[x-unsloth] Enable/disable thinking/reasoning mode for supported models", + ) + reasoning_effort: Optional[ + Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"] + ] = Field( + None, + description = "[x-unsloth] Reasoning effort level ('none'|'minimal'|'low'|'medium'|'high'|'max'|'xhigh'). OpenAI `/v1/responses` accepts model-dependent subsets; Anthropic adaptive thinking uses `max` as the top tier on Claude 4.6 Opus/Sonnet (inbound `xhigh` is mapped to `max`) and `xhigh` on Claude 4.7 Opus; local Harmony/gpt-oss templates support low|medium|high.", + ) + preserve_thinking: Optional[bool] = Field( + None, + description = "[x-unsloth] When true, keep historical blocks from past assistant turns in the prompt (Qwen3.6 templates). Independent of enable_thinking / reasoning_effort.", + ) + thinking: Optional[ThinkingConfig] = Field( + None, + description = "[Anthropic-compatible] Thinking configuration. " + "Use {type: 'disabled'} to disable thinking, {type: 'enabled'} to enable.", + ) + enable_tools: Optional[bool] = Field( + None, + description = "[x-unsloth] Enable tool calling for supported models", + ) + enabled_tools: Optional[list[str]] = Field( + None, + description = ( + "[x-unsloth] List of enabled tool names. Local GGUF/safetensors models " + "accept ['web_search', 'python', 'terminal', 'render_html']. External " + "providers accept ['web_search', 'web_fetch', 'code_execution'] for " + "Anthropic and ['web_search', 'code_execution', 'image_generation'] for " + "OpenAI Responses. If None, all local tools are enabled and no " + "server-side tools are forwarded." + ), + ) + mcp_enabled: Optional[bool] = Field( + None, + description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.", + ) + confirm_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.", + ) + bypass_permissions: Optional[bool] = Field( + False, + description = "[x-unsloth] Bypass Permissions: when true, skip the tool-call confirmation gate AND disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits). Secret env vars are still stripped. Takes precedence over confirm_tool_calls.", + ) + auto_heal_tool_calls: Optional[bool] = Field( + True, + description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", + ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the " + "model emitted a tool signal that healing could not repair, retry ONCE with " + "a short nudge appended (the retry shares the full prompt prefix, so the " + "server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips " + "the process default." + ), + ) + context_overflow: Optional[Literal["error", "truncate_middle"]] = Field( + None, + description = ( + "[x-unsloth] Passthrough behavior when the prompt exceeds the real " + "context window. 'error' (default) returns a 400 with " + "code=context_length_exceeded. 'truncate_middle' drops middle " + "turn-groups (system prompt, first turn, and recent turns kept; " + "tool calls stay paired with their results) and retries." + ), + ) + max_tool_calls_per_message: Optional[int] = Field( + 25, + ge = 0, + description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).", + ) + tool_call_timeout: Optional[int] = Field( + 300, + ge = 1, + description = "[x-unsloth] Timeout in seconds for each tool call execution (9999 = no limit).", + ) + session_id: Optional[str] = Field( + None, + description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.", + ) + rag_scope: Optional[dict] = Field( + None, + description = ( + "[x-unsloth] Hidden RAG retrieval scope for the search_knowledge_base " + "tool: {kb_id?, thread_id?, default_top_k?, mode?, autoinject?, " + "autoinject_min_score?}. Candidate pools and the RRF constant come from " + "server config. The model never sees this; the server resolves which " + "documents to search." + ), + ) + 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( + None, + description = "[x-unsloth] Saved provider config ID. If set with encrypted_api_key, routes to external LLM.", + ) + provider_type: Optional[str] = Field( + None, + description = "[x-unsloth] Provider type (e.g. 'openai', 'mistral'). Used if provider_id is not set.", + ) + external_model: Optional[str] = Field( + None, + description = "[x-unsloth] Model ID at the external provider.", + ) + encrypted_api_key: Optional[str] = Field( + None, + description = "[x-unsloth] RSA-encrypted, base64-encoded API key for the external provider.", + ) + provider_base_url: Optional[str] = Field( + None, + description = "[x-unsloth] Override base URL for the external provider.", + ) + enable_prompt_caching: Optional[Union[bool, str]] = Field( + None, + description = ( + "[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, " + "boolean true attaches cache_control={type:ephemeral} to the system " + "block so the static prefix is reused across turns. On OpenAI cloud, " + "caching is automatic for prompts >=1024 tokens and the boolean is " + "informational. On Gemini, pass a string cache resource name such " + "as `cachedContents/abc123` to attach `cachedContent` on the native " + "request (boolean true is a no-op on Gemini because creating the " + "cache requires a separate POST /cachedContents call). Ignored for " + "every other provider. Treated as enabled when omitted." + ), + ) + + @field_validator("enable_prompt_caching", mode = "before") + @classmethod + def _coerce_enable_prompt_caching(cls, value: Any) -> Any: + """Coerce JSON bool strings back to bool. Widening to Union[bool, str] for + Gemini cache names would let `"false"` read as truthy, so canonical bool + literals are coerced to keep explicit opt-outs working.""" + if isinstance(value, str): + lowered = value.strip().lower() + # Match Pydantic v1's bool coercion table; anything else stays a + # string for Gemini's cachedContent resource path. + if lowered in ("true", "t", "1", "yes", "y", "on"): + return True + if lowered in ("false", "f", "0", "no", "n", "off"): + return False + return value + + prompt_cache_ttl: Optional[str] = Field( + None, + description = ( + "[x-unsloth] Anthropic cache_control TTL. Defaults to the 5-minute " + "ephemeral pool when omitted. Pass `1h` to write into the 1-hour " + "pool instead -- 1h writes are billed at 2x base input vs 1.25x " + "for 5m, but reads stay at 0.1x for both, so 1h pays off the " + "moment a single extra read lands more than 5 minutes after the " + "write. Only `5m` and `1h` are forwarded; any other value is " + "silently ignored downstream so a stale frontend can't make the " + "API 422 on the request. No-op on every non-Anthropic provider." + ), + ) + compaction_threshold: Optional[int] = Field( + None, + ge = 1, + le = 2_000_000, + description = ( + "[x-unsloth] Server-side context compaction trigger, in tokens. " + "Per-provider routing:\n" + " - Anthropic (Opus 4.6+, Sonnet 4.6, Mythos preview): attaches " + "the `compact_20260112` edit and the `compact-2026-01-12` beta " + "header. The upstream floor is 50k; `_stream_anthropic` clamps " + "lower values up.\n" + " - OpenAI cloud (api.openai.com) and Azure OpenAI Foundry " + "(*.openai.azure.com): attaches " + "`context_management:[{type:'compaction', compact_threshold:N}]` " + "to /v1/responses. Effective floor is around 200k (OpenAI's " + "canonical example); values below it surface " + "`compact_threshold is not enabled` 400s upstream.\n" + "Schema floor stays at ge=1 (any positive int) so the field is a " + "silent no-op on non-cloud OpenAI-compatible bases (ollama / " + "llama.cpp / vLLM) and every non-compaction-capable provider " + "rather than returning 422 at request validation time. Per-" + "provider floors are enforced in the corresponding stream helpers." + ), + ) + openai_code_exec_container_id: Optional[str] = Field( + None, + description = ( + "[x-unsloth] OpenAI shell-tool container id from the prior response " + "in the same chat thread. When set and `code_execution` is in " + "`enabled_tools`, the next /v1/responses call uses " + "environment.type='container_reference' so filesystem state " + "persists across turns. Unset → environment.type='container_auto' " + "and OpenAI creates a fresh container. Only meaningful for the " + "OpenAI cloud + gpt-5.5 family path; ignored otherwise." + ), + ) + anthropic_code_exec_container_id: Optional[str] = Field( + None, + description = ( + "[x-unsloth] Anthropic code_execution container id from the prior " + "response in the same chat thread. When set and `code_execution` " + "is in `enabled_tools`, the next /v1/messages call carries a " + "top-level `container` field so the model sees filesystem state " + "from earlier turns. Unset → Anthropic auto-creates a fresh " + "container. Stale ids surface a 4xx with a `container_expired` / " + "`container_not_found` hint; the backend emits a synthetic " + "`container_invalidated` _toolEvent so the next turn falls back " + "to auto-create." + ), + ) + fast_mode: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Anthropic fast-mode toggle. On Claude Opus 4.6 / " + "4.7 adds the `fast-mode-2026-02-01` beta header and sends " + "`speed: 'fast'` for higher OTPS at premium pricing. Silently " + "ignored on every other model + provider. See " + "https://platform.claude.com/docs/en/build-with-claude/fast-mode" + ), + ) + + @model_validator(mode = "after") + def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest": + """Fill missing tool_call_id by walking back to the preceding assistant. + + OpenAI / Anthropic passthrough require the result id to match the + assistant's tool_calls[].id. Prefer function.name match, else first + unconsumed tool_call; synth a random id only if none exists. A user + turn breaks the lookup. + """ + # Pre-mark explicit ids so a missing-id sibling can't steal a claimed one. + consumed: set[tuple[int, int]] = set() + + def _mark_consumed(start_idx: int, tool_call_id: str) -> None: + for asst_idx in range(start_idx - 1, -1, -1): + prev = self.messages[asst_idx] + if prev.role == "user": + break + if prev.role != "assistant" or not prev.tool_calls: + continue + for tc_idx, tc in enumerate(prev.tool_calls): + if isinstance(tc, dict) and tc.get("id") == tool_call_id: + consumed.add((asst_idx, tc_idx)) + return + + for tool_idx, msg in enumerate(self.messages): + if msg.role == "tool" and msg.tool_call_id: + _mark_consumed(tool_idx, msg.tool_call_id) + + for tool_idx, msg in enumerate(self.messages): + if msg.role != "tool" or msg.tool_call_id: + continue + picked: str | None = None + for asst_idx in range(tool_idx - 1, -1, -1): + prev = self.messages[asst_idx] + if prev.role != "assistant" or not prev.tool_calls: + if prev.role == "user": + break + continue + name_match = None + fallback = None + for tc_idx, tc in enumerate(prev.tool_calls): + if (asst_idx, tc_idx) in consumed: + continue + if not isinstance(tc, dict): + continue + tc_id = tc.get("id") + if not tc_id: + continue + function = tc.get("function") + function_name = function.get("name") if isinstance(function, dict) else None + if msg.name and function_name == msg.name: + name_match = (tc_id, asst_idx, tc_idx) + break + if fallback is None: + fallback = (tc_id, asst_idx, tc_idx) + chosen = name_match or fallback + if chosen is not None: + picked, a, t = chosen + consumed.add((a, t)) + break + if picked is None: + import secrets as _secrets + picked = f"call_{_secrets.token_hex(8)}" + msg.tool_call_id = picked + return self + + @model_validator(mode = "after") + def _map_thinking_to_enable_thinking(self) -> "ChatCompletionRequest": + """Map Anthropic-style ``thinking`` parameter to internal ``enable_thinking``. + + ``thinking: {type: 'enabled'}`` sets ``enable_thinking = True`` and + ``thinking: {type: 'disabled'}`` sets ``enable_thinking = False``. + ``enable_thinking`` takes precedence when both are provided so that + callers who already use the internal field are unaffected. Invalid + ``thinking`` shapes are rejected at validation time (422). + """ + if self.thinking is not None and self.enable_thinking is None: + self.enable_thinking = self.thinking.type == "enabled" + return self + + +class ToolConfirmRequest(BaseModel): + session_id: Optional[str] = None + approval_id: Optional[str] = None + decision: Literal["allow", "deny"] = "deny" + + +# ── OpenAI shell-tool container management ───────────────────── + + +class OpenAIContainerRequest(BaseModel): + """Shared body for the OpenAI container endpoints (list / create / delete). + + Carries the encrypted API key + base URL so the route can decrypt and proxy + to the user's account, keeping the key off backend persistent storage. + """ + + encrypted_api_key: str = Field( + ..., + description = "[x-unsloth] RSA-encrypted, base64-encoded OpenAI API key.", + ) + provider_base_url: Optional[str] = Field( + None, + description = "[x-unsloth] OpenAI base URL. Only api.openai.com is supported; non-cloud bases are rejected with 400.", + ) + + +class CreateOpenAIContainerBody(OpenAIContainerRequest): + name: str = Field( + ..., + min_length = 1, + max_length = 256, + description = "Human-readable container name. Surfaces in the picker UI.", + ) + ttl_minutes: int = Field( + 20, + ge = 1, + le = 20, + description = ( + "Idle-timeout TTL the new container will inherit (anchor=" + "last_active_at). OpenAI hard-caps this at 20 minutes and " + "rejects larger values with integer_above_max_value." + ), + ) + + +class DeleteOpenAIContainerBody(OpenAIContainerRequest): + container_id: str = Field( + ..., + description = "OpenAI container id (cntr_...) to delete.", + ) + + +class OpenAIContainerSummary(BaseModel): + """One row from GET /v1/containers, reshaped for the UI.""" + + id: str + name: Optional[str] = None + created_at: Optional[int] = None + last_active_at: Optional[int] = None + expires_after_minutes: Optional[int] = None + status: Optional[str] = None + + +class ListOpenAIContainersResponse(BaseModel): + containers: list[OpenAIContainerSummary] + + +# ── Streaming response chunks ──────────────────────────────────── + + +class ChoiceDelta(BaseModel): + """Delta content for a streaming chunk.""" + + role: Optional[str] = None + content: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[list[dict]] = None + + +OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"] + + +class ChunkChoice(BaseModel): + """A single choice in a streaming chunk.""" + + index: int = 0 + delta: ChoiceDelta + finish_reason: Optional[OpenAIFinishReason] = None + logprobs: Optional[dict] = None + + +class ChatCompletionChunk(BaseModel): + """A single SSE chunk in OpenAI streaming format.""" + + id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}") + object: Literal["chat.completion.chunk"] = "chat.completion.chunk" + created: int = Field(default_factory = lambda: int(time.time())) + model: str = "default" + choices: list[ChunkChoice] + usage: Optional[CompletionUsage] = None + timings: Optional[dict] = None + + +# ── Non-streaming response ─────────────────────────────────────── + + +class CompletionMessage(BaseModel): + """The assistant's complete response message.""" + + role: Literal["assistant"] = "assistant" + # ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise. + content: Optional[str] = None + refusal: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[list[dict]] = None + + +class CompletionChoice(BaseModel): + """A single choice in a non-streaming response.""" + + index: int = 0 + message: CompletionMessage + finish_reason: OpenAIFinishReason = "stop" + logprobs: Optional[dict] = None + + +class CompletionUsage(BaseModel): + """Token usage statistics (approximate).""" + + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + prompt_tokens_details: Optional[dict] = Field( + default_factory = lambda: {"cached_tokens": 0, "audio_tokens": 0} + ) + completion_tokens_details: Optional[dict] = Field( + default_factory = lambda: { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + } + ) + + +class ChatCompletion(BaseModel): + """Non-streaming chat completion response.""" + + id: str = Field(default_factory = lambda: f"chatcmpl-{uuid.uuid4().hex[:12]}") + object: Literal["chat.completion"] = "chat.completion" + created: int = Field(default_factory = lambda: int(time.time())) + model: str = "default" + choices: list[CompletionChoice] + usage: CompletionUsage = Field(default_factory = CompletionUsage) + system_fingerprint: Optional[str] = None + + +# ===================================================================== +# 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", "original"]] = "auto" + + +class ResponsesOutputTextPart(BaseModel): + """Assistant ``output_text`` content part replayed on subsequent turns. + + Clients looping on a stateless Responses endpoint round-trip prior assistant + messages as ``output_text`` parts; we keep the text and ignore the + annotations/logprobs 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 unmodelled content-part types. + + Keeps validation green for newer part types (e.g. ``input_audio``); 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 attaches a `phase` field to assistant messages and requires clients + # to preserve it across turns; we round-trip it, llama-server ignores it. + model_config = {"extra": "allow"} + + +class ResponsesFunctionCallInputItem(BaseModel): + """A prior assistant function_call replayed in a multi-turn Responses input. + + Tool calls are top-level input items (not nested), correlated 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 its + 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 unmodelled Responses input item types. + + Covers ``reasoning`` items and future types. Dropped during normalisation + (GGUFs can't consume them), but kept in the union so unrelated turns don't 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 smart-union matching misreports errors when a strict-``Literal`` + variant doesn't match; an explicit discriminator makes routing deterministic + and falls 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 for the Responses API request. + + Unlike Chat Completions (nested under a ``"function"`` key), this uses a flat + shape with ``type``/``name``/``description``/``parameters``/``strict`` at top level. + """ + + 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 via the Chat Completions + # pass-through. Plain list so built-in tool shapes round-trip without + # validation errors; the translator forwards only ``type=="function"`` entries. + 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 ResponsesOutputReasoningContent(BaseModel): + """A reasoning text content block inside a reasoning output item.""" + + type: Literal["reasoning_text"] = "reasoning_text" + text: str + + +class ResponsesOutputReasoning(BaseModel): + """A top-level reasoning output item in the Responses API response.""" + + type: Literal["reasoning"] = "reasoning" + id: str = Field(default_factory = lambda: f"rs_{uuid.uuid4().hex[:12]}") + status: Literal["completed", "in_progress", "incomplete"] = "completed" + summary: list = Field(default_factory = list) + content: Optional[list[ResponsesOutputReasoningContent]] = None + + +class ResponsesOutputFunctionCall(BaseModel): + """A function-call output item in the Responses API response. + + Each tool call is its own top-level ``output`` item, correlated via ``call_id``. + """ + + 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, + ResponsesOutputReasoning, + 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] = "" + + @field_validator("content", mode = "before") + @classmethod + def _coerce_null_content(cls, v): + # Some clients send null content for an empty tool result; the str|list + # union would 400 on it, so treat null as "". + return "" if v is None else v + + +# Block types the converter translates explicitly. Anything else (thinking / +# redacted_thinking, a provider block a resumed session replays, or a future type) +# is accepted as an unknown block and dropped by the converter, rather than 400-ing +# the whole request on strict validation. +_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"}) + + +class AnthropicUnknownBlock(BaseModel): + type: str + model_config = {"extra": "allow"} + + @field_validator("type") + @classmethod + def _only_unknown_types(cls, v): + # Known types parse as their typed models above (so a malformed known block + # still fails cleanly); this fallback only catches the rest. + if v in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError("known block type handled by its typed model") + return v + + +AnthropicContentBlock = Union[ + AnthropicTextBlock, + AnthropicImageBlock, + AnthropicToolUseBlock, + AnthropicToolResultBlock, + AnthropicUnknownBlock, +] + + +def _anthropic_content_to_system_text(content: Any) -> str: + """Convert misplaced system message content into Anthropic system text.""" + if content is None: # null content must not become the literal "None" + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + text = block.get("text") + if isinstance(text, str): + parts.append(text) + continue + if block is not None: + parts.append(str(block)) + return "\n\n".join(part for part in parts if part) + return str(content) + + +def _merge_anthropic_system(system: Any, additions: list[str]) -> Any: + if not additions: + return system + + addition_blocks = [{"type": "text", "text": text} for text in additions if text.strip()] + if not addition_blocks: + return system + + if system is None: + return addition_blocks[0]["text"] if len(addition_blocks) == 1 else addition_blocks + if isinstance(system, str): + return "\n\n".join([system, *[block["text"] for block in addition_blocks]]) + if isinstance(system, list): + return [*system, *addition_blocks] + return system + + +class AnthropicMessage(BaseModel): + role: Literal["user", "assistant"] + content: Union[str, list[AnthropicContentBlock]] + + @model_validator(mode = "before") + @classmethod + def _normalize_content(cls, data): + # Role-aware leniency that never silently drops real user input: + # - assistant: a resumed tool-only turn's null content -> "" (str|list would + # 400 on null; "" keeps the converter's `for block in content` safe). + # Unknown blocks (thinking / future types) validate via + # AnthropicUnknownBlock and are dropped by the converter. + # - user: keep strict. Null user content stays None so str|list rejects it + # (400) rather than forwarding an empty prompt; and reject block types the + # converter cannot translate, since it silently skips unknown user blocks + # -- a user turn made only of them would validate yet send no content + # (silent data loss). + if not isinstance(data, dict): + return data + content = data.get("content") + if data.get("role") == "assistant": + # Coerce only an explicit null (resumed tool-only turn). A missing + # content key stays malformed so the required-field check still 400s. + if "content" in data and content is None: + return {**data, "content": ""} + return data + if isinstance(content, list): + for block in content: + btype = ( + block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + ) + # Guard the value: a non-string type is unsupported too, and a + # membership test on an unhashable value would raise TypeError + # (escaping as a 500 instead of a clean 400). + if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError(f"unsupported content block type {btype!r} in a user message") + return data + + +class AnthropicTool(BaseModel): + # Client tools have input_schema; server tools may only have type/name. + type: Optional[str] = None + name: Optional[str] = None + description: Optional[str] = None + input_schema: Optional[dict] = None + model_config = {"extra": "allow"} + + +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 mirroring the OpenAI endpoint convenience fields + min_p: Optional[float] = Field( + None, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold" + ) + 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 + bypass_permissions: Optional[bool] = Field( + False, + description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.", + ) + auto_heal_tool_calls: Optional[bool] = Field( + True, + description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).", + ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).", + ) + model_config = {"extra": "allow"} + + @model_validator(mode = "before") + @classmethod + def normalize_system_messages(cls, data: Any) -> Any: + if not isinstance(data, dict): + return data + + messages = data.get("messages") + if not isinstance(messages, list): + return data + + normalized_messages: list[Any] = [] + system_additions: list[str] = [] + changed = False + + for message in messages: + if isinstance(message, dict) and message.get("role") == "system": + system_additions.append( + _anthropic_content_to_system_text(message.get("content", "")) + ) + changed = True + continue + normalized_messages.append(message) + + if not changed: + return data + + normalized = dict(data) + normalized["messages"] = normalized_messages + normalized["system"] = _merge_anthropic_system(normalized.get("system"), system_additions) + return normalized + + +# ── Response models ──────────────────────────────────────────── + + +class AnthropicUsage(BaseModel): + input_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_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) diff --git a/studio/backend/models/mcp_servers.py b/studio/backend/models/mcp_servers.py new file mode 100644 index 0000000..606c242 --- /dev/null +++ b/studio/backend/models/mcp_servers.py @@ -0,0 +1,57 @@ +# 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 typing import Optional + +from pydantic import BaseModel, Field + + +class McpServerCreate(BaseModel): + display_name: str + url: str + headers: Optional[dict[str, str]] = None + is_enabled: bool = True + use_oauth: bool = False + + +class McpServerUpdate(BaseModel): + display_name: Optional[str] = None + url: Optional[str] = None + # Absent in request body = leave as-is; null = drop all headers; dict = set. + headers: Optional[dict[str, str]] = None + is_enabled: Optional[bool] = None + use_oauth: Optional[bool] = None + + +class McpServerResponse(BaseModel): + id: str + display_name: str + url: str + headers: dict[str, str] = Field(default_factory = dict) + is_enabled: bool = True + use_oauth: bool = False + created_at: str + updated_at: str + + +class McpServerTestRequest(BaseModel): + url: str + headers: Optional[dict[str, str]] = None + use_oauth: bool = False + + +class McpServerProbeResult(BaseModel): + ok: bool + tool_count: int = 0 + error: Optional[str] = None + + +class McpServerImportRequest(BaseModel): + # A standard mcpServers JSON config (Claude Desktop / Cursor / Cline / VS Code). + config: dict + + +class McpServerImportResult(BaseModel): + created: list[McpServerResponse] = Field(default_factory = list) + skipped: list[str] = Field(default_factory = list) # display names skipped as duplicates + errors: list[str] = Field(default_factory = list) diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py new file mode 100644 index 0000000..54e88fe --- /dev/null +++ b/studio/backend/models/models.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic schemas for Model Management API""" + +from pydantic import BaseModel, Field +from typing import Optional, List, Dict, Any, Literal + +ModelType = Literal["text", "vision", "audio", "embeddings"] + + +class CheckpointInfo(BaseModel): + """Information about a discovered checkpoint directory.""" + + display_name: str = Field(..., description = "User-friendly checkpoint name (folder name)") + path: str = Field(..., description = "Full path to the checkpoint directory") + loss: Optional[float] = Field(None, description = "Training loss at this checkpoint") + + +class ModelCheckpoints(BaseModel): + """A training run and its associated checkpoints.""" + + name: str = Field(..., description = "Training run folder name") + checkpoints: List[CheckpointInfo] = Field( + default_factory = list, + description = "List of checkpoints for this training run (final + intermediate)", + ) + base_model: Optional[str] = Field( + None, + description = "Base model name from adapter_config.json or config.json", + ) + peft_type: Optional[str] = Field( + None, + description = "PEFT type (e.g. LORA) if adapter training, None for full fine-tune", + ) + lora_rank: Optional[int] = Field( + None, + description = "LoRA rank (r) if applicable", + ) + is_quantized: bool = Field( + False, + description = "Whether the model uses BNB quantization (e.g. bnb-4bit)", + ) + + +class CheckpointListResponse(BaseModel): + """Response for listing available checkpoints in an outputs directory.""" + + outputs_dir: str = Field(..., description = "Directory that was scanned") + models: List[ModelCheckpoints] = Field( + default_factory = list, + description = "List of training runs with their checkpoints", + ) + + +class ExportSizeResponse(BaseModel): + """Model fp16/bf16-equivalent size; size fields are null when unknown.""" + + model: str = Field(..., description = "Model id or path the estimate was computed for") + fp16_bytes: Optional[int] = Field( + None, + description = "Estimated FP16/BF16-equivalent on-disk size in bytes, or null if unknown", + ) + total_params: Optional[int] = Field( + None, + description = "Estimated total parameter count (fp16_bytes // 2), or null if unknown", + ) + source: str = Field( + "unavailable", + description = "How the estimate was derived (e.g. safetensors, config, local, vllm, unavailable)", + ) + + +class ModelDetails(BaseModel): + """Model configuration and metadata; used for both list and detail views""" + + id: str = Field(..., description = "Model identifier") + model_name: Optional[str] = Field( + None, description = "Model identifier (alias for id, for backward compatibility)" + ) + name: Optional[str] = Field(None, description = "Display name for the model") + config: Optional[Dict[str, Any]] = Field(None, description = "Model configuration dictionary") + is_vision: bool = Field(False, description = "Whether model is a vision model") + is_embedding: bool = Field( + False, description = "Whether model is an embedding/sentence-transformer model" + ) + is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") + is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp format)") + is_mlx: bool = Field( + False, description = "Whether model is served via the MLX backend (Apple Silicon)" + ) + is_audio: bool = Field(False, description = "Whether model is a TTS audio model") + audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") + has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") + model_type: Optional[ModelType] = Field( + None, description = "Collapsed model modality: text, vision, audio, or embeddings" + ) + base_model: Optional[str] = Field(None, description = "Base model if this is a LoRA adapter") + max_position_embeddings: Optional[int] = Field( + None, description = "Maximum context length supported by the model" + ) + model_size_bytes: Optional[int] = Field( + None, description = "Total size of model weight files in bytes" + ) + + +class LoRAInfo(BaseModel): + """LoRA adapter or exported model information""" + + display_name: str = Field(..., description = "Display name for the LoRA") + adapter_path: str = Field(..., description = "Path to the LoRA adapter or exported model") + base_model: Optional[str] = Field(None, description = "Base model identifier") + source: Optional[str] = Field(None, description = "'training' or 'exported'") + export_type: Optional[str] = Field( + None, description = "'lora', 'merged', or 'gguf' (for exports)" + ) + + +class LoRAScanResponse(BaseModel): + """Response schema for scanning trained LoRA adapters""" + + loras: List[LoRAInfo] = Field(default_factory = list, description = "List of found LoRA adapters") + outputs_dir: str = Field(..., description = "Directory that was scanned") + + +class ModelListResponse(BaseModel): + """Response schema for listing models""" + + models: List[ModelDetails] = Field(default_factory = list, description = "List of models") + default_models: List[str] = Field(default_factory = list, description = "List of default model IDs") + + +class GgufVariantDetail(BaseModel): + """A single GGUF quantization variant in a HuggingFace repo.""" + + filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')") + quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')") + size_bytes: int = Field(0, description = "File size in bytes") + download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant") + downloaded: bool = Field( + False, description = "Whether this variant is already in the local HF cache" + ) + update_available: bool = Field( + False, description = "Whether a newer version of this variant is available on HF" + ) + + +class GgufVariantsResponse(BaseModel): + """Response for listing GGUF quantization variants in a HuggingFace repo.""" + + repo_id: str = Field(..., description = "HuggingFace repo ID") + variants: List[GgufVariantDetail] = Field( + default_factory = list, description = "Available GGUF variants" + ) + has_vision: bool = Field( + False, description = "Whether the model has vision support (mmproj files)" + ) + default_variant: Optional[str] = Field( + None, description = "Recommended default quantization variant" + ) + context_length: Optional[int] = Field( + None, + description = "Native max context from GGUF metadata; set once a variant is downloaded", + ) + + +class LocalModelInfo(BaseModel): + """Discovered local model candidate.""" + + id: str = Field(..., description = "Identifier to use for loading/training") + display_name: str = Field(..., description = "Display label") + path: str = Field(..., description = "Local path where model data was discovered") + source: Literal["models_dir", "hf_cache", "lmstudio", "custom"] = Field( + ..., + description = "Discovery source", + ) + model_id: Optional[str] = Field( + None, + description = "HF repo id for cached models, e.g. org/model", + ) + model_format: Optional[str] = Field( + None, + description = "Detected weights format ('gguf' when known). Lets the UI " + "classify scanned folders whose name lacks a -GGUF suffix.", + ) + updated_at: Optional[float] = Field( + None, + description = "Unix timestamp of latest observed update", + ) + + +class LocalModelListResponse(BaseModel): + """Response schema for listing local/cached models.""" + + models_dir: str = Field(..., description = "Directory scanned for custom local models") + hf_cache_dir: Optional[str] = Field( + None, + description = "HF cache root that was scanned", + ) + lmstudio_dirs: List[str] = Field( + default_factory = list, + description = "LM Studio model directories that were scanned", + ) + models: List[LocalModelInfo] = Field( + default_factory = list, + description = "Discovered local/cached models", + ) + + +class AddScanFolderRequest(BaseModel): + """Request body for adding a custom scan folder.""" + + path: str = Field(..., description = "Absolute or relative directory path to scan for models") + + +class ScanFolderInfo(BaseModel): + """A registered custom model scan folder.""" + + id: int = Field(..., description = "Database row ID") + path: str = Field(..., description = "Normalized absolute path") + created_at: str = Field(..., description = "ISO 8601 creation timestamp") + + +class 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)." + ), + ) diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py new file mode 100644 index 0000000..5a75246 --- /dev/null +++ b/studio/backend/models/providers.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic schemas for the external LLM providers API.""" + +from typing import Literal, Optional + +from pydantic import BaseModel, Field + + +# ── Registry (static provider info) ─────────────────────────────── + + +class ProviderRegistryEntry(BaseModel): + """A supported provider type with its default configuration.""" + + provider_type: str = Field(..., description = "Provider identifier (e.g. 'openai', 'mistral')") + display_name: str = Field(..., description = "Human-readable provider name") + base_url: str = Field(..., description = "Default API base URL") + default_models: list[str] = Field( + default_factory = list, description = "Well-known model IDs for this provider" + ) + supports_streaming: bool = Field( + True, description = "Whether this provider supports SSE streaming" + ) + supports_vision: bool = Field( + False, description = "Whether this provider supports vision/image input" + ) + supports_tool_calling: bool = Field( + False, description = "Whether this provider supports tool/function calling" + ) + model_list_mode: Literal["remote", "curated"] = Field( + "remote", + description = "remote = fetch /models; curated = huge catalogs — UI uses defaults + manual IDs only", + ) + + +# ── Provider config CRUD ────────────────────────────────────────── + + +class ProviderCreate(BaseModel): + """Request to create a saved provider configuration.""" + + provider_type: str = Field(..., description = "Provider type from the registry") + display_name: str = Field(..., description = "User-chosen label (e.g. 'My OpenAI Key')") + base_url: Optional[str] = Field( + None, + description = "Custom base URL (overrides registry default). Omit to use the default.", + ) + + +class ProviderUpdate(BaseModel): + """Request to update a saved provider configuration.""" + + display_name: Optional[str] = Field(None, description = "New display name") + base_url: Optional[str] = Field(None, description = "New base URL") + is_enabled: Optional[bool] = Field(None, description = "Enable or disable this provider") + + +class ProviderResponse(BaseModel): + """A saved provider configuration (returned by list/get endpoints).""" + + id: str = Field(..., description = "Unique provider config ID") + provider_type: str = Field(..., description = "Provider type (e.g. 'openai')") + display_name: str = Field(..., description = "User-chosen label") + base_url: str = Field(..., description = "API base URL") + is_enabled: bool = Field(True, description = "Whether this provider is enabled") + created_at: str = Field(..., description = "ISO 8601 creation timestamp") + updated_at: str = Field(..., description = "ISO 8601 last-update timestamp") + + +# ── Model listing ───────────────────────────────────────────────── + + +class ProviderModelInfo(BaseModel): + """A model available from an external provider.""" + + id: str = Field(..., description = "Model ID as expected by the provider API") + display_name: str = Field("", description = "Human-readable model name") + context_length: Optional[int] = Field(None, description = "Maximum context length in tokens") + owned_by: Optional[str] = Field(None, description = "Model owner/organization") + + +class ProviderModelsRequest(BaseModel): + """Request to list models from an external provider.""" + + provider_type: str = Field(..., description = "Provider type from the registry") + encrypted_api_key: Optional[str] = Field( + None, + description = "RSA-encrypted, base64-encoded API key (optional for local providers)", + ) + base_url: Optional[str] = Field( + None, description = "Custom base URL (overrides registry default)" + ) + + +# ── Connection testing ──────────────────────────────────────────── + + +class ProviderTestRequest(BaseModel): + """Request to test connectivity to an external provider.""" + + provider_type: str = Field(..., description = "Provider type from the registry") + encrypted_api_key: Optional[str] = Field( + None, + description = "RSA-encrypted, base64-encoded API key (optional for local providers)", + ) + base_url: Optional[str] = Field( + None, description = "Custom base URL (overrides registry default)" + ) + model_id: Optional[str] = Field( + None, description = "Model ID for providers that need a chat probe" + ) + + +class ProviderTestResult(BaseModel): + """Result of a provider connectivity test.""" + + success: bool = Field(..., description = "Whether the test succeeded") + message: str = Field(..., description = "Human-readable result message") + models_count: Optional[int] = Field( + None, description = "Number of models found (if test succeeded)" + ) diff --git a/studio/backend/models/responses.py b/studio/backend/models/responses.py new file mode 100644 index 0000000..cc420c0 --- /dev/null +++ b/studio/backend/models/responses.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic response models for training and model management routes +(previously returned as raw dicts).""" + +from pydantic import BaseModel, Field +from typing import Optional, List + + +# --- Training route response models --- + + +class TrainingStopResponse(BaseModel): + """Response for stopping a training job""" + + status: str = Field(..., description = "Current status: 'stopped' or 'idle'") + message: str = Field(..., description = "Human-readable status message") + + +class TrainingMetricsResponse(BaseModel): + """Response for training metrics history""" + + loss_history: List[float] = Field(default_factory = list, description = "Loss values per step") + lr_history: List[float] = Field(default_factory = list, description = "Learning rate per step") + step_history: List[int] = Field(default_factory = list, description = "Step numbers") + grad_norm_history: List[float] = Field(default_factory = list, description = "Gradient norm values") + grad_norm_step_history: List[int] = Field( + default_factory = list, description = "Step numbers for gradient norm values" + ) + current_loss: Optional[float] = Field(None, description = "Most recent loss value") + current_lr: Optional[float] = Field(None, description = "Most recent learning rate") + current_step: Optional[int] = Field(None, description = "Most recent step number") + + +# --- Model management route response models --- + + +class LoRABaseModelResponse(BaseModel): + """Response for getting a LoRA's base model""" + + lora_path: str = Field(..., description = "Path to the LoRA adapter") + base_model: str = Field(..., description = "Base model identifier") + + +class VisionCheckResponse(BaseModel): + """Response for checking if a model is a vision model""" + + model_name: str = Field(..., description = "Model identifier") + is_vision: bool = Field(..., description = "Whether the model is a vision model") + + +class EmbeddingCheckResponse(BaseModel): + """Response for checking if a model is an embedding model""" + + model_name: str = Field(..., description = "Model identifier") + is_embedding: bool = Field( + ..., description = "Whether the model is an embedding/sentence-transformer model" + ) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py new file mode 100644 index 0000000..ff815a2 --- /dev/null +++ b/studio/backend/models/training.py @@ -0,0 +1,667 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Pydantic schemas for Training API +""" + +import re +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from typing import Any, Optional, List, Dict, Literal + +from utils.training_runs import normalize_project_name + + +# ASCII integer, optional single sign. Rejects "++512" and Unicode digits +# ("５１２") that slip through str.isdigit() + int(). +_INT_RE = re.compile(r"[+-]?[0-9]+") + + +_MAX_BATCH_SIZE = 4096 +_MAX_GRAD_ACCUM = 4096 +_MAX_STEPS = 1_000_000 +_MAX_EPOCHS = 1000 +# 2M is a sanity cap; host RAM runs out long before this. +_MAX_SEQ_LENGTH = 2_000_000 +_MAX_LR_VALUE = 1.0 +_MAX_LORA_R = 16_384 +_MAX_LORA_ALPHA = 32_768 +_MIN_VISION_IMAGE_SIZE = 256 +# 2048 is the highest most llms stay stable at +_MAX_VISION_IMAGE_SIZE = 2048 +# Upper bound for dataset slice indices. Caps `.skip(n)` on streaming datasets so +# an absurd index can't make the loader iterate effectively forever (DoS guard). +# 1e9 is far beyond any realistic fine-tuning dataset row count. +_MAX_DATASET_SLICE_INDEX = 1_000_000_000 + + +class S3Config(BaseModel): + """S3 bucket configuration for loading datasets from AWS S3""" + + # Accept both snake_case and the frontend's camelCase field names. + model_config = ConfigDict(populate_by_name = True) + + bucket: str = Field(..., description = "S3 bucket name") + region: str = Field("us-east-1", description = "AWS region") + prefix: Optional[str] = Field(None, description = "Optional path prefix within bucket") + access_key_id: Optional[str] = Field( + None, + alias = "accessKeyId", + description = "AWS access key ID (optional if using IAM role)", + ) + secret_access_key: Optional[str] = Field( + None, + alias = "secretAccessKey", + description = "AWS secret access key (optional if using IAM role)", + ) + use_iam_role: bool = Field( + False, + alias = "useIamRole", + description = "Use IAM role credentials instead of access keys", + ) + + @model_validator(mode = "after") + def _check_credentials(self) -> "S3Config": + # Require either IAM role auth or a full key pair so credentials are + # never half-configured. + if not self.use_iam_role and not (self.access_key_id and self.secret_access_key): + raise ValueError( + "s3_config requires either use_iam_role=True or both " + "access_key_id and secret_access_key" + ) + return self + + +def _parse_lr(v: Any) -> float: + """Parse learning_rate as a positive float strictly below _MAX_LR_VALUE.""" + if v is None: + raise ValueError("learning_rate is required") + if isinstance(v, bool): + raise ValueError("learning_rate must be a number, not a bool") + try: + lr = float(v) + except (TypeError, ValueError): + raise ValueError(f"learning_rate must be parseable as float (got {v!r})") + if not (lr > 0.0): + raise ValueError(f"learning_rate must be > 0 (got {lr!r}); typical range is 1e-6 .. 1e-3") + if lr >= _MAX_LR_VALUE: + raise ValueError( + f"learning_rate must be < 1.0 (got {lr!r}); " + "values that large always diverge training" + ) + return lr + + +class TrainingStartRequest(BaseModel): + """Request schema for starting training""" + + # Model parameters + model_name: str = Field( + ..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')" + ) + project_name: Optional[str] = Field( + None, + max_length = 80, + description = "Optional user-defined project name appended to run folders and shown in history", + ) + training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = Field( + ..., + description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'", + ) + hf_token: Optional[str] = Field(None, description = "HuggingFace token") + load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") + max_seq_length: int = Field(2048, description = "Maximum sequence length") + vision_image_size: Optional[int] = Field( + None, + description = "Optional maximum image side length for VLM training. Null uses model default.", + ) + trust_remote_code: bool = Field( + False, + description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", + ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) + + # Dataset parameters + hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier") + local_datasets: List[str] = Field( + default_factory = list, description = "List of local dataset paths" + ) + local_eval_datasets: List[str] = Field( + default_factory = list, description = "List of local eval dataset paths" + ) + format_type: str = Field(..., description = "Dataset format type") + subset: Optional[str] = None + train_split: Optional[str] = Field("train", description = "Training split name") + eval_split: Optional[str] = Field(None, description = "Eval split name. None = auto-detect") + dataset_streaming: bool = Field( + False, + description = "Whether to load the Hugging Face dataset in streaming mode", + ) + eval_steps: float = Field(0.00, description = "Fraction of total steps between evals (0-1)") + dataset_slice_start: Optional[int] = Field( + None, + ge = 0, + le = _MAX_DATASET_SLICE_INDEX, + description = "Inclusive start row index for dataset slicing", + ) + dataset_slice_end: Optional[int] = Field( + None, + ge = 0, + le = _MAX_DATASET_SLICE_INDEX, + description = "Inclusive end row index for dataset slicing", + ) + + @model_validator(mode = "before") + @classmethod + def _compat_split(cls, values: Any) -> Any: + """Accept legacy 'split' field as alias for 'train_split'.""" + if isinstance(values, dict) and "split" in values: + values.setdefault("train_split", values.pop("split")) + return values + + @field_validator("project_name") + @classmethod + def _normalize_project_name(cls, value: Optional[str]) -> Optional[str]: + return normalize_project_name(value) + + # NOTE: pydantic runs all `mode="after"` validators in definition order. A + # second one, `_check_steps_or_epochs`, is defined lower in this class; keep + # these cross-field checks order-independent so the two stay decoupled. + @model_validator(mode = "after") + def _validate_dataset_slice(self) -> "TrainingStartRequest": + # Only the ordering is validated here. No upper bound is enforced on the + # indices: the trainer slices via datasets `.take()` / `.select()`, which + # clamp gracefully when the end index exceeds the dataset length. + # start == end is intentionally allowed (deliberate single-row slice, + # e.g. for debugging); the trainer logs a warning for that 1-row case. + if ( + self.dataset_slice_start is not None + and self.dataset_slice_end is not None + and self.dataset_slice_end < self.dataset_slice_start + ): + raise ValueError( + "dataset_slice_end must be greater than or equal to dataset_slice_start" + ) + return self + + @field_validator("hf_dataset") + @classmethod + def _check_hf_dataset(cls, v: Optional[str]) -> Optional[str]: + # Constrain the HF dataset id to a safe charset + length to shrink the + # path-traversal / SSRF surface of `load_dataset(, ...)`. + if v is None: + return v + v = v.strip() + if not v: + return None + if len(v) > 256: + raise ValueError("hf_dataset is too long (max 256 chars)") + if ".." in v: + raise ValueError("hf_dataset must not contain '..'") + if not re.fullmatch(r"[A-Za-z0-9._\-/]+", v): + raise ValueError("hf_dataset may only contain letters, digits, '_', '-', '.', '/'") + return v + + @field_validator("subset") + @classmethod + def _check_subset(cls, v: Optional[str]) -> Optional[str]: + if v is None: + return v + if len(v) > 128: + raise ValueError("subset is too long (max 128 chars)") + if not re.fullmatch(r"[A-Za-z0-9._\-]*", v): + raise ValueError("subset may only contain letters, digits, '_', '-', '.'") + return v + + @field_validator("train_split", "eval_split") + @classmethod + def _check_split_name(cls, v: Optional[str]) -> Optional[str]: + # Split names feed HF slice syntax (e.g. "train[:80%]"), so allow that + # charset but cap length and block path-traversal / NUL bytes. + if v is None: + return v + if len(v) > 128: + raise ValueError("split name is too long (max 128 chars)") + if "\x00" in v or ".." in v or "/" in v or "\\" in v: + raise ValueError("split name contains invalid characters") + if not re.fullmatch(r"[A-Za-z0-9_\-\[\]:%.+ ]*", v): + raise ValueError("split name contains invalid characters") + return v + + @field_validator("learning_rate", mode = "before") + @classmethod + def _check_learning_rate(cls, v): + # Stringify because downstream call sites float() it themselves. + lr = _parse_lr(v) + return str(lr) + + @field_validator("batch_size") + @classmethod + def _check_batch_size(cls, v: int) -> int: + if v is None: + raise ValueError("batch_size is required") + if v < 1 or v > _MAX_BATCH_SIZE: + raise ValueError(f"batch_size must be in [1, {_MAX_BATCH_SIZE}] (got {v!r})") + return v + + @field_validator("gradient_accumulation_steps") + @classmethod + def _check_grad_accum(cls, v: int) -> int: + if v is None: + return 1 + if v < 1 or v > _MAX_GRAD_ACCUM: + raise ValueError( + f"gradient_accumulation_steps must be in [1, {_MAX_GRAD_ACCUM}] " f"(got {v!r})" + ) + return v + + @field_validator("num_epochs") + @classmethod + def _check_num_epochs(cls, v: int) -> int: + # 0 is a sentinel for "use max_steps instead" (frontend toggle). + if v is None: + return 1 + if v < 0 or v > _MAX_EPOCHS: + raise ValueError(f"num_epochs must be in [0, {_MAX_EPOCHS}] (got {v!r})") + return v + + @field_validator("max_steps") + @classmethod + def _check_max_steps(cls, v: Optional[int]) -> Optional[int]: + # 0 is the frontend's sentinel for "use num_epochs instead". + if v is None: + return v + if not isinstance(v, int) or v < 0 or v > _MAX_STEPS: + raise ValueError(f"max_steps must be a non-negative int <= {_MAX_STEPS} (got {v!r})") + return v + + @field_validator("max_seq_length") + @classmethod + def _check_max_seq_length(cls, v: int) -> int: + if v is None or v < 1 or v > _MAX_SEQ_LENGTH: + raise ValueError(f"max_seq_length must be in [1, {_MAX_SEQ_LENGTH}] (got {v!r})") + return v + + @field_validator("vision_image_size", mode = "before") + @classmethod + def _check_vision_image_size(cls, v: Any) -> Optional[int]: + # mode="before" sees True/False as bool (not 1/0) for a precise error. + if v is None: + return v + if isinstance(v, bool): + raise ValueError("vision_image_size must be an integer or null") + if isinstance(v, int): + coerced = v + elif isinstance(v, str) and _INT_RE.fullmatch(v.strip()): + coerced = int(v.strip()) + elif isinstance(v, float) and v.is_integer(): + coerced = int(v) + else: + # numpy ints / Integral subclasses, without a hard numpy import. + try: + import numbers + if isinstance(v, numbers.Integral): + coerced = int(v) + elif isinstance(v, numbers.Real) and float(v).is_integer(): + coerced = int(v) + else: + raise TypeError + except Exception: + raise ValueError("vision_image_size must be an integer or null") + if coerced < _MIN_VISION_IMAGE_SIZE or coerced > _MAX_VISION_IMAGE_SIZE: + raise ValueError( + f"vision_image_size must be in [{_MIN_VISION_IMAGE_SIZE}, " + f"{_MAX_VISION_IMAGE_SIZE}] (got {coerced!r})" + ) + return coerced + + @field_validator("warmup_steps") + @classmethod + def _check_warmup_steps(cls, v: Optional[int]) -> Optional[int]: + if v is None: + return v + if not isinstance(v, int) or v < 0 or v > _MAX_STEPS: + raise ValueError( + f"warmup_steps must be a non-negative int <= {_MAX_STEPS} " f"(got {v!r})" + ) + return v + + @field_validator("warmup_ratio") + @classmethod + def _check_warmup_ratio(cls, v): + if v is None: + return v + try: + r = float(v) + except (TypeError, ValueError): + raise ValueError(f"warmup_ratio must be a number (got {v!r})") + if not (0.0 <= r <= 1.0): + raise ValueError(f"warmup_ratio must be in [0.0, 1.0] (got {r!r})") + return r + + @field_validator("save_steps") + @classmethod + def _check_save_steps(cls, v: int) -> int: + if v is None: + return 100 + if v < 0 or v > _MAX_STEPS: + raise ValueError(f"save_steps must be in [0, {_MAX_STEPS}] (got {v!r})") + return v + + @field_validator("weight_decay") + @classmethod + def _check_weight_decay(cls, v: float) -> float: + if v is None: + return 0.0 + try: + wd = float(v) + except (TypeError, ValueError): + raise ValueError(f"weight_decay must be a number (got {v!r})") + if wd < 0 or wd > 10.0: + raise ValueError(f"weight_decay must be in [0, 10] (got {wd!r}); typical 0..0.1") + return wd + + @field_validator("lora_r") + @classmethod + def _check_lora_r(cls, v: int) -> int: + if v is None: + return 16 + if v < 1 or v > _MAX_LORA_R: + raise ValueError(f"lora_r must be in [1, {_MAX_LORA_R}] (got {v!r})") + return v + + @field_validator("lora_alpha") + @classmethod + def _check_lora_alpha(cls, v: int) -> int: + if v is None: + return 16 + if v < 1 or v > _MAX_LORA_ALPHA: + raise ValueError(f"lora_alpha must be in [1, {_MAX_LORA_ALPHA}] (got {v!r})") + return v + + @field_validator("lora_dropout") + @classmethod + def _check_lora_dropout(cls, v: float) -> float: + if v is None: + return 0.0 + try: + d = float(v) + except (TypeError, ValueError): + raise ValueError(f"lora_dropout must be a number (got {v!r})") + if not (0.0 <= d < 1.0): + raise ValueError(f"lora_dropout must be in [0.0, 1.0) (got {d!r})") + return d + + custom_format_mapping: Optional[Dict[str, Any]] = Field( + None, + description = ( + "User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} " + "for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM. " + "Enhanced format includes __system_prompt, __user_template, " + "__assistant_template, __label_mapping metadata keys." + ), + ) + # Training parameters + num_epochs: int = Field(1, description = "Number of training epochs") + learning_rate: str = Field("2e-4", description = "Learning rate") + batch_size: int = Field(1, description = "Batch size") + gradient_accumulation_steps: int = Field(1, description = "Gradient accumulation steps") + warmup_steps: Optional[int] = Field(None, description = "Warmup steps") + warmup_ratio: Optional[float] = Field(None, description = "Warmup ratio") + max_steps: Optional[int] = Field(None, description = "Maximum training steps") + save_steps: int = Field(100, description = "Steps between checkpoints") + weight_decay: float = Field(0.001, description = "Weight decay") + max_grad_norm: float = Field( + 0.0, + ge = 0, + description = "Global gradient norm clipping threshold. Set 0 to disable.", + ) + max_grad_value: Optional[float] = Field( + None, + ge = 0, + description = ( + "MLX-only elementwise gradient value clipping threshold. " + "If unset, MLX uses its runtime default." + ), + ) + max_grad_leaf_norm: Optional[float] = Field( + None, + ge = 0, + description = ( + "MLX-only proportional per-parameter gradient norm cap. " + "Preserves each tensor's gradient direction without global norm " + "clipping's memory overhead." + ), + ) + cast_norm_output_to_input_dtype: bool = Field( + True, + description = ( + "MLX-only: keep norm parameters in fp32 but cast norm outputs " + "back to the incoming activation dtype." + ), + ) + random_seed: int = Field( + 3407, + description = ( + "Random seed; matches the Studio backend / MLX worker default " + "and unsloth's historical recommended value." + ), + ) + packing: bool = Field(False, description = "Enable sequence packing") + optim: str = Field("adamw_8bit", description = "Optimizer") + lr_scheduler_type: str = Field("linear", description = "Learning rate scheduler type") + embedding_learning_rate: Optional[float] = Field( + None, + gt = 0, + lt = 1.0, + description = "Separate learning rate for embedding matrices (CPT). " + "Must be in (0, 1). Should be 2-10x smaller than the main learning rate.", + ) + + # LoRA parameters + use_lora: bool = Field(True, description = "Use LoRA (derived from training_type)") + lora_r: int = Field(16, description = "LoRA rank") + lora_alpha: int = Field(16, description = "LoRA alpha") + lora_dropout: float = Field(0.0, description = "LoRA dropout") + target_modules: List[str] = Field(default_factory = list, description = "Target modules for LoRA") + gradient_checkpointing: str = Field("", description = "Gradient checkpointing setting") + use_rslora: bool = Field(False, description = "Use RSLoRA") + use_loftq: bool = Field(False, description = "Use LoftQ") + train_on_completions: bool = Field(False, description = "Train on completions only") + + # Vision-specific LoRA parameters + finetune_vision_layers: bool = Field(False, description = "Finetune vision layers") + finetune_language_layers: bool = Field(False, description = "Finetune language layers") + finetune_attention_modules: bool = Field(False, description = "Finetune attention modules") + finetune_mlp_modules: bool = Field(False, description = "Finetune MLP modules") + is_dataset_image: bool = Field(False, description = "Whether the dataset contains image data") + is_dataset_audio: bool = Field(False, description = "Whether the dataset contains audio data") + is_embedding: bool = Field( + False, description = "Whether model is an embedding/sentence-transformer model" + ) + + # Logging parameters + enable_wandb: bool = Field(False, description = "Enable Weights & Biases logging") + wandb_token: Optional[str] = Field(None, description = "W&B token") + 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( + 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.", + ) + + # S3 dataset source configuration + s3_config: Optional[S3Config] = Field( + None, + description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", + ) + + @model_validator(mode = "after") + def _validate_streaming_splits(self) -> "TrainingStartRequest": + # Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]" + # or "train[:20]"). Probe-confirmed: raises ValueError: Bad split. Reject + # early with a clear message so the user knows to use a plain split name. + if self.dataset_streaming: + for field_name, split_val in ( + ("train_split", self.train_split), + ("eval_split", self.eval_split), + ): + if split_val is not None and "[" in split_val: + raise ValueError( + f"dataset_streaming does not support HF slice syntax in {field_name} " + f"(got {split_val!r}); streaming load_dataset raises 'Bad split' on " + "bracket expressions. Use a plain split name (e.g. 'train', 'validation')." + ) + return self + + @model_validator(mode = "after") + def _check_steps_or_epochs(self) -> "TrainingStartRequest": + # Each accepts 0 as "use the other"; both 0 means nothing to train. + if (self.max_steps is None or self.max_steps == 0) and self.num_epochs == 0: + raise ValueError("Either num_epochs or max_steps must be > 0; both cannot be 0.") + return self + + +class TrainingJobResponse(BaseModel): + """Immediate response when training is initiated""" + + job_id: str = Field(..., description = "Unique training job identifier") + status: Literal["queued", "error"] = Field(..., description = "Initial job status") + message: str = Field(..., description = "Human-readable status message") + error: Optional[str] = Field(None, description = "Error details if status is 'error'") + + +class TrainingStatus(BaseModel): + """Current training job status - works for streaming or polling""" + + job_id: str = Field(..., description = "Training job identifier") + phase: Literal[ + "idle", + "loading_model", + "loading_dataset", + "configuring", + "training", + "completed", + "error", + "stopped", + ] = Field(..., description = "Current phase of training pipeline") + is_training_running: bool = Field(..., description = "True if training loop is actively running") + eval_enabled: bool = Field( + False, + description = "True if evaluation dataset is configured for this training run", + ) + message: str = Field(..., description = "Human-readable status message") + error: Optional[str] = Field(None, description = "Error details if phase is 'error'") + details: Optional[dict] = Field( + None, description = "Phase-specific info, e.g. {'model_size': '8B'}" + ) + metric_history: Optional[dict] = Field( + None, + description = "Full metric history arrays for chart recovery after SSE reconnection. " + "Keys: 'steps', 'loss', 'lr', 'grad_norm', 'grad_norm_steps' — each a list of numeric values.", + ) + + +class TrainingProgress(BaseModel): + """Training progress metrics - for streaming or polling""" + + job_id: str = Field(..., description = "Training job identifier") + step: int = Field(..., description = "Current training step") + total_steps: int = Field(..., description = "Total training steps") + loss: Optional[float] = Field(None, description = "Current loss value") + learning_rate: Optional[float] = Field(None, description = "Current learning rate") + progress_percent: float = Field(..., description = "Progress percentage (0.0 to 100.0)") + epoch: Optional[float] = Field(None, description = "Current epoch") + elapsed_seconds: Optional[float] = Field( + None, description = "Time elapsed since training started" + ) + eta_seconds: Optional[float] = Field(None, description = "Estimated time remaining") + grad_norm: Optional[float] = Field( + None, description = "L2 norm of gradients, computed before gradient clipping" + ) + num_tokens: Optional[int] = Field(None, description = "Total number of tokens processed so far") + eval_loss: Optional[float] = Field( + None, description = "Eval loss from the most recent evaluation step" + ) + + +class TrainingRunSummary(BaseModel): + """Summary of a training run for list views.""" + + id: str + status: Literal["running", "completed", "stopped", "error"] + model_name: str + project_name: Optional[str] = None + dataset_name: str + display_name: Optional[str] = None + started_at: str + ended_at: Optional[str] = None + total_steps: Optional[int] = None + final_step: Optional[int] = None + final_loss: Optional[float] = None + output_dir: Optional[str] = None + duration_seconds: Optional[float] = None + error_message: Optional[str] = None + loss_sparkline: Optional[List[float]] = None + can_resume: bool = False + resumed_later: bool = False + has_preview_model: bool = False + preview_ref: Optional[str] = None + # HMAC capability token for the `/p/{preview_ref}` share link; None when not + # previewable. The frontend appends it as `?k=` so a guessed ref can't be used. + preview_sig: Optional[str] = None + + +class TrainingRunUpdateRequest(BaseModel): + """Mutable fields on a training run.""" + + model_config = ConfigDict(extra = "forbid") + + display_name: Optional[str] = Field(None, max_length = 120) + + +class TrainingRunListResponse(BaseModel): + """Response for listing training runs.""" + + runs: List[TrainingRunSummary] + total: int + + +class TrainingRunMetrics(BaseModel): + """Metrics arrays for a training run, using paired step arrays per metric.""" + + step_history: List[int] = Field(default_factory = list) + loss_history: List[float] = Field(default_factory = list) + loss_step_history: List[int] = Field(default_factory = list) + lr_history: List[float] = Field(default_factory = list) + lr_step_history: List[int] = Field(default_factory = list) + grad_norm_history: List[float] = Field(default_factory = list) + grad_norm_step_history: List[int] = Field(default_factory = list) + eval_loss_history: List[float] = Field(default_factory = list) + eval_step_history: List[int] = Field(default_factory = list) + final_epoch: Optional[float] = None + final_num_tokens: Optional[int] = None + + +class TrainingRunDetailResponse(BaseModel): + """Response for a single training run with config and metrics.""" + + run: TrainingRunSummary + config: dict + metrics: TrainingRunMetrics + + +class TrainingRunDeleteResponse(BaseModel): + """Response for deleting a training run.""" + + status: str + message: str diff --git a/studio/backend/models/users.py b/studio/backend/models/users.py new file mode 100644 index 0000000..738c76c --- /dev/null +++ b/studio/backend/models/users.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic models for authentication tokens.""" + +from pydantic import BaseModel, Field + + +class Token(BaseModel): + """Authentication response model for session credentials.""" + + access_token: str = Field( + ..., description = "Session access credential used for authenticated API requests" + ) + refresh_token: str = Field( + ..., + description = "Session refresh credential used to renew an expired access credential", + ) + token_type: str = Field( + ..., description = "Credential type for the Authorization header, always 'bearer'" + ) + must_change_password: bool = Field( + ..., description = "True when the user must change the seeded default password" + ) diff --git a/studio/backend/plugins/__init__.py b/studio/backend/plugins/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/studio/backend/plugins/data-designer-github-repo-seed/README.md b/studio/backend/plugins/data-designer-github-repo-seed/README.md new file mode 100644 index 0000000..346d94b --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/README.md @@ -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. diff --git a/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml new file mode 100644 index 0000000..e232adc --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml @@ -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"] diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py new file mode 100644 index 0000000..62ecb2e --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +# Intentionally empty. Data-designer loads submodules lazily via qualified names +# in plugin.py, so importing this package must not touch data_designer.engine.* +# during Studio bootstrap (circular import). diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py new file mode 100644 index 0000000..6b347c4 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py @@ -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 diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py new file mode 100644 index 0000000..330e917 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py @@ -0,0 +1,78 @@ +# 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: config signature -> JSONL path. A recipe job reads the seed +# multiple times (validation, preview, sampling); memoize to avoid re-scraping +# GitHub on each call. Key uses a short SHA-256 of the token (never the raw value) +# so token never hits memory twice and 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/restart); drop it so the caller re-materializes. + 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) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py new file mode 100644 index 0000000..f87dbd0 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py @@ -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, +) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py new file mode 100644 index 0000000..637193e --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py @@ -0,0 +1,233 @@ +# 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 scraper in `scraper_impl/` per repo, capped via trial_limits +to stop at `limit` items per resource. Then reads the per-resource JSONL shards +and flattens them into one unified JSONL with stable columns (`item_type`, +`repo`, `number`, `title`, `body`, ...). +""" + +from __future__ import annotations + +import json +import os +import sys +import time +import uuid +from dataclasses import dataclass +from pathlib import Path + +# Defer scraper_impl imports until scrape() has 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 + + +@dataclass(frozen = True) +class ResolvedToken: + value: str + source: str + + +def _resolve_token(token: str) -> ResolvedToken: + if token: + return ResolvedToken( + value = token, + source = "explicit token argument (recipe-level field)", + ) + if os.environ.get("GH_TOKEN"): + return ResolvedToken( + value = os.environ["GH_TOKEN"], + source = "GH_TOKEN environment variable", + ) + if os.environ.get("GITHUB_TOKEN"): + return ResolvedToken( + value = os.environ["GITHUB_TOKEN"], + source = "GITHUB_TOKEN environment variable", + ) + raise ValueError( + "GitHub token is required. Set it in the recipe config or the GH_TOKEN / GITHUB_TOKEN env var." + ) + + +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.value, token_source = token.source) + base_dir.mkdir(parents = True, exist_ok = True) + + # Per-resource limits; limit <= 0 means "all" (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 diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py new file mode 100644 index 0000000..3201423 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py @@ -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 diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py new file mode 100644 index 0000000..0ba3394 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py @@ -0,0 +1,317 @@ +# 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 datetime import timezone +from email.utils import parsedate_to_datetime +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 GitHubAuthError(RuntimeError): + """Raised when GitHub returns 401/403 due to invalid or insufficient credentials.""" + + +def _retry_after_seconds(value: str | None) -> int | None: + if not value: + return None + try: + return max(0, int(value)) + except ValueError: + pass + try: + retry_at = parsedate_to_datetime(value) + except (TypeError, ValueError, IndexError, OverflowError): + return None + if retry_at.tzinfo is None: + retry_at = retry_at.replace(tzinfo = timezone.utc) + return max(0, int(retry_at.timestamp() - time.time())) + + +class GitHubClient: + def __init__( + self, + min_remaining_graphql: int = 100, + min_remaining_rest: int = 100, + token: str | None = None, + token_source: str | None = None, + ): + if token: + self._token_source = token_source or "explicit token argument (recipe-level field)" + elif os.environ.get("GH_TOKEN"): + self._token_source = "GH_TOKEN environment variable" + token = os.environ["GH_TOKEN"] + elif os.environ.get("GITHUB_TOKEN"): + self._token_source = "GITHUB_TOKEN environment variable" + token = os.environ["GITHUB_TOKEN"] + else: + raise RuntimeError("GH_TOKEN or GITHUB_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 _is_rate_limit_response(self, r: "requests.Response") -> bool: + if r.headers.get("Retry-After"): + return True + if r.headers.get("X-RateLimit-Remaining") == "0": + return True + body = (r.text or "").lower() + return any( + marker in body + for marker in ( + "api rate limit exceeded", + "rate limit exceeded", + "secondary rate limit", + "secondary limit", + "abuse detection mechanism", + "abuse detection", + ) + ) + + def _is_auth_failure(self, r: "requests.Response") -> bool: + """Tell auth failures apart from rate limiting on 401/403. + + 401 is always auth; 403 is auth unless it carries a rate-limit signal + (Retry-After, X-RateLimit-Remaining: 0, or abuse/secondary text). + """ + if r.status_code == 401: + return True + if r.status_code == 403: + return not self._is_rate_limit_response(r) + return False + + def _raise_auth_error(self, r: "requests.Response", endpoint: str) -> None: + snippet = (r.text or "").strip()[:200] + request_id = r.headers.get("X-GitHub-Request-Id") + request_id_message = f" Request ID: {request_id}." if request_id else "" + raise GitHubAuthError( + f"GitHub {endpoint} returned {r.status_code} {r.reason}. " + f"Token source: {self._token_source}. " + f"The token is invalid, expired, or missing required scopes — " + f"retrying will not recover.{request_id_message} Response: {snippet}" + ) + + 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 + 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 self._is_auth_failure(r): + self._raise_auth_error(r, "GraphQL") + if r.status_code == 403 or r.status_code == 429: + # Secondary/abuse rate limit + retry_after = _retry_after_seconds(r.headers.get("Retry-After")) + if retry_after is not None: + log.warning("Secondary rate limit. Sleep %ds.", retry_after) + time.sleep(retry_after + 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"]: + # Allow partial data; retry on RATE_LIMITED + errs = data["errors"] + 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 self._is_auth_failure(r): + self._raise_auth_error(r, "REST") + if r.status_code in (403, 429): + retry_after = _retry_after_seconds(r.headers.get("Retry-After")) + if retry_after is not None: + log.warning("Secondary rate limit on REST. Sleep %ds.", retry_after) + time.sleep(retry_after + 2) + continue + # Primary rate limit + 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 wrap the list in an "items" field + items = items.get("items", []) + for it in items: + yield it + 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 {} diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py new file mode 100644 index 0000000..ff2a81d --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py @@ -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 with unused fragments, so each query includes +only the fragments it references. +""" + +# ---- Fragments (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 } +} +""" diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py new file mode 100644 index 0000000..a7ddaef --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py @@ -0,0 +1,705 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Scraper orchestration: 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 {} + # light=True uses trimmed GraphQL queries (no reviewThreads/reviews/ + # commits/timelineItems/files) so PR pages can be larger without + # hitting GitHub's node-count ceiling. + self.light = light + self.repo_dir = base_dir / f"{owner}__{name}" + self.repo_dir.mkdir(parents = True, exist_ok = True) + 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/page. Clamp by + # trial_limit so limit=1 asks for first:1, not a full 50-item page. + 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 caps at 3/page (GitHub node-count ceiling); + # light query skips nested fields and goes to 25/page. Clamp by + # trial_limit so limit=1 does not fetch a whole 25-item page. + page_cap = 25 if self.light else 3 + trial_cap = self.trial_limits.get(key) + per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap + 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 have no id; synthesize one + f["_syntheticId"] = f"{self.owner}/{self.name}#{number}:{f.get('path')}" + self.writers[out_key].write(f) + info = ff.get("pageInfo") or {} + 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() diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py new file mode 100644 index 0000000..67107c2 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""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 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 diff --git a/studio/backend/plugins/data-designer-unstructured-seed/__init__.py b/studio/backend/plugins/data-designer-unstructured-seed/__init__.py new file mode 100644 index 0000000..3201423 --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/__init__.py @@ -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 diff --git a/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml new file mode 100644 index 0000000..c9d988b --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "data-designer-unstructured-seed" +version = "0.1.0" +description = "Local Data Designer unstructured seed reader plugin" +requires-python = ">=3.11" +dependencies = [ + "data-designer-engine>=0.5.4,<0.6", + "pandas>=2,<3", + "pymupdf>=1.24.0", + "pymupdf4llm>=0.0.17", + "mammoth>=1.8.0", +] + +[project.entry-points."data_designer.plugins"] +unstructured = "data_designer_unstructured_seed.plugin:unstructured_seed_plugin" + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/__init__.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/__init__.py new file mode 100644 index 0000000..78e2e40 --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from .chunking import ( + DEFAULT_CHUNK_OVERLAP, + DEFAULT_CHUNK_SIZE, + build_unstructured_preview_rows, + materialize_unstructured_seed_dataset, + resolve_chunking, +) +from .config import UnstructuredSeedSource +from .impl import UnstructuredSeedReader +from .plugin import unstructured_seed_plugin + +__all__ = [ + "DEFAULT_CHUNK_OVERLAP", + "DEFAULT_CHUNK_SIZE", + "build_unstructured_preview_rows", + "materialize_unstructured_seed_dataset", + "resolve_chunking", + "UnstructuredSeedSource", + "UnstructuredSeedReader", + "unstructured_seed_plugin", +] diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py new file mode 100644 index 0000000..ee7d872 --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py @@ -0,0 +1,272 @@ +# 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 re +from pathlib import Path +from typing import Any + +import pandas as pd + +from utils.paths import ensure_dir, unstructured_seed_cache_root + +DEFAULT_CHUNK_SIZE = 1200 +DEFAULT_CHUNK_OVERLAP = 200 +MAX_CHUNK_SIZE = 20000 +_MIN_BREAK_RATIO = 0.6 +_CACHE_DIR = unstructured_seed_cache_root() + + +def resolve_chunking(chunk_size: Any, chunk_overlap: Any) -> tuple[int, int]: + size = _to_int(chunk_size, DEFAULT_CHUNK_SIZE) + size = max(1, min(size, MAX_CHUNK_SIZE)) + overlap = _to_int(chunk_overlap, DEFAULT_CHUNK_OVERLAP) + overlap = max(0, min(overlap, max(0, size - 1))) + return size, overlap + + +def build_unstructured_preview_rows( + *, source_path: Path, preview_size: int, chunk_size: Any, chunk_overlap: Any +) -> list[dict[str, str]]: + parquet_path, rows = materialize_unstructured_seed_dataset( + source_path = source_path, + chunk_size = chunk_size, + chunk_overlap = chunk_overlap, + ) + count = max(0, int(preview_size)) + if rows: + return rows[:count] + + try: + import pandas as pd + except ImportError as exc: # pragma: no cover + raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc + + dataframe = pd.read_parquet(parquet_path).head(count) + return [ + {"chunk_text": str(value.get("chunk_text", "")).strip()} + for value in dataframe.to_dict(orient = "records") + if str(value.get("chunk_text", "")).strip() + ] + + +def build_multi_file_preview_rows( + *, + file_entries: list[tuple[Path, str]], + preview_size: int, + chunk_size: int | None, + chunk_overlap: int | None, +) -> list[dict[str, str]]: + cs = _to_int(chunk_size, DEFAULT_CHUNK_SIZE) + co = _to_int(chunk_overlap, DEFAULT_CHUNK_OVERLAP) + _, rows = materialize_multi_file_unstructured_seed( + file_entries = file_entries, + chunk_size = cs, + chunk_overlap = co, + ) + return _round_robin_preview(rows, preview_size) + + +def _round_robin_preview(rows: list[dict[str, str]], preview_size: int) -> list[dict[str, str]]: + """Pick preview rows round-robin across source files so each is represented.""" + if not rows or preview_size <= 0: + return [] + + # Group by source_file, preserving first-appearance order. + from collections import OrderedDict + + grouped: OrderedDict[str, list[dict[str, str]]] = OrderedDict() + for row in rows: + key = row.get("source_file", "") + if key not in grouped: + grouped[key] = [] + grouped[key].append(row) + + result: list[dict[str, str]] = [] + iterators = [iter(chunks) for chunks in grouped.values()] + while len(result) < preview_size and iterators: + exhausted: list[int] = [] + for i, it in enumerate(iterators): + if len(result) >= preview_size: + break + val = next(it, None) + if val is not None: + result.append(val) + else: + exhausted.append(i) + for i in reversed(exhausted): + iterators.pop(i) + + return result + + +def materialize_unstructured_seed_dataset( + *, source_path: Path, chunk_size: Any, chunk_overlap: Any +) -> tuple[Path, list[dict[str, str]]]: + resolved = source_path.expanduser().resolve() + if not resolved.is_file(): + raise FileNotFoundError(f"unstructured seed file not found: {resolved}") + + size, overlap = resolve_chunking(chunk_size, chunk_overlap) + key = _compute_cache_key( + source_path = resolved, + chunk_size = size, + chunk_overlap = overlap, + ) + parquet_path = _CACHE_DIR / f"{key}.parquet" + if parquet_path.exists(): + return parquet_path, [] + + text = load_unstructured_text_file(resolved) + chunks = split_text_into_chunks( + text = text, + chunk_size = size, + chunk_overlap = overlap, + ) + if not chunks: + raise ValueError("No text found in unstructured seed source.") + + rows = [{"chunk_text": chunk} for chunk in chunks] + ensure_dir(_CACHE_DIR) + try: + import pandas as pd + except ImportError as exc: # pragma: no cover + raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc + + tmp_path = _CACHE_DIR / f"{key}.tmp.parquet" + pd.DataFrame(rows).to_parquet(tmp_path, index = False) + tmp_path.replace(parquet_path) + return parquet_path, rows + + +def materialize_multi_file_unstructured_seed( + *, + file_entries: list[tuple[Path, str]], # (extracted_txt_path, original_filename) + chunk_size: int, + chunk_overlap: int, +) -> tuple[Path, list[dict[str, str]]]: + """Chunk multiple files into one parquet dataset with a source_file column.""" + chunk_size, chunk_overlap = resolve_chunking(chunk_size, chunk_overlap) + cache_key = _compute_multi_file_cache_key(file_entries, chunk_size, chunk_overlap) + cached = _CACHE_DIR / f"{cache_key}.parquet" + if cached.exists(): + df = pd.read_parquet(cached) + rows = df.to_dict(orient = "records") + return cached, rows + + all_rows: list[dict[str, str]] = [] + for txt_path, orig_name in file_entries: + text = load_unstructured_text_file(txt_path) + chunks = split_text_into_chunks( + text = text, + chunk_size = chunk_size, + chunk_overlap = chunk_overlap, + ) + for chunk in chunks: + all_rows.append({"chunk_text": chunk, "source_file": orig_name}) + + if not all_rows: + raise ValueError("No text found in any uploaded files.") + + df = pd.DataFrame(all_rows) + ensure_dir(_CACHE_DIR) + tmp = _CACHE_DIR / f"{cache_key}.tmp.parquet" + df.to_parquet(tmp, index = False) + tmp.replace(cached) + return cached, all_rows + + +def load_unstructured_text_file(path: Path) -> str: + ext = path.suffix.lower() + if ext not in {".txt", ".md"}: + raise ValueError(f"Unsupported unstructured seed file type: {ext}") + + raw = path.read_text(encoding = "utf-8", errors = "ignore") + return normalize_unstructured_text(raw) + + +def normalize_unstructured_text(text: str) -> str: + normalized = text.replace("\r\n", "\n").replace("\r", "\n") + return re.sub(r"\n{3,}", "\n\n", normalized).strip() + + +def split_text_into_chunks(*, text: str, chunk_size: int, chunk_overlap: int) -> list[str]: + if not text: + return [] + if chunk_size <= 0: + return [text] + + chunks: list[str] = [] + start = 0 + min_break_index = int(chunk_size * _MIN_BREAK_RATIO) + text_len = len(text) + while start < text_len: + end = min(text_len, start + chunk_size) + if end < text_len: + window = text[start:end] + cut = _find_break_index(window, min_break_index) + if cut is not None and cut > 0: + end = start + cut + + if end <= start: + end = min(text_len, start + chunk_size) + + chunk = text[start:end].strip() + if chunk: + chunks.append(chunk) + if end >= text_len: + break + + next_start = end - chunk_overlap + if next_start <= start: + next_start = end + start = max(0, next_start) + + return chunks + + +def _find_break_index(window: str, min_index: int) -> int | None: + breakpoints = ["\n\n", "\n", " "] + for token in breakpoints: + idx = window.rfind(token) + if idx >= min_index: + return idx + len(token) + return None + + +def _to_int(value: Any, fallback: int) -> int: + if isinstance(value, bool): + return fallback + try: + parsed = int(str(value).strip()) + except (TypeError, ValueError): + return fallback + return parsed + + +def _compute_cache_key(*, source_path: Path, chunk_size: int, chunk_overlap: int) -> str: + stat = source_path.stat() + payload = "|".join( + [ + str(source_path), + str(stat.st_size), + str(stat.st_mtime_ns), + str(chunk_size), + str(chunk_overlap), + ] + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _compute_multi_file_cache_key( + file_entries: list[tuple[Path, str]], chunk_size: int, chunk_overlap: int +) -> str: + parts: list[str] = [] + for path, name in sorted(file_entries, key = lambda e: e[1]): + st = path.stat() + parts.append(f"{path}|{st.st_size}|{st.st_mtime_ns}|{name}") + parts.append(f"cs={chunk_size}|co={chunk_overlap}") + raw = "\n".join(parts) + return hashlib.sha256(raw.encode()).hexdigest() diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py new file mode 100644 index 0000000..9bef34c --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py @@ -0,0 +1,51 @@ +# 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 pathlib import Path +from typing import Literal + +from pydantic import Field, field_validator, model_validator + +from data_designer.config.seed_source import SeedSource + +from .chunking import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE, resolve_chunking + + +class UnstructuredSeedSource(SeedSource): + seed_type: Literal["unstructured"] = "unstructured" + paths: list[str] = Field(min_length = 1) + + @model_validator(mode = "before") + @classmethod + def _normalize_legacy_path(cls, data): + if isinstance(data, dict) and "paths" not in data and data.get("path"): + data = dict(data) + data["paths"] = [data["path"]] + return data + + chunk_size: int = DEFAULT_CHUNK_SIZE + chunk_overlap: int = DEFAULT_CHUNK_OVERLAP + + @field_validator("paths") + @classmethod + def _validate_paths(cls, v: list[str]) -> list[str]: + for p in v: + expanded = Path(p).expanduser() + if not expanded.is_file(): + raise ValueError(f"Seed file does not exist: {expanded}") + return v + + @field_validator("chunk_size") + @classmethod + def _resolve_chunk_size(cls, v: int) -> int: + cs, _ = resolve_chunking(v, 0) + return cs + + @field_validator("chunk_overlap") + @classmethod + def _resolve_chunk_overlap(cls, v: int, info) -> int: + cs = info.data.get("chunk_size", DEFAULT_CHUNK_SIZE) + _, co = resolve_chunking(cs, v) + return co diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py new file mode 100644 index 0000000..7272e42 --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py @@ -0,0 +1,41 @@ +# 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 pathlib import Path + +import data_designer.lazy_heavy_imports as lazy +from data_designer.engine.resources.seed_reader import SeedReader + +from .config import UnstructuredSeedSource + + +class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]): + def create_duckdb_connection(self): + return lazy.duckdb.connect() + + def get_dataset_uri(self) -> str: + from .chunking import materialize_multi_file_unstructured_seed + import json as json_mod + + file_entries: list[tuple[Path, str]] = [] + for p in self.source.paths: + path_obj = Path(p) + file_id = path_obj.name.replace(".extracted.txt", "") + meta_path = path_obj.parent / f"{file_id}.meta.json" + orig_name = path_obj.name + if meta_path.exists(): + try: + meta = json_mod.loads(meta_path.read_text()) + orig_name = meta.get("original_filename", path_obj.name) + except (json_mod.JSONDecodeError, OSError): + pass + file_entries.append((path_obj, orig_name)) + + path, _ = materialize_multi_file_unstructured_seed( + file_entries = file_entries, + chunk_size = self.source.chunk_size, + chunk_overlap = self.source.chunk_overlap, + ) + return str(path) diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py new file mode 100644 index 0000000..6f0d7ff --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py @@ -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 + +unstructured_seed_plugin = Plugin( + impl_qualified_name = "data_designer_unstructured_seed.impl.UnstructuredSeedReader", + config_qualified_name = "data_designer_unstructured_seed.config.UnstructuredSeedSource", + plugin_type = PluginType.SEED_READER, +) diff --git a/studio/backend/requirements/__init__.py b/studio/backend/requirements/__init__.py new file mode 100644 index 0000000..3201423 --- /dev/null +++ b/studio/backend/requirements/__init__.py @@ -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 diff --git a/studio/backend/requirements/base.txt b/studio/backend/requirements/base.txt new file mode 100644 index 0000000..407ae01 --- /dev/null +++ b/studio/backend/requirements/base.txt @@ -0,0 +1,3 @@ +# Core unsloth packages +unsloth-zoo +unsloth diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt new file mode 100644 index 0000000..5830a47 --- /dev/null +++ b/studio/backend/requirements/extras-no-deps.txt @@ -0,0 +1,22 @@ +# Audio extras (installed with --no-deps --no-cache-dir) +descript-audio-codec +descript-audiotools +julius +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 +# executorch>=1.0.1 # 41.5 MB - no imports in unsloth/zoo/studio +torch-c-dlpack-ext +sentence_transformers==5.2.0 +transformers==4.57.6 +pytorch_tokenizers +kernels==0.12.1 +# kernels<3.11 imports tomli as its tomllib fallback; --no-deps skips its own +# marker dep, so list it here (no-op on the 3.12/3.13 default installs). +tomli; python_version < "3.11" diff --git a/studio/backend/requirements/extras.txt b/studio/backend/requirements/extras.txt new file mode 100644 index 0000000..1baf2b6 --- /dev/null +++ b/studio/backend/requirements/extras.txt @@ -0,0 +1,35 @@ +# transitive dep of onnxruntime (via data-designer's pymupdf4llm) +flatbuffers +# Also needed by sentence_transformers (installed with --no-deps in extras-no-deps.txt); +# librosa pulls it in too, but is skipped in no-torch mode. +scikit-learn==1.7.1 + +# Additional extras +jiwer # WER/CER metrics for vision OCR save-merge benchmarks +omegaconf +einx +pyloudnorm +openai-whisper +uroman # 4.0 MB - used for Outetts. +MeCab # 19.9 MB - used for Outetts. +inflect # number-to-words, required by OuteTTS +loguru +flatten_dict +ffmpy +randomname +argbind +tiktoken +ftfy +importlib-resources +librosa +markdown2 +matplotlib==3.10.9 +pystoi +soundfile +tensorboard +torch-stoi +timm +einops +tabulate +openai>=2.7.2 +websockets>=15.0.1 diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt new file mode 100644 index 0000000..de321f8 --- /dev/null +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -0,0 +1,81 @@ +# Runtime dependencies for no-torch (GGUF-only) mode. +# Installed with --no-deps to prevent transitive torch resolution +# from packages like accelerate, peft, trl, sentence-transformers. +# +# Includes unsloth's own direct deps (typer, pydantic, pyyaml, +# nest-asyncio) since unsloth is also installed with --no-deps +# (current PyPI metadata still declares torch as a hard dep). + +# unsloth direct deps (from pyproject.toml [project].dependencies) +typer +# typer's full runtime dep tree. Required explicitly because this +# file is installed with --no-deps. On Linux/Mac CI runners these +# are often cached transitively; on a fresh windows-latest venv they +# are not, and `unsloth studio setup` crashes with +# `ModuleNotFoundError: No module named 'click'`, then 'annotated_doc', +# then 'rich', etc. as each is hit. Pin the full chain so the +# no-torch path works cleanly on every fresh venv. +click>=8.0 +shellingham>=1.5 +annotated-doc>=0.0.3 +rich>=13.0 +markdown-it-py>=3.0 +mdurl>=0.1 +pygments>=2.0 +# pydantic is intentionally NOT pinned here. install.sh / install.ps1 +# / install_python_stack.py run `pip install pydantic` WITH deps just +# before this --no-deps file is applied, so pip resolves pydantic-core +# to the exact version pydantic's _ensure_pydantic_core_version check +# expects. Pinning both under --no-deps used to drift them apart. +pyyaml +nest-asyncio + +# HF ecosystem (from [huggingfacenotorch] extras in pyproject.toml) +wheel>=0.42.0 +packaging +numpy +tqdm +psutil +tyro +protobuf +sentencepiece>=0.2.0 +safetensors>=0.4.3 +datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0 +accelerate>=0.34.1 +peft>=0.18.0,!=0.11.0 +huggingface_hub>=0.34.0 +hf_transfer +diffusers + +# Transitive deps required because this file is installed with --no-deps. +# Without these, `from transformers import AutoConfig` fails at import time. +regex +typing_extensions +filelock +httpx +httpcore +certifi +idna +anyio>=3.0,<4.14.0 # 4.14 asyncio cancel-scope RuntimeError on Py3.13 streaming (#6483); 4.13 unaffected +sniffio +h11 + +# Unpinned resolves to 0.23.1+ which breaks `from transformers import +# AutoConfig`; transformers 4.56..5.3 declares tokenizers<=0.23.0. +tokenizers<=0.23.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.3.0 +trl>=0.18.2,!=0.19.0,<=0.24.0 +sentence-transformers +cut_cross_entropy +pillow + +# RAG store + document parsing, mirroring studio.txt. Pinned here because +# this file installs --no-deps; without them Studio runs with RAG disabled. +sqlite-vec==0.1.9 +pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 +python-docx==1.2.0 + +lxml==6.0.2 diff --git a/studio/backend/requirements/overrides.txt b/studio/backend/requirements/overrides.txt new file mode 100644 index 0000000..8df4089 --- /dev/null +++ b/studio/backend/requirements/overrides.txt @@ -0,0 +1,5 @@ +# torchao is installed by studio/install_python_stack.py, which selects the +# version matching the torch release actually installed in the venv (torchao's +# C++ extensions are built against one exact torch version, so a fixed pin here +# would skip them on a newer torch). See _select_torchao_spec / +# _probe_installed_torch_version in that file. diff --git a/studio/backend/requirements/single-env/constraints.txt b/studio/backend/requirements/single-env/constraints.txt new file mode 100644 index 0000000..0ed2bf8 --- /dev/null +++ b/studio/backend/requirements/single-env/constraints.txt @@ -0,0 +1,23 @@ +# Single-env pins for unsloth + studio + data-designer +# Keep compatible with unsloth transformers bounds. +transformers==4.57.6 +trl==0.23.1 +huggingface-hub==0.36.2 + +# Studio stack +datasets==4.3.0 +pyarrow==23.0.1 + +# FastMCP compat +fastmcp>=3.0.2 +mcp>=1.24,<2 +websockets>=15.0.1 + +# Cap anyio <4.14: 4.14's new asyncio per-task cancel scope (TaskHandle/_run_coro) +# gets exited in the wrong task on Python 3.13 under starlette's collapsing task +# group, raising "RuntimeError: ... exit a cancel scope that isn't the current +# task's" on streaming responses (#6483); 4.13 has no such code. Global cap so +# later with-deps steps can't re-resolve it up. +anyio<4.14.0 + +pandas==2.3.3 diff --git a/studio/backend/requirements/single-env/data-designer-deps.txt b/studio/backend/requirements/single-env/data-designer-deps.txt new file mode 100644 index 0000000..f63c076 --- /dev/null +++ b/studio/backend/requirements/single-env/data-designer-deps.txt @@ -0,0 +1,26 @@ +# Data Designer runtime deps installed explicitly (single-env mode). +# Synced with data-designer-engine==0.5.4 requirements. +anyascii<1,>=0.3.3 +chardet<6,>=3.0.2 +duckdb<2,>=1.5.0 +faker<21,>=20.1.0 +fsspec<2026,>=2025.3.0 +httpx<1,>=0.27.2 +httpx-retries<1,>=0.4.2 +json-repair<1,>=0.48.0 +jsonpath-rust-bindings<2,>=1.0 +jsonschema<5,>=4.0.0 +lxml<7,>=6.0.2 +marko<3,>=2.1.2 +mcp<2,>=1.26.0 +networkx<4,>=3.0 +python-json-logger>=3,<4 +ruff<1,>=0.14.10 +scipy<2,>=1.11.0 +sqlfluff<4,>=3.2.0 +tiktoken<1,>=0.8.0 +# Local seed plugin deps (plugins installed with --no-deps) +requests>=2.31 +pymupdf>=1.24.0 +pymupdf4llm>=0.0.17 +mammoth>=1.8.0 diff --git a/studio/backend/requirements/single-env/data-designer.txt b/studio/backend/requirements/single-env/data-designer.txt new file mode 100644 index 0000000..2e0ba62 --- /dev/null +++ b/studio/backend/requirements/single-env/data-designer.txt @@ -0,0 +1,5 @@ +# Install Data Designer in same env as Unsloth. +data-designer==0.5.4 +data-designer-config==0.5.4 +data-designer-engine==0.5.4 +prompt-toolkit>=3,<4 diff --git a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt new file mode 100644 index 0000000..43f37b3 --- /dev/null +++ b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt @@ -0,0 +1,18 @@ +# mlx-vlm / mlx-lm declare transformers>=5.x which conflicts with the +# main venv's constraints.txt pin transformers==4.57.6 and forces uv to +# backtrack unsloth. Relax to match the pin -- per-model 5.x routing +# happens at runtime via the side-car venvs. +transformers>=4.57.6 + +# mlx-vlm / mlx-lm pull anyio>=4.14, which fights the constraints.txt cap (needed +# for the 4.14 Python-3.13 streaming cancel-scope RuntimeError, #6483). The -c +# constraint loses that fight on macOS-arm, leaving a half-resolved 4.14/4.13 +# anyio that also ImportErrors on TaskHandle and 500s the server. An override +# wins the fight, so force one consistent <4.14 here too. +anyio<4.14.0 + +# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights +# rejects q_norm/k_norm, so those checkpoints fail to load. mlx-lm #1242. +# The override also drops it from transitive resolution; keep the >=0.22.0 floor +# (mirrors mlx_repair.py _MLX_MIN_VERSIONS) or the resolver could go below it. +mlx-lm>=0.22.0,!=0.31.3 diff --git a/studio/backend/requirements/single-env/patch_metadata.py b/studio/backend/requirements/single-env/patch_metadata.py new file mode 100644 index 0000000..6c3ee4a --- /dev/null +++ b/studio/backend/requirements/single-env/patch_metadata.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Relax strict metadata pins so pip check passes on the single-env stack. + +Why: +- data-designer pins huggingface-hub>=1.0.1 and pyarrow<20. +- unsloth/transformers pins huggingface-hub<1. +- studio datasets pins pyarrow>=21. + +Runtime works with hub 0.36.x + pyarrow 23.x; only the metadata conflicts. +""" + +from __future__ import annotations + +import importlib.metadata as im +import re +from pathlib import Path + +TARGETS = ( + "data-designer", + "data-designer-engine", + "data-designer-config", +) + +PATCHES: tuple[tuple[re.Pattern[str], str], ...] = ( + ( + re.compile(r"^Requires-Dist: huggingface-hub<2,>=1\.0\.1$", re.MULTILINE), + "Requires-Dist: huggingface-hub<2,>=0.34.0", + ), + ( + re.compile(r"^Requires-Dist: pyarrow<20,>=19\.0\.1$", re.MULTILINE), + "Requires-Dist: pyarrow>=21.0.0", + ), +) + + +def metadata_path(dist_name: str) -> Path | None: + try: + dist = im.distribution(dist_name) + except im.PackageNotFoundError: + return None + for f in dist.files or []: + sf = str(f) + if sf.endswith(".dist-info/METADATA"): + return Path(dist.locate_file(f)) + return None + + +def patch_file(path: Path) -> bool: + original = path.read_text(encoding = "utf-8") + updated = original + for pattern, repl in PATCHES: + updated = pattern.sub(repl, updated) + if updated == original: + return False + path.write_text(updated, encoding = "utf-8") + return True + + +def main() -> int: + changed = 0 + checked = 0 + for name in TARGETS: + p = metadata_path(name) + if p is None: + continue + checked += 1 + if patch_file(p): + changed += 1 + print(f"single-env metadata patch: checked={checked}, changed={changed}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt new file mode 100644 index 0000000..6f4a5c3 --- /dev/null +++ b/studio/backend/requirements/studio.txt @@ -0,0 +1,28 @@ +# Studio UI backend dependencies +typer +fastapi +uvicorn +pydantic +packaging +matplotlib==3.10.9 +pandas +nest_asyncio +datasets==4.3.0 +pyjwt +# gradio>=4.0.0 # 148 MB - Studio uses React + FastAPI, not Gradio +huggingface-hub==0.36.2 +structlog>=24.1.0 +diceware +ddgs +cryptography>=42.0.0 +boto3>=1.34.0 # optional: S3 dataset loading +httpx>=0.27.0 +fastmcp>=3.0.2 +# RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in +# extras-no-deps.txt; these add the lexical+dense store and document parsing. +sqlite-vec==0.1.9 +pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 +python-docx==1.2.0 diff --git a/studio/backend/requirements/triton-kernels.txt b/studio/backend/requirements/triton-kernels.txt new file mode 100644 index 0000000..17e265b --- /dev/null +++ b/studio/backend/requirements/triton-kernels.txt @@ -0,0 +1,2 @@ +# Triton kernels (installed with --no-deps, from source) +triton_kernels @ git+https://github.com/triton-lang/triton.git@release/3.6.x#subdirectory=python/triton_kernels diff --git a/studio/backend/routes/.gitkeep b/studio/backend/routes/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py new file mode 100644 index 0000000..2a3baac --- /dev/null +++ b/studio/backend/routes/__init__.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +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 +from routes.export import router as export_router +from routes.training_history import router as training_history_router +from routes.chat_history import router as chat_history_router +from routes.providers import router as providers_router +from routes.mcp_servers import router as mcp_servers_router +from routes.rag import router as rag_router + +__all__ = [ + "training_router", + "models_router", + "inference_router", + "inference_studio_router", + "datasets_router", + "auth_router", + "data_recipe_router", + "export_router", + "training_history_router", + "chat_history_router", + "providers_router", + "mcp_servers_router", + "rag_router", +] + +# Bind the re-export so the import-hoist verifier counts it as used. +_ = (rag_router,) diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py new file mode 100644 index 0000000..92ecdbf --- /dev/null +++ b/studio/backend/routes/auth.py @@ -0,0 +1,575 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Authentication API routes.""" + +from fastapi import APIRouter, Depends, HTTPException, Request, Response, status + +import base64 +import ipaddress +import os +import shlex +import sys +import threading +import time +from collections import deque +from datetime import datetime, timedelta, timezone + +from models.auth import ( + ApiKeyListResponse, + ApiKeyResponse, + AuthLoginRequest, + AuthStatusResponse, + ChangePasswordRequest, + CreateApiKeyRequest, + CreateApiKeyResponse, + DesktopLoginRequest, + RefreshTokenRequest, +) +from models.users import Token +from auth import storage, hashing +from auth.authentication import ( + create_access_token, + create_refresh_token, + get_current_subject, + get_current_subject_allow_password_change, + refresh_access_token, +) + +router = APIRouter() + + +def _reset_password_command() -> str: + """Shell command shown in the 'incorrect password' hint. + + Prefer the absolute path to this install's ``unsloth`` launcher (sibling of + the running interpreter) so the hint works even when its dir isn't on PATH. + + POSIX paths are shell-quoted. On Windows we use the bare absolute path only + when it has no spaces (a quoted path differs between cmd and PowerShell); + otherwise, or if the launcher can't be located, fall back to the PATH form. + """ + try: + bin_dir = os.path.dirname(os.path.abspath(sys.executable)) + if os.name == "nt": + exe = os.path.join(bin_dir, "unsloth.exe") + if os.path.isfile(exe) and " " not in exe: + return f"{exe} studio reset-password" + else: + exe = os.path.join(bin_dir, "unsloth") + if os.path.isfile(exe): + return f"{shlex.quote(exe)} studio reset-password" + except Exception: + pass + return "unsloth studio reset-password" + + +# Per-(ip, username) bucket + per-IP aggregate. Account bucket stops one user's +# typos from blocking others; the aggregate stops username-rotation spray. +# Single-process only; multi-worker deployments need a shared store. +_LOGIN_BUCKETS: dict[tuple[str, str], deque] = {} +_LOGIN_IP_BUCKETS: dict[str, deque] = {} +_LOGIN_BUCKETS_LOCK = threading.Lock() +_LOGIN_WINDOW_SECONDS = 60.0 +_LOGIN_MAX_FAILS = 5 +_LOGIN_IP_MAX_FAILS = 30 +_LOGIN_LOCKOUT_SECONDS = 60 +# Bucket-dict cap. On overflow, reclaim expired buckets; a new IP that still can't +# fit falls back to a sharded overflow rather than evicting a hot bucket. +_LOGIN_MAX_BUCKETS = 4096 +# Last full stale-sweep time; rate-limits the O(n) sweep under a burst of new IPs. +_LAST_IP_PRUNE = 0.0 +# Sharded overflow for per-IP failures that can't get their own bucket while the +# dict is saturated. Each shard is a small fixed-capacity dict ``ip -> [count, +# window_start]``: a per-IP count (so a source is throttled, and cleared on +# success, by its own failures -- no cross-IP collateral) with hard-bounded +# memory and O(1) lookups. When a shard is full a new IP evicts the lowest-count +# entry (and starts clean, never inheriting its count) rather than growing without +# bound, so a high-cardinality spray can't blow memory/CPU the way a per-failure +# deque could; a persistent attacker keeps a high count and is never the one +# evicted. +_LOGIN_IP_OVERFLOW_SHARDS = 256 +_LOGIN_IP_OVERFLOW_MAX = 64 # distinct IPs tracked per shard +_LOGIN_IP_OVERFLOW: list[dict] = [dict() for _ in range(_LOGIN_IP_OVERFLOW_SHARDS)] + + +def _overflow_shard(ip: str) -> dict: + return _LOGIN_IP_OVERFLOW[hash(ip) % _LOGIN_IP_OVERFLOW_SHARDS] + + +def _overflow_record(ip: str, now: float) -> int: + """Record an overflow failure for ``ip`` and return its windowed count.""" + shard = _overflow_shard(ip) + entry = shard.get(ip) + if entry is not None: + if now - entry[1] > _LOGIN_WINDOW_SECONDS: + entry[0], entry[1] = 1, now + else: + # Only "at or above the per-IP threshold" matters for blocking, so cap + # the count there. This also keeps the migration into a per-IP bucket + # bounded -- without the cap a saturated source could accrue an + # unbounded count, then materialize one deque entry per failure + # (``[start] * carried``) on the next attempt, allocating an arbitrarily + # large deque while holding the login lock. + entry[0] = min(entry[0] + 1, _LOGIN_IP_MAX_FAILS) + return entry[0] + if len(shard) >= _LOGIN_IP_OVERFLOW_MAX: + # Make room by dropping the lowest-count entry, but the new source starts + # clean -- never inherit the evicted IP's failures, or an unrelated source + # could be 429'd after one attempt. Worst case under a saturated shard is + # that a heavy hitter briefly resets, not that a bystander is blocked. + del shard[min(shard, key = lambda k: shard[k][0])] + shard[ip] = [1, now] + return 1 + + +def _overflow_blocked(ip: str, now: float) -> int: + """Seconds this IP is throttled by its own overflow count, or 0.""" + shard = _overflow_shard(ip) + entry = shard.get(ip) + if entry is None: + return 0 + if now - entry[1] > _LOGIN_WINDOW_SECONDS: + del shard[ip] + return 0 + if entry[0] >= _LOGIN_IP_MAX_FAILS: + return max(1, int(_LOGIN_WINDOW_SECONDS - (now - entry[1]))) + return 0 + + +def _overflow_take(ip: str, now: float) -> tuple[int, float]: + """Pop ip's overflow entry, returning its ``(count, window_start)`` so the + count can migrate into a fresh per-IP bucket. ``(0, now)`` if none/expired.""" + entry = _overflow_shard(ip).pop(ip, None) + if entry is None or now - entry[1] > _LOGIN_WINDOW_SECONDS: + return 0, now + # Cap the carried count so the bucket migration never allocates more than the + # per-IP threshold worth of deque entries (defensive; _overflow_record already + # clamps, but keep the bound at the consumption site too). + return min(entry[0], _LOGIN_IP_MAX_FAILS), entry[1] + + +# Unrepresentable as a real username (leading NUL); folds unknown-user attempts +# into one slot so attacker cardinality can't blow the bucket dict. +_UNKNOWN_LOGIN_USER = "\x00unknown-user" + + +def _trust_forwarded_for() -> bool: + """Honour X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is set. + + Off by default so a direct caller can't spoof the header. + """ + return os.environ.get("UNSLOTH_STUDIO_TRUST_FORWARDED", "").lower() in ( + "1", + "true", + "yes", + ) + + +def _normalize_forwarded_addr(value: str) -> str: + """Parse an XFF / Forwarded `for=` value into a bare IP (port-stripped).""" + value = (value or "").strip().strip('"') + if not value or value.lower() == "unknown": + return "" + if value.startswith("["): + # Bracketed IPv6, optionally with port. + end = value.find("]") + if end <= 0: + return "" + host = value[1:end] + elif value.count(":") == 1: + # IPv4:port. Bare IPv6 has multiple colons → else branch. + head, _, tail = value.rpartition(":") + host = head if tail.isdigit() and head else value + else: + host = value + try: + return str(ipaddress.ip_address(host)) + except ValueError: + return "" + + +def _forwarded_for_from_element(element: str) -> str: + """Pick the `for=` token out of a single ``Forwarded`` element.""" + for tok in element.split(";"): + key, sep, val = tok.strip().partition("=") + if sep and key.lower() == "for": + return _normalize_forwarded_addr(val) + return "" + + +def _client_ip(request: Request | None) -> str: + if request is None: + return "_unknown" + if _trust_forwarded_for(): + xff = request.headers.get("x-forwarded-for", "") + if xff: + # First entry is the originating client. + normalized = _normalize_forwarded_addr(xff.split(",", 1)[0]) + if normalized: + return normalized + fwd = request.headers.get("forwarded", "") + if fwd: + # First element only; multi-element headers can't fork buckets. + normalized = _forwarded_for_from_element(fwd.split(",", 1)[0]) + if normalized: + return normalized + return (request.client.host if request.client else None) or "_unknown" + + +def _bucket_key(request: Request | None, username: str) -> tuple[str, str]: + return (_client_ip(request), (username or "").casefold()) + + +def _unknown_user_key(request: Request | None) -> tuple[str, str]: + return (_client_ip(request), _UNKNOWN_LOGIN_USER) + + +def _prune_bucket(bucket: deque, now: float) -> None: + while bucket and now - bucket[0] > _LOGIN_WINDOW_SECONDS: + bucket.popleft() + + +def _prune_stale_buckets(now: float) -> None: + """Drop empty / expired account buckets to bound memory under spray.""" + stale: list[tuple[str, str]] = [] + for key, bucket in _LOGIN_BUCKETS.items(): + _prune_bucket(bucket, now) + if not bucket: + stale.append(key) + for key in stale: + _LOGIN_BUCKETS.pop(key, None) + + +def _prune_stale_ip_buckets(now: float) -> None: + """Drop empty / expired per-IP buckets to bound memory under spray. + + The dict is otherwise reclaimed only on a successful login, so a failure-only + spray from many (or spoofed) IPs would grow it without bound. + """ + stale: list[str] = [] + for bucket_ip, bucket in _LOGIN_IP_BUCKETS.items(): + _prune_bucket(bucket, now) + if not bucket: + stale.append(bucket_ip) + for bucket_ip in stale: + _LOGIN_IP_BUCKETS.pop(bucket_ip, None) + + +def _record_login_failure(key: tuple[str, str]) -> int: + global _LAST_IP_PRUNE + now = time.monotonic() + ip, _username = key + with _LOGIN_BUCKETS_LOCK: + # Keep the dict bounded without disabling throttling and without letting a + # spray reset a hot bucket: for a new IP at the cap, reclaim expired buckets + # (rate-limited) to make room. + ip_bucket = _LOGIN_IP_BUCKETS.get(ip) + if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS: + if now - _LAST_IP_PRUNE >= 1.0: + _prune_stale_ip_buckets(now) + _LAST_IP_PRUNE = now + if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS: + # Still full -- every bucket is hot. Count this failure in the IP's + # bounded overflow shard instead of evicting a live one, so the spray + # stays throttled but can't push out (and reset) any IP's own counter. + ip_fails = _overflow_record(ip, now) + else: + if ip_bucket is None: + ip_bucket = _LOGIN_IP_BUCKETS[ip] = deque() + # Carry over any overflow failures this IP accrued while the dict + # was saturated, so straddling the overflow -> bucket transition + # can't double the effective per-IP limit. + carried, start = _overflow_take(ip, now) + ip_bucket.extend([start] * carried) + _prune_bucket(ip_bucket, now) + ip_bucket.append(now) + ip_fails = len(ip_bucket) + + if key not in _LOGIN_BUCKETS and len(_LOGIN_BUCKETS) >= _LOGIN_MAX_BUCKETS: + _prune_stale_buckets(now) + if key in _LOGIN_BUCKETS or len(_LOGIN_BUCKETS) < _LOGIN_MAX_BUCKETS: + account_bucket = _LOGIN_BUCKETS.setdefault(key, deque()) + _prune_bucket(account_bucket, now) + account_bucket.append(now) + return len(account_bucket) + # Both dicts at cap (sustained spray): fall back to the per-IP count. + return ip_fails + + +def _blocked_for(bucket: deque | None, now: float, max_fails: int) -> int: + if not bucket: + return 0 + _prune_bucket(bucket, now) + if len(bucket) >= max_fails: + return max(1, int(_LOGIN_WINDOW_SECONDS - (now - bucket[0]))) + return 0 + + +def _login_blocked(key: tuple[str, str]) -> int: + """Return seconds until the next attempt is allowed, or 0.""" + now = time.monotonic() + ip, _username = key + with _LOGIN_BUCKETS_LOCK: + # Honor the IP's overflow shard regardless of current dict capacity: a + # source counted there during saturation must stay throttled until those + # failures age out, even if a bucket later frees up -- otherwise a fresh + # bucket would reset it. Shards are empty outside saturation, so this is a + # no-op in the common case. + ip_blocked = max( + _blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS), + _overflow_blocked(ip, now), + ) + return max(_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), ip_blocked) + + +def _clear_login_bucket(key: tuple[str, str]) -> None: + ip, _username = key + with _LOGIN_BUCKETS_LOCK: + _LOGIN_BUCKETS.pop(key, None) + _LOGIN_IP_BUCKETS.pop(ip, None) + # A successful login resets the IP's throttle, including any overflow it + # accumulated during saturation (drop only this IP's entry, so a + # shard-mate's throttle is untouched). + _overflow_shard(ip).pop(ip, None) + + +# Sync def (not async): compute_identity_proof touches SQLite on the first call, +# so FastAPI runs it in the threadpool rather than blocking the event loop. +@router.get("/identity") +def identity(nonce: str, request: Request) -> dict: + """Challenge-response proof this is the real local Studio: caller sends a nonce, + gets HMAC(install identity secret, nonce, connection address + port). + Unauthenticated and side-effect free; a process that can't read the same-user + secret can't forge a proof, and binding to the address/port the connection + landed on stops a squatter relaying a proof from the real Studio elsewhere.""" + try: + raw = base64.urlsafe_b64decode(nonce) + except Exception: + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, detail = "nonce must be base64url" + ) + if not 16 <= len(raw) <= 128: + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, detail = "nonce must decode to 16-128 bytes" + ) + # The address + port the connection actually landed on, from the socket + # (request.scope is getsockname, so it is the real local address even when + # bound to 0.0.0.0), never the client-controlled Host header. + server = request.scope.get("server") or ("", 0) + host = server[0] or "" + port = server[1] if server[1] is not None else 0 + return {"proof": storage.compute_identity_proof(raw, host, port)} + + +@router.get("/status", response_model = AuthStatusResponse) +async def auth_status() -> AuthStatusResponse: + """Auth initialization state; ``default_username`` is exposed for first-boot UI prefill only.""" + return AuthStatusResponse( + initialized = storage.is_initialized(), + default_username = storage.DEFAULT_ADMIN_USERNAME, + requires_password_change = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME) + if storage.is_initialized() + else True, + ) + + +@router.post("/login", response_model = Token) +async def login(payload: AuthLoginRequest, request: Request) -> Token: + """Login with username/password. Per-account + per-IP rate-limited.""" + key = _bucket_key(request, payload.username) + unknown_key = _unknown_user_key(request) + blocked_for = max(_login_blocked(key), _login_blocked(unknown_key)) + if blocked_for > 0: + raise HTTPException( + status_code = status.HTTP_429_TOO_MANY_REQUESTS, + # IP not interpolated into the body; behind a proxy/NAT it's + # misleading or an info leak. + detail = (f"Too many failed login attempts. " f"Try again in {blocked_for} seconds."), + headers = {"Retry-After": str(blocked_for)}, + ) + + record = storage.get_user_and_secret(payload.username) + if record is None: + # Record under one sentinel key per IP so attacker-controlled username + # cardinality can't allocate unbounded buckets. + _record_login_failure(unknown_key) + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = f"Incorrect password. To reset it, run this in your terminal: {_reset_password_command()}", + ) + + salt, pwd_hash, _jwt_secret, must_change_password = record + if not hashing.verify_password(payload.password, salt, pwd_hash): + _record_login_failure(key) + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = f"Incorrect password. To reset it, run this in your terminal: {_reset_password_command()}", + ) + + _clear_login_bucket(key) + _clear_login_bucket(unknown_key) + access_token = create_access_token(subject = payload.username) + refresh_token = create_refresh_token(subject = payload.username) + return Token( + access_token = access_token, + refresh_token = refresh_token, + token_type = "bearer", + must_change_password = must_change_password, + ) + + +@router.post("/logout", status_code = status.HTTP_204_NO_CONTENT) +async def logout( + request: Request, current_subject: str = Depends(get_current_subject_allow_password_change) +) -> Response: + """Revoke refresh tokens for the subject; the access token is stateless and expires on its own.""" + try: + storage.revoke_user_refresh_tokens(current_subject) + except Exception: + pass + try: + request.app.state.bootstrap_password = None + except AttributeError: + pass + return Response(status_code = status.HTTP_204_NO_CONTENT) + + +@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: + """Exchange a refresh token for a new access+refresh pair (single-use).""" + consumed = storage.consume_refresh_token(payload.refresh_token) + if consumed is None: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid or expired refresh token", + ) + username, is_desktop = consumed + new_access_token = create_access_token(subject = username, desktop = is_desktop) + new_refresh_token = create_refresh_token(subject = username, desktop = is_desktop) + + return Token( + access_token = new_access_token, + refresh_token = new_refresh_token, + token_type = "bearer", + must_change_password = False if is_desktop else storage.requires_password_change(username), + ) + + +@router.post("/change-password", response_model = Token) +async def change_password( + payload: ChangePasswordRequest, + request: Request, + current_subject: str = Depends(get_current_subject_allow_password_change), +) -> Token: + """Allow the authenticated user to replace the default password.""" + record = storage.get_user_and_secret(current_subject) + if record is None: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "User session is invalid", + ) + + salt, pwd_hash, _jwt_secret, _must_change_password = record + if not hashing.verify_password(payload.current_password, salt, pwd_hash): + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Current password is incorrect", + ) + if payload.current_password == payload.new_password: + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, + detail = "New password must be different from the current password", + ) + + storage.update_password(current_subject, payload.new_password) + storage.revoke_user_refresh_tokens(current_subject) + try: + request.app.state.bootstrap_password = None + except AttributeError: + pass + access_token = create_access_token(subject = current_subject) + refresh_token = create_refresh_token(subject = current_subject) + return Token( + access_token = access_token, + refresh_token = refresh_token, + 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"} diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py new file mode 100644 index 0000000..7a27a58 --- /dev/null +++ b/studio/backend/routes/chat_history.py @@ -0,0 +1,622 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Chat history API routes backed by studio.db. +""" + +from typing import Any, Literal, Optional + +from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from auth.authentication import get_current_subject +from loggers import get_logger +from utils.utils import safe_curated_detail, log_and_http_error +from storage.studio_db import ( + ChatMessageConflictError, + CorruptSettingsError, + clear_chat_history, + count_chat_threads, + count_forks_for_message, + delete_chat_threads, + delete_chat_project, + ensure_chat_project_workspace, + fork_chat_thread, + get_chat_project, + get_chat_thread, + get_chat_message, + list_chat_projects, + list_chat_legacy_imports, + list_chat_settings, + list_chat_messages, + list_chat_messages_for_threads, + list_chat_threads, + sync_chat_messages, + update_chat_project, + update_chat_thread, + upsert_chat_project, + upsert_chat_legacy_imports, + upsert_chat_message, + upsert_chat_settings_merge, + upsert_chat_thread, +) + +router = APIRouter() + +logger = get_logger(__name__) + + +class ChatThread(BaseModel): + id: str + title: str = "New Chat" + modelType: Literal["base", "lora", "model1", "model2"] + modelId: str = "" + pairId: Optional[str] = None + projectId: Optional[str] = None + archived: bool = False + createdAt: int + updatedAt: Optional[int] = None + openaiCodeExecContainerId: Optional[str] = None + anthropicCodeExecContainerId: Optional[str] = None + forkedFromThreadId: Optional[str] = None + forkedFromMessageId: Optional[str] = None + + +class ChatThreadPatch(BaseModel): + title: Optional[str] = None + modelType: Optional[Literal["base", "lora", "model1", "model2"]] = None + modelId: Optional[str] = None + pairId: Optional[str] = None + projectId: Optional[str] = None + archived: Optional[bool] = None + createdAt: Optional[int] = None + updatedAt: Optional[int] = None + openaiCodeExecContainerId: Optional[str] = None + anthropicCodeExecContainerId: Optional[str] = None + + +class ChatMessage(BaseModel): + id: str + threadId: str + parentId: Optional[str] = None + role: str + content: Any = Field(default_factory = list) + attachments: Optional[Any] = None + metadata: Optional[dict[str, Any]] = None + createdAt: int + + +class ChatProject(BaseModel): + id: str + name: str + instructions: str = "" + rootPath: Optional[str] = None + sandboxPath: Optional[str] = None + archived: bool = False + createdAt: int + updatedAt: int + + +class ChatProjectPatch(BaseModel): + name: Optional[str] = None + instructions: Optional[str] = None + archived: Optional[bool] = None + createdAt: Optional[int] = None + updatedAt: Optional[int] = None + + +class ChatThreadListResponse(BaseModel): + threads: list[ChatThread] + + +class ChatProjectListResponse(BaseModel): + projects: list[ChatProject] + + +class ChatMessageListResponse(BaseModel): + messages: list[ChatMessage] + + +class ChatMessageSyncRequest(BaseModel): + messages: list[ChatMessage] + pruneMissing: bool = False + + +class ChatDeleteRequest(BaseModel): + ids: list[str] + + +class ChatCountResponse(BaseModel): + count: int + + +class ChatExportResponse(BaseModel): + exportedAt: str + version: int + threadCount: int + projects: list[ChatProject] = Field(default_factory = list) + threads: list[ChatThread] + messages: list[ChatMessage] + + +class ChatInferenceSettings(BaseModel): + model_config = ConfigDict(extra = "forbid") + + temperature: Optional[float] = None + topP: Optional[float] = None + topK: Optional[float] = None + minP: Optional[float] = None + repetitionPenalty: Optional[float] = None + presencePenalty: Optional[float] = None + maxSeqLength: Optional[float] = None + maxTokens: Optional[float] = None + systemPrompt: Optional[str] = None + systemVariables: Optional[str] = None + trustRemoteCode: Optional[bool] = None + fastMode: Optional[bool] = None + + +class ChatPreset(BaseModel): + model_config = ConfigDict(extra = "forbid") + + name: str + params: ChatInferenceSettings + + +class ChatSettingsPayload(BaseModel): + model_config = ConfigDict(extra = "forbid") + + inferenceParams: Optional[ChatInferenceSettings] = None + customPresets: Optional[list[ChatPreset]] = None + activePreset: Optional[str] = None + activePresetSource: Optional[Literal["builtin-default", "custom", "modified"]] = None + autoTitle: Optional[bool] = None + reasoningEffort: Optional[ + Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"] + ] = None + preserveThinking: Optional[bool] = None + collapseHtmlArtifacts: Optional[bool] = None + allowArtifactNetworkAccess: Optional[bool] = None + autoHealToolCalls: Optional[bool] = None + nudgeToolCalls: Optional[bool] = None + maxToolCallsPerMessage: Optional[int] = Field(default = None, ge = 1) + toolCallTimeout: Optional[int] = Field(default = None, ge = 1) + + +class ChatSettingsResponse(BaseModel): + settings: dict[str, Any] + + +class ChatMessagesBatchRequest(BaseModel): + threadIds: list[str] + + +class ChatMessagesBatchResponse(BaseModel): + messagesByThreadId: dict[str, list[ChatMessage]] + + +class ChatImportLedgerResponse(BaseModel): + threadIds: list[str] + + +class ChatImportLedgerRecordRequest(BaseModel): + # 10k cap bounds the request body; real users have << 1k threads. + threadIds: list[str] = Field(default_factory = list, max_length = 10_000) + + +class ChatImportLedgerRecordResponse(BaseModel): + # accepted: deduped non-empty input count. inserted: rows actually new + # (ON CONFLICT DO NOTHING skips already-recorded ids). + accepted: int + inserted: int + + +@router.get("/threads", response_model = ChatThreadListResponse) +async def list_threads( + model_type: Optional[str] = Query(None), + pair_id: Optional[str] = Query(None), + project_id: Optional[str] = Query(None), + include_archived: bool = Query(True), + current_subject: str = Depends(get_current_subject), +): + threads = list_chat_threads( + model_type = model_type, + pair_id = pair_id, + project_id = project_id, + include_archived = include_archived, + ) + return ChatThreadListResponse(threads = [ChatThread(**t) for t in threads]) + + +@router.post("/threads", response_model = ChatThread) +async def save_thread(payload: ChatThread, current_subject: str = Depends(get_current_subject)): + if payload.projectId and get_chat_project(payload.projectId) is None: + raise HTTPException( + status_code = 404, + detail = f"Project {payload.projectId} not found", + ) + return ChatThread(**upsert_chat_thread(payload.model_dump())) + + +@router.get("/threads/{thread_id}", response_model = ChatThread) +async def get_thread(thread_id: str, current_subject: str = Depends(get_current_subject)): + thread = get_chat_thread(thread_id) + if thread is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + return ChatThread(**thread) + + +@router.patch("/threads/{thread_id}", response_model = ChatThread) +async def patch_thread( + thread_id: str, + payload: ChatThreadPatch, + current_subject: str = Depends(get_current_subject), +): + patch = payload.model_dump(exclude_unset = True) + for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"): + if field in patch and patch[field] is None: + raise HTTPException(status_code = 400, detail = f"{field} cannot be null") + if patch.get("projectId") and get_chat_project(patch["projectId"]) is None: + raise HTTPException( + status_code = 404, + detail = f"Project {patch['projectId']} not found", + ) + thread = update_chat_thread( + thread_id, + patch, + ) + if thread is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + return ChatThread(**thread) + + +@router.delete("/threads") +async def delete_threads( + payload: ChatDeleteRequest, current_subject: str = Depends(get_current_subject) +): + delete_chat_threads(payload.ids) + return {"status": "deleted"} + + +@router.get("/projects", response_model = ChatProjectListResponse) +async def list_projects( + include_archived: bool = Query(False), current_subject: str = Depends(get_current_subject) +): + return ChatProjectListResponse( + projects = [ + ChatProject(**(ensure_chat_project_workspace(project["id"]) or project)) + for project in list_chat_projects(include_archived = include_archived) + ] + ) + + +@router.post("/projects", response_model = ChatProject) +async def save_project(payload: ChatProject, current_subject: str = Depends(get_current_subject)): + return ChatProject(**upsert_chat_project(payload.model_dump())) + + +@router.get("/projects/{project_id}", response_model = ChatProject) +async def get_project(project_id: str, current_subject: str = Depends(get_current_subject)): + project = ensure_chat_project_workspace(project_id) + if project is None: + raise HTTPException( + status_code = 404, + detail = f"Project {project_id} not found", + ) + return ChatProject(**project) + + +@router.patch("/projects/{project_id}", response_model = ChatProject) +async def patch_project( + project_id: str, + payload: ChatProjectPatch, + current_subject: str = Depends(get_current_subject), +): + patch = payload.model_dump(exclude_unset = True) + for field in ("name", "archived", "createdAt", "updatedAt"): + if field in patch and patch[field] is None: + raise HTTPException(status_code = 400, detail = f"{field} cannot be null") + project = update_chat_project(project_id, patch) + if project is not None: + project = ensure_chat_project_workspace(project_id) + if project is None: + raise HTTPException( + status_code = 404, + detail = f"Project {project_id} not found", + ) + return ChatProject(**project) + + +@router.delete("/projects/{project_id}", response_model = ChatProject) +async def delete_project( + project_id: str, + delete_files: bool = Query(False), + current_subject: str = Depends(get_current_subject), +): + project = delete_chat_project(project_id, delete_files = delete_files) + if project is None: + raise HTTPException( + status_code = 404, + detail = f"Project {project_id} not found", + ) + # Best-effort: drop the project's RAG sources (lazy import keeps RAG optional). + try: + import os + + from storage import rag_db + if rag_db.RAG_AVAILABLE: + from core.rag import store as rag_store + from utils.paths import rag_uploads_root + + uploads = os.path.realpath(str(rag_uploads_root())) + conn = rag_db.get_connection() + try: + scope = rag_store.project_scope(project_id) + for doc in rag_store.list_documents(conn, scope): + full = rag_store.get_document(conn, doc["id"]) or {} + rag_store.delete_document(conn, doc["id"]) + stored = full.get("stored_path") + # Also remove the uploaded file; confined to the uploads root. + if stored: + target = os.path.realpath(stored) + if ( + os.path.isfile(target) + and os.path.commonpath([uploads, target]) == uploads + ): + os.remove(target) + finally: + conn.close() + except Exception: # noqa: BLE001 - source cleanup must not block project deletion + logger.warning("failed to delete RAG sources for project %s", project_id, exc_info = True) + return ChatProject(**project) + + +@router.get("/threads/{thread_id}/messages", response_model = ChatMessageListResponse) +async def get_thread_messages(thread_id: str, current_subject: str = Depends(get_current_subject)): + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + return ChatMessageListResponse( + messages = [ChatMessage(**m) for m in list_chat_messages(thread_id)] + ) + + +@router.post("/messages:batch", response_model = ChatMessagesBatchResponse) +async def batch_thread_messages( + payload: ChatMessagesBatchRequest, current_subject: str = Depends(get_current_subject) +): + """One round-trip per sidebar/search rebuild instead of N. Unknown thread ids return empty lists.""" + by_thread: dict[str, list[ChatMessage]] = {tid: [] for tid in payload.threadIds} + for m in list_chat_messages_for_threads(payload.threadIds): + tid = m["threadId"] + if tid in by_thread: + by_thread[tid].append(ChatMessage(**m)) + return ChatMessagesBatchResponse(messagesByThreadId = by_thread) + + +@router.get("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage) +async def get_thread_message( + thread_id: str, + message_id: str, + current_subject: str = Depends(get_current_subject), +): + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + message = get_chat_message(thread_id, message_id) + if message is None: + raise HTTPException(status_code = 404, detail = f"Message {message_id} not found") + return ChatMessage(**message) + + +@router.put("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage) +async def save_thread_message( + thread_id: str, + message_id: str, + payload: ChatMessage, + current_subject: str = Depends(get_current_subject), +): + if thread_id != payload.threadId or message_id != payload.id: + raise HTTPException(status_code = 400, detail = "Message id mismatch") + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + try: + return ChatMessage(**upsert_chat_message(payload.model_dump())) + except ChatMessageConflictError as exc: + raise log_and_http_error( + exc, + 409, + safe_curated_detail(exc), + event = "chat_history.save_message_conflict", + log = logger, + ) from exc + + +@router.put("/threads/{thread_id}/messages", response_model = ChatMessageListResponse) +async def replace_thread_messages( + thread_id: str, + payload: ChatMessageSyncRequest, + current_subject: str = Depends(get_current_subject), +): + mismatched_ids = [message.id for message in payload.messages if message.threadId != thread_id] + if mismatched_ids: + preview = ", ".join(mismatched_ids[:5]) + suffix = "" if len(mismatched_ids) <= 5 else f" (+{len(mismatched_ids) - 5} more)" + raise HTTPException( + status_code = 400, + detail = f"Message threadId mismatch: {preview}{suffix}", + ) + if get_chat_thread(thread_id) is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + messages = [message.model_dump() for message in payload.messages] + try: + return ChatMessageListResponse( + messages = [ + ChatMessage(**m) + for m in sync_chat_messages( + thread_id, + messages, + prune_missing = payload.pruneMissing, + ) + ] + ) + except ChatMessageConflictError as exc: + raise log_and_http_error( + exc, + 409, + safe_curated_detail(exc), + event = "chat_history.replace_messages_conflict", + log = logger, + ) from exc + + +@router.get("/count", response_model = ChatCountResponse) +async def count_threads(current_subject: str = Depends(get_current_subject)): + return ChatCountResponse(count = count_chat_threads()) + + +@router.get("/import-ledger", response_model = ChatImportLedgerResponse) +async def get_import_ledger(current_subject: str = Depends(get_current_subject)): + """Legacy-Dexie import ledger: legacy thread ids already copied into chat tables. + + The frontend checks this on tab open to decide whether to re-run the Dexie -> studio.db import. + """ + return ChatImportLedgerResponse(threadIds = list_chat_legacy_imports()) + + +@router.post("/import-ledger", response_model = ChatImportLedgerRecordResponse) +async def record_import_ledger( + payload: ChatImportLedgerRecordRequest, current_subject: str = Depends(get_current_subject) +): + """Mark each legacy thread id as imported. Idempotent.""" + accepted, inserted = upsert_chat_legacy_imports(payload.threadIds) + return ChatImportLedgerRecordResponse(accepted = accepted, inserted = inserted) + + +@router.delete("") +async def clear_history(current_subject: str = Depends(get_current_subject)): + clear_chat_history() + return {"status": "deleted"} + + +@router.get("/settings", response_model = ChatSettingsResponse) +async def get_settings(current_subject: str = Depends(get_current_subject)): + return ChatSettingsResponse(settings = list_chat_settings()) + + +@router.put("/settings", response_model = ChatSettingsResponse) +async def put_settings( + payload: dict[str, Any], current_subject: str = Depends(get_current_subject) +): + try: + parsed = ChatSettingsPayload.model_validate(payload) + except ValidationError as exc: + raise HTTPException(status_code = 400, detail = exc.errors()) from exc + # Atomic read + deep-merge + write in one BEGIN IMMEDIATE so concurrent updates don't clobber. + try: + return ChatSettingsResponse( + settings = upsert_chat_settings_merge(parsed.model_dump(exclude_unset = True)) + ) + except CorruptSettingsError as exc: + raise log_and_http_error( + exc, + 409, + safe_curated_detail(exc), + event = "chat_history.put_settings_conflict", + log = logger, + ) from exc + + +class ChatForkRequest(BaseModel): + messageId: str + newThreadId: str + createdAt: int + + +class ChatForkResponse(BaseModel): + thread: ChatThread + messages: list[ChatMessage] + containerSnapshotWarning: Optional[str] = None + + +class ChatForkCountResponse(BaseModel): + count: int + + +@router.post("/threads/{thread_id}/fork", response_model = ChatForkResponse) +async def fork_thread( + thread_id: str, + payload: ChatForkRequest, + current_subject: str = Depends(get_current_subject), +): + """Fork a thread at `messageId` -- creates a new thread with + ancestor msgs [root..messageId] copied with fresh ids. Both + code-exec container ids reset on the fork. OpenAI snapshot is a + best-effort enhancement; failure surfaces as + `containerSnapshotWarning` and the fork still succeeds with a + clean sandbox. + """ + import uuid + + source = get_chat_thread(thread_id) + if source is None: + raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") + if get_chat_message(thread_id, payload.messageId) is None: + raise HTTPException( + status_code = 404, + detail = f"Message {payload.messageId} not found in thread {thread_id}", + ) + base_title = source.get("title") or "New Chat" + new_title = f"fork · {base_title}" + forked = fork_chat_thread( + source_thread_id = thread_id, + branch_message_id = payload.messageId, + new_thread_id = payload.newThreadId, + new_title = new_title, + created_at = payload.createdAt, + id_factory = lambda: str(uuid.uuid4()), + ) + if forked is None: + raise HTTPException(status_code = 500, detail = "Fork failed") + messages = list_chat_messages(payload.newThreadId) + # Best-effort OpenAI container snapshot. Stub: a follow-up patch can + # call /v1/containers list+download / create+upload here and patch + # the new openaiCodeExecContainerId. For v1 we always start clean + # and surface the same warning regardless of provider so the UI can + # show a consistent "sandbox starts fresh" toast. + warning: Optional[str] = None + if source.get("openaiCodeExecContainerId") or source.get("anthropicCodeExecContainerId"): + warning = "Sandbox starts fresh in fork; files from parent are not carried over." + return ChatForkResponse( + thread = ChatThread(**forked), + messages = [ChatMessage(**m) for m in messages], + containerSnapshotWarning = warning, + ) + + +@router.get( + "/threads/{thread_id}/messages/{message_id}/forks", + response_model = ChatForkCountResponse, +) +async def get_fork_count( + thread_id: str, + message_id: str, + current_subject: str = Depends(get_current_subject), +): + return ChatForkCountResponse(count = count_forks_for_message(thread_id, message_id)) + + +@router.get("/export", response_model = ChatExportResponse) +async def export_history(current_subject: str = Depends(get_current_subject)): + from datetime import datetime, timezone + + threads = list_chat_threads(include_archived = True) + projects = list_chat_projects(include_archived = True) + messages = list_chat_messages_for_threads([thread["id"] for thread in threads]) + return ChatExportResponse( + exportedAt = datetime.now(timezone.utc).isoformat(), + version = 1, + threadCount = len(threads), + projects = [ChatProject(**project) for project in projects], + threads = [ChatThread(**thread) for thread in threads], + messages = [ChatMessage(**message) for message in messages], + ) diff --git a/studio/backend/routes/data_recipe/__init__.py b/studio/backend/routes/data_recipe/__init__.py new file mode 100644 index 0000000..c596f18 --- /dev/null +++ b/studio/backend/routes/data_recipe/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Data Recipe route package.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +from fastapi import APIRouter, Depends + +from auth.authentication import get_current_subject + +backend_path = Path(__file__).parent.parent.parent +if str(backend_path) not in sys.path: + sys.path.insert(0, str(backend_path)) + +from .jobs import router as jobs_router +from .mcp import router as mcp_router +from .seed import router as seed_router +from .validate import router as validate_router + +router = APIRouter(dependencies = [Depends(get_current_subject)]) +router.include_router(seed_router) +router.include_router(validate_router) +router.include_router(jobs_router) +router.include_router(mcp_router) + +__all__ = ["router"] diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py new file mode 100644 index 0000000..5971438 --- /dev/null +++ b/studio/backend/routes/data_recipe/jobs.py @@ -0,0 +1,627 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Job lifecycle endpoints for data recipe.""" + +from __future__ import annotations + +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 +from pydantic import ValidationError + +from core.data_recipe.huggingface import ( + RecipeDatasetPublishError, + publish_recipe_dataset, +) +from core.data_recipe.jobs import get_job_manager +from loggers import get_logger +from models.data_recipe import ( + JobCreateResponse, + PublishDatasetRequest, + PublishDatasetResponse, + RecipePayload, +) +from utils.utils import safe_error_detail, safe_curated_detail, log_and_http_error + +logger = get_logger(__name__) +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`` (run.py, post-bind) - survives proxies/tunnels. + 2. ``request.scope["server"]`` - when Studio starts outside ``run_server``. + 3. parsed ``request.base_url`` - last resort for test fixtures. + """ + 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 model_aliases actually referenced by an LLM column. + + Narrows the "Chat model loaded" gate so orphan model_config nodes don't + block unrelated runs. The ``llm-`` prefix matches + ``service.py::_recipe_has_llm_columns`` and covers all LLM column types. + """ + 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 _used_local_model_selections( + recipe: dict[str, Any], local_provider_names: set[str] +) -> dict[tuple[str, str], list[str]]: + used_aliases = _used_llm_model_aliases(recipe) + selections: dict[tuple[str, str], list[str]] = {} + for mc in recipe.get("model_configs", []): + if not isinstance(mc, dict): + continue + alias = mc.get("alias") + if not isinstance(alias, str) or alias not in used_aliases: + continue + provider = mc.get("provider") + if not isinstance(provider, str) or provider not in local_provider_names: + continue + model = mc.get("model") + target = model.strip() if isinstance(model, str) else "" + if not target or target.lower() == "local": + continue + variant = mc.get("gguf_variant") + gguf_variant = variant.strip() if isinstance(variant, str) else "" + selections.setdefault((target, gguf_variant), []).append(alias) + return selections + + +def _single_used_local_model_selection( + recipe: dict[str, Any], local_provider_names: set[str] +) -> tuple[str, str] | None: + selections = _used_local_model_selections(recipe, local_provider_names) + if not selections: + return None + if len(selections) > 1: + aliases = ", ".join(alias for values in selections.values() for alias in values) + raise ValueError( + "Recipes supports one active local model per run. " + f"Select the same local model and GGUF variant for: {aliases}." + ) + return next(iter(selections)) + + +def _loaded_local_model_identity() -> tuple[bool, str, str]: + from routes.inference import get_llama_cpp_backend + from core.inference import get_inference_backend + + llama = get_llama_cpp_backend() + if llama.is_loaded: + model = str(getattr(llama, "model_identifier", "") or "").strip() + variant = str(getattr(llama, "hf_variant", "") or "").strip() + return True, model, variant + + backend = get_inference_backend() + active_model = str(getattr(backend, "active_model_name", "") or "").strip() + if active_model: + return True, active_model, "" + return False, "", "" + + +def _ensure_selected_local_model_loaded( + recipe: dict[str, Any], local_provider_names: set[str] +) -> None: + model_loaded, active_model, active_variant = _loaded_local_model_identity() + if not model_loaded: + raise ValueError("No model loaded in Chat. Load a model first, then run the recipe.") + + selection = _single_used_local_model_selection(recipe, local_provider_names) + if selection is None: + return + + target, gguf_variant = selection + variant_matches = not gguf_variant or active_variant == gguf_variant + if active_model.lower() != target.lower() or not variant_matches: + selected = f"{target} ({gguf_variant})" if gguf_variant else target + active = f"{active_model} ({active_variant})" if active_variant else active_model + raise ValueError( + "Selected local model is not loaded. " + f"Selected {selected}; active {active or 'none'}. " + "Load the selected model again, then run the recipe." + ) + + +def _inject_local_structured_response_format( + recipe: dict[str, Any], local_provider_names: set[str] +) -> None: + """Inject an OpenAI ``response_format`` for each local llm-structured column. + + Clones the model_config and repoints the column at the clone so llm-text / + llm-judge columns sharing the alias keep free-form sampling. Without this, + data_designer only adds a prompt-level "return JSON" hint, which small GGUFs + often break. Forwarding ``response_format`` lets llama-server apply + grammar-constrained sampling, guaranteeing parseable output. + """ + 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 sharing the + # 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 + # BaseInferenceParams is extra="forbid", so response_format can't sit at + # the top level. Its `extra_body` passthrough is spread into the request + # body top level by the OpenAI client, where llama-server reads + # response_format. Per tools/server/README.md the schema sits directly + # under response_format (not nested in a json_schema object as OpenAI + # expects) and is converted to a GBNF grammar for 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, + } + # The OpenAI chat endpoint now returns raw JSON by default for + # response_format requests (spec compliance for public clients). This + # internal opt-in flag rides through the OpenAI SDK's extra_body + # passthrough alongside response_format and re-enables the ```json + # markdown fence that data_designer's structured-output parser expects. + extra_body["_unsloth_guided_fence"] = True + 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 in-place: point is_local providers at this server and mint + a short-lived internal sk-unsloth-* key for workflow auth. + + Returns the minted key's row id (for the caller to revoke on completion), or + ``None`` when no local provider is reachable from an LLM column. + """ + providers = recipe.get("model_providers") + if not providers: + return None + + # Collect local providers and pop is_local from ALL dicts. Strict `is True` + # guard so malformed payloads (1, "true") don't 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 reachable from an LLM + # column via a model_config. Orphan model_config nodes shouldn't block runs; + # the recipe never calls /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 the selected local model is loaded before minting a key. Still + # a point-in-time check (TOCTOU); the /v1 endpoint returns a clear 400 if + # the model is unloaded or swapped before the subprocess calls it. + _ensure_selected_local_model_loaded(recipe, local_names) + + from auth import storage # deferred: avoids circular import + + # Mint an internal sk-unsloth-* key scoped to this run via the unified + # API-key path. Marked internal so it's hidden from the user's key list; + # the caller revokes it when the job terminates. + expires_at = (datetime.now(timezone.utc) + timedelta(hours = 24)).isoformat() + token, row = storage.create_api_key( + username = "unsloth", + name = "data-recipe workflow", + expires_at = expires_at, + internal = True, + ) + internal_key_id = int(row["id"]) + + # Strip stale "external"-only fields (extra_headers/extra_body/api_key_env) + # the frontend may have serialized; a provider flipped from external to local + # could otherwise carry invalid JSON or rogue auth headers into the /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 local model_configs. llama-server's /v1/models + # response can differ from the selected id (cache aliases, GGUF variants), + # and we already gated on a loaded backend, so the health check would be + # redundant and could reject valid local selections. + for mc in recipe.get("model_configs", []): + if not isinstance(mc, dict): + continue + if mc.get("provider") in local_names: + mc["skip_health_check"] = True + # Disable thinking for local data-recipe inference. The + # ... preamble roughly doubles tokens per row and + # pushes answers past data_designer's json-fence regex. Forward + # chat_template_kwargs={enable_thinking: False} via extra_body so + # llama-server renders the template without it: llm-text columns get + # the latency cut, structured columns stop leaking think tags. + 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["enable_thinking"] = False + extra_body["chat_template_kwargs"] = tpl_kwargs + params["extra_body"] = extra_body + + # Forward each llm-structured column's output_format as a response_format so + # llama-server uses grammar-constrained sampling instead of 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 + if not isinstance(value, str): + raise HTTPException(status_code = 400, detail = "invalid run_name: must be a string") + trimmed = value.strip() + if not trimmed: + return None + return trimmed[:120] + + +@router.post("/jobs", response_class = JSONResponse, response_model = JobCreateResponse) +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.") + + run: dict[str, Any] = payload.run or {} + run.pop("artifact_path", None) + run.pop("dataset_name", None) + execution_type = str(run.get("execution_type") or "full").strip().lower() + if execution_type not in {"preview", "full"}: + raise HTTPException( + status_code = 400, + detail = "invalid execution_type: must be 'preview' or 'full'", + ) + run["execution_type"] = execution_type + run["run_name"] = _normalize_run_name(run.get("run_name")) + run_config_raw = run.get("run_config") + if run_config_raw is not None: + try: + from data_designer.config.run_config import RunConfig + RunConfig.model_validate(run_config_raw) + except (ImportError, ValidationError, TypeError, ValueError) as exc: + raise log_and_http_error( + exc, + 400, + "invalid run_config", + event = "data_recipe.jobs.run_config_invalid", + log = logger, + ) from exc + + try: + internal_api_key_id = _inject_local_providers(recipe, request) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_curated_detail(exc), + event = "data_recipe.jobs.inject_local_providers_failed", + log = logger, + ) from exc + + # Single try over get_job_manager() AND mgr.start() so a minted key never + # outlives the request on an unexpected exception; without the bare except it + # would 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 log_and_http_error( + exc, + 409, + safe_curated_detail(exc), + event = "data_recipe.jobs.start_conflict", + log = logger, + ) from exc + except ValueError as exc: + if internal_api_key_id is not None: + _revoke_internal_api_key_safe(internal_api_key_id) + raise log_and_http_error( + exc, + 400, + safe_curated_detail(exc), + event = "data_recipe.jobs.start_failed", + log = logger, + ) 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; never mask the caller's error.""" + 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() + state = mgr.get_status(job_id) + if state is None: + raise HTTPException(status_code = 404, detail = "job not found") + return state + + +@router.get("/jobs/current") +def current_job(): + mgr = get_job_manager() + state = mgr.get_current_status() + if state is None: + raise HTTPException(status_code = 404, detail = "no job") + return state + + +@router.post("/jobs/{job_id}/cancel") +def cancel_job(job_id: str): + mgr = get_job_manager() + ok = mgr.cancel(job_id) + if not ok: + raise HTTPException(status_code = 404, detail = "job not found") + return mgr.get_status(job_id) + + +@router.get("/jobs/{job_id}/analysis") +def job_analysis(job_id: str): + mgr = get_job_manager() + analysis = mgr.get_analysis(job_id) + if analysis is None: + raise HTTPException(status_code = 404, detail = "analysis not ready") + return analysis + + +@router.get("/jobs/{job_id}/dataset") +def job_dataset( + job_id: str, + limit: int = Query(default = 20, ge = 1, le = 500), + offset: int = Query(default = 0, ge = 0), +): + mgr = get_job_manager() + result = mgr.get_dataset(job_id, limit = limit, offset = offset) + if result is None: + raise HTTPException(status_code = 404, detail = "dataset not ready") + if "error" in result: + raise HTTPException(status_code = 422, detail = result["error"]) + return { + "dataset": result["dataset"], + "total": result["total"], + "limit": limit, + "offset": offset, + } + + +@router.post( + "/jobs/{job_id}/publish", + response_class = JSONResponse, + response_model = PublishDatasetResponse, +) +def publish_job_dataset(job_id: str, payload: PublishDatasetRequest): + repo_id = payload.repo_id.strip() + description = payload.description.strip() + hf_token = payload.hf_token.strip() if isinstance(payload.hf_token, str) else None + artifact_path = ( + payload.artifact_path.strip() if isinstance(payload.artifact_path, str) else None + ) + + if not repo_id: + raise HTTPException(status_code = 400, detail = "repo_id is required") + if not description: + raise HTTPException(status_code = 400, detail = "description is required") + + mgr = get_job_manager() + status = mgr.get_status(job_id) + if status is not None: + if status.get("status") != "completed" or status.get("execution_type") != "full": + raise HTTPException( + status_code = 409, + detail = "Only completed full runs can be published.", + ) + status_artifact = status.get("artifact_path") + if isinstance(status_artifact, str) and status_artifact.strip(): + artifact_path = status_artifact.strip() + + if not artifact_path: + raise HTTPException( + status_code = 400, + detail = "This execution does not have publishable dataset artifacts.", + ) + + try: + url = publish_recipe_dataset( + artifact_path = artifact_path, + repo_id = repo_id, + description = description, + hf_token = hf_token or None, + private = payload.private, + ) + except RecipeDatasetPublishError as exc: + raise log_and_http_error( + exc, + 400, + safe_curated_detail(exc), + event = "data_recipe.jobs.publish_failed", + log = logger, + ) from exc + except Exception as exc: + raise log_and_http_error( + exc, + 500, + safe_error_detail(exc), + event = "data_recipe.jobs.publish_error", + log = logger, + ) from exc + + return { + "success": True, + "url": url, + "message": f"Published dataset to {repo_id}.", + } + + +@router.get("/jobs/{job_id}/events") +async def job_events(request: Request, job_id: str): + mgr = get_job_manager() + last_id = request.headers.get("last-event-id") + after_seq: int | None = None + if last_id: + try: + after_seq = int(str(last_id).strip()) + except (TypeError, ValueError): + after_seq = None + + after_q = request.query_params.get("after") + if after_q: + try: + after_seq = int(str(after_q).strip()) + except (TypeError, ValueError): + pass + + sub = mgr.subscribe(job_id, after_seq = after_seq) + if sub is None: + raise HTTPException(status_code = 404, detail = "job not found") + + async def gen(): + try: + for event in sub.replay: + yield sub.format_sse(event) + + while True: + if await request.is_disconnected(): + break + event = await sub.next_event(timeout_sec = 1.0) + if event is None: + continue + yield sub.format_sse(event) + finally: + mgr.unsubscribe(sub) + + return StreamingResponse(gen(), media_type = "text/event-stream") diff --git a/studio/backend/routes/data_recipe/mcp.py b/studio/backend/routes/data_recipe/mcp.py new file mode 100644 index 0000000..78a3987 --- /dev/null +++ b/studio/backend/routes/data_recipe/mcp.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""MCP helper endpoints for data recipe.""" + +from __future__ import annotations + +from collections import defaultdict + +from fastapi import APIRouter + +from core.data_recipe.service import build_mcp_providers +from loggers import get_logger +from models.data_recipe import ( + McpToolsListRequest, + McpToolsListResponse, + McpToolsProviderResult, +) +from utils.utils import safe_error_detail + +logger = get_logger(__name__) +router = APIRouter() + + +@router.post("/mcp/tools", response_model = McpToolsListResponse) +def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse: + try: + from data_designer.engine.mcp import io as mcp_io + except ImportError as exc: + logger.error( + "data_recipe.mcp.dependencies_unavailable", + error = str(exc), + exc_info = True, + ) + return McpToolsListResponse( + providers = [ + McpToolsProviderResult( + name = "", + error = "MCP dependencies unavailable.", + ) + ] + ) + + providers: list[McpToolsProviderResult] = [] + tool_to_providers: dict[str, list[str]] = defaultdict(list) + + from core.inference.mcp_client import stdio_mcp_enabled + + for provider_payload in payload.mcp_providers: + provider_name = str(provider_payload.get("name", "")).strip() + if provider_payload.get("provider_type") == "stdio" and not stdio_mcp_enabled(): + providers.append( + McpToolsProviderResult( + name = provider_name, + error = "Local (stdio) MCP servers are disabled on this host.", + ) + ) + continue + built = build_mcp_providers({"mcp_providers": [provider_payload]}) + if len(built) != 1: + providers.append( + McpToolsProviderResult( + name = provider_name, + error = "Unsupported MCP provider config.", + ) + ) + continue + + provider = built[0] + try: + tools = mcp_io.list_tools(provider, timeout_sec = payload.timeout_sec) + tool_names = sorted({tool.name for tool in tools if getattr(tool, "name", "")}) + for tool_name in tool_names: + tool_to_providers[tool_name].append(provider.name) + providers.append( + McpToolsProviderResult( + name = provider.name, + tools = tool_names, + ) + ) + except Exception as exc: + logger.error( + "data_recipe.mcp.list_tools_failed", + error = str(exc), + exc_info = True, + ) + providers.append( + McpToolsProviderResult( + name = provider.name or provider_name, + error = safe_error_detail(exc, fallback = "Failed to load tools."), + ) + ) + + duplicate_tools = { + tool_name: provider_names + for tool_name, provider_names in sorted(tool_to_providers.items()) + if len(provider_names) > 1 + } + + return McpToolsListResponse( + providers = providers, + duplicate_tools = duplicate_tools, + ) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py new file mode 100644 index 0000000..a5b75b7 --- /dev/null +++ b/studio/backend/routes/data_recipe/seed.py @@ -0,0 +1,717 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Seed inspect endpoints for data recipe.""" + +from __future__ import annotations + +import base64 +import binascii +import json +import os +import re +import shutil +from itertools import islice +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from fastapi import APIRouter, HTTPException, UploadFile, File as FastAPIFile, Form + +try: + from data_designer_unstructured_seed.chunking import ( + build_multi_file_preview_rows, + build_unstructured_preview_rows, + normalize_unstructured_text, + resolve_chunking, + ) +except ImportError: + build_multi_file_preview_rows = None + build_unstructured_preview_rows = None + normalize_unstructured_text = None + resolve_chunking = None +from core.data_recipe.jsonable import to_preview_jsonable +from loggers import get_logger +from utils.paths import ensure_dir, seed_uploads_root, unstructured_uploads_root +from utils.utils import log_and_http_error +from utils.upload_limits import ( + LOCAL_SEED_UPLOAD_MAX_BYTES, + LOCAL_SEED_UPLOAD_MAX_LABEL, + UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES, + UNSTRUCTURED_RECIPE_UPLOAD_MAX_LABEL, + UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES, + UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_LABEL, +) + +from models.data_recipe import ( + SeedInspectRequest, + SeedInspectResponse, + SeedInspectUploadRequest, + UnstructuredFileUploadResponse, +) + +logger = get_logger(__name__) +router = APIRouter() + +DATA_EXTS = (".parquet", ".jsonl", ".json", ".csv") +DEFAULT_SPLIT = "train" +LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl"} +UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"} +SEED_UPLOAD_DIR = seed_uploads_root() +UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root() +_SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$") +# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...) +# never match: those directories can be shared by several recipes. +_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$") + + +def _validate_safe_id(value: str, label: str) -> str: + if not value or not _SAFE_ID_RE.match(value): + raise HTTPException(400, f"Invalid {label}: must be alphanumeric/dash/underscore only") + return value + + +def _serialize_preview_value(value: Any) -> Any: + return to_preview_jsonable(value) + + +def _serialize_preview_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + {str(key): _serialize_preview_value(value) for key, value in row.items()} for row in rows + ] + + +def _normalize_optional_text(value: str | None) -> str | None: + if value is None: + return None + trimmed = value.strip() + return trimmed if trimmed else None + + +def _list_hf_data_files(*, dataset_name: str, token: str | None) -> list[str]: + try: + from huggingface_hub import HfApi + from huggingface_hub.utils import HfHubHTTPError + except ImportError: + return [] + try: + api = HfApi() + repo_files = api.list_repo_files(dataset_name, repo_type = "dataset", token = token) + return [file for file in repo_files if file.lower().endswith(DATA_EXTS)] + except (HfHubHTTPError, OSError, ValueError): + return [] + + +def _select_best_file(data_files: list[str], split: str = DEFAULT_SPLIT) -> str | None: + if not data_files: + return None + split_lower = split.lower() + + def score(path: str) -> tuple[int, int]: + name = path.lower() + if f"/{split_lower}/" in name: + return (0, len(path)) + if ( + f"_{split_lower}." in name + or f"-{split_lower}." in name + or f"/{split_lower}." in name + or f"/{split_lower}_" in name + or f"/{split_lower}-" in name + ): + return (1, len(path)) + return (2, len(path)) + + return sorted(data_files, key = score)[0] + + +def _resolve_seed_hf_path( + dataset_name: str, + data_files: list[str], + split: str = DEFAULT_SPLIT, +) -> str | None: + selected = _select_best_file(data_files, split) + if not selected: + return None + + ext = Path(selected).suffix.lower() + if ext not in DATA_EXTS: + return f"datasets/{dataset_name}/{selected}" + + parent = Path(selected).parent.as_posix() + if not parent or parent == ".": + return f"datasets/{dataset_name}/**/*{ext}" + return f"datasets/{dataset_name}/{parent}/**/*{ext}" + + +def _build_stream_load_kwargs( + *, + dataset_name: str, + split: str, + subset: str | None, + token: str | None, + data_file: str | None = None, +) -> dict[str, Any]: + kwargs: dict[str, Any] = { + "path": dataset_name, + "split": split, + "streaming": True, + "trust_remote_code": False, + } + if data_file: + kwargs["data_files"] = [data_file] + if subset: + kwargs["name"] = subset + if token: + kwargs["token"] = token + return kwargs + + +def _load_preview_rows( + *, load_dataset_fn, load_kwargs: dict[str, Any], preview_size: int +) -> list[dict[str, Any]]: + streamed_ds = load_dataset_fn(**load_kwargs) + return [row for row in islice(streamed_ds, preview_size)] + + +def _extract_columns(rows: list[dict[str, Any]]) -> list[str]: + columns_seen: dict[str, None] = {} + for row in rows: + for key in row.keys(): + columns_seen[str(key)] = None + return list(columns_seen.keys()) + + +def _sanitize_filename(filename: str) -> str: + name = Path(filename).name.strip().replace("\x00", "") + if not name: + return "seed_upload" + return name + + +def _decode_base64_payload(content_base64: str) -> bytes: + raw = content_base64.strip() + if "," in raw and raw.lower().startswith("data:"): + raw = raw.split(",", 1)[1] + try: + return base64.b64decode(raw, validate = True) + except binascii.Error as exc: + raise HTTPException(status_code = 400, detail = "invalid base64 payload") from exc + + +def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[dict[str, Any]]: + try: + import pandas as pd + except ImportError as exc: + raise log_and_http_error( + exc, + 500, + "seed inspect dependencies unavailable", + event = "data_recipe.seed.dependencies_unavailable", + log = logger, + ) from exc + + ext = path.suffix.lower() + try: + if ext == ".csv": + df = pd.read_csv(path, nrows = preview_size, encoding = "utf-8-sig") + df.columns = df.columns.str.strip() + unnamed = [c for c in df.columns if c == "" or c.startswith("Unnamed:")] + if unnamed: + df = df.drop(columns = unnamed) + full_df = pd.read_csv(path, encoding = "utf-8-sig") + full_df.columns = full_df.columns.str.strip() + full_df = full_df.drop(columns = unnamed) + tmp_csv = path.with_suffix(".tmp.csv") + full_df.to_csv(tmp_csv, index = False, encoding = "utf-8") + tmp_csv.replace(path) + elif ext == ".jsonl": + df = pd.read_json(path, lines = True).head(preview_size) + elif ext == ".json": + try: + df = pd.read_json(path).head(preview_size) + except ValueError: + df = pd.read_json(path, lines = True).head(preview_size) + else: + raise HTTPException(status_code = 422, detail = f"unsupported file type: {ext}") + except HTTPException: + raise + except (ValueError, OSError) as exc: + raise log_and_http_error( + exc, + 422, + "seed inspect failed", + event = "data_recipe.seed.local_preview_failed", + log = logger, + ) from exc + + rows = df.to_dict(orient = "records") + return _serialize_preview_rows(rows) + + +def _read_preview_rows_from_unstructured_file( + *, path: Path, preview_size: int, chunk_size: int | None, chunk_overlap: int | None +) -> list[dict[str, Any]]: + if resolve_chunking is None or build_unstructured_preview_rows is None: + raise HTTPException( + 500, + "Unstructured seed support not available (missing data_designer_unstructured_seed)", + ) + size, overlap = resolve_chunking(chunk_size, chunk_overlap) + try: + rows = build_unstructured_preview_rows( + source_path = path, + preview_size = preview_size, + chunk_size = size, + chunk_overlap = overlap, + ) + except (FileNotFoundError, RuntimeError, ValueError, OSError) as exc: + raise log_and_http_error( + exc, + 422, + "seed inspect failed", + event = "data_recipe.seed.unstructured_preview_failed", + log = logger, + ) from exc + return _serialize_preview_rows(rows) + + +def _read_preview_rows_from_multi_files( + *, + block_id: str, + file_ids: list[str], + file_names: list[str], + preview_size: int, + chunk_size: int | None, + chunk_overlap: int | None, +) -> list[dict[str, str]]: + if build_multi_file_preview_rows is None: + raise HTTPException( + 500, + "Unstructured seed support not available (missing data_designer_unstructured_seed)", + ) + + _validate_safe_id(block_id, "block_id") + block_dir = UNSTRUCTURED_UPLOAD_ROOT / block_id + file_entries: list[tuple[Path, str]] = [] + for fid, fname in zip(file_ids, file_names): + extracted = block_dir / f"{fid}.extracted.txt" + if not extracted.exists(): + raise HTTPException(404, f"Extracted text not found for file: {fname} (id: {fid})") + file_entries.append((extracted, fname)) + + return build_multi_file_preview_rows( + file_entries = file_entries, + preview_size = preview_size, + chunk_size = chunk_size, + chunk_overlap = chunk_overlap, + ) + + +@router.post("/seed/inspect", response_model = SeedInspectResponse) +def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse: + dataset_name = payload.dataset_name.strip() + if not dataset_name or dataset_name.count("/") < 1: + raise HTTPException( + status_code = 400, + detail = "dataset_name must be a Hugging Face repo id like org/repo", + ) + + try: + from datasets import load_dataset + except ImportError as exc: + raise log_and_http_error( + exc, + 500, + "seed inspect dependencies unavailable", + event = "data_recipe.seed.dependencies_unavailable", + log = logger, + ) from exc + + split = _normalize_optional_text(payload.split) or DEFAULT_SPLIT + subset = _normalize_optional_text(payload.subset) + token = _normalize_optional_text(payload.hf_token) + preview_size = int(payload.preview_size) + + preview_rows: list[dict[str, Any]] = [] + data_files = _list_hf_data_files(dataset_name = dataset_name, token = token) + + selected_file = _select_best_file(data_files, split) + if selected_file: + try: + single_file_kwargs = _build_stream_load_kwargs( + dataset_name = dataset_name, + split = split, + subset = subset, + token = token, + data_file = selected_file, + ) + preview_rows = _load_preview_rows( + load_dataset_fn = load_dataset, + load_kwargs = single_file_kwargs, + preview_size = preview_size, + ) + except (ValueError, OSError, RuntimeError): + preview_rows = [] + + if not preview_rows: + try: + split_kwargs = _build_stream_load_kwargs( + dataset_name = dataset_name, + split = split, + subset = subset, + token = token, + ) + preview_rows = _load_preview_rows( + load_dataset_fn = load_dataset, + load_kwargs = split_kwargs, + preview_size = preview_size, + ) + except (ValueError, OSError, RuntimeError) as exc: + raise log_and_http_error( + exc, + 422, + "seed inspect failed", + event = "data_recipe.seed.hf_preview_failed", + log = logger, + ) from exc + + if not preview_rows: + raise HTTPException(status_code = 422, detail = "dataset appears empty or unreadable") + preview_rows = _serialize_preview_rows(preview_rows) + columns = _extract_columns(preview_rows) + + if not data_files: + resolved_path = f"datasets/{dataset_name}/**/*.parquet" + else: + resolved_path = _resolve_seed_hf_path(dataset_name, data_files, split) + if not resolved_path: + raise HTTPException(status_code = 422, detail = "unable to resolve seed dataset path") + + return SeedInspectResponse( + dataset_name = dataset_name, + resolved_path = resolved_path, + columns = columns, + preview_rows = preview_rows, + split = split, + subset = subset, + ) + + +def _extract_text_from_file(file_path: Path, ext: str) -> str: + """Extract text from an uploaded file by extension, to markdown where possible.""" + if ext in {".txt", ".md"}: + raw = file_path.read_text(encoding = "utf-8", errors = "ignore") + elif ext == ".pdf": + import pymupdf4llm + raw = pymupdf4llm.to_markdown( + str(file_path), write_images = False, show_progress = False, use_ocr = False + ) + elif ext == ".docx": + import mammoth + with open(str(file_path), "rb") as f: + result = mammoth.convert_to_markdown(f) + raw = result.value + else: + raise ValueError(f"Unsupported file type: {ext}") + + if normalize_unstructured_text is None: + return raw + return normalize_unstructured_text(raw) + + +def _get_block_total_size(block_dir: Path) -> int: + """Sum raw upload sizes for the whole block from server-owned files.""" + if not block_dir.exists(): + return 0 + total = 0 + for f in block_dir.iterdir(): + if not f.is_file(): + continue + if f.name.endswith(".extracted.txt") or f.name.endswith(".meta.json"): + continue + total += f.stat().st_size + return total + + +@router.post("/seed/upload-unstructured-file") +async def upload_unstructured_file( + file: UploadFile = FastAPIFile(...), block_id: str = Form(...) +) -> UnstructuredFileUploadResponse: + _validate_safe_id(block_id, "block_id") + + original_filename = file.filename or "upload" + ext = Path(original_filename).suffix.lower() + if ext not in UNSTRUCTURED_ALLOWED_EXTS: + raise HTTPException( + 400, + f"Unsupported file type: {ext}. Allowed: {', '.join(sorted(UNSTRUCTURED_ALLOWED_EXTS))}", + ) + + content = await file.read() + size_bytes = len(content) + + if size_bytes == 0: + raise HTTPException(400, "Empty file not allowed") + + if size_bytes > UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES: + raise HTTPException( + 413, + f"File too large ({size_bytes} bytes). Maximum is {UNSTRUCTURED_RECIPE_UPLOAD_MAX_LABEL}.", + ) + + block_dir = UNSTRUCTURED_UPLOAD_ROOT / block_id + ensure_dir(block_dir) + current_total = _get_block_total_size(block_dir) + if current_total + size_bytes > UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES: + raise HTTPException( + 413, + f"Total upload limit ({UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_LABEL}) exceeded", + ) + + file_id = uuid4().hex + raw_path = block_dir / f"{file_id}{ext}" + raw_path.write_bytes(content) + + extracted_path = block_dir / f"{file_id}.extracted.txt" + try: + extracted_text = _extract_text_from_file(raw_path, ext) + if not extracted_text or not extracted_text.strip(): + raw_path.unlink(missing_ok = True) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = "No extractable text found in file", + ) + extracted_path.write_text(extracted_text, encoding = "utf-8") + except ImportError as e: + raw_path.unlink(missing_ok = True) + extracted_path.unlink(missing_ok = True) + missing = getattr(e, "name", None) + expected_missing = {".pdf": "pymupdf4llm", ".docx": "mammoth"}.get(ext) + if isinstance(e, ModuleNotFoundError) and missing == expected_missing: + logger.error( + "data_recipe.seed.text_extraction_dependency_missing", + error = str(e), + missing = missing, + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = f"Cannot read {ext} files: the '{missing}' package is not installed.", + ) + logger.error( + "data_recipe.seed.text_extraction_failed", + error = str(e), + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = "Text extraction failed.", + ) + except Exception as e: + raw_path.unlink(missing_ok = True) + extracted_path.unlink(missing_ok = True) + logger.error( + "data_recipe.seed.text_extraction_failed", + error = str(e), + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = "Text extraction failed.", + ) + + try: + meta_path = block_dir / f"{file_id}.meta.json" + meta_path.write_text( + json.dumps({"original_filename": original_filename, "size_bytes": size_bytes}), + encoding = "utf-8", + ) + except OSError: + raw_path.unlink(missing_ok = True) + extracted_path.unlink(missing_ok = True) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = "Failed to save file metadata", + ) + + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "ok", + ) + + +@router.delete("/seed/unstructured-file/{block_id}/{file_id}") +async def remove_unstructured_file(block_id: str, file_id: str): + _validate_safe_id(block_id, "block_id") + _validate_safe_id(file_id, "file_id") + + block_dir = UNSTRUCTURED_UPLOAD_ROOT / block_id + if not block_dir.exists(): + raise HTTPException(404, "Block not found") + + deleted = False + for f in block_dir.iterdir(): + stem = f.name.split(".")[0] + if stem == file_id: + f.unlink(missing_ok = True) + deleted = True + + if not deleted: + raise HTTPException(404, "File not found") + try: + if not any(block_dir.iterdir()): + block_dir.rmdir() + except OSError: + pass + + return {"status": "ok"} + + +@router.delete("/seed/unstructured-block/{block_id}") +async def remove_unstructured_block(block_id: str): + """Delete a block's upload directory; files on disk still count toward its quota. + + Only uid-namespaced directories may be bulk-deleted: they have exactly one + owning block. Legacy node-id directories (n1, ...) can be shared by other + recipes, so they are managed file-by-file instead. + """ + _validate_safe_id(block_id, "block_id") + if not _UPLOAD_UID_RE.match(block_id): + raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted") + + block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve() + if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()): + raise HTTPException(400, "Invalid block_id: outside upload root") + if not block_dir.exists(): + return {"status": "ok", "deleted": False} + + try: + shutil.rmtree(block_dir) + except OSError as exc: + raise log_and_http_error( + exc, + 500, + "failed to delete uploaded files", + event = "data_recipe.seed.unstructured_block_delete_failed", + log = logger, + ) from exc + if block_dir.exists(): + raise HTTPException(500, "failed to delete uploaded files") + return {"status": "ok", "deleted": True} + + +@router.post("/seed/inspect-upload", response_model = SeedInspectResponse) +def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: + if payload.file_ids is not None: + if len(payload.file_ids) == 0: + raise HTTPException(400, "file_ids must not be empty") + _validate_safe_id(payload.block_id, "block_id") + for fid in payload.file_ids: + _validate_safe_id(fid, "file_id") + preview_rows = _read_preview_rows_from_multi_files( + block_id = payload.block_id, + file_ids = payload.file_ids, + file_names = payload.file_names, + preview_size = payload.preview_size, + chunk_size = payload.unstructured_chunk_size, + chunk_overlap = payload.unstructured_chunk_overlap, + ) + columns = ["chunk_text", "source_file"] if preview_rows else [] + resolved_paths = [ + str(UNSTRUCTURED_UPLOAD_ROOT / payload.block_id / f"{fid}.extracted.txt") + for fid in payload.file_ids + ] + return SeedInspectResponse( + dataset_name = "unstructured_seed", + resolved_path = resolved_paths[0] if resolved_paths else "", + resolved_paths = resolved_paths, + columns = columns, + preview_rows = _serialize_preview_rows(preview_rows), + ) + + seed_source_type = _normalize_optional_text(payload.seed_source_type) or "local" + filename = _sanitize_filename(payload.filename) + ext = Path(filename).suffix.lower() + # Legacy single-file path is .txt/.md only; PDF/DOCX use multi-file upload + _LEGACY_UNSTRUCTURED_EXTS = {".txt", ".md"} + if seed_source_type == "unstructured": + if ext not in _LEGACY_UNSTRUCTURED_EXTS: + allowed = ", ".join(sorted(_LEGACY_UNSTRUCTURED_EXTS)) + raise HTTPException( + status_code = 400, + detail = f"unsupported file type: {ext}. allowed: {allowed}", + ) + else: + if ext not in LOCAL_UPLOAD_EXTS: + allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS)) + raise HTTPException( + status_code = 400, + detail = f"unsupported file type: {ext}. allowed: {allowed}", + ) + + file_bytes = _decode_base64_payload(payload.content_base64) + if not file_bytes: + raise HTTPException(status_code = 400, detail = "empty upload payload") + if len(file_bytes) > LOCAL_SEED_UPLOAD_MAX_BYTES: + raise HTTPException( + status_code = 413, + detail = f"file too large (max {LOCAL_SEED_UPLOAD_MAX_LABEL})", + ) + + ensure_dir(SEED_UPLOAD_DIR) + stored_name = f"{uuid4().hex}_{filename}" + stored_path = SEED_UPLOAD_DIR / stored_name + stored_path.write_bytes(file_bytes) + + if seed_source_type == "unstructured": + preview_rows = _read_preview_rows_from_unstructured_file( + path = stored_path, + preview_size = int(payload.preview_size), + chunk_size = payload.unstructured_chunk_size, + chunk_overlap = payload.unstructured_chunk_overlap, + ) + else: + preview_rows = _read_preview_rows_from_local_file( + stored_path, + int(payload.preview_size), + ) + if not preview_rows: + raise HTTPException(status_code = 422, detail = "dataset appears empty or unreadable") + columns = _extract_columns(preview_rows) + + return SeedInspectResponse( + dataset_name = filename, + resolved_path = str(stored_path), + columns = columns, + preview_rows = preview_rows, + 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} diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py new file mode 100644 index 0000000..ffe36ba --- /dev/null +++ b/studio/backend/routes/data_recipe/validate.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Validation endpoints for data recipe.""" + +from __future__ import annotations + +from typing import Any + +from fastapi import APIRouter + +from core.data_recipe.service import ( + build_config_builder, + create_data_designer, + validate_recipe, +) +from loggers import get_logger +from models.data_recipe import RecipePayload, ValidateError, ValidateResponse +from utils.utils import safe_error_detail, safe_curated_detail, log_and_http_error + +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: + from data_designer.engine.compiler import ( + _add_internal_row_id_column_if_needed, + _get_allowed_references, + _resolve_and_add_seed_columns, + ) + from data_designer.engine.validation import ( + ViolationLevel, + validate_data_designer_config, + ) + except ImportError: + return [] + + try: + builder = build_config_builder(recipe) + designer = create_data_designer(recipe) + resource_provider = designer._create_resource_provider( # type: ignore[attr-defined] + "validate-configuration", + builder, + ) + config = builder.build() + _resolve_and_add_seed_columns(config, resource_provider.seed_reader) + _add_internal_row_id_column_if_needed(config) + violations = validate_data_designer_config( + columns = config.columns, + processor_configs = config.processors or [], + allowed_references = _get_allowed_references(config), + ) + except (TypeError, ValueError, AttributeError): + return [] + + errors: list[ValidateError] = [] + for violation in violations: + if violation.level != ViolationLevel.ERROR: + continue + code = getattr(violation.type, "value", None) + path = violation.column if violation.column else None + message = str(violation.message).strip() or "Validation failed." + errors.append( + ValidateError( + message = message, + path = path, + code = code, + ) + ) + return errors + + +def _patch_local_providers(recipe: dict[str, Any]) -> None: + """Strip is_local and fill a dummy endpoint so validation doesn't choke. + + Strict `is True` matches _inject_local_providers: truthy non-boolean values + aren't 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 + if not recipe.get("columns"): + return ValidateResponse( + valid = False, + 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 passed + # and full validation is deferred to run start, so a missing import + # shouldn't block the recipe. Restrict the bypass to data_designer so + # other ImportErrors still surface as failures. + 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: + logger.error( + "data_recipe.validate.github_config_failed", + error = str(exc), + exc_info = True, + ) + detail = safe_error_detail(exc, fallback = "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: + raise log_and_http_error( + exc, + 503, + safe_error_detail(exc), + event = "data_recipe.validate.service_unavailable", + log = logger, + ) from exc + except Exception as exc: + logger.error( + "data_recipe.validate.recipe_failed", + error = str(exc), + exc_info = True, + ) + detail = safe_curated_detail(exc, fallback = "Validation failed.") + parsed_errors = _collect_validation_errors(recipe) + return ValidateResponse( + valid = False, + errors = parsed_errors or [ValidateError(message = detail)], + raw_detail = detail, + ) + + return ValidateResponse(valid = True) diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py new file mode 100644 index 0000000..46319ca --- /dev/null +++ b/studio/backend/routes/datasets.py @@ -0,0 +1,817 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Datasets API routes.""" + +import base64 +import io +import json +import sys +from contextlib import suppress +from pathlib import Path +from uuid import uuid4 +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]: + """Resolved realpath for a HF cache repo dir: most-recent snapshot, else cache root. + + Mirrors routes/models.py; duplicated here to keep this module 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 + + +backend_path = Path(__file__).parent.parent.parent +if str(backend_path) not in sys.path: + sys.path.insert(0, str(backend_path)) + +from utils.datasets import check_dataset_format +from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label +from auth.authentication import get_current_subject + +router = APIRouter() +logger = get_logger(__name__) + + +from models.datasets import ( + AiAssistMappingRequest, + AiAssistMappingResponse, + CheckFormatRequest, + CheckFormatResponse, + LocalDatasetItem, + LocalDatasetsResponse, + UploadDatasetResponse, +) +from utils.paths import ( + dataset_uploads_root, + ensure_dir, + recipe_datasets_root, + resolve_dataset_path, +) + + +def _serialize_preview_value(value): + """Make a value JSON-safe for the client preview.""" + if value is None or isinstance(value, (str, int, float, bool)): + return value + + try: + from PIL.Image import Image as PILImage + if isinstance(value, PILImage): + buffer = io.BytesIO() + value.convert("RGB").save(buffer, format = "JPEG", quality = 85) + return { + "type": "image", + "mime": "image/jpeg", + "width": value.width, + "height": value.height, + "data": base64.b64encode(buffer.getvalue()).decode("ascii"), + } + except Exception: + pass + + if isinstance(value, dict): + return {str(key): _serialize_preview_value(item) for key, item in value.items()} + + if isinstance(value, (list, tuple)): + return [_serialize_preview_value(item) for item in value] + + return str(value) + + +def _serialize_preview_rows(rows): + return [ + {str(key): _serialize_preview_value(value) for key, value in dict(row).items()} + for row in rows + ] + + +# Data-file extensions for single-file preview. Tier 1 only uses tabular +# files; archives/text/config fall through to full load_dataset. +_COLUMNAR_EXTS = (".parquet", ".arrow") +_RECORD_EXTS = (".jsonl", ".csv", ".tsv") +_JSON_EXTS = (".json",) +_TABULAR_EXTS = _COLUMNAR_EXTS + _RECORD_EXTS + _JSON_EXTS +LOCAL_FILE_EXTS = (".json", ".jsonl", ".csv", ".parquet") +LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"} +# sync: training dataset upload limits are exposed by /api/settings/upload-limit +LOCAL_DATASETS_ROOT = recipe_datasets_root() +DATASET_UPLOAD_DIR = dataset_uploads_root() + + +def _safe_read_metadata(path: Path) -> dict | None: + try: + payload = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, ValueError, TypeError): + return None + if not isinstance(payload, dict): + return None + return payload + + +_HF_PREVIEW_EXT_PRIORITY = { + ".parquet": 0, + ".arrow": 0, + ".jsonl": 1, + ".csv": 2, + ".tsv": 3, + ".json": 4, +} +_HF_NON_DATA_EXACT_FILENAMES = { + ".gitattributes", + "builder_config.json", + "config.json", + "dataset_info.json", + "dataset_infos.json", + "metadata.json", +} +_HF_NON_DATA_CARD_FILENAMES = {"card.json", "dataset_card.json"} + + +def _normalize_hf_repo_path(path: str) -> str: + return path.strip().replace("\\", "/").lstrip("./") + + +def _hf_preview_extension(path: str) -> str | None: + lower = path.lower() + for ext in _HF_PREVIEW_EXT_PRIORITY: + if lower.endswith(ext): + return ext + return None + + +def _is_known_hf_non_data_file(path: str) -> bool: + name = Path(path).name.lower() + if name in _HF_NON_DATA_EXACT_FILENAMES: + return True + if name in _HF_NON_DATA_CARD_FILENAMES: + return True + if name == "readme" or name.startswith("readme."): + return True + if name.endswith("_config.json") or name.endswith("-config.json"): + return True + if name.endswith("_card.json") or name.endswith("-card.json"): + return True + return False + + +def _is_hf_preview_data_file(path: str) -> bool: + normalized = _normalize_hf_repo_path(path) + if not normalized or _is_known_hf_non_data_file(normalized): + return False + return _hf_preview_extension(normalized) is not None + + +def _extract_hf_metadata_data_paths(metadata: dict | None) -> list[str]: + if not metadata: + return [] + file_paths = metadata.get("file_paths") + if not isinstance(file_paths, dict): + return [] + raw_data_paths = file_paths.get("data") + if isinstance(raw_data_paths, str): + values = [raw_data_paths] + elif isinstance(raw_data_paths, list): + values = raw_data_paths + else: + return [] + + paths: list[str] = [] + for value in values: + if not isinstance(value, str): + continue + normalized = _normalize_hf_repo_path(value) + if normalized: + paths.append(normalized) + return paths + + +def _select_best_hf_preview_candidate( + candidates: list[str], *, subset: str | None, split: str | None +) -> str | None: + if not candidates: + return None + subset_lower = subset.lower() if subset else None + split_lower = split.lower() if split else None + + def score(path: str) -> tuple[int, int, int, int, str]: + ext = _hf_preview_extension(path) + ext_priority = _HF_PREVIEW_EXT_PRIORITY[ext] if ext else 99 + stem = Path(path).stem.lower() + path_lower = path.lower() + + subset_miss = 0 + if subset_lower: + subset_miss = 0 if subset_lower in stem or subset_lower in path_lower else 1 + + split_miss = 0 + if split_lower: + split_hit = ( + stem == split_lower + or stem.startswith(f"{split_lower}_") + or stem.startswith(f"{split_lower}-") + or f"/{split_lower}/" in path_lower + or f"_{split_lower}." in path_lower + or f"-{split_lower}." in path_lower + or f"/{split_lower}." in path_lower + or f"/{split_lower}_" in path_lower + or f"/{split_lower}-" in path_lower + ) + split_miss = 0 if split_hit else 1 + + return (subset_miss, split_miss, ext_priority, len(path), path) + + return sorted(candidates, key = score)[0] + + +def _select_hf_preview_file( + repo_files: list[str], *, metadata: dict | None, subset: str | None, split: str | None +) -> str | None: + normalized_repo_files = [_normalize_hf_repo_path(path) for path in repo_files] + repo_file_set = set(normalized_repo_files) + + metadata_candidates = [ + path + for path in _extract_hf_metadata_data_paths(metadata) + if path in repo_file_set and _is_hf_preview_data_file(path) + ] + if metadata_candidates: + return _select_best_hf_preview_candidate(metadata_candidates, subset = subset, split = split) + + data_candidates = [path for path in normalized_repo_files if _is_hf_preview_data_file(path)] + return _select_best_hf_preview_candidate(data_candidates, subset = subset, split = split) + + +def _download_hf_metadata(*, repo_id: str, repo_files: list[str], token: str | None) -> dict | None: + metadata_file = next( + ( + path + for path in repo_files + if Path(_normalize_hf_repo_path(path)).name.lower() == "metadata.json" + ), + None, + ) + if not metadata_file: + return None + + try: + from huggingface_hub import hf_hub_download + local_path = hf_hub_download( + repo_id = repo_id, + filename = metadata_file, + repo_type = "dataset", + token = token, + ) + except Exception as exc: + logger.warning(f"Could not read HF dataset metadata for {repo_id}: {exc}") + return None + + return _safe_read_metadata(Path(local_path)) + + +def _safe_read_rows_from_metadata(payload: dict | None) -> int | None: + if not payload: + return None + for key in ("actual_num_records", "target_num_records"): + value = payload.get(key) + if isinstance(value, int): + return value + return None + + +def _safe_read_metadata_summary(payload: dict | None) -> dict | None: + if not payload: + return None + + actual_num_records = ( + payload.get("actual_num_records") + if isinstance(payload.get("actual_num_records"), int) + else None + ) + target_num_records = ( + payload.get("target_num_records") + if isinstance(payload.get("target_num_records"), int) + else actual_num_records + ) + + columns: list[str] | None = None + schema = payload.get("schema") + if isinstance(schema, dict): + columns = [str(key) for key in schema.keys()] + if not columns: + stats = payload.get("column_statistics") + if isinstance(stats, list): + derived = [ + str(item.get("column_name")) + for item in stats + if isinstance(item, dict) and item.get("column_name") + ] + columns = derived or None + + parquet_files_count = None + file_paths = payload.get("file_paths") + if isinstance(file_paths, dict): + parquet_files = file_paths.get("parquet-files") + if isinstance(parquet_files, list): + parquet_files_count = len(parquet_files) + + total_num_batches = ( + payload.get("total_num_batches") + if isinstance(payload.get("total_num_batches"), int) + else parquet_files_count + ) + num_completed_batches = ( + payload.get("num_completed_batches") + if isinstance(payload.get("num_completed_batches"), int) + else total_num_batches + ) + + return { + "actual_num_records": actual_num_records, + "target_num_records": target_num_records, + "total_num_batches": total_num_batches, + "num_completed_batches": num_completed_batches, + "columns": columns, + } + + +def _build_local_dataset_items() -> list[LocalDatasetItem]: + if not LOCAL_DATASETS_ROOT.exists(): + return [] + + items: list[LocalDatasetItem] = [] + for entry in LOCAL_DATASETS_ROOT.iterdir(): + if not entry.is_dir() or not entry.name.startswith("recipe_"): + continue + parquet_dir = entry / "parquet-files" + if not parquet_dir.exists() or not any(parquet_dir.glob("*.parquet")): + continue + + rows = None + metadata_summary = None + metadata_path = entry / "metadata.json" + if metadata_path.exists(): + metadata_payload = _safe_read_metadata(metadata_path) + rows = _safe_read_rows_from_metadata(metadata_payload) + metadata_summary = _safe_read_metadata_summary(metadata_payload) + + try: + updated_at = entry.stat().st_mtime + except OSError: + updated_at = None + + items.append( + LocalDatasetItem( + id = entry.name, + label = entry.name, + path = str(parquet_dir.resolve()), + rows = rows, + updated_at = updated_at, + metadata = metadata_summary, + ) + ) + + items.sort(key = lambda item: item.updated_at or 0, reverse = True) + return items + + +def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int): + # Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper. + from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset + + if dataset_path.is_dir(): + parquet_dir = ( + dataset_path / "parquet-files" + if (dataset_path / "parquet-files").exists() + else dataset_path + ) + parquet_files = sorted(parquet_dir.glob("*.parquet")) + if parquet_files: + dataset = load_dataset( + "parquet", + data_files = [str(path) for path in parquet_files], + split = train_split, + ) + total_rows = len(dataset) + preview_slice = dataset.select(range(min(preview_size, total_rows))) + return preview_slice, total_rows + else: + candidate_files: list[Path] = [] + for ext in LOCAL_FILE_EXTS: + candidate_files.extend(sorted(dataset_path.glob(f"*{ext}"))) + if not candidate_files: + raise HTTPException( + status_code = 400, + detail = "Unsupported local dataset directory (expected parquet/json/jsonl/csv files)", + ) + dataset_path = candidate_files[0] + + if dataset_path.suffix in [".json", ".jsonl"]: + dataset = load_dataset("json", data_files = str(dataset_path), split = train_split) + elif dataset_path.suffix == ".csv": + dataset = load_dataset("csv", data_files = str(dataset_path), split = train_split) + elif dataset_path.suffix == ".parquet": + dataset = load_dataset("parquet", data_files = str(dataset_path), split = train_split) + else: + raise HTTPException( + status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}" + ) + + total_rows = len(dataset) + preview_slice = dataset.select(range(min(preview_size, total_rows))) + return preview_slice, total_rows + + +def _sanitize_filename(filename: str) -> str: + name = Path(filename).name.strip().replace("\x00", "") + if not name: + return "dataset_upload" + return name + + +@router.post("/upload", response_model = UploadDatasetResponse) +async def upload_dataset( + file: UploadFile, current_subject: str = Depends(get_current_subject) +) -> UploadDatasetResponse: + filename = _sanitize_filename(file.filename or "dataset_upload") + ext = Path(filename).suffix.lower() + if ext not in LOCAL_UPLOAD_EXTS: + allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS)) + raise HTTPException( + status_code = 400, + detail = f"Unsupported file type: {ext}. Allowed: {allowed}", + ) + + ensure_dir(DATASET_UPLOAD_DIR) + stem = Path(filename).stem + stored_name = f"{uuid4().hex}_{stem}{ext}" + stored_path = DATASET_UPLOAD_DIR / stored_name + + # Stream to disk in chunks to avoid holding the whole file in memory. The + # route-level cap gives a clear training-dataset error and avoids leaving + # oversized partial files in the Studio uploads directory. + upload_limit_bytes = get_upload_limit_bytes() + total_bytes = 0 + upload_complete = False + try: + with open(stored_path, "wb") as f: + while chunk := await file.read(1024 * 1024): + total_bytes += len(chunk) + if total_bytes > upload_limit_bytes: + raise HTTPException( + status_code = 413, + detail = ( + "Training dataset upload too large. " + f"Maximum is {get_upload_limit_label()}." + ), + ) + f.write(chunk) + upload_complete = True + finally: + if not upload_complete: + with suppress(OSError): + stored_path.unlink(missing_ok = True) + + if stored_path.stat().st_size == 0: + stored_path.unlink(missing_ok = True) + raise HTTPException(status_code = 400, detail = "Empty upload payload") + + return UploadDatasetResponse(filename = filename, stored_path = str(stored_path)) + + +@router.get("/local", response_model = LocalDatasetsResponse) +def list_local_datasets( + current_subject: str = Depends(get_current_subject), +) -> LocalDatasetsResponse: + 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 dir under HF_HUB_CACHE, where in-progress + download bytes are visible. Returns ``cache_path`` so the UI can show it. + """ + _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, + } + + # 95% threshold (as in the model endpoint): HF blob dedup makes + # completed_bytes drift under expected_bytes; inter-file gaps look "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, current_subject: str = Depends(get_current_subject)): + """Check if a dataset requires manual column mapping. + + HuggingFace strategy: + 1. list_repo_files -> select one tabular data file -> load_dataset + (avoids resolving thousands of files; ~2-4 s). + 2. Full streaming load_dataset as a last-resort fallback. + + Local files load directly. Plain `def` (not async) so FastAPI runs it in a + thread-pool, keeping blocking IO off the event loop. + """ + try: + from itertools import islice + from datasets import Dataset, load_dataset + from utils.datasets import format_dataset + + PREVIEW_SIZE = 10 + + logger.info(f"Checking format for dataset: {request.dataset_name}") + + dataset_path = resolve_dataset_path(request.dataset_name) + total_rows = None + + if dataset_path.exists(): + # ── Local file ────────────────────────────────────────── + train_split = request.train_split or "train" + preview_slice, total_rows = _load_local_preview_slice( + dataset_path = dataset_path, + train_split = train_split, + preview_size = PREVIEW_SIZE, + ) + else: + # ── HuggingFace dataset ───────────────────────────────── + # Tier 1: list_repo_files -> load one selected tabular data file + preview_slice = None + + try: + from huggingface_hub import HfApi + + api = HfApi() + repo_files = api.list_repo_files( + request.dataset_name, + repo_type = "dataset", + token = request.hf_token or None, + ) + metadata = _download_hf_metadata( + repo_id = request.dataset_name, + repo_files = repo_files, + token = request.hf_token or None, + ) + selected_file = _select_hf_preview_file( + repo_files, + metadata = metadata, + subset = request.subset, + split = request.train_split or "train", + ) + + if selected_file: + logger.info(f"Tier 1: loading single file {selected_file}") + load_kwargs = { + "path": request.dataset_name, + "data_files": [selected_file], + "split": "train", + "streaming": True, + } + if request.hf_token: + load_kwargs["token"] = request.hf_token + + streamed_ds = load_dataset(**load_kwargs) + rows = list(islice(streamed_ds, PREVIEW_SIZE)) + if rows: + preview_slice = Dataset.from_list(rows) + except Exception as e: + logger.warning(f"Tier 1 (single-file) failed: {e}") + + if preview_slice is None: + # Tier 2: full streaming (resolves all files; slow for large repos) + logger.info("Tier 2: falling back to full streaming load_dataset") + load_kwargs = { + "path": request.dataset_name, + "split": request.train_split, + "streaming": True, + } + if request.subset: + load_kwargs["name"] = request.subset + if request.hf_token: + load_kwargs["token"] = request.hf_token + + streamed_ds = load_dataset(**load_kwargs) + + rows = list(islice(streamed_ds, PREVIEW_SIZE)) + if not rows: + raise HTTPException( + status_code = 400, + detail = "Dataset appears to be empty or could not be streamed", + ) + + preview_slice = Dataset.from_list(rows) + total_rows = None + + result = check_dataset_format(preview_slice, is_vlm = request.is_vlm) + + logger.info( + f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}" + ) + + preview_samples = None + if not result["requires_manual_mapping"]: + if result.get("suggested_mapping"): + # Heuristic-detected: show raw data so columns match the API response. + # Column stripping happens at training time, not preview. + preview_samples = _serialize_preview_rows(preview_slice) + else: + try: + format_result = format_dataset( + preview_slice, + format_type = "auto", + num_proc = None, # Only 10 preview rows + ) + processed = format_result["dataset"] + preview_samples = _serialize_preview_rows(processed) + except Exception as e: + logger.warning(f"Processed preview generation failed (non-fatal): {e}") + preview_samples = _serialize_preview_rows(preview_slice) + else: + preview_samples = _serialize_preview_rows(preview_slice) + + # Warnings from check_dataset_format plus URL-based image detection. + warning = result.get("warning") + image_col = result.get("detected_image_column") + if image_col and image_col in (result.get("columns") or []): + try: + sample_val = preview_slice[0][image_col] + if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")): + url_warning = ( + "This dataset contains image URLs instead of embedded images. " + "Images will be downloaded during training, which may be slow for large datasets." + ) + logger.info(f"URL-based image column detected: {image_col}") + warning = f"{warning} {url_warning}" if warning else url_warning + except Exception: + pass + + return CheckFormatResponse( + requires_manual_mapping = result["requires_manual_mapping"], + detected_format = result["detected_format"], + columns = result["columns"], + is_image = result.get("is_image", False), + is_audio = result.get("is_audio", False), + multimodal_columns = result.get("multimodal_columns"), + suggested_mapping = result.get("suggested_mapping"), + detected_image_column = result.get("detected_image_column"), + detected_audio_column = result.get("detected_audio_column"), + detected_text_column = result.get("detected_text_column"), + detected_speaker_column = result.get("detected_speaker_column"), + chat_column = result.get("chat_column"), + preview_samples = preview_samples, + total_rows = total_rows, + warning = warning, + ) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error checking dataset format: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = "Failed to check dataset format") + + +@router.post("/ai-assist-mapping", response_model = AiAssistMappingResponse) +def ai_assist_mapping( + request: AiAssistMappingRequest, current_subject: str = Depends(get_current_subject) +): + """Run LLM-assisted dataset conversion advisor (user-triggered). + + Multi-pass analysis with a 7B helper model: classify dataset type, generate + conversion strategy, validate quality. Falls back to simple column + classification if the advisor fails. + """ + try: + from utils.datasets.llm_assist import llm_conversion_advisor + + # Truncate sample values for the LLM prompt. + truncated = [ + {col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5] + ] + + result = llm_conversion_advisor( + column_names = request.columns, + samples = truncated, + dataset_name = request.dataset_name, + hf_token = request.hf_token, + model_name = request.model_name, + model_type = request.model_type, + ) + + if result and result.get("success"): + return AiAssistMappingResponse( + success = True, + suggested_mapping = result.get("suggested_mapping"), + system_prompt = result.get("system_prompt"), + user_template = result.get("user_template"), + assistant_template = result.get("assistant_template"), + label_mapping = result.get("label_mapping"), + dataset_type = result.get("dataset_type"), + is_conversational = result.get("is_conversational"), + user_notification = result.get("user_notification"), + ) + + return AiAssistMappingResponse( + success = False, + warning = "AI could not determine column roles. Please assign them manually.", + ) + + except Exception as e: + logger.error(f"AI assist mapping failed: {e}", exc_info = True) + raise HTTPException(status_code = 500, detail = "AI assist failed") diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py new file mode 100644 index 0000000..a7fd7cb --- /dev/null +++ b/studio/backend/routes/export.py @@ -0,0 +1,577 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Export API routes: checkpoint discovery and model export operations.""" + +import asyncio +import json +import os +import sys +import time +from pathlib import Path +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 + +backend_path = Path(__file__).parent.parent.parent +if str(backend_path) not in sys.path: + sys.path.insert(0, str(backend_path)) + +from auth.authentication import get_current_subject + +from utils.utils import safe_error_detail + +try: + from core.export import get_export_backend +except ImportError: + parent_backend = backend_path.parent / "backend" + if str(parent_backend) not in sys.path: + sys.path.insert(0, str(parent_backend)) + from core.export import get_export_backend + +from models import ( + LoadCheckpointRequest, + ExportStatusResponse, + ExportOperationResponse, + ExportMergedModelRequest, + ExportBaseModelRequest, + ExportGGUFRequest, + ExportLoRAAdapterRequest, +) + +router = APIRouter() +logger = get_logger(__name__) + + +def _ensure_export_supported() -> None: + """Reject a mutating export request up front (HTTP 400) when the host can't export. + + Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints + (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + """ + from utils.hardware import export_capability + + cap = export_capability() + if not cap.get("export_supported", True): + raise HTTPException( + status_code = 400, + detail = cap.get("export_unsupported_message") + or "Export is not supported on this platform.", + ) + + +@router.post("/load-checkpoint", response_model = ExportOperationResponse) +async def load_checkpoint( + request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject) +): + """Load a checkpoint into the export backend (ExportBackend.load_checkpoint). + + Export runs in its own subprocess and is allowed to run in parallel with + training and inference. We deliberately do NOT stop training or unload the + chat model here -- if the GPU runs out of memory the load/export fails with + a clear error instead of tearing down the user's other running workloads. + """ + try: + _ensure_export_supported() + backend = get_export_backend() + # Run in a worker thread (spawns and waits on a subprocess, can take + # minutes) so the event loop stays free to serve the live log SSE stream. + 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, + trust_remote_code = request.trust_remote_code, + approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, + hf_token = request.hf_token, + subject = current_subject, + ) + + if not success: + raise HTTPException(status_code = 400, detail = message) + + return ExportOperationResponse(success = True, message = message) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error loading checkpoint: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to load checkpoint", + ) + + +@router.post("/cleanup", response_model = ExportOperationResponse) +async def cleanup_export_memory(current_subject: str = Depends(get_current_subject)): + """Cleanup export-related models from memory (ExportBackend.cleanup_memory).""" + try: + backend = get_export_backend() + success = await asyncio.to_thread(backend.cleanup_memory) + + if not success: + raise HTTPException( + status_code = 500, + detail = "Memory cleanup failed. See server logs for details.", + ) + + return ExportOperationResponse( + success = True, + message = "Memory cleanup completed successfully", + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error during export memory cleanup: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to cleanup export memory", + ) + + +@router.post("/cancel", response_model = ExportOperationResponse) +async def cancel_export(current_subject: str = Depends(get_current_subject)): + """Cancel the in-flight export by terminating its worker subprocess. + + Only the export subprocess is killed; training and inference run in their + own subprocesses and keep going. + """ + try: + backend = get_export_backend() + cancelled = await asyncio.to_thread(backend.cancel_export) + return ExportOperationResponse( + success = True, + message = "Export cancelled" if cancelled else "No active export to cancel", + ) + except Exception as e: + logger.error(f"Error cancelling export: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to cancel export", + ) + + +@router.get("/status", response_model = ExportStatusResponse) +async def get_export_status(current_subject: str = Depends(get_current_subject)): + """Get export backend status (loaded checkpoint, model type, PEFT flag).""" + try: + backend = get_export_backend() + last_op = backend.get_last_op() + # Relativise the recovered output path the same way the per-op POST response + # does, so the success banner shows an identical path on either route. + last_op_output_path = None + if last_op and last_op.get("output_path"): + details = _export_details(last_op["output_path"]) + last_op_output_path = (details or {}).get("output_path") + return ExportStatusResponse( + current_checkpoint = backend.current_checkpoint, + is_vision = bool(getattr(backend, "is_vision", False)), + is_peft = bool(getattr(backend, "is_peft", False)), + is_export_active = bool(backend.is_export_active()), + active_op_kind = backend.get_active_op_kind(), + last_op_seq = int(last_op["seq"]) if last_op else 0, + last_op_kind = last_op.get("kind") if last_op else None, + last_op_status = last_op.get("status") if last_op else None, + last_op_output_path = last_op_output_path, + last_op_error = last_op.get("error") if last_op else None, + ) + except Exception as e: + logger.error(f"Error getting export status: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to get export status", + ) + + +@router.get("/logs") +async def get_export_logs( + since: Optional[int] = Query( + None, + description = "Return log entries with seq strictly greater than this cursor.", + ), + current_subject: str = Depends(get_current_subject), +): + """Tunnel-safe JSON fallback for the live export log stream. + + The SSE endpoint (`/logs/stream`) is the low-latency path, but some reverse + proxies -- notably Cloudflare quick tunnels (`*.trycloudflare.com`) used by + `--secure` mode -- buffer `text/event-stream` responses and only flush when + the stream closes, so over the tunnel the browser sees nothing for the whole + export ("connecting..." with no logs). This endpoint returns the same + ring-buffer lines as a short, complete JSON response that no proxy buffers, + so the frontend can poll it and still show logs in near real time. + + Shares the orchestrator's monotonic `seq` cursor with the SSE stream, so the + two transports can run together and the client de-dupes by seq. + """ + try: + backend = get_export_backend() + # No cursor on the first poll of a run: start from the run-start snapshot + # so the client gets every line since the run began (matches the SSE + # default), not the entire historical ring buffer. + if since is None: + cursor = backend.get_run_start_seq() + else: + cursor = max(0, int(since)) + + entries, new_cursor = backend.get_logs_since(cursor) + return { + "entries": [ + { + "seq": int(entry.get("seq", 0)), + "stream": entry.get("stream", "stdout"), + "line": entry.get("line", ""), + "ts": entry.get("ts"), + } + for entry in entries + ], + "cursor": new_cursor, + "active": bool(backend.is_export_active()), + } + except Exception as e: + logger.error(f"Error getting export logs: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to get export logs", + ) + + +def _try_register_external_export(path: Path) -> tuple[bool, Optional[str]]: + """Best-effort registration so absolute exports show up in local scans.""" + try: + from storage.studio_db import add_scan_folder + folder = add_scan_folder(str(path)) + return True, str(folder.get("path") or path) + except Exception as exc: + logger.warning("Could not register export scan folder %s: %s", path, exc) + return False, None + + +def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]: + """Return relative export paths, keeping external absolute paths visible.""" + if not output_path: + return None + try: + from utils.paths.storage_roots import exports_root + + path = Path(output_path) + # If it's outside exports_root, return the full absolute path + # so users can find their files on a different drive. + if path.is_absolute(): + try: + path.resolve().relative_to(exports_root().resolve()) + except ValueError: + registered, registered_path = _try_register_external_export(path) + return { + "output_path": str(path), + "scan_folder_registered": registered, + "scan_folder_path": registered_path, + } + rel = os.path.relpath(output_path, exports_root()) + return {"output_path": rel} + except Exception: + return {"output_path": output_path} + + +@router.post("/export/merged", response_model = ExportOperationResponse) +async def export_merged_model( + request: ExportMergedModelRequest, current_subject: str = Depends(get_current_subject) +): + """Export a merged PEFT model (16-bit or 4-bit), optionally pushing to Hub. + + Wraps ExportBackend.export_merged_model. + """ + try: + _ensure_export_supported() + backend = get_export_backend() + 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, + repo_id = request.repo_id, + hf_token = request.hf_token, + private = request.private, + compressed_method = request.compressed_method, + ) + + if not success: + raise HTTPException(status_code = 400, detail = message) + + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting merged model: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to export merged model", + ) + + +@router.post("/export/base", response_model = ExportOperationResponse) +async def export_base_model( + request: ExportBaseModelRequest, current_subject: str = Depends(get_current_subject) +): + """Export a non-PEFT base model, optionally pushing to Hub. + + Wraps ExportBackend.export_base_model. + """ + try: + _ensure_export_supported() + backend = get_export_backend() + 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, + hf_token = request.hf_token, + private = request.private, + base_model_id = request.base_model_id, + ) + + if not success: + raise HTTPException(status_code = 400, detail = message) + + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting base model: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to export base model", + ) + + +@router.post("/export/gguf", response_model = ExportOperationResponse) +async def export_gguf( + request: ExportGGUFRequest, current_subject: str = Depends(get_current_subject) +): + """Export the current model to GGUF format, optionally pushing to Hub. + + Wraps ExportBackend.export_gguf. + """ + try: + _ensure_export_supported() + backend = get_export_backend() + # A custom path wins; otherwise the imatrix toggle requests the upstream auto-download. + imatrix_file = request.imatrix_path or (True if request.imatrix else None) + 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, + repo_id = request.repo_id, + hf_token = request.hf_token, + imatrix_file = imatrix_file, + ) + + if not success: + raise HTTPException(status_code = 400, detail = message) + + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting GGUF model: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to export GGUF model", + ) + + +@router.post("/export/lora", response_model = ExportOperationResponse) +async def export_lora_adapter( + request: ExportLoRAAdapterRequest, current_subject: str = Depends(get_current_subject) +): + """Export only the LoRA adapter (if the loaded model is PEFT). + + Wraps ExportBackend.export_lora_adapter. + """ + try: + _ensure_export_supported() + backend = get_export_backend() + 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, + hf_token = request.hf_token, + private = request.private, + gguf = request.gguf, + gguf_outtype = request.gguf_outtype, + ) + + if not success: + raise HTTPException(status_code = 400, detail = message) + + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to export LoRA adapter", + ) + + +# Live export log stream (Server-Sent Events). +# +# The export worker's stdout/stderr is piped to the orchestrator as log +# entries (core/export/worker.py, orchestrator.py); this endpoint streams +# them to the browser for a live terminal panel during export operations. +# +# Shape follows routes/training.py::stream_training_progress: each event +# carries id/event/data, 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 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` : once the worker is idle and no new lines arrived for + ~1 second. Clients should close. + - `error` : unrecoverable server-side error + + Each event's `id:` field is the log entry's monotonic seq number so the + browser can resume via `Last-Event-ID` on reconnect. + """ + backend = get_export_backend() + + # Starting cursor: explicit `since` wins, then Last-Event-ID on reconnect, + # else the run-start snapshot so the client sees every line since the run + # began even if the SSE connection opened after the export-kickoff POST. + 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 + # 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(): + # Let the reader thread drain trailing lines printed just + # before the worker signalled 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: end 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": safe_error_detail(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", + }, + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py new file mode 100644 index 0000000..ec5d309 --- /dev/null +++ b/studio/backend/routes/inference.py @@ -0,0 +1,14184 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Inference API routes for model loading and text generation. +""" + +import os +import sys +import time +import uuid +from pathlib import Path +from fastapi import APIRouter, Depends, HTTPException, Request, status +from fastapi.responses import StreamingResponse, JSONResponse, Response +from starlette.requests import ClientDisconnect +from typing import Any, Callable, List, Optional, Union +import json +import httpx +from loggers import get_logger +import asyncio +import threading +import weakref + + +import re as _re + +# Model size extraction (shared with core/inference/llama_cpp.py) +from utils.models import extract_model_size_b as _extract_model_size_b + +from utils.api_errors import openai_error_body, anthropic_error_body +from core.inference.llama_admission import ( + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + LlamaAdmissionLease, + LlamaAdmissionQueueFull, + LlamaAdmissionReservation, + LlamaAdmissionTimeout, + get_llama_admission_queue, + llama_admission_config_from_env, +) + + +def _positive_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int > 0 else None + + +def _nonnegative_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int >= 0 else None + + +_MLX_MPI_DISTRIBUTED_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) + + +def _mlx_distributed_launch_detected() -> bool: + if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None: + world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE")) + if world_size is not None and world_size > 1: + return True + return bool( + os.environ.get("MLX_HOSTFILE") + or os.environ.get("MLX_IBV_DEVICES") + or os.environ.get("MLX_JACCL_COORDINATOR") + or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT")) + ) + return any( + _nonnegative_int_or_none(os.environ.get(rank_env)) is not None + and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1 + for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS + ) + + +def _install_httpcore_asyncgen_silencer() -> None: + """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. + + When Studio proxies a llama-server stream via httpx, the innermost + ``HTTP11ConnectionByteStream.__aiter__`` async generator is finalised by + the asyncgen GC hook on a task different from the one that opened it. Its + ``aclose`` calls ``anyio.Lock.acquire`` → ``cancel_shielded_checkpoint``, + entering a ``CancelScope`` on the finaliser task; Python 3.13 flags the + cross-task exit as ``"Attempted to exit cancel scope in a different task"`` + and prints ``"async generator ignored GeneratorExit"`` as an unraisable + warning. + + Known httpx + httpcore + anyio interaction (MCP SDK python-sdk#831, agno + #3556, chainlit #2361, langchain-mcp-adapters #254). Benign: the 200 + response is already delivered. The streaming pass-throughs + (``/v1/chat/completions``, ``/v1/messages``, ``/v1/responses``, + ``/v1/completions``) manage their httpx lifecycle in one task with explicit + ``aclose()``; we don't hold a reference to the errant generator and can't + close it ourselves. + + Install one process-wide unraisable hook that swallows only this + interaction -- identified by (RuntimeError mentioning cancel scope / + GeneratorExit) + (object repr referencing HTTP11ConnectionByteStream) -- + and defers to the default hook otherwise. Idempotent. + """ + prior_hook = sys.unraisablehook + if getattr(prior_hook, "_unsloth_httpcore_silencer", False): + return + + def _hook(unraisable): + exc_value = getattr(unraisable, "exc_value", None) + obj = getattr(unraisable, "object", None) + obj_repr = repr(obj) if obj is not None else "" + if ( + isinstance(exc_value, RuntimeError) + and "HTTP11ConnectionByteStream" in obj_repr + and ( + "cancel scope" in str(exc_value) + or "GeneratorExit" in str(exc_value) + or "no running event loop" in str(exc_value) + ) + ): + return + prior_hook(unraisable) + + _hook._unsloth_httpcore_silencer = True # type: ignore[attr-defined] + sys.unraisablehook = _hook + + +_install_httpcore_asyncgen_silencer() + + +def _loaded_chat_template() -> Optional[str]: + """Chat template of the currently loaded GGUF model, if any.""" + try: + return get_llama_cpp_backend().chat_template + except Exception: + return None + + +def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Optional[str]: + """A chat-template raise_exception message to surface, but only when it appears + verbatim in chat_template (simple substring check), so we never leak arbitrary + llama-server text. Anchors on llama.cpp's "Jinja Exception:" prefix.""" + if not chat_template: + return None + marker = "Jinja Exception:" + idx = error_text.find(marker) + if idx == -1: + return None + candidate = error_text[idx + len(marker) :] + # llama-server appends JSON after the message; cut at the first boundary. + for stop in ('"', "\n"): + cut = candidate.find(stop) + if cut != -1: + candidate = candidate[:cut] + candidate = candidate.strip() + return candidate if candidate and candidate in chat_template else None + + +_LOST_CONNECTION_MSG = ( + "Lost connection to the model server. It may have crashed -- try reloading the model." +) + + +def _friendly_error(exc: Exception) -> str: + """Extract a user-friendly message from known llama-server errors.""" + if isinstance(exc, httpx.ReadTimeout): + if "stopped producing tokens" in str(exc).lower(): + return ( + "The model stopped producing tokens before the response " + "completed. Try stopping and retrying, or reduce max tokens." + ) + return ( + "The model is still processing the prompt but did not produce a " + "first token within 20 minutes. Try reducing context length, " + "using more GPU offload, or loading a smaller model." + ) + if isinstance(exc, httpx.TimeoutException): + return "Timed out communicating with the model server. Try again shortly." + # httpx transport failures from the async pass-through helpers. Any + # RequestError subclass (ConnectError, ReadError, RemoteProtocolError, + # WriteError, PoolTimeout, ...) means the llama-server subprocess is + # unreachable -- crashed or still coming up. + if isinstance(exc, httpx.RequestError): + return _LOST_CONNECTION_MSG + msg = str(exc) + m = _re.search( + r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)", + msg, + ) + if m: + return ( + f"Message too long: {m.group(1)} tokens exceeds the {m.group(2)}-token " + f"context window. Try increasing the Context Length in Model settings, " + f"or shorten the conversation." + ) + if "Lost connection to llama-server" in msg: + return _LOST_CONNECTION_MSG + template_msg = _template_raise_message(msg, _loaded_chat_template()) + if template_msg: + return f"An internal error occurred: {template_msg}" + return "An internal error occurred" + + +def _friendly_upstream_error(text: str) -> str: + """Rewrite a raw llama-server error body into an actionable message where we can. + + The main case is a tool-calling grammar that llama-server can't compile ("failed to + parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as + a hard 400 on every tool-bearing turn. It is a llama-server limitation with some + model/quant + tool-schema combinations, and recent llama.cpp builds handle the common + coding-agent tools, so point the user at updating Studio rather than the raw body. + """ + lowered = text.lower() + if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered: + return ( + "The model couldn't compile a tool-calling grammar for this request. This is a " + "llama-server limitation with some model/quant and tool-schema combinations. " + "Update Studio (it installs the latest llama.cpp, which handles the common " + "coding-agent tools) or try a different GGUF model." + ) + return f"llama-server error: {text}" + + +def _clamp_finish_reason(value) -> str: + """Coerce an upstream finish_reason into OpenAI's known chat values. + + Unknown values (including ``None``) become ``"stop"`` so local upstream + quirks do not leak into the public API shape. + """ + return ( + value + if value + in ( + "stop", + "length", + "tool_calls", + "content_filter", + "function_call", + ) + else "stop" + ) + + +def _normalize_stop_sequences(raw): + """Coerce an OpenAI/Anthropic ``stop`` value into the list-of-non-empty-strings + shape llama-server expects, or ``None`` when absent. A bare string becomes a + single-element list; empty strings are dropped (an empty stop sequence would + terminate generation immediately at position 0).""" + if isinstance(raw, str): + return [raw] if raw else None + if isinstance(raw, list): + return [s for s in raw if isinstance(s, str) and s] or None + return None + + +def _effective_max_tokens(payload): + """Resolve the generation cap, preferring OpenAI's replacement field. + + ``max_tokens`` is deprecated in favor of ``max_completion_tokens``; honor + either for compatibility, but let the replacement field win when both are + supplied. + """ + return ( + payload.max_completion_tokens + if payload.max_completion_tokens is not None + else payload.max_tokens + ) + + +_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT" + + +def _positive_float_env(env_name: str, default): + """Parse a positive float from an env var. A parseable non-positive value + returns ``None`` (0 disables the guarded feature); only unparseable or unset + values fall back to ``default``.""" + raw_value = os.environ.get(env_name) + if raw_value is None or not raw_value.strip(): + return default + try: + value = float(raw_value.strip()) + except ValueError: + return default + return value if value > 0 else None + + +def _effective_openai_max_tokens_from_values(max_tokens, max_completion_tokens = None): + """Resolve the OpenAI-compatible generation cap from raw request values. + + Prefers ``max_completion_tokens`` over the deprecated ``max_tokens``, and + returns ``None`` when both are omitted so callers keep their context-window + default (OpenAI treats an omitted cap as bounded only by the context + window). Explicit client caps pass through unchanged. + """ + + def _validate_explicit(value, param: str): + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be an integer.", + status = 400, + code = "invalid_type", + param = param, + ), + ) + # The legacy completions spec declares ``minimum: 0`` for max_tokens, + # so 0 is a valid (if degenerate) cap and only negatives are rejected. + # The chat fields never reach here with 0 (pydantic enforces ge=1). + if value < 0: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be at least 0.", + status = 400, + code = "invalid_value", + param = param, + ), + ) + return value + + max_tokens = _validate_explicit(max_tokens, "max_tokens") + max_completion_tokens = _validate_explicit(max_completion_tokens, "max_completion_tokens") + return max_completion_tokens if max_completion_tokens is not None else max_tokens + + +def _effective_openai_max_tokens(payload): + return _effective_openai_max_tokens_from_values( + getattr(payload, "max_tokens", None), + getattr(payload, "max_completion_tokens", None), + ) + + +def _wants_multiple_choices(payload) -> bool: + return (payload.n or 1) > 1 + + +def _has_openai_tool_history(messages) -> bool: + for message in messages or []: + if isinstance(message, dict): + if message.get("role") == "tool" or message.get("tool_calls"): + return True + continue + if getattr(message, "role", None) == "tool" or getattr(message, "tool_calls", None): + return True + return False + + +def _raise_unsupported_openai_parameter(param: str, message: str) -> None: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + message, + status = 400, + code = "unsupported_parameter", + param = param, + ), + ) + + +def _raise_unsupported_n(path_label: str) -> None: + _raise_unsupported_openai_parameter("n", f"n > 1 is not supported for {path_label}.") + + +def _sse_streaming_response(content) -> StreamingResponse: + """A ``text/event-stream`` response with the standard SSE headers used by + every streaming path here: no client/proxy caching, no proxy buffering, and + a one-shot connection. Two callers build their response inline instead: the + external-provider proxy omits ``Connection: close``, and the OpenAI + passthrough returns an empty ``keep-alive`` stream when the request is + cancelled before the upstream response starts. + + Built on ``_SameTaskStreamingResponse`` (not Starlette's stock + ``StreamingResponse``) so the SSE generator runs in the request task. The + legacy AnyIO task-group wrapper trips "Attempted to exit a cancel scope in a + different task" on Python 3.13 + httpx, which surfaced as a mid-stream + ``response.failed``. The streaming paths that take their response inline use + ``_SameTaskStreamingResponse`` directly for the same reason.""" + return _SameTaskStreamingResponse( + content, + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) + + +def _openai_stream_error_chunk(exc) -> dict: + """Build an in-band OpenAI error chunk for a mid-stream failure. Once the + stream's 200 headers are flushed the status can't change, so the error must + ride in the SSE body. An upstream context-window overflow is mapped to + code=context_length_exceeded so client compaction/trim loops can detect it + (a code-less error hides it).""" + _cls = _classify_llama_generation_error(exc) + if _cls: + return openai_error_body( + _friendly_error(exc), + status = 400, + code = "context_length_exceeded", + ) + if _cls is False: + return openai_error_body(_friendly_error(exc), status = 400) + return openai_error_body(_friendly_error(exc), status = 500) + + +def _openai_stream_error_sse(error: dict) -> str: + return f"data: {json.dumps(error)}\n\ndata: [DONE]\n\n" + + +def _openai_stream_error_sse_bytes(error: dict) -> bytes: + return _openai_stream_error_sse(error).encode("utf-8") + + +def _openai_passthrough_error(status_code, text) -> "HTTPException": + """HTTPException for a non-200 upstream response on the OpenAI passthrough + (tools / response_format). An over-context upstream error is mapped to a 400 + with code="context_length_exceeded" so these paths deliver the same signal as + the non-passthrough path; a tool-grammar compile failure gets the same actionable + guidance as the Anthropic passthrough; any other upstream error stays verbatim.""" + if _classify_llama_generation_error(Exception(text)): + return HTTPException( + status_code = 400, + detail = openai_error_body( + _friendly_error(Exception(text)), + status = 400, + code = "context_length_exceeded", + param = "messages", + ), + ) + return HTTPException( + status_code = status_code, + detail = _friendly_upstream_error(text[:500]), + ) + + +_OVERFLOW_TRUNCATE_MAX_RETRIES = 3 +# Truncated-prompt share of the real window; the rest is generation headroom +# so a near-full prompt cannot cut a tool call mid-JSON at the wall. +_OVERFLOW_PROMPT_TARGET_FRACTION = 0.75 + + +def _overflow_truncation_requested(payload) -> bool: + """True when the request (or the UNSLOTH_CONTEXT_OVERFLOW server default, + for clients that cannot send custom fields) opted into truncation.""" + requested = getattr(payload, "context_overflow", None) + if requested is not None: + return requested == "truncate_middle" + return os.environ.get("UNSLOTH_CONTEXT_OVERFLOW", "").strip().lower() == "truncate_middle" + + +def _parse_overflow_counts(err_text: str): + """(n_prompt_tokens, n_ctx) from an exceed_context_size_error body, or + None. Tolerates \\" around keys (body may be a re-wrapped JSON string).""" + m_prompt = _re.search(r'n_prompt_tokens\\?"?\s*:\s*(\d+)', err_text) + m_ctx = _re.search(r'n_ctx\\?"?\s*:\s*(\d+)', err_text) + if m_prompt and m_ctx: + return int(m_prompt.group(1)), int(m_ctx.group(1)) + return None + + +def _estimate_message_tokens(msg: dict) -> int: + try: + return max(1, len(json.dumps(msg, ensure_ascii = False)) // 4) + except Exception: + return 1 + + +def _truncate_middle_messages(messages: list, keep_ratio: float): + """Drop whole turn-groups from the middle of an OpenAI message list. + + Always kept: leading system message(s), the first group (task anchor), + and the trailing groups. A group is a user message, or an assistant + message plus its following tool results, so surviving tool_calls stay + paired with their results as chat templates require. + Returns (new_messages, dropped_message_count). + """ + if not messages or keep_ratio >= 1.0: + return messages, 0 + + head: list = [] + idx = 0 + while idx < len(messages) and messages[idx].get("role") in ("system", "developer"): + head.append(messages[idx]) + idx += 1 + + groups: list[list] = [] + for msg in messages[idx:]: + role = msg.get("role") + if role == "tool" and groups: + groups[-1].append(msg) + elif role == "tool": + groups.append([msg]) # orphan tool result; treat as its own group + else: + groups.append([msg]) + + # Anchor group plus the last 3 groups stay. + protected_tail = min(3, max(1, len(groups) - 1)) + if len(groups) <= 1 + protected_tail: + return messages, 0 + + total_est = sum(_estimate_message_tokens(m) for m in messages) + target_est = int(total_est * keep_ratio) + + anchor = groups[0] + middle = groups[1:-protected_tail] + tail = groups[-protected_tail:] + + current_est = total_est + kept_middle: list[list] = list(middle) + dropped = 0 + # Drop oldest-first until the estimate fits the target. + while kept_middle and current_est > target_est: + victim = kept_middle.pop(0) + dropped += len(victim) + current_est -= sum(_estimate_message_tokens(m) for m in victim) + + if dropped == 0: + return messages, 0 + + new_messages = head + anchor + for grp in kept_middle: + new_messages.extend(grp) + for grp in tail: + new_messages.extend(grp) + return new_messages, dropped + + +_CLIP_MARKER = "\n[... truncated by context_overflow=truncate_middle ...]\n" +# Generous head+tail first; cut harder if the estimate still misses the target. +_CLIP_KEEP_CHARS = (1500, 400) + + +def _clip_long_contents(messages: list, target_est: int) -> int: + """Clip oversized string contents middle-out until ``target_est`` is met. + + Tool results first, then earlier user turns, the final message last. + Message count and roles never change, so tool pairing holds even when + group-dropping could not free enough. Returns messages clipped. + """ + + def _candidates(): + tools = [m for m in messages if m.get("role") == "tool"] + users = [m for m in messages[:-1] if m.get("role") == "user"] + last = [messages[-1]] if messages else [] + return tools + users + last + + clipped = 0 + for keep in _CLIP_KEEP_CHARS: + for msg in _candidates(): + if sum(_estimate_message_tokens(m) for m in messages) <= target_est: + return clipped + content = msg.get("content") + if not isinstance(content, str) or len(content) <= 2 * keep + len(_CLIP_MARKER): + continue + msg["content"] = content[:keep] + _CLIP_MARKER + content[-keep:] + clipped += 1 + return clipped + + +def _apply_overflow_truncation(body: dict, err_text: str) -> bool: + """Shrink a passthrough body after an upstream context overflow: drop + middle turn-groups, clip still-oversized contents, clamp ``max_tokens`` + to the generation headroom. Returns False when nothing could shrink.""" + counts = _parse_overflow_counts(err_text) + messages = body.get("messages") or [] + total_est = sum(_estimate_message_tokens(m) for m in messages) + if counts: + n_prompt, n_ctx = counts + keep_ratio = min(0.95, (_OVERFLOW_PROMPT_TARGET_FRACTION * n_ctx) / max(1, n_prompt)) + else: + n_ctx = None + keep_ratio = 0.6 # no counts in the error; cut conservatively + # Scale the server-token target into char-estimate units. + target_est = int(total_est * keep_ratio) + + new_messages, dropped = _truncate_middle_messages(messages, keep_ratio) + if dropped: + body["messages"] = new_messages + clipped = 0 + if sum(_estimate_message_tokens(m) for m in body.get("messages") or []) > target_est: + clipped = _clip_long_contents(body.get("messages") or [], target_est) + if not dropped and not clipped: + return False + if n_ctx: + headroom = max(1024, int(n_ctx * (1.0 - _OVERFLOW_PROMPT_TARGET_FRACTION))) + cur_max = body.get("max_tokens") + body["max_tokens"] = min(cur_max, headroom) if cur_max else headroom + logger.warning( + "context_overflow=truncate_middle: dropped %d middle messages, clipped " + "%d contents (keep_ratio %.2f); retrying within the real window", + dropped, + clipped, + keep_ratio, + ) + return True + + +def _anthropic_stream_error_event(exc, *, force: bool = False): + """Return an Anthropic in-band stream error event when one is useful.""" + _cls = _classify_llama_generation_error(exc) + if _cls is None and not force: + return None + status = 400 if _cls is not None else 500 + return build_anthropic_sse_event( + "error", + anthropic_error_body(_friendly_error(exc), status = status), + ) + + +def _drop_parallel_tool_call_deltas(chunk) -> bool: + """In-place: drop tool_call deltas whose index >= 1 from a parsed OpenAI + streaming chunk so only the first tool call survives (parallel_tool_calls=false + / disable_parallel_tool_use, best-effort). Returns True if anything changed.""" + if not isinstance(chunk, dict): + return False + changed = False + for ch in chunk.get("choices") or []: + delta = ch.get("delta") or {} + tcs = delta.get("tool_calls") + if isinstance(tcs, list): + kept = [tc for tc in tcs if isinstance(tc, dict) and (tc.get("index") or 0) == 0] + if len(kept) != len(tcs): + delta["tool_calls"] = kept + changed = True + return changed + + +def _add_empty_content_to_reasoning_deltas(chunk: dict) -> bool: + """Make reasoning-only deltas palatable to strict OpenAI adapters. + + Some clients built on OpenAI-compatible streams ignore or reject chunks whose + delta only contains non-standard ``reasoning_content``. Preserve that field, + but add an empty standard ``content`` member so the chunk is still a valid + text-delta shape and downstream parsers keep the stream alive. + """ + changed = False + choices = chunk.get("choices") + if not isinstance(choices, list): + return False + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") + if not isinstance(delta, dict): + continue + if "reasoning_content" in delta and "content" not in delta: + delta["content"] = "" + changed = True + return changed + + +def _normalize_openai_passthrough_sse_line( + raw_line: str, *, cap_parallel_tool_calls: bool = False +) -> str: + """Normalize one passthrough OpenAI SSE ``data:`` line before relaying. + + The function is intentionally narrow: it leaves comments, blank events, + ``[DONE]``, and unparseable upstream bytes untouched; parsed chunks are + re-serialized only when a compatibility mutation is actually required. + """ + if not raw_line.startswith("data:"): + return raw_line + # Both mutations key off JSON object keys, so a line without either quoted + # key can never change; skip the parse on the per-token common case. + if '"reasoning_content"' not in raw_line and not ( + cap_parallel_tool_calls and '"tool_calls"' in raw_line + ): + return raw_line + payload = raw_line[len("data:") :].lstrip() + if payload.strip() in ("", "[DONE]"): + return raw_line + try: + obj = json.loads(payload) + except Exception: + return raw_line + if not isinstance(obj, dict): + return raw_line + changed = _add_empty_content_to_reasoning_deltas(obj) + if cap_parallel_tool_calls and _drop_parallel_tool_call_deltas(obj): + changed = True + if not changed: + return raw_line + return "data: " + json.dumps(obj, separators = (",", ":"), ensure_ascii = False) + + +def _prompt_tokens_details(upstream): + """Surface llama-server's real ``cached_tokens`` (KV-cache prompt hits) while + keeping the full OpenAI ``prompt_tokens_details`` shape. Defaults to zero when + the upstream usage doesn't carry it, so the field is always present.""" + out = {"cached_tokens": 0, "audio_tokens": 0} + if isinstance(upstream, dict): + out.update({k: v for k, v in upstream.items() if v is not None}) + return out + + +def _wants_stream_usage(payload) -> bool: + return bool((payload.stream_options or {}).get("include_usage")) + + +_OPENAI_PASSTHROUGH_TERMINAL_GRACE_S = 2.0 +_SSE_DONE_LINE = "data: [DONE]" + + +def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]: + """Classify OpenAI-compatible chat stream terminal markers. + + Some llama-server builds can emit the logical final chunk (``finish_reason``) + and optional usage chunk, then keep the HTTP stream open without sending the + OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close + the client stream promptly while preserving an optional trailing usage chunk. + """ + if not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + if data_str == "[DONE]": + return "done" + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + return _openai_passthrough_terminal_state_from_data(data) + + +def _openai_passthrough_terminal_state_from_data(data) -> Optional[str]: + """Dict-level core of ``_openai_passthrough_sse_line_terminal_state`` for + callers that already parsed the chunk (avoids a re-parse per relayed line).""" + if not isinstance(data, dict): + return None + if _monitor_openai_error_message(data): + return "error" + choices = data.get("choices") + if isinstance(choices, list): + if not choices and isinstance(data.get("usage"), dict): + return "usage" + for choice in choices: + if isinstance(choice, dict) and choice.get("finish_reason") is not None: + return "finish" + elif isinstance(data.get("usage"), dict): + return "usage" + return None + + +def _openai_stream_usage_chunk( + payload, completion_id, created, model_name, stream_usage, stream_timings +): + """Build the final OpenAI-standard usage chunk (choices=[], usage populated) + for a chat stream. Returns the SSE ``data:`` line, or None when the client + did not opt in via ``stream_options.include_usage`` (or no usage exists).""" + if not _wants_stream_usage(payload): + return None + if not (stream_usage or stream_timings): + return None + _usage = stream_usage or {} + _prompt_tokens = _usage.get("prompt_tokens") or 0 + _completion_tokens = _usage.get("completion_tokens") or 0 + _total_tokens = _usage.get("total_tokens") or (_prompt_tokens + _completion_tokens) + usage_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [], + usage = CompletionUsage( + prompt_tokens = _prompt_tokens, + completion_tokens = _completion_tokens, + total_tokens = _total_tokens, + prompt_tokens_details = _prompt_tokens_details(_usage.get("prompt_tokens_details")), + ), + timings = stream_timings, + ) + return f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" + + +def _chat_chunk_sse(completion_id, created, model_name, *, delta, finish_reason) -> str: + """One ``ChatCompletionChunk`` as an SSE ``data:`` line. The role / content / + final chunks every in-process streamer emits differ only in their ``delta`` + and ``finish_reason``.""" + chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ChunkChoice(delta = delta, finish_reason = finish_reason)], + ) + return f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + + +def _chat_role_chunk(completion_id, created, model_name) -> str: + """Opening assistant-role chunk for a chat stream.""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(role = "assistant"), + finish_reason = None, + ) + + +def _chat_content_chunk(completion_id, created, model_name, text) -> str: + """A content-delta chunk carrying ``text``.""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(content = text), + finish_reason = None, + ) + + +def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str: + """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block). + + Carries ``content: ""`` alongside, like the GGUF and passthrough paths, so + strict OpenAI adapters don't drop the reasoning-only delta. + """ + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(content = "", reasoning_content = text), + finish_reason = None, + ) + + +def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: + """Terminal stop chunk (empty delta) carrying the finish reason.""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(), + finish_reason = finish_reason, + ) + + +def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> str: + """Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``).""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(tool_calls = tool_calls), + finish_reason = None, + ) + + +def _sf_heal_events_to_sse( + events, + completion_id, + created, + model_name, + state, + parallel_tool_calls, + monitor_id = None, +): + """Serialize ``StreamToolCallHealer`` events into chat SSE lines. + + ``state["idx"]`` tracks the call index across ``feed``/``finalize``; + ``parallel_tool_calls is False`` caps promotion to one call (GGUF parity). + The monitor is fed from the same events the client receives, never the + healed-away markup.""" + lines = [] + for kind, value in events: + if kind == "text": + if value: + lines.append(_chat_content_chunk(completion_id, created, model_name, value)) + api_monitor.append_reply(monitor_id, value) + continue + if parallel_tool_calls is False and state["idx"] >= 1: + continue + lines.append( + _chat_tool_calls_chunk( + completion_id, + created, + model_name, + [ + { + "index": state["idx"], + "id": value["id"], + "type": "function", + "function": value["function"], + } + ], + ) + ) + _fn = value.get("function") or {} + api_monitor.append_reply( + monitor_id, + ("[tool_calls] " if state["idx"] == 0 else "; ") + + f"{_fn.get('name', '')}({_fn.get('arguments', '')})", + ) + state["idx"] += 1 + return lines + + +def _rewrite_cmpl_id(raw: bytes) -> bytes: + """Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-`` + prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key + (both spacing variants) so the rest of the body stays byte-exact.""" + return raw.replace(b'"id":"chatcmpl-', b'"id":"cmpl-').replace( + b'"id": "chatcmpl-', b'"id": "cmpl-' + ) + + +def _cmpl_stream_event_out(event: bytes, include_usage: bool) -> Optional[bytes]: + """Process one legacy /v1/completions SSE event (text between blank-line + separators). + + Always rewrites the ``chatcmpl-`` -> ``cmpl-`` id prefix. When the client + did NOT request ``stream_options.include_usage``, also removes the usage + statistics so the stream matches OpenAI's contract. + + Shape note: on /v1/completions, llama-server attaches ``usage`` to the + FINAL content chunk (the ``finish_reason`` chunk, which has a populated + ``choices`` array) -- unlike the chat stream, which emits a standalone + ``choices: []`` usage chunk. Both shapes are handled: a standalone + usage-only chunk is dropped; an inline ``usage`` field is stripped from a + content chunk while keeping ``choices``/``finish_reason`` intact. + + Returns the event bytes to emit, or ``None`` to drop the event. Only a + usage-bearing event is re-serialized; every other event keeps exact bytes. + """ + if include_usage: + return _rewrite_cmpl_id(event) + lines = event.split(b"\n") + changed = False + for i, ln in enumerate(lines): + if not ln.startswith(b"data:"): + continue + payload = ln[len(b"data:") :].strip() + if not payload or payload == b"[DONE]": + continue + try: + obj = json.loads(payload) + except Exception: + continue + if not isinstance(obj, dict) or obj.get("usage") is None: + continue + # Standalone usage-only chunk (chat-style) -> drop the whole event. + if obj.get("choices") == []: + return None + # Usage on a content/finish chunk (completions-style) -> strip it. + obj.pop("usage", None) + lines[i] = b"data: " + json.dumps(obj, separators = (",", ":")).encode("utf-8") + changed = True + return _rewrite_cmpl_id(b"\n".join(lines) if changed else event) + + +def _classify_llama_generation_error(exc: Exception) -> Optional[bool]: + """Classify an error raised while consuming the GGUF generator. + + Returns True for a context-window overflow, False for any other upstream + 4xx (a client error), or None when it should stay a 500. Distinguishes a + real client error from a genuine crash by the explicit "llama-server + returned 4xx" marker, not a bare "tokens"/"exceed" substring. + """ + msg = str(exc) + msg_l = msg.lower() + if "n_ctx" in msg_l or ( + "context" in msg_l and any(t in msg_l for t in ("exceed", "length", "window", "too long")) + ): + return True + if _re.search(r"llama-server returned (4\d\d)", msg): + return False + return None + + +# Add backend directory to path +backend_path = Path(__file__).parent.parent.parent +if str(backend_path) not in sys.path: + sys.path.insert(0, str(backend_path)) + +try: + from core.inference import get_inference_backend + from core.inference.llama_cpp import ( + LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + _DEFAULT_MAX_TOKENS_FLOOR, + _DEFAULT_STREAM_STALL_TIMEOUT_S, + _canonicalize_spec_mode, + _extra_args_set_spec_type, + _hf_offline_if_dns_dead, + detect_reasoning_flags, + ) + from core.inference.llama_server_args import ( + _effective_tensor_parallel, + _tensor_parallel_matches_loaded, + parse_split_mode_override, + resolve_tensor_parallel, + strip_shadowing_flags, + validate_extra_args, + ) + from core.inference.tensor_fallback import load_with_tensor_fallback + from utils.models import ModelConfig + from utils.inference import load_inference_config + from utils.models.model_config import ( + detect_mtp_file, + load_model_defaults, + ) + from utils.native_path_leases import ( + NativePathLeaseError, + display_label_for_native_path, + is_registered_native_path_label, + redact_native_paths, + verify_native_path_lease, + ) +except ImportError: + parent_backend = backend_path.parent / "backend" + if str(parent_backend) not in sys.path: + sys.path.insert(0, str(parent_backend)) + from core.inference import get_inference_backend + from core.inference.llama_cpp import ( + LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + _DEFAULT_MAX_TOKENS_FLOOR, + _DEFAULT_STREAM_STALL_TIMEOUT_S, + _canonicalize_spec_mode, + _extra_args_set_spec_type, + _hf_offline_if_dns_dead, + detect_reasoning_flags, + ) + from core.inference.llama_server_args import ( + _effective_tensor_parallel, + _tensor_parallel_matches_loaded, + parse_split_mode_override, + resolve_tensor_parallel, + strip_shadowing_flags, + validate_extra_args, + ) + from core.inference.tensor_fallback import load_with_tensor_fallback + from utils.models import ModelConfig + from utils.inference import load_inference_config + from utils.models.model_config import ( + detect_mtp_file, + load_model_defaults, + ) + from utils.native_path_leases import ( + NativePathLeaseError, + display_label_for_native_path, + is_registered_native_path_label, + redact_native_paths, + verify_native_path_lease, + ) + + +def _llama_non_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout(_DEFAULT_FIRST_TOKEN_TIMEOUT_S) + + +def _llama_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout(_DEFAULT_FIRST_TOKEN_TIMEOUT_S) + + +def _set_stream_response_read_timeout( + response: httpx.Response, read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S +) -> None: + # ``read_timeout_s = None`` clears httpx's read timeout (wait indefinitely), + # used when the stall guard is disabled so a stale first-token deadline + # can't keep timing out post-first-chunk gaps. + try: + timeout_ext = response.request.extensions.get("timeout") + if isinstance(timeout_ext, dict): + timeout_ext["read"] = read_timeout_s + except Exception: + pass + + +_STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 +_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 +_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0 +_OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n" +_OPENAI_LLAMA_ADMISSION_POLL_S = 0.25 + + +def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int: + """Serving slots available for one local llama-server backend. + + The loaded backend is the source of truth because it may have reduced + ``--parallel`` at load time to keep the model on GPU. The app state is a + launch-intent fallback for tests and for the short window before a backend + reports its committed runtime slots. + """ + slots = _positive_int_or_none(getattr(llama_backend, "effective_parallel_slots", None)) + if slots is not None: + return slots + try: + slots = getattr(request.app.state, "llama_parallel_slots", None) + except Exception: + slots = None + return _positive_int_or_none(slots) or 1 + + +def _openai_llama_admission_reserve( + *, request: Optional[Request], llama_backend +) -> tuple[LlamaAdmissionReservation, LlamaAdmissionConfig]: + config = llama_admission_config_from_env() + capacity = _openai_llama_admission_capacity(request, llama_backend) + key = str(getattr(llama_backend, "base_url", "llama-server")) + reservation = get_llama_admission_queue(key).reserve( + capacity = capacity, + config = config, + ) + return reservation, config + + +def _openai_admission_request_path(request: Optional[Request]) -> Optional[str]: + try: + return str(request.url.path) if request is not None else None + except Exception: + return None + + +def _openai_admission_log( + event: str, + reservation: Optional[LlamaAdmissionReservation] = None, + *, + snapshot = None, + request: Optional[Request], + mode: str, + wait_started_at: Optional[float] = None, + completion_id: Optional[str] = None, + level: str = "debug", +) -> None: + if snapshot is None and reservation is not None: + snapshot = reservation.snapshot_now() + wait_ms = None + if wait_started_at is not None: + wait_ms = int(max(0.0, time.monotonic() - wait_started_at) * 1000) + log = getattr(logger, level, logger.debug) + log( + "openai admission %s: mode=%s path=%s completion_id=%s capacity=%s active=%s queued=%s wait_ms=%s", + event, + mode, + _openai_admission_request_path(request), + completion_id, + getattr(snapshot, "capacity", None), + getattr(snapshot, "active", None), + getattr(snapshot, "queued", None), + wait_ms, + ) + + +def _openai_admission_error_body(exc: Exception, *, status_code: int) -> dict: + snapshot = getattr(exc, "snapshot", None) + message = str(exc) + if snapshot is not None: + message = ( + f"{message} " + f"(active={snapshot.active}, queued={snapshot.queued}, capacity={snapshot.capacity})" + ) + return openai_error_body(message, status = status_code) + + +def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: + return HTTPException( + status_code = status_code, + detail = _openai_admission_error_body(exc, status_code = status_code), + ) + + +def _openai_admission_timeout_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionTimeout: + return LlamaAdmissionTimeout( + "Timed out waiting for an available local llama-server generation slot", + snapshot = reservation.snapshot_now(), + ) + + +def _openai_admission_cancelled_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionCancelled: + return LlamaAdmissionCancelled( + "Client disconnected before an upstream llama-server generation slot was available", + snapshot = reservation.snapshot_now(), + ) + + +async def _raise_if_openai_admission_cancelled( + reservation: LlamaAdmissionReservation, *, request: Optional[Request], cancel_event +) -> None: + if reservation.is_cancelled: + raise _openai_admission_cancelled_error(reservation) + if await _preheader_cancelled(cancel_event, request): + reservation.cancel() + raise _openai_admission_cancelled_error(reservation) + + +async def _wait_for_openai_admission_non_streaming( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +) -> LlamaAdmissionLease: + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + wait_s = _OPENAI_LLAMA_ADMISSION_POLL_S + if deadline is not None: + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + continue + if lease is not None: + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _openai_admission_wait_stream_chunks( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +): + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + keepalive_interval_s = max(0.001, config.keepalive_interval_s) + next_keepalive_at = time.monotonic() + keepalive_interval_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + now = time.monotonic() + wait_s = min(_OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001)) + if deadline is not None: + remaining_s = deadline - now + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + lease = None + if lease is not None: + yield lease + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + now = time.monotonic() + if now >= next_keepalive_at: + next_keepalive_at = now + keepalive_interval_s + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _close_openai_admitted_stream_iterator(iterator, *, cancelled: bool) -> None: + if iterator is None: + return + if cancelled: + athrow = getattr(iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + return + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + + +def _openai_compat_stream_stall_timeout(): + """Max silent gap after an OpenAI passthrough stream has produced data. + + If the socket goes silent after valid SSE data, this bounds how long the + client is kept open. Defaults to the backend-wide stall timeout so this + path stalls out like every sibling stream; set the env var to tighten it + for local serving, or to 0 to disable the guard. + """ + return _positive_float_env( + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, + _DEFAULT_STREAM_STALL_TIMEOUT_S, + ) + + +def _openai_passthrough_upstream_headers(*, llama_backend = None) -> dict: + headers = {} + auth_headers = getattr(llama_backend, "_auth_headers", None) + if isinstance(auth_headers, dict): + headers.update(auth_headers) + headers["Connection"] = "close" + return headers + + +class _CompatSameTaskTimeout: + """Same-task timeout fallback for Python versions before asyncio.timeout.""" + + def __init__(self, timeout_s: float): + self.timeout_s = timeout_s + self._task = None + self._handle = None + self._timed_out = False + self._cancelling = 0 + + async def __aenter__(self): + self._task = asyncio.current_task() + if self._task is None: + return self + if hasattr(self._task, "cancelling"): + self._cancelling = self._task.cancelling() + loop = asyncio.get_running_loop() + self._handle = loop.call_later(max(self.timeout_s, 0), self._cancel_task) + return self + + async def __aexit__(self, exc_type, exc, tb): + if self._handle is not None: + self._handle.cancel() + if exc_type is not None and issubclass(exc_type, asyncio.CancelledError): + if self._timed_out: + if self._task is not None and hasattr(self._task, "uncancel"): + if self._task.uncancel() > self._cancelling: + return None + raise asyncio.TimeoutError from exc + return None + + def _cancel_task(self) -> None: + self._timed_out = True + if self._task is not None: + self._task.cancel() + + +def _same_task_timeout(timeout_s: float): + timeout_ctx = getattr(asyncio, "timeout", None) + if timeout_ctx is not None: + return timeout_ctx(timeout_s) + return _CompatSameTaskTimeout(timeout_s) + + +class _SameTaskStreamingResponse(StreamingResponse): + """StreamingResponse without Starlette's legacy AnyIO task-group wrapper.""" + + def __init__( + self, + *args, + unstarted_cleanup = None, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + # Released when the client disconnects before the body iterator starts: + # its try/finally never runs, so a stream that opens resources before the + # first yield (the passthrough's upstream httpx stream) passes this. + self._unstarted_cleanup = unstarted_cleanup + + async def __call__(self, scope, receive, send) -> None: + # send() emits a body message only after the first chunk, so no body + # message means the generator never entered its try/finally. + body_started = False + + async def _tracking_send(message) -> None: + nonlocal body_started + if message.get("type") == "http.response.body": + body_started = True + await send(message) + + try: + await self.stream_response(_tracking_send) + except OSError: # client disconnected mid-send + if body_started: + # Generator is suspended in its try/finally: throw CancelledError + # (not aclose's GeneratorExit) so its handler finishes the + # api_monitor entry. Fall back to aclose() without athrow. + athrow = getattr(self.body_iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + pass + else: + aclose = getattr(self.body_iterator, "aclose", None) + if aclose is not None: + await aclose() + else: + # Generator never started; aclose()/athrow() are no-ops on it, so + # release eager resources via the hook. getattr guards a response + # built through __new__ without __init__ (tests, pickling). + aclose = getattr(self.body_iterator, "aclose", None) + if aclose is not None: + await aclose() + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: + try: + await cleanup() + except Exception: + pass + raise ClientDisconnect() + if self.background is not None: + await self.background() + + +def _tracked_cancel_unstarted_cleanup(tracker): + """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when + the generator's finally (which normally exits it) never runs.""" + + async def _cleanup() -> None: + tracker.__exit__(None, None, None) + + return _cleanup + + +async def _aclose_stream_resources( + *, + watchers = (), + iterator = None, + resp = None, + client = None, +) -> None: + """Tear down an httpx streaming generator's resources in the required order: + cancel + await each watcher task, then aclose() the byte/line iterator, the + response, and the client. Each step swallows its own exceptions so teardown + always completes; a close-time CancelledError is re-raised only after every + step has run. See _anthropic_passthrough_stream for the ordering rationale.""" + for watcher in watchers: + if watcher is not None: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + close_cancelled = False + if iterator is not None: + try: + await iterator.aclose() + except asyncio.CancelledError: + close_cancelled = True + except Exception: + pass + if resp is not None: + try: + await resp.aclose() + except asyncio.CancelledError: + close_cancelled = True + except Exception: + pass + if client is not None: + try: + await client.aclose() + except asyncio.CancelledError: + close_cancelled = True + except Exception: + pass + if close_cancelled: + raise asyncio.CancelledError() + + +async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: + if cancel_event is not None and cancel_event.is_set(): + return True + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return True + return False + + +async def _wait_preheader_cancel(cancel_event = None, request: Optional[Request] = None) -> None: + while not await _preheader_cancelled(cancel_event, request): + await asyncio.sleep(0.05) + + +async def _send_stream_with_preheader_cancel( + client: httpx.AsyncClient, + req: httpx.Request, + cancel_event = None, + request: Optional[Request] = None, + mark_cancel_on_cancel: bool = True, +) -> Optional[httpx.Response]: + if cancel_event is None and request is None: + return await client.send(req, stream = True) + if await _preheader_cancelled(cancel_event, request): + return None + + send_task = asyncio.create_task(client.send(req, stream = True)) + cancel_task = asyncio.create_task(_wait_preheader_cancel(cancel_event, request)) + + async def _stop_send_task() -> None: + try: + await client.aclose() + except Exception: + pass + send_task.cancel() + try: + await send_task + except (asyncio.CancelledError, Exception): + pass + + try: + done, _pending = await asyncio.wait( + {send_task, cancel_task}, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + return await send_task + + await _stop_send_task() + return None + except asyncio.CancelledError: + if mark_cancel_on_cancel and cancel_event is not None: + cancel_event.set() + await _stop_send_task() + raise + finally: + cancel_task.cancel() + try: + await cancel_task + except (asyncio.CancelledError, Exception): + pass + + +async def _aiter_llama_stream_items( + async_iter, + *, + cancel_event = None, + request: Optional[Request] = None, + first_token_deadline: Optional[float] = None, + response: Optional[httpx.Response] = None, + post_first_item_read_timeout_s: Optional[ + Union[float, Callable[[], Optional[float]]] + ] = _DEFAULT_STREAM_STALL_TIMEOUT_S, +): + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + last_item_at: Optional[float] = None + + def _post_first_timeout_s() -> Optional[float]: + if callable(post_first_item_read_timeout_s): + return post_first_item_read_timeout_s() + return post_first_item_read_timeout_s + + while True: + if cancel_event is not None and cancel_event.is_set(): + return + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return + waiting_first_item = last_item_at is None + try: + if waiting_first_item: + remaining_s = first_token_deadline - time.monotonic() + if remaining_s <= 0: + raise httpx.ReadTimeout("The model did not produce a first token in time.") + if response is not None: + _set_stream_response_read_timeout(response, remaining_s) + # Keep httpx/httpcore's AnyIO cancel scope in this task. + # asyncio.wait_for would drive __anext__ in a child task. + async with _same_task_timeout(remaining_s): + item = await async_iter.__anext__() + else: + timeout_s = _post_first_timeout_s() + if ( + request is not None + and response is not None + and timeout_s is not None + and last_item_at is not None + ): + stall_remaining_s = timeout_s - (time.monotonic() - last_item_at) + if stall_remaining_s <= 0: + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + _set_stream_response_read_timeout(response, stall_remaining_s) + item = await async_iter.__anext__() + except asyncio.TimeoutError as exc: + if waiting_first_item: + raise httpx.ReadTimeout("The model did not produce a first token in time.") from exc + raise + except StopAsyncIteration: + return + except httpx.ReadTimeout: + now = time.monotonic() + if last_item_at is None: + if now >= first_token_deadline: + raise + continue + timeout_s = _post_first_timeout_s() + if request is not None and timeout_s is not None and now - last_item_at < timeout_s: + continue + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + if last_item_at is None and response is not None: + # The first-token read deadline no longer applies once a chunk has + # arrived: switch to the stall timeout, or clear the read timeout + # entirely when the stall guard is disabled (callable returns None) + # so a long gap can't trip the stale first-token deadline. + _set_stream_response_read_timeout(response, _post_first_timeout_s()) + last_item_at = time.monotonic() + yield item + + +from models.inference import ( + LoadRequest, + UnloadRequest, + GenerateRequest, + LoadResponse, + LoadProgressResponse, + UnloadResponse, + InferenceStatusResponse, + ChatCompletionRequest, + ChatCompletionChunk, + ChatCompletion, + ToolConfirmRequest, + ChatMessage, + ChunkChoice, + ChoiceDelta, + CompletionChoice, + CompletionMessage, + CompletionUsage, + ValidateModelRequest, + ValidateModelResponse, + TextContentPart, + ImageContentPart, + ImageUrl, + ResponsesRequest, + ResponsesInputTextPart, + ResponsesInputImagePart, + ResponsesOutputTextPart, + ResponsesUnknownInputItem, + ResponsesFunctionCallInputItem, + ResponsesFunctionCallOutputInputItem, + ResponsesOutputTextContent, + ResponsesOutputMessage, + ResponsesOutputReasoning, + ResponsesOutputReasoningContent, + ResponsesOutputFunctionCall, + ResponsesUsage, + ResponsesResponse, + AnthropicMessagesRequest, + AnthropicMessagesResponse, + AnthropicResponseTextBlock, + AnthropicResponseToolUseBlock, + AnthropicUsage, + CreateOpenAIContainerBody, + DeleteOpenAIContainerBody, + ListOpenAIContainersResponse, + OpenAIContainerRequest, + OpenAIContainerSummary, +) +from core.inference.anthropic_compat import ( + anthropic_messages_to_openai, + anthropic_tools_to_openai, + anthropic_tool_choice_to_openai, + openai_finish_to_anthropic_stop, + anthropic_tool_use_id, + build_anthropic_sse_event, + AnthropicStreamEmitter, + AnthropicPassthroughEmitter, +) +from auth.authentication import get_current_subject +from state.tool_approvals import resolve_tool_decision + +from core.inference.key_exchange import decrypt_api_key +from core.inference.model_ids import public_model_id +from core.inference.api_monitor import api_monitor +from core.inference.llama_http import nonstreaming_client +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) +from core.inference.tool_call_parser import TOOL_XML_SIGNALS as _PARSER_TOOL_SIGNALS +from core.inference.passthrough_healing import ( + StreamToolCallHealer, + heal_gate, + heal_openai_message, + heal_openai_message_events, + nudge_enabled, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) +from core.inference.providers import get_base_url +from core.inference.external_provider import ExternalProviderClient +from core.inference.chat_templates import resolve_effective_chat_template_override +from storage import providers_db +from utils.utils import safe_error_detail, log_and_http_error + +import io +import base64 +import numpy as np +from datetime import date as _date + +router = APIRouter() +# Studio-only router (not mounted on /v1 OpenAI-compat). +studio_router = APIRouter() + + +_ARTIFACT_PREVIEW_FRAME_ANCESTORS = "'self' tauri://localhost http://tauri.localhost" +_ARTIFACT_PREVIEW_FRAME_STRICT_CSP = ( + "default-src 'none'; " + "script-src 'unsafe-inline'; " + "style-src 'unsafe-inline'; " + "img-src data: blob:; " + "font-src data:; " + "media-src data: blob:; " + "connect-src 'none'; " + "object-src 'none'; " + "base-uri 'none'; " + "form-action 'none'; " + f"frame-ancestors {_ARTIFACT_PREVIEW_FRAME_ANCESTORS}; " + "sandbox allow-scripts" +) +_ARTIFACT_PREVIEW_FRAME_NETWORK_CSP = ( + "default-src http: https: data: blob:; " + "script-src 'unsafe-inline' 'unsafe-eval' http: https: data: blob:; " + "script-src-elem 'unsafe-inline' http: https: data: blob:; " + "style-src 'unsafe-inline' http: https: data: blob:; " + "style-src-elem 'unsafe-inline' http: https: data: blob:; " + "img-src http: https: data: blob:; " + "font-src http: https: data: blob:; " + "media-src http: https: data: blob:; " + "connect-src http: https: ws: wss: data: blob:; " + "worker-src http: https: blob:; " + "object-src 'none'; " + "base-uri 'none'; " + "form-action 'none'; " + f"frame-ancestors {_ARTIFACT_PREVIEW_FRAME_ANCESTORS}; " + "sandbox allow-scripts" +) +_ARTIFACT_PREVIEW_FRAME_HTML = """ + + + + + +""" + + +async def _authenticate_header_or_query(request: Request, token: Optional[str]) -> str: + """Resolve the bearer token from the Authorization header or the ``?token=`` + query param (needed for /