Merge branch 'main' into pip
This commit is contained in:
commit
d459f60458
385 changed files with 84538 additions and 6698 deletions
7
.github/CODEOWNERS
vendored
7
.github/CODEOWNERS
vendored
|
|
@ -53,3 +53,10 @@
|
|||
/studio/backend/tests/ @rolandtannous @danielhanchen
|
||||
/tests/ @rolandtannous @danielhanchen
|
||||
/scripts/ @rolandtannous @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
|
||||
|
|
|
|||
87
.github/dependabot.yml
vendored
87
.github/dependabot.yml
vendored
|
|
@ -5,23 +5,96 @@ updates:
|
|||
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: ["*"]
|
||||
|
||||
- package-ecosystem: "bun"
|
||||
directory: "/studio/frontend"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
groups:
|
||||
bun-frontend:
|
||||
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: ["*"]
|
||||
...
|
||||
|
|
|
|||
107
.github/scripts/hf-download-with-retry.sh
vendored
Executable file
107
.github/scripts/hf-download-with-retry.sh
vendored
Executable file
|
|
@ -0,0 +1,107 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# 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
|
||||
2281
.github/workflows/consolidated-tests-ci.yml
vendored
Normal file
2281
.github/workflows/consolidated-tests-ci.yml
vendored
Normal file
File diff suppressed because it is too large
Load diff
321
.github/workflows/lint-ci.yml
vendored
Normal file
321
.github/workflows/lint-ci.yml
vendored
Normal file
|
|
@ -0,0 +1,321 @@
|
|||
# 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: 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
|
||||
430
.github/workflows/mlx-ci.yml
vendored
Normal file
430
.github/workflows/mlx-ci.yml
vendored
Normal file
|
|
@ -0,0 +1,430 @@
|
|||
# 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 \
|
||||
'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
|
||||
|
||||
# Studio prebuilt llama.cpp install + GGUF inference. Drives the
|
||||
# exact path Studio's setup.sh takes on macOS: invokes
|
||||
# studio/install_llama_prebuilt.py with --published-repo
|
||||
# ggml-org/llama.cpp and --published-release-tag b9049 (the
|
||||
# latest llama.cpp release at the time this step was added; bump
|
||||
# via UNSLOTH_LLAMA_TAG / DEFAULT_LLAMA_TAG when refreshing).
|
||||
# The installer downloads llama-b9049-bin-macos-arm64.tar.gz,
|
||||
# which is the universal Apple Silicon (arm64) build -- the
|
||||
# same artifact works on M1/M2/M3/M4 because llama.cpp compiles
|
||||
# against the ARMv8.2 baseline.
|
||||
#
|
||||
# The b9049 release also publishes:
|
||||
# - llama-b9049-bin-macos-arm64-kleidiai.tar.gz
|
||||
# KleidiAI dispatches at runtime; on M1 it falls back where
|
||||
# ISA features (e.g. I8MM) are missing, so this asset also
|
||||
# runs on M1 -- Studio just doesn't choose it by default.
|
||||
# - llama-b9049-bin-macos-x64.tar.gz
|
||||
# Intel-only; would only run on M1 via Rosetta 2 emulation,
|
||||
# which we explicitly avoid.
|
||||
# - iOS XCFramework
|
||||
# iOS-app build artifact, unrelated to a macOS desktop CI.
|
||||
#
|
||||
# After install, downloads a small published GGUF
|
||||
# (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) from HuggingFace and
|
||||
# runs the prebuilt llama-cli on it. Asserts the prompt echo
|
||||
# appears in stdout. If the install fails OR the binary exits
|
||||
# non-zero, that's an Unsloth/Studio bug.
|
||||
- name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1)
|
||||
env:
|
||||
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
||||
# install_llama_prebuilt.py hits the GitHub releases API to
|
||||
# resolve the asset URL. Anonymous calls share the runner-IP
|
||||
# rate-limit bucket and 403 quickly -- pass the workflow's
|
||||
# automatic GITHUB_TOKEN to bump us to the 5000/hr authenticated
|
||||
# bucket.
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
|
||||
rm -rf "$INSTALL_DIR"
|
||||
# --simple-policy is required when --published-repo points
|
||||
# at upstream ggml-org/llama.cpp; that repo doesn't ship the
|
||||
# llama-prebuilt-manifest.json asset Studio's default policy
|
||||
# expects, so the simple platform-specific policy maps
|
||||
# Darwin+arm64 -> bin-macos-arm64 directly. studio/setup.sh
|
||||
# passes both --published-repo ggml-org/llama.cpp AND
|
||||
# --simple-policy automatically on macOS, so this CI step
|
||||
# exercises the same code path users hit when they run
|
||||
# `curl -fsSL https://unsloth.ai/install.sh | sh`.
|
||||
python studio/install_llama_prebuilt.py \
|
||||
--install-dir "$INSTALL_DIR" \
|
||||
--published-repo ggml-org/llama.cpp \
|
||||
--published-release-tag b9049 \
|
||||
--simple-policy
|
||||
|
||||
# Studio bundles only llama-server + llama-quantize from the
|
||||
# prebuilt (not llama-cli) -- inference goes through
|
||||
# llama-server's HTTP /completion endpoint. Validate both:
|
||||
# llama-quantize --help proves the dynamic libs link, then
|
||||
# spin up llama-server and POST a /completion request on a
|
||||
# tiny published GGUF.
|
||||
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"
|
||||
|
||||
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
|
||||
|
||||
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"
|
||||
|
||||
# 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 ("<<HELLO!!>> 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:
|
||||
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
||||
UNSLOTH_COMPILE_DISABLE: '1'
|
||||
run: |
|
||||
mkdir -p mlx_workdir
|
||||
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:
|
||||
HF_TOKEN: ${{ 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:
|
||||
HF_TOKEN: ${{ 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:
|
||||
HF_TOKEN: ${{ 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
|
||||
440
.github/workflows/notebooks-ci.yml
vendored
Normal file
440
.github/workflows/notebooks-ci.yml
vendored
Normal file
|
|
@ -0,0 +1,440 @@
|
|||
# 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 \
|
||||
'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.
|
||||
pip install --no-deps unsloth_zoo
|
||||
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
|
||||
740
.github/workflows/release-desktop.yml
vendored
740
.github/workflows/release-desktop.yml
vendored
|
|
@ -3,16 +3,306 @@ 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'
|
||||
description: 'Create as draft release; draft runs do not advance desktop-latest updater channel'
|
||||
type: boolean
|
||||
default: true
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
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
|
||||
|
|
@ -32,14 +322,31 @@ jobs:
|
|||
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:
|
||||
- uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5
|
||||
# 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
|
||||
|
|
@ -50,12 +357,18 @@ jobs:
|
|||
|
||||
# ── Node.js ──
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020
|
||||
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e
|
||||
with:
|
||||
node-version: 24
|
||||
|
||||
- name: Install pinned Tauri CLI
|
||||
run: npm install --save-dev --prefix studio @tauri-apps/cli@2.10.1
|
||||
# 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
|
||||
|
|
@ -67,38 +380,152 @@ jobs:
|
|||
exit 1
|
||||
fi
|
||||
|
||||
- name: Install frontend dependencies
|
||||
working-directory: studio/frontend
|
||||
run: npm install
|
||||
|
||||
- name: Verify backend package is published
|
||||
- name: Verify desktop updater and Linux package config
|
||||
shell: bash
|
||||
run: |
|
||||
node <<'JS'
|
||||
const { readFileSync } = require('node:fs');
|
||||
|
||||
(async () => {
|
||||
const cargo = readFileSync('studio/src-tauri/Cargo.toml', 'utf8');
|
||||
const match = cargo.match(/^version\s*=\s*"([^"]+)"/m);
|
||||
if (!match) throw new Error('Could not read desktop app version');
|
||||
const 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 appVersion = match[1];
|
||||
const response = await fetch(`https://pypi.org/pypi/unsloth/${appVersion}/json`);
|
||||
if (!response.ok) {
|
||||
const message = 'Publish unsloth=={app_version} to PyPI before the desktop release';
|
||||
throw new Error(`${message.replace('{app_version}', appVersion)} (HTTP ${response.status})`);
|
||||
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');
|
||||
}
|
||||
|
||||
const workflow = readFileSync('.github/workflows/release-desktop.yml', 'utf8');
|
||||
const lines = workflow.split(/\r?\n/);
|
||||
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');
|
||||
}
|
||||
}
|
||||
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@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@42dc69e1aa15d09112580998cf2ef0119e2e91ae
|
||||
uses: swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32
|
||||
with:
|
||||
workspaces: 'studio/src-tauri -> target'
|
||||
|
||||
|
|
@ -146,8 +573,8 @@ jobs:
|
|||
with:
|
||||
projectPath: studio
|
||||
tauriScript: npx --prefix . tauri
|
||||
tagName: desktop-v__VERSION__
|
||||
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
|
||||
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
|
||||
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
|
||||
releaseBody: |
|
||||
Desktop app for Unsloth Studio.
|
||||
|
||||
|
|
@ -159,7 +586,7 @@ jobs:
|
|||
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
|
||||
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
|
||||
releaseDraft: ${{ inputs.draft }}
|
||||
prerelease: false
|
||||
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
|
||||
args: -v ${{ matrix.args }}
|
||||
|
||||
# ── macOS: build + sign + notarize + upload ──
|
||||
|
|
@ -177,8 +604,8 @@ jobs:
|
|||
with:
|
||||
projectPath: studio
|
||||
tauriScript: npx --prefix . tauri
|
||||
tagName: desktop-v__VERSION__
|
||||
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
|
||||
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
|
||||
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
|
||||
releaseBody: |
|
||||
Desktop app for Unsloth Studio.
|
||||
|
||||
|
|
@ -190,7 +617,7 @@ jobs:
|
|||
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
|
||||
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
|
||||
releaseDraft: ${{ inputs.draft }}
|
||||
prerelease: false
|
||||
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
|
||||
args: -v ${{ matrix.args }}
|
||||
|
||||
# ── Windows: build + sign + upload ──
|
||||
|
|
@ -209,8 +636,8 @@ jobs:
|
|||
with:
|
||||
projectPath: studio
|
||||
tauriScript: npx --prefix . tauri
|
||||
tagName: desktop-v__VERSION__
|
||||
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
|
||||
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
|
||||
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
|
||||
releaseBody: |
|
||||
Desktop app for Unsloth Studio.
|
||||
|
||||
|
|
@ -222,5 +649,254 @@ jobs:
|
|||
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
|
||||
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
|
||||
releaseDraft: ${{ inputs.draft }}
|
||||
prerelease: false
|
||||
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
|
||||
|
|
|
|||
1126
.github/workflows/security-audit.yml
vendored
Normal file
1126
.github/workflows/security-audit.yml
vendored
Normal file
File diff suppressed because it is too large
Load diff
2
.github/workflows/stale.yml
vendored
2
.github/workflows/stale.yml
vendored
|
|
@ -11,7 +11,7 @@ jobs:
|
|||
issues: write
|
||||
|
||||
steps:
|
||||
- uses: actions/stale@v10
|
||||
- uses: actions/stale@b5d41d4e1d5dceea10e7104786b73624c18a190f # v10.2.0
|
||||
with:
|
||||
# The message to post on stale issues.
|
||||
# This message will ping the issue author.
|
||||
|
|
|
|||
166
.github/workflows/studio-api-smoke.yml
vendored
Normal file
166
.github/workflows/studio-api-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
# 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 }}-v1
|
||||
|
||||
- name: Prime HF_HOME with the GGUF
|
||||
id: prime-hf
|
||||
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
|
||||
env:
|
||||
HF_TOKEN: ${{ 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 }}-v1
|
||||
|
||||
- name: Install Studio (--local, --no-torch)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_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
|
||||
221
.github/workflows/studio-backend-ci.yml
vendored
Normal file
221
.github/workflows/studio-backend-ci.yml
vendored
Normal file
|
|
@ -0,0 +1,221 @@
|
|||
# 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, etc.):
|
||||
pip install \
|
||||
python-multipart aiofiles sqlalchemy cryptography \
|
||||
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 '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 \
|
||||
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 \
|
||||
'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 even though
|
||||
# it is an optional dep of unsloth.
|
||||
pip install 'unsloth_zoo>=2026.5.1'
|
||||
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_tauri_install_exit_order.sh \
|
||||
tests/sh/test_torch_constraint.sh; do
|
||||
echo "::group::$s"
|
||||
bash "$s"
|
||||
echo "::endgroup::"
|
||||
done
|
||||
|
||||
151
.github/workflows/studio-frontend-ci.yml
vendored
Normal file
151
.github/workflows/studio-frontend-ci.yml
vendored
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
# 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'
|
||||
- '.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'
|
||||
|
||||
# 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
|
||||
|
||||
- name: Lockfile must agree with package.json (npm ci is strict)
|
||||
# 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
|
||||
|
||||
- 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/<base> is reliably created (a bare
|
||||
# `git fetch origin <ref>` 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
|
||||
887
.github/workflows/studio-inference-smoke.yml
vendored
Normal file
887
.github/workflows/studio-inference-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,887 @@
|
|||
# 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
|
||||
# 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: 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 }}-v1
|
||||
|
||||
- name: Prime HF_HOME with the GGUF
|
||||
id: prime-hf
|
||||
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
|
||||
env:
|
||||
HF_TOKEN: ${{ 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 }}-v1
|
||||
|
||||
- name: Install Studio (--local, --no-torch)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_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()
|
||||
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"
|
||||
assert a == b, (
|
||||
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()
|
||||
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 IQ3_XXS = ~890 MiB). 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.
|
||||
GGUF_REPO: unsloth/Qwen3.5-2B-GGUF
|
||||
GGUF_FILE: Qwen3.5-2B-UD-IQ3_XXS.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:
|
||||
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 }}
|
||||
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 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",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout = timeout) as resp:
|
||||
return resp.status, json.loads(resp.read().decode())
|
||||
|
||||
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"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
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})")
|
||||
|
||||
# ── 2. Server-side python tool ───────────────────────────────
|
||||
# 123 * 456 = 56088. The agentic loop streams SSE; we
|
||||
# accumulate the assistant text and look for the answer. We
|
||||
# accept "56088" or "56,088" since the model may format it.
|
||||
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": 0.0,
|
||||
"seed": SEED,
|
||||
"max_tokens": 600,
|
||||
})
|
||||
assert "56088" in content or "56,088" in content, (
|
||||
f"expected 56088 in python-tool answer, got: {content!r}"
|
||||
)
|
||||
print(f"[tools] PASS python tool ({len(content)} chars)")
|
||||
|
||||
# ── 3. Server-side bash (terminal) tool ──────────────────────
|
||||
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": 0.0,
|
||||
"seed": SEED,
|
||||
"max_tokens": 600,
|
||||
})
|
||||
assert "hello-bash-tool" in content, (
|
||||
f"expected 'hello-bash-tool' in terminal-tool answer, got: {content!r}"
|
||||
)
|
||||
print(f"[tools] PASS bash/terminal tool ({len(content)} chars)")
|
||||
|
||||
# ── 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.
|
||||
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": 0.0,
|
||||
"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 ─────────────────────────────────────
|
||||
# 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 <think> 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 = ("<think>" in on_text) or len(on_text) > 80
|
||||
had_think_off = ("<think>" 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 <think> marker.
|
||||
assert "<think>" not in off_text, (
|
||||
f"enable_thinking=False but <think> 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()
|
||||
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: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF
|
||||
GGUF_VARIANT: UD-IQ3_XXS
|
||||
GGUF_FILE: gemma-4-E2B-it-UD-IQ3_XXS.gguf
|
||||
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 }}-v1
|
||||
|
||||
- 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:
|
||||
HF_TOKEN: ${{ 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"
|
||||
|
||||
- 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 }}-v1
|
||||
|
||||
- name: Install Studio (--local, --no-torch)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_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"
|
||||
# Load the GGUF (mmproj is auto-detected via the HF repo
|
||||
# lookup, the cached file is pulled out of HF_HOME).
|
||||
curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \
|
||||
-H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \
|
||||
--max-time 900 \
|
||||
-d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}" \
|
||||
| jq '{status, display_name, is_vision}'
|
||||
|
||||
- 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 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",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout = timeout) as resp:
|
||||
return resp.status, json.loads(resp.read().decode())
|
||||
|
||||
# ── 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()
|
||||
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
|
||||
with:
|
||||
name: json-images-log
|
||||
path: |
|
||||
logs/studio.log
|
||||
logs/install.log
|
||||
retention-days: 7
|
||||
153
.github/workflows/studio-mac-api-smoke.yml
vendored
Normal file
153
.github/workflows/studio-mac-api-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
# 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 }}-v1
|
||||
|
||||
- name: Prime HF_HOME with the GGUF
|
||||
id: prime-hf
|
||||
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
|
||||
env:
|
||||
HF_TOKEN: ${{ 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 }}-v1
|
||||
|
||||
- name: Install Studio (--local, --no-torch)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
mkdir -p logs
|
||||
set -o pipefail
|
||||
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
|
||||
|
||||
- name: Assert install.sh used the Mac llama.cpp prebuilt
|
||||
run: |
|
||||
if grep -q "falling back to source build" logs/install.log; then
|
||||
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
|
||||
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- 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
|
||||
1042
.github/workflows/studio-mac-inference-smoke.yml
vendored
Normal file
1042
.github/workflows/studio-mac-inference-smoke.yml
vendored
Normal file
File diff suppressed because it is too large
Load diff
345
.github/workflows/studio-mac-ui-smoke.yml
vendored
Normal file
345
.github/workflows/studio-mac-ui-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,345 @@
|
|||
# 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 }}-v1
|
||||
|
||||
- name: Prime HF_HOME with the GGUF
|
||||
id: prime-hf
|
||||
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
|
||||
env:
|
||||
HF_TOKEN: ${{ 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 }}-v1
|
||||
|
||||
- name: Install Studio (--local, --no-torch)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
mkdir -p logs
|
||||
set -o pipefail
|
||||
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
|
||||
|
||||
- name: Assert install.sh used the Mac llama.cpp prebuilt
|
||||
run: |
|
||||
if grep -q "falling back to source build" logs/install.log; then
|
||||
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
|
||||
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- 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, and (2) Chromium net::ERR_NO_BUFFER_SPACE
|
||||
# when the runner's kernel briefly runs out of socket buffers.
|
||||
# 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 either 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; } \
|
||||
&& [ "$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 and ERR_NO_BUFFER_SPACE.
|
||||
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; } \
|
||||
&& [ "$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
|
||||
184
.github/workflows/studio-mac-update-smoke.yml
vendored
Normal file
184
.github/workflows/studio-mac-update-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
# 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'
|
||||
- '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 }}
|
||||
run: |
|
||||
mkdir -p logs
|
||||
set -o pipefail
|
||||
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
|
||||
|
||||
- name: Assert install.sh used the Mac llama.cpp prebuilt
|
||||
run: |
|
||||
# Mac install must take the prebuilt path. Source-build
|
||||
# fallback here is an Unsloth bug.
|
||||
if grep -q "falling back to source build" logs/install.log; then
|
||||
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
|
||||
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
|
||||
exit 1
|
||||
fi
|
||||
if ! grep -qE "prebuilt installed and validated|prebuilt up to date and validated|bin-macos-arm64" logs/install.log; then
|
||||
echo "::error::no Mac prebuilt llama.cpp marker in install.log."
|
||||
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
|
||||
exit 1
|
||||
fi
|
||||
echo "install.sh installed the Mac prebuilt llama.cpp"
|
||||
|
||||
- name: First update should be a no-op (prebuilt already validated)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_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 }}
|
||||
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 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
|
||||
# uninstall.sh has not landed yet (lets this workflow merge
|
||||
# before #5497).
|
||||
run: |
|
||||
set -o pipefail
|
||||
if [ ! -f uninstall.sh ]; then
|
||||
echo "uninstall.sh not present in this tree; skipping round-trip"
|
||||
: > logs/uninstall.log
|
||||
exit 0
|
||||
fi
|
||||
sh 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 uninstall.sh 2>&1 | tail -5
|
||||
sh 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
|
||||
128
.github/workflows/studio-tauri-smoke.yml
vendored
Normal file
128
.github/workflows/studio-tauri-smoke.yml
vendored
Normal file
|
|
@ -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 libayatana-appindicator3-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
|
||||
293
.github/workflows/studio-ui-smoke.yml
vendored
Normal file
293
.github/workflows/studio-ui-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
# 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 }}-v1
|
||||
|
||||
- name: Prime HF_HOME with the GGUF
|
||||
id: prime-hf
|
||||
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
|
||||
env:
|
||||
HF_TOKEN: ${{ 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 }}-v1
|
||||
|
||||
- name: Install Studio (--local, --no-torch)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_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
|
||||
|
||||
- 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/playwright
|
||||
logs/playwright_extra
|
||||
logs/playwright_ime
|
||||
retention-days: 7
|
||||
191
.github/workflows/studio-update-smoke.yml
vendored
Normal file
191
.github/workflows/studio-update-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
# 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'
|
||||
- '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 }}
|
||||
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 }}
|
||||
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 }}
|
||||
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 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 uninstall.sh hasn't caught up.
|
||||
# Skips gracefully if uninstall.sh has not landed yet (lets this
|
||||
# workflow merge before #5497).
|
||||
run: |
|
||||
set -o pipefail
|
||||
if [ ! -f uninstall.sh ]; then
|
||||
echo "uninstall.sh not present in this tree; skipping round-trip"
|
||||
: > logs/uninstall.log
|
||||
exit 0
|
||||
fi
|
||||
sh 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 uninstall.sh 2>&1 | tail -5
|
||||
sh 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
|
||||
246
.github/workflows/studio-windows-api-smoke.yml
vendored
Normal file
246
.github/workflows/studio-windows-api-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,246 @@
|
|||
# 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 }}-v1
|
||||
|
||||
- name: Prime HF_HOME with the GGUF
|
||||
id: prime-hf
|
||||
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
|
||||
env:
|
||||
HF_TOKEN: ${{ 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 }}-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 }}
|
||||
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: Patch Studio venv with full typer / pydantic dep trees
|
||||
# Belt-and-suspenders: install.ps1's --no-deps install of
|
||||
# no-torch-runtime.txt drops typer's and pydantic's runtime
|
||||
# deps unless explicitly pinned. Re-install the ones whose
|
||||
# deps don't pull torch.
|
||||
run: |
|
||||
STUDIO_PY=~/.unsloth/studio/unsloth_studio/Scripts/python.exe
|
||||
if [ ! -f "$STUDIO_PY" ]; then
|
||||
echo "::error::Studio venv python not at $STUDIO_PY"
|
||||
ls -la ~/.unsloth/studio/ || true
|
||||
exit 1
|
||||
fi
|
||||
"$STUDIO_PY" -m pip install --upgrade typer pydantic huggingface_hub
|
||||
|
||||
- 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
|
||||
1244
.github/workflows/studio-windows-inference-smoke.yml
vendored
Normal file
1244
.github/workflows/studio-windows-inference-smoke.yml
vendored
Normal file
File diff suppressed because it is too large
Load diff
342
.github/workflows/studio-windows-ui-smoke.yml
vendored
Normal file
342
.github/workflows/studio-windows-ui-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,342 @@
|
|||
# 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 }}-v1
|
||||
|
||||
- name: Prime HF_HOME with the GGUF
|
||||
id: prime-hf
|
||||
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
|
||||
env:
|
||||
HF_TOKEN: ${{ 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 }}-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)
|
||||
# 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 }}
|
||||
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: 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: Patch Studio venv with full typer / pydantic dep trees
|
||||
# Belt-and-suspenders: install.ps1's --no-deps install of
|
||||
# no-torch-runtime.txt drops typer's and pydantic's runtime
|
||||
# deps unless explicitly pinned. Re-install the ones whose
|
||||
# deps don't pull torch.
|
||||
run: |
|
||||
STUDIO_PY=~/.unsloth/studio/unsloth_studio/Scripts/python.exe
|
||||
if [ ! -f "$STUDIO_PY" ]; then
|
||||
echo "::error::Studio venv python not at $STUDIO_PY"
|
||||
ls -la ~/.unsloth/studio/ || true
|
||||
exit 1
|
||||
fi
|
||||
"$STUDIO_PY" -m pip install --upgrade typer pydantic huggingface_hub
|
||||
|
||||
- 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
|
||||
314
.github/workflows/studio-windows-update-smoke.yml
vendored
Normal file
314
.github/workflows/studio-windows-update-smoke.yml
vendored
Normal file
|
|
@ -0,0 +1,314 @@
|
|||
# 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 (llama-bNNNN-bin-win-cpu-
|
||||
# x64 from ggml-org/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'
|
||||
- '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 line 1109-1145 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 npm
|
||||
# check fails and setup.ps1 falls through to the
|
||||
# "winget install Node.js LTS" branch -- a ~35 s reinstall
|
||||
# of Node we don't need. `npm install -g npm@^11` updates
|
||||
# the bundled npm in-place in ~5 s, which makes setup.ps1
|
||||
# short-circuit on the existing Node.
|
||||
#
|
||||
# (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 }}
|
||||
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: Patch Studio venv with full typer / pydantic dep trees
|
||||
# install.ps1 runs `uv pip install --no-deps -r
|
||||
# no-torch-runtime.txt` to keep torch out of transitive
|
||||
# resolution from accelerate/peft/trl. That also drops
|
||||
# typer's and pydantic's runtime deps unless they're
|
||||
# explicitly pinned in no-torch-runtime.txt. We pin the
|
||||
# known ones (click, shellingham, annotated-doc, rich,
|
||||
# pydantic-core, annotated-types, typing-inspection, ...)
|
||||
# but typer / pydantic minor versions can introduce new
|
||||
# transitive deps that are NOT in our pin list.
|
||||
#
|
||||
# Belt-and-suspenders: re-install typer + pydantic +
|
||||
# huggingface_hub WITH their deps into the Studio venv.
|
||||
# `pip install --upgrade` only adds missing packages; it
|
||||
# never down-shifts an installed version. Cannot pull
|
||||
# torch (none of typer / pydantic / huggingface_hub depend
|
||||
# on it).
|
||||
run: |
|
||||
STUDIO_PY=~/.unsloth/studio/unsloth_studio/Scripts/python.exe
|
||||
if [ ! -f "$STUDIO_PY" ]; then
|
||||
echo "::error::Studio venv python not at $STUDIO_PY"
|
||||
ls -la ~/.unsloth/studio/ || true
|
||||
exit 1
|
||||
fi
|
||||
"$STUDIO_PY" -m pip install --upgrade typer pydantic huggingface_hub
|
||||
|
||||
- name: First update should be a no-op (prebuilt already validated)
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_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 }}
|
||||
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 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 uninstall.ps1 has not been updated
|
||||
# to match. Skips gracefully if 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\uninstall.ps1")) {
|
||||
Write-Host "uninstall.ps1 not present in this tree; skipping round-trip"
|
||||
"" | Set-Content logs/uninstall.log
|
||||
exit 0
|
||||
}
|
||||
pwsh -NoProfile -File "$PWD\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\uninstall.ps1" *>&1 | Select-Object -Last 5
|
||||
pwsh -NoProfile -File "$PWD\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
|
||||
312
.github/workflows/version-compat-ci.yml
vendored
Normal file
312
.github/workflows/version-compat-ci.yml
vendored
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
# 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 \
|
||||
'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.<module>` 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
|
||||
|
||||
# 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
|
||||
136
.github/workflows/wheel-smoke.yml
vendored
Normal file
136
.github/workflows/wheel-smoke.yml
vendored
Normal file
|
|
@ -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-<version>-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
|
||||
7
.gitignore
vendored
7
.gitignore
vendored
|
|
@ -3,6 +3,8 @@ __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/
|
||||
|
|
@ -24,8 +26,8 @@ dist/
|
|||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
/lib/
|
||||
/lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
|
|
@ -228,3 +230,4 @@ setup_leo.sh
|
|||
server.pid
|
||||
*.log
|
||||
package-lock.json
|
||||
llama.cpp/
|
||||
|
|
|
|||
|
|
@ -218,10 +218,12 @@ unsloth studio -p 8888
|
|||
```
|
||||
|
||||
#### Uninstall
|
||||
You can uninstall Unsloth Studio by deleting its install folder usually located under `$HOME/.unsloth/studio` on Mac/Linux/WSL and `%USERPROFILE%\.unsloth\studio` on Windows. Using the `rm -rf` commands will **delete everything**, including your history, cache:
|
||||
The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows):
|
||||
|
||||
* **MacOS, WSL, Linux:** `rm -rf ~/.unsloth/studio`
|
||||
* **Windows (PowerShell):** `Remove-Item -Recurse -Force "$HOME\.unsloth\studio"`
|
||||
* **MacOS, WSL, Linux:** `curl -fsSL https://unsloth.ai/uninstall.sh | sh`
|
||||
* **Windows (PowerShell):** `irm https://unsloth.ai/uninstall.ps1 | iex`
|
||||
|
||||
If you only want to drop the install dir and keep the launcher/shortcut for a later reinstall, you can instead run `rm -rf ~/.unsloth/studio` (Mac/Linux/WSL) or `Remove-Item -Recurse -Force "$HOME\.unsloth\studio"` (Windows). The model cache at `~/.cache/huggingface` is not touched by any of these.
|
||||
|
||||
For more info, [see our docs](https://unsloth.ai/docs/new/studio/install#uninstall).
|
||||
|
||||
|
|
|
|||
31
build.sh
31
build.sh
|
|
@ -2,6 +2,10 @@
|
|||
|
||||
set -euo pipefail
|
||||
|
||||
# PyPI/Studio release publishing must use `./build.sh publish` (or an
|
||||
# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Studio
|
||||
# artifacts include the display-only Studio release version.
|
||||
|
||||
# 1. Build frontend (Vite outputs to dist/)
|
||||
cd studio/frontend
|
||||
|
||||
|
|
@ -70,10 +74,33 @@ cd ../..
|
|||
# 2. Clean old artifacts
|
||||
rm -rf build dist *.egg-info
|
||||
|
||||
# 3. Build wheel
|
||||
# 3. Stamp display-only Studio release metadata for packaged builds.
|
||||
_STUDIO_BUILD_INFO="studio/backend/utils/_studio_release_build.py"
|
||||
_STUDIO_BUILD_INFO_BACKUP="$(mktemp)"
|
||||
cp "$_STUDIO_BUILD_INFO" "$_STUDIO_BUILD_INFO_BACKUP"
|
||||
_restore_studio_build_info() {
|
||||
cp "$_STUDIO_BUILD_INFO_BACKUP" "$_STUDIO_BUILD_INFO" 2>/dev/null || true
|
||||
rm -f "$_STUDIO_BUILD_INFO_BACKUP"
|
||||
}
|
||||
trap _restore_studio_build_info EXIT
|
||||
|
||||
if [ "${1:-}" = "publish" ]; then
|
||||
STUDIO_STAMPED_VERSION="$(python scripts/stamp_studio_release.py --require-release)"
|
||||
else
|
||||
STUDIO_STAMPED_VERSION="$(python scripts/stamp_studio_release.py)"
|
||||
fi
|
||||
|
||||
# 4. Build wheel/sdist
|
||||
python -m build
|
||||
|
||||
# 4. Optionally publish
|
||||
if [ "${1:-}" = "publish" ]; then
|
||||
python scripts/stamp_studio_release.py --verify-dist dist --expected "$STUDIO_STAMPED_VERSION"
|
||||
fi
|
||||
|
||||
_restore_studio_build_info
|
||||
trap - EXIT
|
||||
|
||||
# 5. Optionally publish
|
||||
if [ "${1:-}" = "publish" ]; then
|
||||
python -m twine upload dist/*
|
||||
fi
|
||||
|
|
|
|||
Binary file not shown.
|
Before Width: | Height: | Size: 13 KiB After Width: | Height: | Size: 15 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 12 KiB After Width: | Height: | Size: 12 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 59 KiB After Width: | Height: | Size: 31 KiB |
400
install.ps1
400
install.ps1
|
|
@ -3,6 +3,11 @@
|
|||
# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local
|
||||
# NoTorch: .\install.ps1 --no-torch (skip PyTorch, GGUF-only mode)
|
||||
# Test: .\install.ps1 --package roland-sloth
|
||||
#
|
||||
# Env vars (priority: UNSLOTH_STUDIO_HOME > STUDIO_HOME > USERPROFILE-redirect > default):
|
||||
# UNSLOTH_STUDIO_HOME / STUDIO_HOME = path -> install under that path
|
||||
# (DataDir nests inside; user PATH not modified persistently).
|
||||
# Default ($USERPROFILE\.unsloth\studio) is preserved when no env var is set.
|
||||
|
||||
function Install-UnslothStudio {
|
||||
$ErrorActionPreference = "Stop"
|
||||
|
|
@ -126,7 +131,94 @@ function Install-UnslothStudio {
|
|||
}
|
||||
|
||||
$PythonVersion = "3.13"
|
||||
$StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
|
||||
|
||||
# Resolve install destinations. Priority: UNSLOTH_STUDIO_HOME, then
|
||||
# STUDIO_HOME alias, then USERPROFILE-redirect, then default.
|
||||
# Reject whitespace-only values so " " is treated as unset (matches the
|
||||
# Python resolvers' .strip()), preventing install/runtime layout drift.
|
||||
$envOverrideVar = $null
|
||||
$envOverride = $null
|
||||
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) {
|
||||
$envOverrideVar = "UNSLOTH_STUDIO_HOME"
|
||||
$envOverride = $env:UNSLOTH_STUDIO_HOME.Trim()
|
||||
} elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) {
|
||||
$envOverrideVar = "STUDIO_HOME"
|
||||
$envOverride = $env:STUDIO_HOME.Trim()
|
||||
}
|
||||
|
||||
# Custom Studio roots are not supported with --tauri (desktop app still
|
||||
# resolves %USERPROFILE%\.unsloth\studio). Pass through if override == legacy.
|
||||
if ($TauriMode -and $envOverride) {
|
||||
$_tauriOverride = $envOverride
|
||||
if ($_tauriOverride -eq "~" -or $_tauriOverride -like "~/*" -or $_tauriOverride -like "~\*") {
|
||||
$_tauriOverride = (Join-Path $env:USERPROFILE $_tauriOverride.Substring(1).TrimStart('/','\'))
|
||||
}
|
||||
try {
|
||||
$_tauriOverride = [System.IO.Path]::GetFullPath($_tauriOverride)
|
||||
} catch {}
|
||||
$_legacyTauriRoot = Join-Path $env:USERPROFILE ".unsloth\studio"
|
||||
try {
|
||||
$_legacyTauriRoot = [System.IO.Path]::GetFullPath($_legacyTauriRoot)
|
||||
} catch {}
|
||||
# Strip trailing separators so ".../studio\" matches ".../studio".
|
||||
$_trimSeps = @(
|
||||
[System.IO.Path]::DirectorySeparatorChar,
|
||||
[System.IO.Path]::AltDirectorySeparatorChar
|
||||
)
|
||||
$_tauriOverride = $_tauriOverride.TrimEnd($_trimSeps)
|
||||
$_legacyTauriRoot = $_legacyTauriRoot.TrimEnd($_trimSeps)
|
||||
if ($_tauriOverride -ne $_legacyTauriRoot) {
|
||||
Write-Host "ERROR: $envOverrideVar is not supported with --tauri." -ForegroundColor Red
|
||||
Write-Host " The desktop app still uses the legacy %USERPROFILE%\.unsloth\studio root." -ForegroundColor Red
|
||||
Write-Host " Run install.ps1 without --tauri for custom-root shell installs," -ForegroundColor Yellow
|
||||
Write-Host " or unset the env var for default desktop installs." -ForegroundColor Yellow
|
||||
throw "$envOverrideVar is not supported with --tauri."
|
||||
}
|
||||
}
|
||||
|
||||
$defaultProfile = $null
|
||||
try { $defaultProfile = [Environment]::GetFolderPath("UserProfile") } catch {}
|
||||
|
||||
# LOCALAPPDATA may be unset in service / CI contexts; Join-Path would abort
|
||||
# under ErrorActionPreference=Stop without this guard.
|
||||
$defaultDataDir = if ($env:LOCALAPPDATA -and -not [string]::IsNullOrWhiteSpace($env:LOCALAPPDATA)) {
|
||||
Join-Path $env:LOCALAPPDATA "Unsloth Studio"
|
||||
} else { $null }
|
||||
|
||||
if ($envOverride) {
|
||||
# Tilde expansion: env vars aren't subject to it when quoted on assignment.
|
||||
if ($envOverride -eq "~" -or $envOverride -like "~/*" -or $envOverride -like "~\*") {
|
||||
$envOverride = (Join-Path $env:USERPROFILE $envOverride.Substring(1).TrimStart('/','\'))
|
||||
}
|
||||
try {
|
||||
# .NET API: New-Item -Path treats brackets as wildcards and has no
|
||||
# -LiteralPath in PS 5.1, so a root like C:\studio[abc] would fail.
|
||||
[System.IO.Directory]::CreateDirectory($envOverride) | Out-Null
|
||||
$StudioHome = (Resolve-Path -LiteralPath $envOverride).Path
|
||||
} catch {
|
||||
Write-Host "ERROR: $envOverrideVar=$envOverride cannot be created or accessed." -ForegroundColor Red
|
||||
throw "$envOverrideVar=$envOverride cannot be created or accessed."
|
||||
}
|
||||
$probe = Join-Path $StudioHome (".unsloth-write-probe-" + [guid]::NewGuid())
|
||||
try {
|
||||
# WriteAllText: literal-path safe + closes handle so Remove-Item works.
|
||||
[System.IO.File]::WriteAllText($probe, "")
|
||||
Remove-Item -LiteralPath $probe -Force -ErrorAction SilentlyContinue
|
||||
} catch {
|
||||
Write-Host "ERROR: $envOverrideVar=$StudioHome is not writable." -ForegroundColor Red
|
||||
throw "$envOverrideVar=$StudioHome is not writable."
|
||||
}
|
||||
$StudioDataDir = Join-Path $StudioHome "share"
|
||||
$StudioRedirectMode = 'env'
|
||||
} elseif ($defaultProfile -and $env:USERPROFILE -and ($env:USERPROFILE -ne $defaultProfile)) {
|
||||
$StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
|
||||
$StudioDataDir = $defaultDataDir
|
||||
$StudioRedirectMode = 'profile'
|
||||
} else {
|
||||
$StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
|
||||
$StudioDataDir = $defaultDataDir
|
||||
$StudioRedirectMode = 'default'
|
||||
}
|
||||
$VenvDir = Join-Path $StudioHome "unsloth_studio"
|
||||
|
||||
$Rule = [string]::new([char]0x2500, 52)
|
||||
|
|
@ -378,24 +470,24 @@ function Install-UnslothStudio {
|
|||
[Parameter(Mandatory = $true)][string]$UnslothExePath
|
||||
)
|
||||
|
||||
if (-not (Test-Path $UnslothExePath)) {
|
||||
if (-not (Test-Path -LiteralPath $UnslothExePath)) {
|
||||
substep "cannot create shortcuts, unsloth.exe not found at $UnslothExePath" "Yellow"
|
||||
return
|
||||
}
|
||||
try {
|
||||
# Persist an absolute path in launcher scripts so shortcut working
|
||||
# directory changes do not break process startup.
|
||||
$UnslothExePath = (Resolve-Path $UnslothExePath).Path
|
||||
$UnslothExePath = (Resolve-Path -LiteralPath $UnslothExePath).Path
|
||||
# Escape for single-quoted embedding in generated launcher script.
|
||||
# This prevents runtime variable expansion for paths containing '$'.
|
||||
$SingleQuotedExePath = $UnslothExePath -replace "'", "''"
|
||||
|
||||
$localAppDataDir = $env:LOCALAPPDATA
|
||||
if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) {
|
||||
substep "LOCALAPPDATA path unavailable; skipped shortcut creation" "Yellow"
|
||||
# $StudioDataDir = LOCALAPPDATA\Unsloth Studio, or $StudioHome\share in env-mode.
|
||||
if (-not $StudioDataDir -or [string]::IsNullOrWhiteSpace($StudioDataDir)) {
|
||||
substep "DataDir path unavailable; skipped shortcut creation" "Yellow"
|
||||
return
|
||||
}
|
||||
$appDir = Join-Path $localAppDataDir "Unsloth Studio"
|
||||
$appDir = $StudioDataDir
|
||||
$launcherPs1 = Join-Path $appDir "launch-studio.ps1"
|
||||
$launcherVbs = Join-Path $appDir "launch-studio.vbs"
|
||||
$desktopDir = [Environment]::GetFolderPath("Desktop")
|
||||
|
|
@ -427,23 +519,89 @@ function Install-UnslothStudio {
|
|||
}
|
||||
$iconUrl = "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth.ico"
|
||||
|
||||
if (-not (Test-Path $appDir)) {
|
||||
New-Item -ItemType Directory -Path $appDir -Force | Out-Null
|
||||
if (-not (Test-Path -LiteralPath $appDir)) {
|
||||
[System.IO.Directory]::CreateDirectory($appDir) | Out-Null
|
||||
}
|
||||
|
||||
# Same-install discriminator: per-install opaque id written once at
|
||||
# install time and read by both this launcher and the backend
|
||||
# (/api/health). Replaces the older sha256(resolved $StudioHome)
|
||||
# scheme to (a) avoid leaking the install path on -H 0.0.0.0
|
||||
# deployments and (b) sidestep launcher/backend canonicalization
|
||||
# drift (Resolve-Path vs Path.resolve() junction handling). Lives
|
||||
# at $StudioHome\share\ (not $appDir) so the backend can find it
|
||||
# via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id"
|
||||
# regardless of mode. 32 bytes of crypto random -> 64 hex chars.
|
||||
$_studioIdDir = Join-Path $StudioHome "share"
|
||||
if (-not (Test-Path -LiteralPath $_studioIdDir)) {
|
||||
[System.IO.Directory]::CreateDirectory($_studioIdDir) | Out-Null
|
||||
}
|
||||
$_studioIdFile = Join-Path $_studioIdDir "studio_install_id"
|
||||
$_studioRootId = ""
|
||||
if ((Test-Path -LiteralPath $_studioIdFile) -and `
|
||||
((Get-Item -LiteralPath $_studioIdFile).Length -gt 0)) {
|
||||
$_studioRootId = ([System.IO.File]::ReadAllText($_studioIdFile)).Trim()
|
||||
}
|
||||
if (-not $_studioRootId) {
|
||||
$_idBytes = New-Object byte[] 32
|
||||
[Security.Cryptography.RandomNumberGenerator]::Create().GetBytes($_idBytes)
|
||||
$_studioRootId = -join ($_idBytes | ForEach-Object { $_.ToString('x2') })
|
||||
# Atomic write: write to a temp sibling then rename, so a partial
|
||||
# install cannot leave a half-written id.
|
||||
$_idTmp = $_studioIdFile + ".$PID.tmp"
|
||||
[System.IO.File]::WriteAllText($_idTmp, $_studioRootId)
|
||||
Move-Item -LiteralPath $_idTmp -Destination $_studioIdFile -Force
|
||||
}
|
||||
|
||||
# Env-mode: persist UNSLOTH_STUDIO_HOME (and llama path) so fresh
|
||||
# shells don't need to re-export, and bake per-install $portFile /
|
||||
# $mutexName so concurrent custom-root launchers cannot serialize
|
||||
# through one global mutex on 8888..8908. Default installs get an
|
||||
# empty prefix to match pre-PR behavior.
|
||||
$studioHomeExport = if ($StudioRedirectMode -eq 'env') {
|
||||
# When override == legacy default, llama.cpp stays at
|
||||
# ~/.unsloth/llama.cpp (one shared build). Canonicalize the
|
||||
# legacy side so the comparison survives path normalization.
|
||||
$_legacyStudio = Join-Path $env:USERPROFILE ".unsloth\studio"
|
||||
if (Test-Path -LiteralPath $_legacyStudio -PathType Container) {
|
||||
$_legacyStudio = (Resolve-Path -LiteralPath $_legacyStudio).Path
|
||||
}
|
||||
$_llamaPath = if ($StudioHome -eq $_legacyStudio) {
|
||||
Join-Path $env:USERPROFILE ".unsloth\llama.cpp"
|
||||
} else {
|
||||
Join-Path $StudioHome "llama.cpp"
|
||||
}
|
||||
$_sq = $StudioHome -replace "'", "''"
|
||||
$_llama = $_llamaPath -replace "'", "''"
|
||||
$_appDirSq = $appDir -replace "'", "''"
|
||||
$_appBytes = [Text.Encoding]::UTF8.GetBytes($appDir)
|
||||
$_appHash = ([BitConverter]::ToString(
|
||||
[Security.Cryptography.SHA256]::Create().ComputeHash($_appBytes)
|
||||
) -replace '-', '').Substring(0, 16)
|
||||
# UNSLOTH_LLAMA_CPP_PATH is a pre-existing user override; only default if unset.
|
||||
"`$env:UNSLOTH_STUDIO_HOME = '$_sq'`nif (-not `$env:UNSLOTH_LLAMA_CPP_PATH) {`n `$env:UNSLOTH_LLAMA_CPP_PATH = '$_llama'`n}`n`$portFile = '$_appDirSq\studio.port'`n`$mutexName = 'Local\UnslothStudioLauncher-$_appHash'`n"
|
||||
} else {
|
||||
"`$portFile = `$null`n`$mutexName = 'Local\UnslothStudioLauncher'`n"
|
||||
}
|
||||
|
||||
$launcherContent = @"
|
||||
`$ErrorActionPreference = 'Stop'
|
||||
$studioHomeExport`$ErrorActionPreference = 'Stop'
|
||||
`$basePort = 8888
|
||||
`$maxPortOffset = 20
|
||||
`$timeoutSec = 60
|
||||
`$pollIntervalMs = 1000
|
||||
`$_ExpectedStudioRootId = '$_studioRootId'
|
||||
|
||||
function Test-StudioHealth {
|
||||
param([Parameter(Mandatory = `$true)][int]`$Port)
|
||||
try {
|
||||
`$url = "http://127.0.0.1:`$Port/api/health"
|
||||
`$resp = Invoke-RestMethod -Uri `$url -TimeoutSec 1 -Method Get
|
||||
return (`$resp -and `$resp.status -eq 'healthy' -and `$resp.service -eq 'Unsloth UI Backend')
|
||||
if (-not (`$resp -and `$resp.status -eq 'healthy' -and `$resp.service -eq 'Unsloth UI Backend')) { return `$false }
|
||||
# why: verify the backend belongs to THIS install via the install-time
|
||||
# hex digest; raw path is not leaked over /api/health.
|
||||
if (`$_ExpectedStudioRootId -and `$resp.studio_root_id -ne `$_ExpectedStudioRootId) { return `$false }
|
||||
return `$true
|
||||
} catch {
|
||||
return `$false
|
||||
}
|
||||
|
|
@ -469,6 +627,17 @@ function Get-CandidatePorts {
|
|||
}
|
||||
|
||||
function Find-HealthyStudioPort {
|
||||
if (`$portFile) {
|
||||
if (Test-Path -LiteralPath `$portFile) {
|
||||
`$cached = Get-Content -LiteralPath `$portFile -ErrorAction SilentlyContinue | Select-Object -First 1
|
||||
if (`$cached -match '^\d+`$') {
|
||||
`$cachedPort = [int]`$cached
|
||||
if (Test-StudioHealth -Port `$cachedPort) { return `$cachedPort }
|
||||
Remove-Item -LiteralPath `$portFile -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
}
|
||||
return `$null
|
||||
}
|
||||
foreach (`$candidate in (Get-CandidatePorts)) {
|
||||
if (Test-StudioHealth -Port `$candidate) {
|
||||
return `$candidate
|
||||
|
|
@ -522,7 +691,7 @@ if (`$existingPort) {
|
|||
exit 0
|
||||
}
|
||||
|
||||
`$launchMutex = [System.Threading.Mutex]::new(`$false, 'Local\UnslothStudioLauncher')
|
||||
`$launchMutex = [System.Threading.Mutex]::new(`$false, `$mutexName)
|
||||
`$haveMutex = `$false
|
||||
try {
|
||||
try {
|
||||
|
|
@ -552,7 +721,9 @@ try {
|
|||
} catch {}
|
||||
exit 1
|
||||
}
|
||||
`$studioCommand = '& "' + `$studioExe + '" studio -p ' + `$launchPort
|
||||
# Single-quote the path in the child -Command so `$` / backtick in custom
|
||||
# roots don't get reparsed; double any apostrophes so 'O''Brien' survives.
|
||||
`$studioCommand = "& '" + (`$studioExe -replace "'", "''") + "' studio -p " + `$launchPort
|
||||
`$launchArgs = @(
|
||||
'-NoExit',
|
||||
'-NoProfile',
|
||||
|
|
@ -576,9 +747,13 @@ try {
|
|||
`$browserOpened = `$false
|
||||
`$deadline = (Get-Date).AddSeconds(`$timeoutSec)
|
||||
while ((Get-Date) -lt `$deadline) {
|
||||
`$healthyPort = Find-HealthyStudioPort
|
||||
if (`$healthyPort) {
|
||||
Start-Process "http://localhost:`$healthyPort"
|
||||
if (Test-StudioHealth -Port `$launchPort) {
|
||||
if (`$portFile) {
|
||||
try {
|
||||
[System.IO.File]::WriteAllText(`$portFile, "`$launchPort`n")
|
||||
} catch {}
|
||||
}
|
||||
Start-Process "http://localhost:`$launchPort"
|
||||
`$browserOpened = `$true
|
||||
break
|
||||
}
|
||||
|
|
@ -613,19 +788,19 @@ cmd = "powershell -NoProfile -ExecutionPolicy Bypass -WindowStyle Hidden -File "
|
|||
shell.Run cmd, 0, False
|
||||
"@
|
||||
# WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths.
|
||||
Set-Content -Path $launcherVbs -Value $vbsContent -Encoding Unicode -Force
|
||||
Set-Content -LiteralPath $launcherVbs -Value $vbsContent -Encoding Unicode -Force
|
||||
|
||||
# Prefer bundled icon from local clone/dev installs.
|
||||
# If not available, best-effort download from raw GitHub.
|
||||
# We only attach the icon if the resulting file has a valid ICO header.
|
||||
$hasValidIcon = $false
|
||||
if ($bundledIcon -and (Test-Path $bundledIcon)) {
|
||||
if ($bundledIcon -and (Test-Path -LiteralPath $bundledIcon)) {
|
||||
try {
|
||||
Copy-Item -Path $bundledIcon -Destination $iconPath -Force
|
||||
Copy-Item -LiteralPath $bundledIcon -Destination $iconPath -Force
|
||||
} catch {
|
||||
Write-Host "[DEBUG] Error copying bundled icon: $($_.Exception.Message)" -ForegroundColor DarkGray
|
||||
}
|
||||
} elseif (-not (Test-Path $iconPath)) {
|
||||
} elseif (-not (Test-Path -LiteralPath $iconPath)) {
|
||||
try {
|
||||
Invoke-WebRequest -Uri $iconUrl -OutFile $iconPath -UseBasicParsing
|
||||
} catch {
|
||||
|
|
@ -633,7 +808,7 @@ shell.Run cmd, 0, False
|
|||
}
|
||||
}
|
||||
|
||||
if (Test-Path $iconPath) {
|
||||
if (Test-Path -LiteralPath $iconPath) {
|
||||
try {
|
||||
$bytes = [System.IO.File]::ReadAllBytes($iconPath)
|
||||
if (
|
||||
|
|
@ -645,14 +820,21 @@ shell.Run cmd, 0, False
|
|||
) {
|
||||
$hasValidIcon = $true
|
||||
} else {
|
||||
Remove-Item $iconPath -Force -ErrorAction SilentlyContinue
|
||||
Remove-Item -LiteralPath $iconPath -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
} catch {
|
||||
Write-Host "[DEBUG] Error validating or removing icon: $($_.Exception.Message)" -ForegroundColor DarkGray
|
||||
Remove-Item $iconPath -Force -ErrorAction SilentlyContinue
|
||||
Remove-Item -LiteralPath $iconPath -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
}
|
||||
|
||||
# Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts
|
||||
# that may point at a deleted workspace; launcher + icon stay.
|
||||
if ($StudioRedirectMode -eq 'env') {
|
||||
substep "wrote launcher at $launcherPs1 (persistent shortcuts skipped in env-override mode)"
|
||||
return
|
||||
}
|
||||
|
||||
$wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe"
|
||||
$shortcutArgs = "//B //Nologo `"$launcherVbs`""
|
||||
|
||||
|
|
@ -850,8 +1032,9 @@ shell.Run cmd, 0, False
|
|||
# Pass the resolved executable path to uv so it does not re-resolve
|
||||
# a version string back to a conda interpreter.
|
||||
Write-TauriLog "STEP" "Creating virtual environment"
|
||||
if (-not (Test-Path $StudioHome)) {
|
||||
New-Item -ItemType Directory -Path $StudioHome -Force | Out-Null
|
||||
if (-not (Test-Path -LiteralPath $StudioHome)) {
|
||||
# .NET API: New-Item -Path treats brackets as wildcards.
|
||||
[System.IO.Directory]::CreateDirectory($StudioHome) | Out-Null
|
||||
}
|
||||
|
||||
$VenvPython = Join-Path $VenvDir "Scripts\python.exe"
|
||||
|
|
@ -865,11 +1048,13 @@ shell.Run cmd, 0, False
|
|||
$stamp = Get-Date -Format "yyyyMMddHHmmss"
|
||||
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID"
|
||||
$suffix = 0
|
||||
while (Test-Path $candidate) {
|
||||
# -LiteralPath: a custom $StudioHome may contain [ ] * ? which
|
||||
# plain Test-Path / Move-Item would interpret as wildcards.
|
||||
while (Test-Path -LiteralPath $candidate) {
|
||||
$suffix++
|
||||
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix"
|
||||
}
|
||||
Move-Item -Path $ExistingDir -Destination $candidate -ErrorAction Stop
|
||||
Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop
|
||||
$script:StudioVenvRollbackDir = $candidate
|
||||
$script:StudioVenvRollbackTarget = $ExistingDir
|
||||
$script:StudioVenvRollbackActive = $true
|
||||
|
|
@ -880,16 +1065,16 @@ shell.Run cmd, 0, False
|
|||
if (-not $script:StudioVenvRollbackActive) { return }
|
||||
$backup = $script:StudioVenvRollbackDir
|
||||
$target = $script:StudioVenvRollbackTarget
|
||||
if (-not $backup -or -not (Test-Path $backup)) {
|
||||
if (-not $backup -or -not (Test-Path -LiteralPath $backup)) {
|
||||
$script:StudioVenvRollbackActive = $false
|
||||
return
|
||||
}
|
||||
substep "restoring previous environment after failed install..." "Yellow"
|
||||
try {
|
||||
if (Test-Path $target) {
|
||||
Remove-Item -Recurse -Force $target -ErrorAction SilentlyContinue
|
||||
if (Test-Path -LiteralPath $target) {
|
||||
Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
Move-Item -Path $backup -Destination $target -Force -ErrorAction Stop
|
||||
Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop
|
||||
substep "restored previous environment"
|
||||
$script:StudioVenvRollbackActive = $false
|
||||
$script:StudioVenvRollbackDir = $null
|
||||
|
|
@ -902,14 +1087,29 @@ shell.Run cmd, 0, False
|
|||
function Complete-StudioVenvRollback {
|
||||
if (-not $script:StudioVenvRollbackActive) { return }
|
||||
$backup = $script:StudioVenvRollbackDir
|
||||
if ($backup -and (Test-Path $backup)) {
|
||||
Remove-Item -Recurse -Force $backup -ErrorAction SilentlyContinue
|
||||
if ($backup -and (Test-Path -LiteralPath $backup)) {
|
||||
Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
$script:StudioVenvRollbackActive = $false
|
||||
$script:StudioVenvRollbackDir = $null
|
||||
}
|
||||
|
||||
if (Test-Path $VenvPython) {
|
||||
if (Test-Path -LiteralPath $VenvPython) {
|
||||
# why: matching guard to the .venv branch below -- in env-mode
|
||||
# $StudioHome is a user-chosen workspace, so refuse to nuke an
|
||||
# existing $StudioHome\unsloth_studio that lacks Studio sentinels.
|
||||
# -PathType Leaf rejects a directory at the sentinel path. Accept the
|
||||
# in-VENV ownership marker so partial-install retries are not blocked.
|
||||
if (
|
||||
$StudioRedirectMode -eq 'env' -and
|
||||
-not (Test-Path -LiteralPath (Join-Path $VenvDir ".unsloth-studio-owned") -PathType Leaf) -and
|
||||
-not (Test-Path -LiteralPath (Join-Path $StudioHome "share\studio.conf") -PathType Leaf) -and
|
||||
-not (Test-Path -LiteralPath (Join-Path $StudioHome "bin\unsloth.exe") -PathType Leaf)
|
||||
) {
|
||||
Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red
|
||||
Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow
|
||||
throw "Refusing to delete non-Studio venv at $VenvDir"
|
||||
}
|
||||
# New layout already exists -- replace only after preserving rollback copy.
|
||||
substep "preserving existing environment for rollback..."
|
||||
try {
|
||||
|
|
@ -918,8 +1118,13 @@ shell.Run cmd, 0, False
|
|||
Write-Host "[ERROR] Could not prepare existing environment for reinstall: $($_.Exception.Message)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Could not prepare existing environment for reinstall")
|
||||
}
|
||||
} elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) {
|
||||
# Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating
|
||||
} elseif (
|
||||
$StudioRedirectMode -ne 'env' `
|
||||
-and (Test-Path -LiteralPath (Join-Path $StudioHome ".venv\Scripts\python.exe"))
|
||||
) {
|
||||
# Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating.
|
||||
# Skip in env-mode so we don't blow away an unrelated .venv at the
|
||||
# workspace root (e.g. user's existing project Python venv).
|
||||
$OldVenv = Join-Path $StudioHome ".venv"
|
||||
$OldPy = Join-Path $OldVenv "Scripts\python.exe"
|
||||
substep "found legacy Studio environment, validating..."
|
||||
|
|
@ -936,24 +1141,29 @@ shell.Run cmd, 0, False
|
|||
$ErrorActionPreference = $prevEAP2
|
||||
if ($legacyOk) {
|
||||
substep "legacy environment is healthy -- migrating..."
|
||||
Move-Item -Path $OldVenv -Destination $VenvDir -Force
|
||||
Move-Item -LiteralPath $OldVenv -Destination $VenvDir -Force
|
||||
substep "moved .venv -> unsloth_studio"
|
||||
$_Migrated = $true
|
||||
} else {
|
||||
substep "legacy environment failed validation -- creating fresh environment" "Yellow"
|
||||
$invalidVenv = Join-Path $StudioHome (".venv.invalid.{0}.{1}" -f (Get-Date -Format "yyyyMMddHHmmss"), $PID)
|
||||
Move-Item -Path $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue
|
||||
Move-Item -LiteralPath $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
} elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) {
|
||||
# CWD-relative venv from old install.ps1 -- migrate to absolute path
|
||||
} elseif (
|
||||
$StudioRedirectMode -ne 'env' `
|
||||
-and (Test-Path -LiteralPath (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe"))
|
||||
) {
|
||||
# CWD-relative venv from old install.ps1 -> migrate to absolute path.
|
||||
# Skip in env-mode so we don't relocate the default-install venv into
|
||||
# the workspace root.
|
||||
$CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio"
|
||||
substep "found CWD-relative Studio environment, migrating to $VenvDir..."
|
||||
Move-Item -Path $CwdVenv -Destination $VenvDir -Force
|
||||
Move-Item -LiteralPath $CwdVenv -Destination $VenvDir -Force
|
||||
substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio"
|
||||
$_Migrated = $true
|
||||
}
|
||||
|
||||
if (-not (Test-Path $VenvPython)) {
|
||||
if (-not (Test-Path -LiteralPath $VenvPython)) {
|
||||
step "venv" "creating Python $($DetectedPython.Version) virtual environment"
|
||||
substep "$VenvDir"
|
||||
$venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" }
|
||||
|
|
@ -966,6 +1176,13 @@ shell.Run cmd, 0, False
|
|||
substep "$VenvDir"
|
||||
}
|
||||
|
||||
# Mark the freshly-created venv as Studio-owned so a partial install can be
|
||||
# repaired by re-running install.ps1; the env-mode deletion guard above
|
||||
# accepts this marker as the primary sentinel.
|
||||
if (Test-Path -LiteralPath $VenvDir -PathType Container) {
|
||||
try { [System.IO.File]::WriteAllText((Join-Path $VenvDir ".unsloth-studio-owned"), "") } catch {}
|
||||
}
|
||||
|
||||
# ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ──
|
||||
$HasNvidiaSmi = $false
|
||||
$NvidiaSmiExe = $null
|
||||
|
|
@ -1054,7 +1271,7 @@ shell.Run cmd, 0, False
|
|||
if ($StudioLocalInstall -and (Test-Path (Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"))) {
|
||||
return Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"
|
||||
}
|
||||
$installed = Get-ChildItem -Path $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue |
|
||||
$installed = Get-ChildItem -LiteralPath $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue |
|
||||
Where-Object { $_.FullName -like "*studio*backend*requirements*no-torch-runtime.txt" } |
|
||||
Select-Object -ExpandProperty FullName -First 1
|
||||
return $installed
|
||||
|
|
@ -1068,7 +1285,7 @@ shell.Run cmd, 0, False
|
|||
if ($SkipTorch) {
|
||||
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
|
||||
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.2" unsloth-zoo }
|
||||
if ($baseInstallExit -eq 0) {
|
||||
$NoTorchReq = Find-NoTorchRuntimeFile
|
||||
if ($NoTorchReq) {
|
||||
|
|
@ -1076,7 +1293,7 @@ shell.Run cmd, 0, False
|
|||
}
|
||||
}
|
||||
} else {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.2" unsloth-zoo }
|
||||
}
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
|
|
@ -1114,7 +1331,7 @@ shell.Run cmd, 0, False
|
|||
if ($SkipTorch) {
|
||||
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
|
||||
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.2" unsloth-zoo }
|
||||
if ($baseInstallExit -eq 0) {
|
||||
$NoTorchReq = Find-NoTorchRuntimeFile
|
||||
if ($NoTorchReq) {
|
||||
|
|
@ -1122,7 +1339,7 @@ shell.Run cmd, 0, False
|
|||
}
|
||||
}
|
||||
} elseif ($StudioLocalInstall) {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.1" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.2" unsloth-zoo }
|
||||
} else {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
|
||||
}
|
||||
|
|
@ -1150,7 +1367,7 @@ shell.Run cmd, 0, False
|
|||
Write-TauriLog "STEP" "Installing unsloth"
|
||||
substep "installing unsloth (this may take a few minutes)..."
|
||||
if ($StudioLocalInstall) {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.1" --torch-backend=auto }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.2" --torch-backend=auto }
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
|
||||
|
|
@ -1192,23 +1409,25 @@ shell.Run cmd, 0, False
|
|||
foreach ($rel in $overlayMap.Keys) {
|
||||
$src = Join-Path $scriptDir $rel
|
||||
$dst = Join-Path $VenvDir $overlayMap[$rel]
|
||||
if (-not (Test-Path $src)) { continue }
|
||||
# -LiteralPath: $VenvDir derives from $StudioHome which may
|
||||
# contain [ ] * ? when the user overrode UNSLOTH_STUDIO_HOME.
|
||||
if (-not (Test-Path -LiteralPath $src)) { continue }
|
||||
$dstParent = Split-Path -Parent $dst
|
||||
if (-not (Test-Path $dstParent)) {
|
||||
if (-not (Test-Path -LiteralPath $dstParent)) {
|
||||
Write-Host "[WARN] Overlay target dir missing: $dstParent; studio setup may use stale bundled file" -ForegroundColor Yellow
|
||||
continue
|
||||
}
|
||||
try {
|
||||
if (-not (Test-Path $dst)) {
|
||||
if (-not (Test-Path -LiteralPath $dst)) {
|
||||
# Backfill: target file missing but parent dir exists.
|
||||
Copy-Item $src $dst -Force
|
||||
Copy-Item -LiteralPath $src -Destination $dst -Force
|
||||
substep ("backfilled bundled " + (Split-Path -Leaf $rel))
|
||||
} else {
|
||||
# Hash-compare so re-runs are no-ops when files already match.
|
||||
$srcHash = (Get-FileHash $src -Algorithm SHA256).Hash
|
||||
$dstHash = (Get-FileHash $dst -Algorithm SHA256).Hash
|
||||
$srcHash = (Get-FileHash -LiteralPath $src -Algorithm SHA256).Hash
|
||||
$dstHash = (Get-FileHash -LiteralPath $dst -Algorithm SHA256).Hash
|
||||
if ($srcHash -ne $dstHash) {
|
||||
Copy-Item $src $dst -Force
|
||||
Copy-Item -LiteralPath $src -Destination $dst -Force
|
||||
substep ("applied bundled " + (Split-Path -Leaf $rel))
|
||||
}
|
||||
}
|
||||
|
|
@ -1225,7 +1444,8 @@ shell.Run cmd, 0, False
|
|||
Write-TauriLog "STEP" "Running studio setup"
|
||||
step "setup" "running unsloth studio setup..."
|
||||
$UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe"
|
||||
if (-not (Test-Path $UnslothExe)) {
|
||||
if (-not (Test-Path -LiteralPath $UnslothExe)) {
|
||||
Write-TauriLog "ERROR" "unsloth CLI was not installed correctly"
|
||||
Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red
|
||||
Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow
|
||||
Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow
|
||||
|
|
@ -1250,6 +1470,15 @@ shell.Run cmd, 0, False
|
|||
# Use 'studio setup' (not 'studio update') because 'update' pops
|
||||
# SKIP_STUDIO_BASE, which would cause redundant package reinstallation
|
||||
# and bypass the fast-path version check from PR #4667.
|
||||
# Propagate UNSLOTH_STUDIO_HOME only for env-override installs; otherwise
|
||||
# an inherited value would put llama.cpp in the wrong place.
|
||||
$previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME
|
||||
$hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome)
|
||||
if ($StudioRedirectMode -eq 'env') {
|
||||
$env:UNSLOTH_STUDIO_HOME = $StudioHome
|
||||
} else {
|
||||
Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue
|
||||
}
|
||||
$studioArgs = @('studio', 'setup')
|
||||
if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
|
||||
$env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1"
|
||||
|
|
@ -1257,6 +1486,11 @@ shell.Run cmd, 0, False
|
|||
& $UnslothExe @studioArgs
|
||||
$setupExit = $LASTEXITCODE
|
||||
} finally {
|
||||
if ($hadPreviousUnslothStudioHome) {
|
||||
$env:UNSLOTH_STUDIO_HOME = $previousUnslothStudioHome
|
||||
} else {
|
||||
Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue
|
||||
}
|
||||
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
|
||||
}
|
||||
if ($setupExit -ne 0) {
|
||||
|
|
@ -1301,20 +1535,32 @@ shell.Run cmd, 0, False
|
|||
}
|
||||
} catch { }
|
||||
$ShimDir = Join-Path $StudioHome "bin"
|
||||
New-Item -ItemType Directory -Force -Path $ShimDir | Out-Null
|
||||
[System.IO.Directory]::CreateDirectory($ShimDir) | Out-Null
|
||||
$ShimExe = Join-Path $ShimDir "unsloth.exe"
|
||||
# Fatal preflight outside the lock-handling try/catch -- a directory at
|
||||
# the shim path must not be downgraded to "Continuing with the existing
|
||||
# launcher", or the install finishes with no usable shim.
|
||||
if (Test-Path -LiteralPath $ShimExe -PathType Container) {
|
||||
Write-Host "[ERROR] Cannot create unsloth launcher: $ShimExe is a directory." -ForegroundColor Red
|
||||
Write-Host " Move or remove it manually, then re-run the installer." -ForegroundColor Yellow
|
||||
throw "Cannot create unsloth launcher: $ShimExe is a directory."
|
||||
}
|
||||
# try/catch: if unsloth.exe is locked (Studio running), keep the old shim.
|
||||
$shimUpdated = $false
|
||||
try {
|
||||
if (Test-Path $ShimExe) { Remove-Item $ShimExe -Force -ErrorAction Stop }
|
||||
if (Test-Path -LiteralPath $ShimExe) { Remove-Item -LiteralPath $ShimExe -Force -ErrorAction Stop }
|
||||
try {
|
||||
# New-Item -ItemType HardLink does NOT accept -LiteralPath in any
|
||||
# PowerShell version, so use -Path. Wildcards in $ShimExe (e.g.
|
||||
# brackets in custom roots) glob-expand here and fall through to
|
||||
# the Copy-Item -LiteralPath fallback below.
|
||||
New-Item -ItemType HardLink -Path $ShimExe -Target $UnslothExe -ErrorAction Stop | Out-Null
|
||||
} catch {
|
||||
Copy-Item -Path $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy
|
||||
Copy-Item -LiteralPath $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy
|
||||
}
|
||||
$shimUpdated = $true
|
||||
} catch {
|
||||
if (Test-Path $ShimExe) {
|
||||
if (Test-Path -LiteralPath $ShimExe) {
|
||||
Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow
|
||||
Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow
|
||||
Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow
|
||||
|
|
@ -1325,10 +1571,13 @@ shell.Run cmd, 0, False
|
|||
Write-Host " Launch unsloth studio directly via '$UnslothExe' until the next successful install." -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
# Only add to PATH when the launcher actually exists on disk.
|
||||
# Add to PATH only when launcher exists. Env-mode: session-only export,
|
||||
# no registry change (workspace path may be deleted later).
|
||||
$pathAdded = $false
|
||||
if (Test-Path $ShimExe) {
|
||||
$pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend'
|
||||
if (Test-Path -LiteralPath $ShimExe) {
|
||||
if ($StudioRedirectMode -ne 'env') {
|
||||
$pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend'
|
||||
}
|
||||
}
|
||||
if ($shimUpdated -and $pathAdded) {
|
||||
step "path" "added unsloth launcher to PATH"
|
||||
|
|
@ -1336,12 +1585,20 @@ shell.Run cmd, 0, False
|
|||
Refresh-SessionPath # sync current session with registry
|
||||
Complete-StudioVenvRollback
|
||||
|
||||
# Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy
|
||||
# User PATH entry (Machine > User > current $env:Path) would win.
|
||||
if ($StudioRedirectMode -eq 'env' -and (Test-Path -LiteralPath $ShimExe)) {
|
||||
$env:Path = "$ShimDir;$env:Path"
|
||||
step "path" "exported $ShimDir for this session (no registry PATH change in env-override mode)"
|
||||
}
|
||||
|
||||
# ── Tauri mode: done, skip shortcuts and auto-launch ──
|
||||
if ($TauriMode) {
|
||||
Write-TauriLog "DONE" ""
|
||||
return
|
||||
}
|
||||
|
||||
# New-StudioShortcuts gates the .lnk shortcuts on env-mode internally.
|
||||
New-StudioShortcuts -UnslothExePath $UnslothExe
|
||||
|
||||
# In interactive terminals, ask the user before starting Studio.
|
||||
|
|
@ -1360,8 +1617,21 @@ shell.Run cmd, 0, False
|
|||
}
|
||||
} else {
|
||||
step "launch" "manual commands:"
|
||||
substep "& `"$VenvDir\Scripts\Activate.ps1`""
|
||||
substep "unsloth studio -p 8888"
|
||||
# Single-quote the printed paths so $-vars / backticks in custom roots
|
||||
# do not reparse when the user pastes the command.
|
||||
$_actLiteral = "'" + ((Join-Path $VenvDir "Scripts\Activate.ps1") -replace "'", "''") + "'"
|
||||
if ($StudioRedirectMode -eq 'env') {
|
||||
# Env-mode skips registry PATH; print the absolute shim path.
|
||||
$_shim = Join-Path $StudioHome "bin\unsloth.exe"
|
||||
$_shimLiteral = "'" + ($_shim -replace "'", "''") + "'"
|
||||
substep "& $_shimLiteral studio -p 8888"
|
||||
substep "or activate env first:"
|
||||
substep "& $_actLiteral"
|
||||
substep "unsloth studio -p 8888"
|
||||
} else {
|
||||
substep "& $_actLiteral"
|
||||
substep "unsloth studio -p 8888"
|
||||
}
|
||||
substep "(add -H 0.0.0.0 to allow network / cloud access)"
|
||||
Write-Host ""
|
||||
}
|
||||
|
|
|
|||
506
install.sh
506
install.sh
|
|
@ -6,6 +6,12 @@
|
|||
# Usage (no-torch): ./install.sh --no-torch (skip PyTorch, GGUF-only mode)
|
||||
# Usage (test): ./install.sh --package roland-sloth (install a different package name)
|
||||
# Usage (py): ./install.sh --python 3.12 (override auto-detected Python version)
|
||||
#
|
||||
# Env vars (priority: UNSLOTH_STUDIO_HOME > STUDIO_HOME > HOME-redirect > default):
|
||||
# UNSLOTH_STUDIO_HOME=/abs/path -> install under that path
|
||||
# STUDIO_HOME=/abs/path -> alias, same effect (UNSLOTH_STUDIO_HOME wins)
|
||||
# (DATA_DIR + unsloth CLI shim nest inside; no shell rc-file append.)
|
||||
# Default ($HOME/.unsloth/studio) is preserved when no env var is set.
|
||||
set -e
|
||||
|
||||
# ── Output style (aligned with studio/setup.sh) ──
|
||||
|
|
@ -66,6 +72,56 @@ if [ "$_VERBOSE" = true ]; then
|
|||
export UNSLOTH_VERBOSE=1
|
||||
fi
|
||||
|
||||
# Custom Studio roots are not supported with --tauri (desktop app still
|
||||
# resolves ~/.unsloth/studio). Pass through if the override == legacy default.
|
||||
if [ "$TAURI_MODE" = true ]; then
|
||||
_tauri_override_var=""
|
||||
_tauri_override="${UNSLOTH_STUDIO_HOME:-}"
|
||||
if [ -n "$_tauri_override" ]; then
|
||||
_tauri_override_var="UNSLOTH_STUDIO_HOME"
|
||||
else
|
||||
_tauri_override="${STUDIO_HOME:-}"
|
||||
[ -n "$_tauri_override" ] && _tauri_override_var="STUDIO_HOME"
|
||||
fi
|
||||
# Strip whitespace so " " is treated as unset (matches Python .strip()).
|
||||
_tauri_override=$(printf '%s' "$_tauri_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')
|
||||
if [ -n "$_tauri_override" ]; then
|
||||
case "$_tauri_override" in
|
||||
"~") _tauri_override="$HOME" ;;
|
||||
"~/"*) _tauri_override="$HOME/${_tauri_override#'~/'}" ;;
|
||||
esac
|
||||
# Canonicalize both sides (CDPATH=, -P) so a CDPATH-set env or
|
||||
# symlinked $HOME doesn't break the legacy-equality comparison.
|
||||
if [ -d "$_tauri_override" ]; then
|
||||
_tauri_override_abs=$(CDPATH= cd -P -- "$_tauri_override" 2>/dev/null && pwd -P) \
|
||||
|| _tauri_override_abs="$_tauri_override"
|
||||
else
|
||||
_tauri_override_abs="$_tauri_override"
|
||||
fi
|
||||
# Strip trailing separators so ".../studio/" matches ".../studio".
|
||||
while [ "$_tauri_override_abs" != "/" ] \
|
||||
&& [ "${_tauri_override_abs%/}" != "$_tauri_override_abs" ]; do
|
||||
_tauri_override_abs=${_tauri_override_abs%/}
|
||||
done
|
||||
_tauri_legacy_root="$HOME/.unsloth/studio"
|
||||
if [ -d "$_tauri_legacy_root" ]; then
|
||||
_tauri_legacy_root=$(CDPATH= cd -P -- "$_tauri_legacy_root" 2>/dev/null && pwd -P) \
|
||||
|| _tauri_legacy_root="$HOME/.unsloth/studio"
|
||||
fi
|
||||
while [ "$_tauri_legacy_root" != "/" ] \
|
||||
&& [ "${_tauri_legacy_root%/}" != "$_tauri_legacy_root" ]; do
|
||||
_tauri_legacy_root=${_tauri_legacy_root%/}
|
||||
done
|
||||
if [ "$_tauri_override_abs" != "$_tauri_legacy_root" ]; then
|
||||
echo "ERROR: $_tauri_override_var is not supported with --tauri." >&2
|
||||
echo " The desktop app still uses the legacy ~/.unsloth/studio root." >&2
|
||||
echo " Run install.sh without --tauri for custom-root shell installs," >&2
|
||||
echo " or unset the env var for default desktop installs." >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
_is_verbose() {
|
||||
[ "${UNSLOTH_VERBOSE:-0}" = "1" ]
|
||||
}
|
||||
|
|
@ -219,7 +275,67 @@ _tauri_gpu_branch() {
|
|||
}
|
||||
|
||||
PYTHON_VERSION="" # resolved after platform detection
|
||||
STUDIO_HOME="$HOME/.unsloth/studio"
|
||||
|
||||
# Resolve install destinations: env override, HOME-redirect (best-effort
|
||||
# via getent/dscl), or default. Env-var priority: UNSLOTH_STUDIO_HOME wins
|
||||
# over STUDIO_HOME (the more specific signal beats the generic alias).
|
||||
_resolve_studio_destinations() {
|
||||
_override_var=""
|
||||
_override="${UNSLOTH_STUDIO_HOME:-}"
|
||||
if [ -n "$_override" ]; then
|
||||
_override_var="UNSLOTH_STUDIO_HOME"
|
||||
else
|
||||
_override="${STUDIO_HOME:-}"
|
||||
[ -n "$_override" ] && _override_var="STUDIO_HOME"
|
||||
fi
|
||||
# Strip surrounding whitespace so " " is treated as unset (matches the
|
||||
# Python resolvers' .strip()), preventing install/runtime layout drift.
|
||||
_override=$(printf '%s' "$_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')
|
||||
# Tilde expansion: env vars are not subject to it when quoted on assignment.
|
||||
case "$_override" in
|
||||
"~") _override="$HOME" ;;
|
||||
"~/"*) _override="$HOME/${_override#'~/'}" ;;
|
||||
esac
|
||||
if [ -n "$_override" ]; then
|
||||
mkdir -p -- "$_override" 2>/dev/null || { echo "ERROR: $_override_var=$_override cannot be created." >&2; exit 1; }
|
||||
[ -w "$_override" ] || { echo "ERROR: $_override_var=$_override is not writable." >&2; exit 1; }
|
||||
STUDIO_HOME="$(CDPATH= cd -P -- "$_override" && pwd -P)" || exit 1
|
||||
DATA_DIR="$STUDIO_HOME/share"
|
||||
_LOCAL_BIN="$STUDIO_HOME/bin"
|
||||
_STUDIO_HOME_REDIRECT=env
|
||||
substep "custom $_override_var=$STUDIO_HOME"
|
||||
return 0
|
||||
fi
|
||||
_default_home=""
|
||||
if command -v getent >/dev/null 2>&1; then
|
||||
_default_home=$(getent passwd "${USER:-$(whoami)}" 2>/dev/null | cut -d: -f6)
|
||||
elif [ "$(uname)" = "Darwin" ] && command -v dscl >/dev/null 2>&1; then
|
||||
_default_home=$(dscl . -read "/Users/${USER:-$(whoami)}" NFSHomeDirectory 2>/dev/null | awk '{print $2}')
|
||||
fi
|
||||
# Canonicalize both sides so a trailing slash on $HOME (or symlink mismatch
|
||||
# with passwd-DB output) doesn't misfire the redirection branch.
|
||||
_home_canon="$HOME"
|
||||
if [ -d "$_home_canon" ]; then
|
||||
_home_canon=$(CDPATH= cd -P -- "$_home_canon" 2>/dev/null && pwd -P) || _home_canon="$HOME"
|
||||
fi
|
||||
_default_home_canon="$_default_home"
|
||||
if [ -n "$_default_home_canon" ] && [ -d "$_default_home_canon" ]; then
|
||||
_default_home_canon=$(CDPATH= cd -P -- "$_default_home_canon" 2>/dev/null && pwd -P) || _default_home_canon="$_default_home"
|
||||
fi
|
||||
if [ -n "$_default_home_canon" ] && [ "$_home_canon" != "$_default_home_canon" ]; then
|
||||
STUDIO_HOME="$HOME/.unsloth/studio"
|
||||
DATA_DIR="$HOME/.local/share/unsloth"
|
||||
_LOCAL_BIN="$HOME/.local/bin"
|
||||
_STUDIO_HOME_REDIRECT=home
|
||||
substep "HOME redirected ($HOME); install follows \$HOME"
|
||||
return 0
|
||||
fi
|
||||
STUDIO_HOME="$HOME/.unsloth/studio"
|
||||
DATA_DIR="$HOME/.local/share/unsloth"
|
||||
_LOCAL_BIN="$HOME/.local/bin"
|
||||
_STUDIO_HOME_REDIRECT=default
|
||||
}
|
||||
_resolve_studio_destinations
|
||||
VENV_DIR="$STUDIO_HOME/unsloth_studio"
|
||||
_VENV_ROLLBACK_DIR=""
|
||||
_VENV_ROLLBACK_TARGET="$VENV_DIR"
|
||||
|
|
@ -383,23 +499,65 @@ create_studio_shortcuts() {
|
|||
_css_exe_dir=$(cd "$(dirname "$_css_exe")" && pwd)
|
||||
_css_exe="$_css_exe_dir/$(basename "$_css_exe")"
|
||||
|
||||
_css_data_dir="$HOME/.local/share/unsloth"
|
||||
_css_data_dir="$DATA_DIR"
|
||||
_css_launcher="$_css_data_dir/launch-studio.sh"
|
||||
_css_icon_png="$_css_data_dir/unsloth-studio.png"
|
||||
_css_gem_png="$_css_data_dir/unsloth-gem.png"
|
||||
|
||||
mkdir -p "$_css_data_dir"
|
||||
|
||||
# Same-install discriminator: per-install opaque id written once at install
|
||||
# time and read by both this launcher and the backend (/api/health). Replaces
|
||||
# the older sha256(canonical $STUDIO_HOME) scheme to (a) avoid leaking the
|
||||
# install path on -H 0.0.0.0 deployments and (b) sidestep launcher/backend
|
||||
# canonicalization drift (cd -P vs Path.resolve() symlink/junction handling).
|
||||
# Lives at $STUDIO_HOME/share/ (not $DATA_DIR) so the backend can find it
|
||||
# via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" regardless of
|
||||
# mode (in env-mode $STUDIO_HOME/share == $DATA_DIR; in default mode they
|
||||
# diverge but the backend only knows the studio_root). 32 bytes of urandom
|
||||
# -> 64 hex chars, byte-compatible with the prior digest so launcher
|
||||
# placeholder, _check_health, and tests stay length-agnostic.
|
||||
_css_id_dir="$STUDIO_HOME/share"
|
||||
mkdir -p "$_css_id_dir"
|
||||
_css_id_file="$_css_id_dir/studio_install_id"
|
||||
if [ ! -s "$_css_id_file" ]; then
|
||||
if [ -r /dev/urandom ]; then
|
||||
_css_new_id=$(od -An -N32 -tx1 /dev/urandom 2>/dev/null | tr -d ' \n')
|
||||
fi
|
||||
if [ -z "${_css_new_id:-}" ] && command -v python3 >/dev/null 2>&1; then
|
||||
_css_new_id=$(python3 -c 'import secrets; print(secrets.token_hex(32))' 2>/dev/null)
|
||||
fi
|
||||
if [ -z "${_css_new_id:-}" ]; then
|
||||
echo "[WARN] Cannot create launcher: no entropy source for studio_install_id" >&2
|
||||
return 1
|
||||
fi
|
||||
# Atomic write so a partial install can't leave a half-written id.
|
||||
_css_id_tmp="$_css_id_file.$$.tmp"
|
||||
printf '%s' "$_css_new_id" > "$_css_id_tmp" \
|
||||
&& mv "$_css_id_tmp" "$_css_id_file"
|
||||
chmod 600 "$_css_id_file" 2>/dev/null || true
|
||||
unset _css_new_id _css_id_tmp
|
||||
fi
|
||||
_css_studio_root_id=$(cat "$_css_id_file" 2>/dev/null)
|
||||
if [ -z "$_css_studio_root_id" ]; then
|
||||
echo "[WARN] Cannot create launcher: failed to read $_css_id_file" >&2
|
||||
return 1
|
||||
fi
|
||||
_css_is_env_mode=false
|
||||
[ "$_STUDIO_HOME_REDIRECT" = "env" ] && _css_is_env_mode=true
|
||||
|
||||
# ── Write launcher script ──
|
||||
# The launcher is Bash (not POSIX sh).
|
||||
# We write it with a placeholder and substitute the exe path via sed.
|
||||
# Single-quoted heredoc; @@DATA_DIR@@, @@STUDIO_ROOT_ID@@, and
|
||||
# @@INSTALLED_IS_ENV_MODE@@ are substituted via sed below.
|
||||
cat > "$_css_launcher" << 'LAUNCHER_EOF'
|
||||
#!/usr/bin/env bash
|
||||
# Unsloth Studio Launcher
|
||||
# Auto-generated by install.sh -- do not edit manually.
|
||||
set -euo pipefail
|
||||
|
||||
DATA_DIR="$HOME/.local/share/unsloth"
|
||||
DATA_DIR='@@DATA_DIR@@'
|
||||
_EXPECTED_STUDIO_ROOT_ID='@@STUDIO_ROOT_ID@@'
|
||||
_INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@'
|
||||
|
||||
# Read exe path from config written at install time.
|
||||
# Sourcing is safe: the config file is written by install.sh, not user input.
|
||||
|
|
@ -414,9 +572,25 @@ fi
|
|||
BASE_PORT=8888
|
||||
MAX_PORT_OFFSET=20
|
||||
TIMEOUT_SEC=60
|
||||
POLL_INTERVAL_SEC=1
|
||||
POLL_INTERVAL_SEC=0.25
|
||||
LOG_FILE="$DATA_DIR/studio.log"
|
||||
# why: in env-override mode multiple installs share an OS user; namespace the
|
||||
# lock and remember our own healthy port so we never attach to an unrelated
|
||||
# Studio listening on the global 8888..8908 range.
|
||||
LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock"
|
||||
PORT_FILE=""
|
||||
# why: gate on the install-time mode (baked above) instead of the runtime env
|
||||
# var; sourcing a custom-root studio.conf in shell must not flip a default-mode
|
||||
# launcher into env-mode behavior with stale state.
|
||||
if [ "$_INSTALLED_IS_ENV_MODE" = "true" ]; then
|
||||
if command -v cksum >/dev/null 2>&1; then
|
||||
_LOCK_KEY=$(printf '%s' "$DATA_DIR" | cksum | awk '{print $1}')
|
||||
else
|
||||
_LOCK_KEY=""
|
||||
fi
|
||||
[ -n "$_LOCK_KEY" ] && LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u)-${_LOCK_KEY}.lock"
|
||||
PORT_FILE="$DATA_DIR/studio.port"
|
||||
fi
|
||||
|
||||
# ── HTTP GET helper (supports curl and wget) ──
|
||||
_http_get() {
|
||||
|
|
@ -435,10 +609,20 @@ _check_health() {
|
|||
_port=$1
|
||||
_resp=$(_http_get "http://127.0.0.1:$_port/api/health") || return 1
|
||||
case "$_resp" in
|
||||
*'"status"'*'"healthy"'*'"service"'*'"Unsloth UI Backend"'*) return 0 ;;
|
||||
*'"service"'*'"Unsloth UI Backend"'*'"status"'*'"healthy"'*) return 0 ;;
|
||||
*'"status"'*'"healthy"'*'"service"'*'"Unsloth UI Backend"'*) ;;
|
||||
*'"service"'*'"Unsloth UI Backend"'*'"status"'*'"healthy"'*) ;;
|
||||
*) return 1 ;;
|
||||
esac
|
||||
return 1
|
||||
# why: verify the backend belongs to THIS install. Baked hex digest avoids
|
||||
# JSON-escape mismatches on paths with `\`/`"` and avoids leaking the raw
|
||||
# install path to unauthenticated callers.
|
||||
if [ -n "$_EXPECTED_STUDIO_ROOT_ID" ]; then
|
||||
case "$_resp" in
|
||||
*"\"studio_root_id\":\"$_EXPECTED_STUDIO_ROOT_ID\""*|*"\"studio_root_id\": \"$_EXPECTED_STUDIO_ROOT_ID\""*) return 0 ;;
|
||||
*) return 1 ;;
|
||||
esac
|
||||
fi
|
||||
return 0
|
||||
}
|
||||
|
||||
# ── Port scanning ──
|
||||
|
|
@ -461,6 +645,25 @@ _candidate_ports() {
|
|||
}
|
||||
|
||||
_find_healthy_port() {
|
||||
if [ -n "$PORT_FILE" ] && [ -f "$PORT_FILE" ]; then
|
||||
# why: env-mode installs only attach to a port we previously launched
|
||||
# ourselves; never to a sibling Studio that happens to be healthy.
|
||||
_p=$(cat "$PORT_FILE" 2>/dev/null || true)
|
||||
case "$_p" in
|
||||
''|*[!0-9]*) ;;
|
||||
*)
|
||||
if _check_health "$_p"; then
|
||||
echo "$_p"
|
||||
return 0
|
||||
fi
|
||||
rm -f "$PORT_FILE"
|
||||
;;
|
||||
esac
|
||||
return 1
|
||||
fi
|
||||
if [ -n "$PORT_FILE" ]; then
|
||||
return 1
|
||||
fi
|
||||
for _p in $(_candidate_ports | sort -un); do
|
||||
if _check_health "$_p"; then
|
||||
echo "$_p"
|
||||
|
|
@ -524,9 +727,65 @@ _spawn_terminal() {
|
|||
_cmd="$1"
|
||||
_os=$(uname)
|
||||
if [ "$_os" = "Darwin" ]; then
|
||||
# Escape backslashes and double-quotes for AppleScript string
|
||||
_cmd_escaped=$(printf '%s' "$_cmd" | sed 's/\\/\\\\/g; s/"/\\"/g')
|
||||
osascript -e "tell application \"Terminal\" to do script \"$_cmd_escaped\"" >/dev/null 2>&1 && return 0
|
||||
# AppleEvents are TCC-denied from unsigned .app bundles; spawn
|
||||
# Terminal via a .command file + Launch Services instead. Server
|
||||
# is nohup'd so warm relaunches hit the fast-path; watcher + trap
|
||||
# in the .command couple Terminal close <-> server shutdown.
|
||||
# `exec` keeps the recorded PID equal to the studio process so
|
||||
# signals reach studio directly rather than a wrapper shell.
|
||||
nohup sh -c "exec $_cmd" >> "$LOG_FILE" 2>&1 &
|
||||
_server_pid=$!
|
||||
_pid_file="$DATA_DIR/studio-$_launch_port.pid"
|
||||
printf '%d\n' "$_server_pid" > "$_pid_file" 2>/dev/null || true
|
||||
|
||||
_cmd_file="$DATA_DIR/launch-terminal.command"
|
||||
_logfile_q=$(printf '%s' "$LOG_FILE" | sed "s/'/'\\\\''/g")
|
||||
_pidfile_q=$(printf '%s' "$_pid_file" | sed "s/'/'\\\\''/g")
|
||||
if {
|
||||
{
|
||||
printf '#!/bin/bash\n'
|
||||
printf "SERVER_PID=%s\n" "$_server_pid"
|
||||
printf "PID_FILE='%s'\n" "$_pidfile_q"
|
||||
# Wait up to 12s for graceful shutdown before SIGKILL.
|
||||
printf 'shutdown_studio() {\n'
|
||||
printf ' kill -TERM "$SERVER_PID" 2>/dev/null\n'
|
||||
printf ' _i=0\n'
|
||||
printf ' while kill -0 "$SERVER_PID" 2>/dev/null && [ "$_i" -lt 24 ]; do\n'
|
||||
printf ' sleep 0.5\n'
|
||||
printf ' _i=$((_i + 1))\n'
|
||||
printf ' done\n'
|
||||
printf ' kill -0 "$SERVER_PID" 2>/dev/null && kill -KILL "$SERVER_PID" 2>/dev/null\n'
|
||||
printf ' rm -f "$PID_FILE" 2>/dev/null\n'
|
||||
printf '}\n'
|
||||
printf "tail -n 100 -F '%s' &\n" "$_logfile_q"
|
||||
printf 'TAIL_PID=$!\n'
|
||||
# Server gone -> kill tail so bash exits cleanly.
|
||||
printf '(\n'
|
||||
printf ' while kill -0 "$SERVER_PID" 2>/dev/null; do sleep 1; done\n'
|
||||
printf ' kill "$TAIL_PID" 2>/dev/null\n'
|
||||
printf ') &\n'
|
||||
printf 'WATCHER_PID=$!\n'
|
||||
printf "trap 'shutdown_studio; kill \"\$WATCHER_PID\" \"\$TAIL_PID\" 2>/dev/null; exit' HUP INT TERM\n"
|
||||
printf "trap 'rm -f \"\$PID_FILE\" 2>/dev/null' EXIT\n"
|
||||
printf 'wait "$TAIL_PID" 2>/dev/null\n'
|
||||
} > "$_cmd_file" 2>/dev/null \
|
||||
&& chmod +x "$_cmd_file" 2>/dev/null \
|
||||
&& open -a Terminal "$_cmd_file" 2>/dev/null
|
||||
}; then
|
||||
# Foreground Terminal (Launch Services spawns us backgrounded).
|
||||
osascript -e 'tell application "Terminal" to activate' >/dev/null 2>&1 || true
|
||||
return 0
|
||||
fi
|
||||
# .command/open failed: kill orphan, fall through to generic fallback.
|
||||
kill -TERM "$_server_pid" 2>/dev/null || true
|
||||
_i=0
|
||||
while kill -0 "$_server_pid" 2>/dev/null && [ "$_i" -lt 6 ]; do
|
||||
sleep 0.5
|
||||
_i=$((_i + 1))
|
||||
done
|
||||
kill -0 "$_server_pid" 2>/dev/null && kill -KILL "$_server_pid" 2>/dev/null || true
|
||||
rm -f "$_pid_file" 2>/dev/null || true
|
||||
echo "[WARN] Could not open Terminal; falling back to background launch" >&2
|
||||
else
|
||||
for _term in gnome-terminal konsole xfce4-terminal mate-terminal lxterminal xterm; do
|
||||
if command -v "$_term" >/dev/null 2>&1; then
|
||||
|
|
@ -611,6 +870,7 @@ if [ -t 1 ]; then
|
|||
_obwr_deadline=$(($(date +%s) + TIMEOUT_SEC))
|
||||
while [ "$(date +%s)" -lt "$_obwr_deadline" ]; do
|
||||
if _check_health "$_launch_port"; then
|
||||
[ -n "$PORT_FILE" ] && printf '%s\n' "$_launch_port" > "$PORT_FILE" 2>/dev/null || true
|
||||
_release_lock
|
||||
_open_browser "http://localhost:$_launch_port"
|
||||
exit 0
|
||||
|
|
@ -634,6 +894,7 @@ else
|
|||
_deadline=$(($(date +%s) + TIMEOUT_SEC))
|
||||
while [ "$(date +%s)" -lt "$_deadline" ]; do
|
||||
if _check_health "$_launch_port"; then
|
||||
[ -n "$PORT_FILE" ] && printf '%s\n' "$_launch_port" > "$PORT_FILE" 2>/dev/null || true
|
||||
_open_browser "http://localhost:$_launch_port"
|
||||
exit 0
|
||||
fi
|
||||
|
|
@ -646,13 +907,62 @@ else
|
|||
fi
|
||||
LAUNCHER_EOF
|
||||
|
||||
# why: bake non-user-controlled placeholders FIRST so a literal
|
||||
# `@@STUDIO_ROOT_ID@@` inside $DATA_DIR cannot be rewritten below.
|
||||
sed -e "s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g" \
|
||||
-e "s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g" \
|
||||
"$_css_launcher" > "$_css_launcher.tmp" \
|
||||
&& mv "$_css_launcher.tmp" "$_css_launcher"
|
||||
|
||||
# Env-mode bakes an absolute DATA_DIR (root fixed at install time);
|
||||
# default / HOME-redirect keeps the literal $HOME/.local/share/unsloth
|
||||
# so behavior is byte-identical to pre-override.
|
||||
if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then
|
||||
# Two-stage escape: (1) `'` -> `'\''` for shell single-quote embedding,
|
||||
# (2) backslash/&/| escape so the value survives the s|...|VALUE| sed
|
||||
# below. Verified end-to-end with apostrophes, spaces, &, |, $.
|
||||
_sq_escaped=$(printf '%s' "$DATA_DIR" | sed "s/'/'\\\\''/g")
|
||||
_sed_safe=$(printf '%s' "$_sq_escaped" | sed 's/[\\&|]/\\&/g')
|
||||
sed "s|@@DATA_DIR@@|$_sed_safe|g" "$_css_launcher" > "$_css_launcher.tmp" \
|
||||
&& mv "$_css_launcher.tmp" "$_css_launcher"
|
||||
else
|
||||
sed "s|DATA_DIR='@@DATA_DIR@@'|DATA_DIR=\"\$HOME/.local/share/unsloth\"|" \
|
||||
"$_css_launcher" > "$_css_launcher.tmp" \
|
||||
&& mv "$_css_launcher.tmp" "$_css_launcher"
|
||||
fi
|
||||
|
||||
chmod +x "$_css_launcher"
|
||||
|
||||
# Write the exe path to a separate conf file sourced by the launcher.
|
||||
# Using single-quote wrapping with the standard '\'' escape for any
|
||||
# embedded apostrophes. This avoids all sed metacharacter issues.
|
||||
# studio.conf: exe path + (env-mode only) persisted env vars so fresh
|
||||
# shells launch the right install without re-exporting.
|
||||
_css_quoted_exe=$(printf '%s' "$_css_exe" | sed "s/'/'\\\\''/g")
|
||||
printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" > "$_css_data_dir/studio.conf"
|
||||
{
|
||||
printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'"
|
||||
if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then
|
||||
# When an override resolves to the legacy default, llama.cpp
|
||||
# still lives at ~/.unsloth/llama.cpp (one shared build).
|
||||
# Canonicalize the legacy side so a symlinked $HOME doesn't
|
||||
# break the comparison.
|
||||
_css_legacy_studio="$HOME/.unsloth/studio"
|
||||
if [ -d "$_css_legacy_studio" ]; then
|
||||
_css_legacy_studio=$(CDPATH= cd -P -- "$_css_legacy_studio" 2>/dev/null && pwd -P) \
|
||||
|| _css_legacy_studio="$HOME/.unsloth/studio"
|
||||
fi
|
||||
if [ "$STUDIO_HOME" = "$_css_legacy_studio" ]; then
|
||||
_css_llama_path="$HOME/.unsloth/llama.cpp"
|
||||
else
|
||||
_css_llama_path="$STUDIO_HOME/llama.cpp"
|
||||
fi
|
||||
_css_quoted_home=$(printf '%s' "$STUDIO_HOME" | sed "s/'/'\\\\''/g")
|
||||
_css_quoted_llama=$(printf '%s' "$_css_llama_path" | sed "s/'/'\\\\''/g")
|
||||
printf '%s\n' "export UNSLOTH_STUDIO_HOME='$_css_quoted_home'"
|
||||
# UNSLOTH_LLAMA_CPP_PATH is a pre-existing user-controlled
|
||||
# llama.cpp dir override; only default it if unset.
|
||||
printf '%s\n' 'if [ -z "${UNSLOTH_LLAMA_CPP_PATH:-}" ]; then'
|
||||
printf '%s\n' " export UNSLOTH_LLAMA_CPP_PATH='$_css_quoted_llama'"
|
||||
printf '%s\n' 'fi'
|
||||
fi
|
||||
} > "$_css_data_dir/studio.conf"
|
||||
|
||||
# ── Icon: try bundled, then download ──
|
||||
# rounded-512.png used for both Linux and macOS icons
|
||||
|
|
@ -698,6 +1008,14 @@ LAUNCHER_EOF
|
|||
fi
|
||||
|
||||
# ── Platform-specific shortcuts ──
|
||||
# Env-mode installs are workspace-scoped: skip persistent desktop /
|
||||
# Start-Menu / dock launchers that may point at a deleted workspace.
|
||||
# Runtime launcher + studio.conf + icon are still written above.
|
||||
if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then
|
||||
substep "wrote launcher at $_css_launcher (persistent shortcuts skipped in env-override mode)"
|
||||
return 0
|
||||
fi
|
||||
|
||||
_css_created=0
|
||||
|
||||
if [ "$_css_os" = "linux" ]; then
|
||||
|
|
@ -743,6 +1061,17 @@ DESKTOP_EOF
|
|||
_css_contents="$_css_app/Contents"
|
||||
_css_macos_dir="$_css_contents/MacOS"
|
||||
_css_res_dir="$_css_contents/Resources"
|
||||
# Recreate bundle if root or any subpath is a symlink (mkdir -p follows them).
|
||||
if [ -L "$_css_app" ] || [ -L "$_css_contents" ] \
|
||||
|| [ -L "$_css_macos_dir" ] || [ -L "$_css_res_dir" ]; then
|
||||
rm -rf "$_css_app" 2>/dev/null || {
|
||||
echo "[ERROR] $_css_app contains a symlinked bundle path; remove manually and re-run install" >&2
|
||||
return 1
|
||||
}
|
||||
elif [ -e "$_css_app" ] && [ ! -d "$_css_app" ]; then
|
||||
echo "[ERROR] $_css_app exists but is not a directory; remove manually and re-run install" >&2
|
||||
return 1
|
||||
fi
|
||||
mkdir -p "$_css_macos_dir" "$_css_res_dir"
|
||||
|
||||
# Info.plist
|
||||
|
|
@ -775,11 +1104,18 @@ DESKTOP_EOF
|
|||
</plist>
|
||||
PLIST_EOF
|
||||
|
||||
# Executable stub
|
||||
cat > "$_css_macos_dir/launch-studio" << STUB_EOF
|
||||
# Executable stub: same single-quoted-heredoc + sed-substitute
|
||||
# pattern as launch-studio.sh so $-vars in $_css_data_dir don't
|
||||
# expand at .app launch time.
|
||||
_css_sq_dir=$(printf '%s' "$_css_data_dir" | sed "s/'/'\\\\''/g")
|
||||
_css_sed_dir=$(printf '%s' "$_css_sq_dir" | sed 's/[\\&|]/\\&/g')
|
||||
cat > "$_css_macos_dir/launch-studio" << 'STUB_EOF'
|
||||
#!/bin/sh
|
||||
exec "$HOME/.local/share/unsloth/launch-studio.sh" "\$@"
|
||||
exec '@@DATA_DIR@@/launch-studio.sh' "$@"
|
||||
STUB_EOF
|
||||
sed "s|@@DATA_DIR@@|$_css_sed_dir|g" "$_css_macos_dir/launch-studio" \
|
||||
> "$_css_macos_dir/launch-studio.tmp" \
|
||||
&& mv "$_css_macos_dir/launch-studio.tmp" "$_css_macos_dir/launch-studio"
|
||||
chmod +x "$_css_macos_dir/launch-studio"
|
||||
|
||||
# Build AppIcon.icns from unsloth-gem.png (2240x2240)
|
||||
|
|
@ -1079,11 +1415,28 @@ mkdir -p "$STUDIO_HOME"
|
|||
_MIGRATED=false
|
||||
|
||||
if [ -x "$VENV_DIR/bin/python" ]; then
|
||||
# why: matching guard to the .venv branch below -- in env-mode
|
||||
# $STUDIO_HOME is a user-chosen workspace, so refuse to nuke an
|
||||
# existing $STUDIO_HOME/unsloth_studio that lacks Studio sentinels.
|
||||
# Accept the in-VENV ownership marker so partial-install retries are
|
||||
# not blocked. Sentinels must be regular files: -f follows symlinks
|
||||
# to files (the legitimate ln -s shim shape) but rejects directories
|
||||
# and broken/dir-targeted symlinks.
|
||||
if [ "$_STUDIO_HOME_REDIRECT" = "env" ] \
|
||||
&& [ ! -f "$VENV_DIR/.unsloth-studio-owned" ] \
|
||||
&& [ ! -f "$STUDIO_HOME/share/studio.conf" ] \
|
||||
&& [ ! -f "$STUDIO_HOME/bin/unsloth" ]; then
|
||||
echo "ERROR: $VENV_DIR already exists but does not look like an Unsloth Studio install." >&2
|
||||
echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2
|
||||
exit 1
|
||||
fi
|
||||
# New layout already exists — replace only after preserving rollback copy.
|
||||
substep "preserving existing environment for rollback..."
|
||||
_start_studio_venv_replacement "$VENV_DIR"
|
||||
elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then
|
||||
elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/python" ]; then
|
||||
# Old layout exists — validate before migrating.
|
||||
# Skip in env-mode so we don't rm -rf an unrelated .venv at the
|
||||
# workspace root (e.g. user's existing project Python venv).
|
||||
# In no-torch mode, a missing torch package is expected; validate Python only.
|
||||
substep "found legacy Studio environment, validating..."
|
||||
_legacy_ok=false
|
||||
|
|
@ -1132,6 +1485,13 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then
|
|||
run_install_cmd "create venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
|
||||
fi
|
||||
|
||||
# Mark the freshly-created venv as Studio-owned so a partial install can be
|
||||
# repaired by re-running install.sh; the env-mode deletion guard above accepts
|
||||
# this marker as the primary sentinel.
|
||||
if [ -x "$VENV_DIR/bin/python" ]; then
|
||||
: > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true
|
||||
fi
|
||||
|
||||
# Guard against Python 3.13.8 torch import bug on Apple Silicon
|
||||
# (skip when the user explicitly chose a version via --python)
|
||||
if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
|
||||
|
|
@ -1143,6 +1503,9 @@ if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
|
|||
rm -rf "$VENV_DIR"
|
||||
PYTHON_VERSION="3.12"
|
||||
run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
|
||||
if [ -x "$VENV_DIR/bin/python" ]; then
|
||||
: > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
|
|
@ -1486,7 +1849,7 @@ if [ "$_MIGRATED" = true ]; then
|
|||
# to prevent transitive torch resolution.
|
||||
run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
|
||||
--reinstall-package unsloth --reinstall-package unsloth-zoo \
|
||||
"unsloth>=2026.5.1" unsloth-zoo
|
||||
"unsloth>=2026.5.2" unsloth-zoo
|
||||
_NO_TORCH_RT="$(_find_no_torch_runtime)"
|
||||
if [ -n "$_NO_TORCH_RT" ]; then
|
||||
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
|
||||
|
|
@ -1494,7 +1857,7 @@ if [ "$_MIGRATED" = true ]; then
|
|||
else
|
||||
run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
|
||||
--reinstall-package unsloth --reinstall-package unsloth-zoo \
|
||||
"unsloth>=2026.5.1" unsloth-zoo
|
||||
"unsloth>=2026.5.2" unsloth-zoo
|
||||
fi
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
substep "overlaying local repo (editable)..."
|
||||
|
|
@ -1662,7 +2025,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
|
||||
run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
|
||||
--upgrade-package unsloth --upgrade-package unsloth-zoo \
|
||||
"unsloth>=2026.5.1" unsloth-zoo
|
||||
"unsloth>=2026.5.2" unsloth-zoo
|
||||
_NO_TORCH_RT="$(_find_no_torch_runtime)"
|
||||
if [ -n "$_NO_TORCH_RT" ]; then
|
||||
run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
|
||||
|
|
@ -1677,7 +2040,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
fi
|
||||
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
|
||||
--upgrade-package unsloth "unsloth>=2026.5.1" unsloth-zoo
|
||||
--upgrade-package unsloth "unsloth>=2026.5.2" unsloth-zoo
|
||||
substep "overlaying local repo (editable)..."
|
||||
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
|
||||
substep "overlaying unsloth-zoo from git main..."
|
||||
|
|
@ -1709,7 +2072,7 @@ else
|
|||
tauri_log "STEP" "Installing Unsloth"
|
||||
substep "installing unsloth (this may take a few minutes)..."
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.1" --torch-backend=auto
|
||||
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.2" --torch-backend=auto
|
||||
substep "overlaying local repo (editable)..."
|
||||
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
|
||||
substep "overlaying unsloth-zoo from git main..."
|
||||
|
|
@ -1721,6 +2084,12 @@ else
|
|||
fi
|
||||
fi
|
||||
|
||||
# ── Install mlx-vlm on Apple Silicon (optional, for VLM training) ──
|
||||
if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
|
||||
substep "installing mlx-vlm (VLM training support)..."
|
||||
run_install_cmd "install mlx-vlm" uv pip install --python "$_VENV_PY" mlx-vlm
|
||||
fi
|
||||
|
||||
# ── Run studio setup ──
|
||||
tauri_log "STEP" "Running Studio setup"
|
||||
# When --local, use the repo's own setup.sh directly.
|
||||
|
|
@ -1768,7 +2137,17 @@ _SKIP_FRONTEND=0
|
|||
if [ "$TAURI_MODE" = true ]; then
|
||||
_SKIP_FRONTEND=1
|
||||
fi
|
||||
# Prepend UNSLOTH_STUDIO_HOME=$STUDIO_HOME to "$@" for env-override installs
|
||||
# without word-splitting on whitespace paths.
|
||||
_run_setup_with_studio_home() {
|
||||
if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then
|
||||
UNSLOTH_STUDIO_HOME="$STUDIO_HOME" "$@"
|
||||
else
|
||||
"$@"
|
||||
fi
|
||||
}
|
||||
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
||||
_run_setup_with_studio_home env \
|
||||
SKIP_STUDIO_BASE="$_SKIP_BASE" \
|
||||
SKIP_STUDIO_FRONTEND="$_SKIP_FRONTEND" \
|
||||
STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
|
||||
|
|
@ -1782,6 +2161,7 @@ else
|
|||
# the same session) does not silently flip a normal install onto the
|
||||
# local-dev path in setup.sh and install_python_stack.py. Mirrors the
|
||||
# reset already done in install.ps1 for PowerShell.
|
||||
_run_setup_with_studio_home env \
|
||||
SKIP_STUDIO_BASE="$_SKIP_BASE" \
|
||||
SKIP_STUDIO_FRONTEND="$_SKIP_FRONTEND" \
|
||||
STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
|
||||
|
|
@ -1791,36 +2171,53 @@ else
|
|||
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
|
||||
fi
|
||||
|
||||
# ── Make 'unsloth' available globally via ~/.local/bin ──
|
||||
mkdir -p "$HOME/.local/bin"
|
||||
ln -sf "$VENV_DIR/bin/unsloth" "$HOME/.local/bin/unsloth"
|
||||
# ── Make 'unsloth' available via $_LOCAL_BIN (resolved earlier) ──
|
||||
# Env-mode: $_LOCAL_BIN is $STUDIO_HOME/bin; skip shell-rc PATH append so we
|
||||
# don't pollute the user's profile with a workspace-scoped path.
|
||||
mkdir -p "$_LOCAL_BIN"
|
||||
# ln -sf into an existing dir creates link inside it. Refuse to delete a
|
||||
# real directory at the shim path -- that could destroy unrelated user data.
|
||||
_shim_path="$_LOCAL_BIN/unsloth"
|
||||
if [ -d "$_shim_path" ] && [ ! -L "$_shim_path" ]; then
|
||||
echo "ERROR: $_shim_path is a directory; refusing to delete it." >&2
|
||||
echo " Move or remove it manually, then re-run the installer." >&2
|
||||
exit 1
|
||||
fi
|
||||
# why: -sfn is atomic and -n prevents descent into a symlink-to-directory at
|
||||
# the shim path (the directory guard above already rejects a real directory).
|
||||
ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path"
|
||||
|
||||
_LOCAL_BIN="$HOME/.local/bin"
|
||||
case ":$PATH:" in
|
||||
*":$_LOCAL_BIN:"*) ;; # already on PATH
|
||||
*)
|
||||
_SHELL_PROFILE=""
|
||||
if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then
|
||||
_SHELL_PROFILE="$HOME/.zshrc"
|
||||
elif [ -f "$HOME/.bashrc" ]; then
|
||||
_SHELL_PROFILE="$HOME/.bashrc"
|
||||
elif [ -f "$HOME/.profile" ]; then
|
||||
_SHELL_PROFILE="$HOME/.profile"
|
||||
fi
|
||||
|
||||
if [ -n "$_SHELL_PROFILE" ]; then
|
||||
if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then
|
||||
echo '' >> "$_SHELL_PROFILE"
|
||||
echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE"
|
||||
echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE"
|
||||
step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE"
|
||||
if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then
|
||||
export PATH="$_LOCAL_BIN:$PATH"
|
||||
step "path" "exported $_LOCAL_BIN for this session (no rc-file append in env-override mode)"
|
||||
else
|
||||
_SHELL_PROFILE=""
|
||||
if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then
|
||||
_SHELL_PROFILE="$HOME/.zshrc"
|
||||
elif [ -f "$HOME/.bashrc" ]; then
|
||||
_SHELL_PROFILE="$HOME/.bashrc"
|
||||
elif [ -f "$HOME/.profile" ]; then
|
||||
_SHELL_PROFILE="$HOME/.profile"
|
||||
fi
|
||||
if [ -n "$_SHELL_PROFILE" ]; then
|
||||
if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then
|
||||
echo '' >> "$_SHELL_PROFILE"
|
||||
echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE"
|
||||
echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE"
|
||||
step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE"
|
||||
fi
|
||||
fi
|
||||
export PATH="$_LOCAL_BIN:$PATH"
|
||||
fi
|
||||
export PATH="$_LOCAL_BIN:$PATH"
|
||||
;;
|
||||
esac
|
||||
|
||||
# Non-Tauri installs keep shortcuts even if setup reports failure.
|
||||
# create_studio_shortcuts gates persistent menu shortcuts on env-mode;
|
||||
# launcher + studio.conf + icon are always written.
|
||||
if [ "$TAURI_MODE" != true ]; then
|
||||
create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS"
|
||||
fi
|
||||
|
|
@ -1883,10 +2280,21 @@ if [ -t 1 ]; then
|
|||
esac
|
||||
else
|
||||
step "launch" "manual commands:"
|
||||
substep "unsloth studio -p 8888"
|
||||
substep "or activate env first:"
|
||||
substep "source ${VENV_DIR}/bin/activate"
|
||||
substep "unsloth studio -p 8888"
|
||||
# Single-quote-escape so paths with spaces / apostrophes copy-paste cleanly.
|
||||
_li_shim_q="'$(printf '%s' "${_LOCAL_BIN}/unsloth" | sed "s/'/'\\\\''/g")'"
|
||||
_li_act_q="'$(printf '%s' "${VENV_DIR}/bin/activate" | sed "s/'/'\\\\''/g")'"
|
||||
if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then
|
||||
# Env-mode skips the rc PATH append, so print the absolute shim path.
|
||||
substep "$_li_shim_q studio -p 8888"
|
||||
substep "or activate env first:"
|
||||
substep "source $_li_act_q"
|
||||
substep "unsloth studio -p 8888"
|
||||
else
|
||||
substep "unsloth studio -p 8888"
|
||||
substep "or activate env first:"
|
||||
substep "source $_li_act_q"
|
||||
substep "unsloth studio -p 8888"
|
||||
fi
|
||||
substep "(add -H 0.0.0.0 to allow network / cloud access)"
|
||||
echo ""
|
||||
fi
|
||||
|
|
|
|||
1195
scripts/check_frontend_dep_removal.py
Normal file
1195
scripts/check_frontend_dep_removal.py
Normal file
File diff suppressed because it is too large
Load diff
292
scripts/check_new_install_scripts.py
Normal file
292
scripts/check_new_install_scripts.py
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""Diff two `package-lock.json` files and flag NEW install-script deps.
|
||||
|
||||
A package with `"hasInstallScript": true` runs `preinstall` / `install` /
|
||||
`postinstall` lifecycle hooks every time `npm ci` lays it down. Every
|
||||
npm supply-chain compromise of the last 18 months (Shai-Hulud,
|
||||
TanStack, axios-style, ArmorCode hijacks) leveraged exactly this lever:
|
||||
the attacker publishes a new malicious version of a dep we already
|
||||
trust, and the post-install hook runs the next time CI installs.
|
||||
|
||||
This scanner refuses to allow a newly-introduced install-script dep to
|
||||
land without a maintainer eyeball on the lifecycle script body.
|
||||
Existing install-script deps are NOT re-flagged -- if `node-gyp` has
|
||||
been in the lockfile since day one, it's not part of this PR's threat
|
||||
model. Only new entries are surfaced.
|
||||
|
||||
Supports lockfileVersion 1 (`dependencies` key, recursive), 2 and 3
|
||||
(flat `packages` key with `node_modules/<a>/node_modules/<b>` nesting
|
||||
for transitive entries). For each NEW install-script package we
|
||||
attempt a stdlib-only fetch of
|
||||
`https://registry.npmjs.org/<name>/<version>` to recover the actual
|
||||
postinstall command body. If the network is blocked we still emit the
|
||||
finding -- the lifecycle command body is informational, not
|
||||
load-bearing.
|
||||
|
||||
Exit codes
|
||||
==========
|
||||
0 no newly-added install-script deps
|
||||
1 one or more newly-added install-script deps; listed on stderr
|
||||
2 internal error (missing lockfile, malformed JSON, etc.)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
REGISTRY_BASE = "https://registry.npmjs.org/"
|
||||
REGISTRY_TIMEOUT_SECS = 5
|
||||
|
||||
CRITICAL = "CRITICAL"
|
||||
HIGH = "HIGH"
|
||||
|
||||
|
||||
class Finding:
|
||||
__slots__ = ("severity", "name", "version", "kind", "detail")
|
||||
|
||||
def __init__(
|
||||
self, severity: str, name: str, version: str, kind: str, detail: str
|
||||
) -> None:
|
||||
self.severity = severity
|
||||
self.name = name
|
||||
self.version = version
|
||||
self.kind = kind
|
||||
self.detail = detail
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
f" [{self.severity}] {self.name}@{self.version}\n"
|
||||
f" kind: {self.kind}\n"
|
||||
f" detail: {self.detail}"
|
||||
)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Lockfile parsing.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _strip_nm_prefix(key: str) -> str:
|
||||
"""Convert a v2/v3 `packages` key into a bare package name.
|
||||
|
||||
`node_modules/foo` -> `foo`; `node_modules/foo/node_modules/bar` ->
|
||||
`bar`. The empty key (`""`) is the project root and returns "".
|
||||
"""
|
||||
if not key:
|
||||
return ""
|
||||
# Use the LAST `node_modules/` segment so transitives map to their
|
||||
# leaf name, matching how npm install resolves a postinstall.
|
||||
marker = "node_modules/"
|
||||
idx = key.rfind(marker)
|
||||
if idx == -1:
|
||||
return key
|
||||
return key[idx + len(marker) :]
|
||||
|
||||
|
||||
def _collect_install_script_entries(lock: dict) -> dict[str, str]:
|
||||
"""Walk a parsed lockfile and return {package_name: version} for
|
||||
every entry with `hasInstallScript: true` (v2/v3) OR a
|
||||
non-empty `scripts.preinstall|install|postinstall` (v1).
|
||||
|
||||
The same package may appear at multiple versions in a single
|
||||
lockfile (de-duplicated copies under different parents); we key by
|
||||
`name@version` so we don't lose either copy. Returns a dict keyed
|
||||
by `name@version` -> the same string for convenience.
|
||||
"""
|
||||
seen: dict[str, str] = {}
|
||||
version = lock.get("lockfileVersion")
|
||||
|
||||
# v2 / v3: flat `packages` map.
|
||||
packages = lock.get("packages") or {}
|
||||
for key, entry in packages.items():
|
||||
if key == "" or not isinstance(entry, dict):
|
||||
continue
|
||||
if entry.get("link"):
|
||||
continue
|
||||
if not entry.get("hasInstallScript"):
|
||||
continue
|
||||
name = _strip_nm_prefix(key)
|
||||
if not name:
|
||||
continue
|
||||
ver = entry.get("version") or "<unversioned>"
|
||||
seen[f"{name}@{ver}"] = name
|
||||
|
||||
# v1 also embeds a `dependencies` tree; v2/v3 carry both for
|
||||
# backwards-compat but `packages` is canonical for them. For v1
|
||||
# there is no `hasInstallScript` flag, so look for a non-empty
|
||||
# `scripts.preinstall|install|postinstall` directly.
|
||||
def _walk_v1(deps: dict, depth: int = 0) -> None:
|
||||
if depth > 64 or not isinstance(deps, dict):
|
||||
return
|
||||
for name, entry in deps.items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
scripts = entry.get("scripts") or {}
|
||||
lifecycle = any(
|
||||
isinstance(scripts, dict) and scripts.get(hook)
|
||||
for hook in ("preinstall", "install", "postinstall")
|
||||
)
|
||||
# v1 also sets `requires` only on the parent, no flag, so
|
||||
# the lifecycle-script presence is the only signal.
|
||||
if lifecycle:
|
||||
ver = entry.get("version") or "<unversioned>"
|
||||
seen[f"{name}@{ver}"] = name
|
||||
_walk_v1(entry.get("dependencies"), depth = depth + 1)
|
||||
|
||||
if version == 1 or "dependencies" in lock:
|
||||
_walk_v1(lock.get("dependencies") or {})
|
||||
|
||||
return seen
|
||||
|
||||
|
||||
def _load_lockfile(path: Path) -> dict:
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"lockfile not found: {path}")
|
||||
try:
|
||||
return json.loads(path.read_text(encoding = "utf-8"))
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"{path}: not valid JSON: {exc}") from exc
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Registry lookup for the postinstall command body (best-effort).
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _fetch_registry_scripts(name: str, version: str) -> dict[str, str] | None:
|
||||
"""Return {hook: command} for any of preinstall / install /
|
||||
postinstall published in the registry metadata for this name@ver.
|
||||
|
||||
Returns None on any error (network blocked, 404, malformed JSON).
|
||||
Never raises; the caller treats absence as "could not enrich, emit
|
||||
finding anyway".
|
||||
"""
|
||||
safe_name = urllib.parse.quote(name, safe = "@/")
|
||||
url = f"{REGISTRY_BASE}{safe_name}/{urllib.parse.quote(version)}"
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout = REGISTRY_TIMEOUT_SECS) as resp:
|
||||
body = resp.read()
|
||||
except (urllib.error.URLError, OSError, ValueError, TimeoutError):
|
||||
return None
|
||||
try:
|
||||
meta = json.loads(body)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
scripts = meta.get("scripts") or {}
|
||||
if not isinstance(scripts, dict):
|
||||
return None
|
||||
keep = {}
|
||||
for hook in ("preinstall", "install", "postinstall"):
|
||||
cmd = scripts.get(hook)
|
||||
if isinstance(cmd, str) and cmd.strip():
|
||||
keep[hook] = cmd
|
||||
return keep or None
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Diff.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
|
||||
base = _collect_install_script_entries(base_lock)
|
||||
head = _collect_install_script_entries(head_lock)
|
||||
findings: list[Finding] = []
|
||||
for key in sorted(head):
|
||||
if key in base:
|
||||
continue # pre-existing install-script dep; not in scope
|
||||
name = head[key]
|
||||
# key is "name@version"; rsplit("@", 1) handles scoped names.
|
||||
version = (
|
||||
key[len(name) + 1 :] if key.startswith(name + "@") else "<unversioned>"
|
||||
)
|
||||
scripts = _fetch_registry_scripts(name, version)
|
||||
if scripts:
|
||||
detail = "; ".join(f"{h}={cmd!r}" for h, cmd in scripts.items())
|
||||
else:
|
||||
detail = (
|
||||
"newly added with hasInstallScript=true; registry "
|
||||
"metadata unreachable -- inspect the package's "
|
||||
"scripts.{preinstall,install,postinstall} manually"
|
||||
)
|
||||
findings.append(
|
||||
Finding(
|
||||
severity = CRITICAL,
|
||||
name = name,
|
||||
version = version,
|
||||
kind = "new-install-script",
|
||||
detail = detail,
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# CLI.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description = (
|
||||
"Diff two package-lock.json files and refuse any newly-"
|
||||
"added install-script dep."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--base",
|
||||
required = True,
|
||||
help = "Path to the BASE package-lock.json (e.g. main branch).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--head",
|
||||
required = True,
|
||||
help = "Path to the HEAD package-lock.json (this PR).",
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
try:
|
||||
base_lock = _load_lockfile(Path(args.base))
|
||||
head_lock = _load_lockfile(Path(args.head))
|
||||
except (FileNotFoundError, ValueError) as exc:
|
||||
print(f"[install-script-diff] ERROR: {exc}", file = sys.stderr)
|
||||
return 2
|
||||
|
||||
findings = diff_new_install_scripts(base_lock, head_lock)
|
||||
if not findings:
|
||||
print(
|
||||
"[install-script-diff] OK: no newly-added install-script "
|
||||
"dependencies between base and head",
|
||||
flush = True,
|
||||
)
|
||||
return 0
|
||||
|
||||
print(
|
||||
f"\n[install-script-diff] FAIL: {len(findings)} newly-added "
|
||||
f"install-script dependency(ies):\n",
|
||||
file = sys.stderr,
|
||||
)
|
||||
for f in findings:
|
||||
print(str(f), file = sys.stderr)
|
||||
print(file = sys.stderr)
|
||||
print(
|
||||
"[install-script-diff] Refusing to proceed. Every new "
|
||||
"install-script dep is a postinstall lifecycle hook that "
|
||||
"would run on the next `npm ci`. Review each finding above, "
|
||||
"confirm the maintainer + version, and re-run.",
|
||||
file = sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
1142
scripts/data/colab_apt_list.gpu.txt
Normal file
1142
scripts/data/colab_apt_list.gpu.txt
Normal file
File diff suppressed because it is too large
Load diff
9
scripts/data/colab_os_info.gpu.txt
Normal file
9
scripts/data/colab_os_info.gpu.txt
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
# Do not modify this file directly; it is generated by extract_colabx_testing_tarballs.sh via
|
||||
# $ (lsb_release -ds;python --version;) > os-info-gpu.txt
|
||||
# Be aware that this list does not necessarily reflect the current state of the
|
||||
# staging or production container, but rather the state as of the most recent
|
||||
# submitted CL where extract_colabx_testing_tarballs.sh was run.
|
||||
Ubuntu 22.04.5 LTS
|
||||
Python 3.12.13
|
||||
R version 4.5.3 (2026-03-11) -- "Reassured Reassurer"
|
||||
julia version 1.12.6
|
||||
731
scripts/data/colab_pip_freeze.gpu.txt
Normal file
731
scripts/data/colab_pip_freeze.gpu.txt
Normal file
|
|
@ -0,0 +1,731 @@
|
|||
# Do not modify this file directly; it is generated by extract_colabx_testing_tarballs.sh via
|
||||
# $ python3 -m pip freeze
|
||||
# Be aware that this list does not necessarily reflect the current state of the
|
||||
# staging or production container, but rather the state as of the most recent
|
||||
# submitted CL where extract_colabx_testing_tarballs.sh was run.
|
||||
absl-py==1.4.0
|
||||
accelerate==1.13.0
|
||||
access==1.1.10.post3
|
||||
affine==2.4.0
|
||||
aiofiles==24.1.0
|
||||
aiohappyeyeballs==2.6.1
|
||||
aiohttp==3.13.5
|
||||
aiosignal==1.4.0
|
||||
aiosqlite==0.22.1
|
||||
alabaster==1.0.0
|
||||
albucore==0.0.24
|
||||
albumentations==2.0.8
|
||||
ale-py==0.11.2
|
||||
alembic==1.18.4
|
||||
altair==5.5.0
|
||||
annotated-doc==0.0.4
|
||||
annotated-types==0.7.0
|
||||
antlr4-python3-runtime==4.9.3
|
||||
anyio==4.13.0
|
||||
anywidget==0.9.21
|
||||
apsw==3.53.0.0
|
||||
apswutils==0.1.2
|
||||
argon2-cffi==25.1.0
|
||||
argon2-cffi-bindings==25.1.0
|
||||
array_record==0.8.3
|
||||
arrow==1.4.0
|
||||
arviz==0.22.0
|
||||
astropy==7.2.0
|
||||
astropy-iers-data==0.2026.4.20.0.58.15
|
||||
astunparse==1.6.3
|
||||
atpublic==5.1
|
||||
attrs==26.1.0
|
||||
audioread==3.1.0
|
||||
Authlib==1.6.11
|
||||
autograd==1.8.0
|
||||
babel==2.18.0
|
||||
backcall==0.2.0
|
||||
beartype==0.22.9
|
||||
beautifulsoup4==4.13.5
|
||||
betterproto==2.0.0b6
|
||||
bigframes==2.39.0
|
||||
bigquery-magics==0.14.0
|
||||
bleach==6.3.0
|
||||
blinker==1.9.0
|
||||
blis==1.3.3
|
||||
blobfile==3.2.0
|
||||
blosc2==4.1.2
|
||||
bokeh==3.8.2
|
||||
Bottleneck==1.4.2
|
||||
bqplot==0.12.45
|
||||
branca==0.8.2
|
||||
brotli==1.2.0
|
||||
CacheControl==0.14.4
|
||||
cachetools==6.2.6
|
||||
catalogue==2.0.10
|
||||
certifi==2026.4.22
|
||||
cffi==2.0.0
|
||||
chardet==5.2.0
|
||||
charset-normalizer==3.4.7
|
||||
clarabel==0.11.1
|
||||
click==8.3.3
|
||||
click-plugins==1.1.1.2
|
||||
cligj==0.7.2
|
||||
cloudpathlib==0.23.0
|
||||
cloudpickle==3.1.2
|
||||
cmake==3.31.10
|
||||
cmdstanpy==1.3.0
|
||||
colorcet==3.1.0
|
||||
colorlover==0.3.0
|
||||
community==1.0.0b1
|
||||
confection==1.3.3
|
||||
cons==0.4.7
|
||||
contourpy==1.3.3
|
||||
cramjam==2.11.0
|
||||
cryptography==43.0.3
|
||||
cucim-cu12 @ https://pypi.nvidia.com/cucim-cu12/cucim_cu12-26.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
|
||||
cuda-bindings==12.9.4
|
||||
cuda-core==0.3.2
|
||||
cuda-pathfinder==1.5.3
|
||||
cuda-python==12.9.4
|
||||
cuda-toolkit==12.8.1
|
||||
cudf-cu12==26.2.1
|
||||
cudf-polars-cu12==26.2.1
|
||||
cufflinks==0.17.3
|
||||
cuml-cu12==26.2.0
|
||||
cupy-cuda12x==14.0.1
|
||||
curl_cffi==0.15.0
|
||||
cuvs-cu12 @ https://pypi.nvidia.com/cuvs-cu12/cuvs_cu12-26.2.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
|
||||
cvxopt==1.3.2
|
||||
cvxpy==1.6.7
|
||||
cycler==0.12.1
|
||||
cyipopt==1.5.0
|
||||
cymem==2.0.13
|
||||
Cython==3.0.12
|
||||
dask==2026.1.1
|
||||
dask-cuda==26.2.0
|
||||
dask-cudf-cu12==26.2.1
|
||||
dataproc-spark-connect==1.1.0
|
||||
datasets==4.0.0
|
||||
db-dtypes==1.5.1
|
||||
dbus-python==1.2.18
|
||||
debugpy==1.8.15
|
||||
decorator==4.4.2
|
||||
defusedxml==0.7.1
|
||||
deprecation==2.1.0
|
||||
diffusers==0.37.1
|
||||
dill==0.3.8
|
||||
distributed==2026.1.1
|
||||
distributed-ucxx-cu12==0.48.0
|
||||
distro==1.9.0
|
||||
dlib==19.24.6
|
||||
dm-tree==0.1.10
|
||||
docstring_parser==0.18.0
|
||||
docutils==0.21.2
|
||||
dopamine_rl==4.1.2
|
||||
duckdb==1.3.2
|
||||
earthengine-api==1.7.22
|
||||
easydict==1.13
|
||||
editdistance==0.8.1
|
||||
eerepr==0.1.2
|
||||
einops==0.8.2
|
||||
en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl#sha256=1932429db727d4bff3deed6b34cfc05df17794f4a52eeb26cf8928f7c1a0fb85
|
||||
entrypoints==0.4
|
||||
esda==2.9.0
|
||||
et_xmlfile==2.0.0
|
||||
etils==1.14.0
|
||||
etuples==0.3.10
|
||||
Farama-Notifications==0.0.4
|
||||
fastai==2.8.7
|
||||
fastapi==0.136.1
|
||||
fastcore==1.12.42
|
||||
fastdownload==0.0.7
|
||||
fastjsonschema==2.21.2
|
||||
fastlite==0.2.4
|
||||
fastprogress==1.1.5
|
||||
fasttransform==0.0.2
|
||||
ffmpy==1.0.0
|
||||
filelock==3.29.0
|
||||
fiona==1.10.1
|
||||
firebase-admin==6.9.0
|
||||
Flask==3.1.3
|
||||
flatbuffers==25.12.19
|
||||
flax==0.11.2
|
||||
folium==0.20.0
|
||||
fonttools==4.62.1
|
||||
fqdn==1.5.1
|
||||
frozendict==2.4.7
|
||||
frozenlist==1.8.0
|
||||
fsspec==2025.3.0
|
||||
future==1.0.0
|
||||
gast==0.7.0
|
||||
gcsfs==2025.3.0
|
||||
GDAL==3.8.4
|
||||
gdown==5.2.2
|
||||
geemap==0.37.2
|
||||
geocoder==1.38.1
|
||||
geographiclib==2.1
|
||||
geopandas==1.1.3
|
||||
geopy==2.4.1
|
||||
giddy==2.3.6
|
||||
gin-config==0.5.0
|
||||
gitdb==4.0.12
|
||||
GitPython==3.1.47
|
||||
glob2==0.7
|
||||
google==3.0.0
|
||||
google-adk==1.29.0
|
||||
google-ai-generativelanguage==0.6.15
|
||||
google-api-core==2.30.3
|
||||
google-api-python-client==2.194.0
|
||||
google-auth==2.47.0
|
||||
google-auth-httplib2==0.3.1
|
||||
google-auth-oauthlib==1.3.1
|
||||
google-cloud-aiplatform==1.148.1
|
||||
google-cloud-appengine-logging==1.9.0
|
||||
google-cloud-audit-log==0.5.0
|
||||
google-cloud-bigquery==3.41.0
|
||||
google-cloud-bigquery-connection==1.21.0
|
||||
google-cloud-bigquery-storage==2.37.0
|
||||
google-cloud-bigtable==2.36.0
|
||||
google-cloud-core==2.5.1
|
||||
google-cloud-dataplex==2.18.0
|
||||
google-cloud-dataproc==5.27.0
|
||||
google-cloud-datastore==2.24.0
|
||||
google-cloud-discoveryengine==0.13.12
|
||||
google-cloud-firestore==2.27.0
|
||||
google-cloud-functions==1.23.0
|
||||
google-cloud-iam==2.22.0
|
||||
google-cloud-language==2.20.0
|
||||
google-cloud-logging==3.15.0
|
||||
google-cloud-monitoring==2.30.0
|
||||
google-cloud-pubsub==2.37.0
|
||||
google-cloud-resource-manager==1.17.0
|
||||
google-cloud-secret-manager==2.27.0
|
||||
google-cloud-spanner==3.65.0
|
||||
google-cloud-speech==2.38.0
|
||||
google-cloud-storage==3.10.1
|
||||
google-cloud-trace==1.19.0
|
||||
google-cloud-translate==3.26.0
|
||||
google-colab @ file:///colabtools/dist/google_colab-1.0.0.tar.gz
|
||||
google-crc32c==1.8.0
|
||||
google-genai==1.68.0
|
||||
google-generativeai==0.8.6
|
||||
google-pasta==0.2.0
|
||||
google-resumable-media==2.8.2
|
||||
googleapis-common-protos==1.74.0
|
||||
googledrivedownloader==1.1.0
|
||||
gradio==5.50.0
|
||||
gradio_client==1.14.0
|
||||
grain==0.2.16
|
||||
graphviz==0.21
|
||||
greenlet==3.4.0
|
||||
groovy==0.1.2
|
||||
grpc-google-iam-v1==0.14.4
|
||||
grpc-interceptor==0.15.4
|
||||
grpcio==1.80.0
|
||||
grpcio-status==1.71.2
|
||||
grpclib==0.4.9
|
||||
gspread==6.2.1
|
||||
gspread-dataframe==4.0.0
|
||||
gym==0.25.2
|
||||
gym-notices==0.1.0
|
||||
gymnasium==1.3.0
|
||||
h11==0.16.0
|
||||
h2==4.3.0
|
||||
h5netcdf==1.8.1
|
||||
h5py==3.16.0
|
||||
hdbscan==0.8.42
|
||||
hf-xet==1.4.3
|
||||
highspy==1.14.0
|
||||
holidays==0.95
|
||||
holoviews==1.22.1
|
||||
hpack==4.1.0
|
||||
html5lib==1.1
|
||||
httpcore==1.0.9
|
||||
httpimport==1.4.1
|
||||
httplib2==0.31.2
|
||||
httptools==0.7.1
|
||||
httpx==0.28.1
|
||||
httpx-sse==0.4.3
|
||||
huggingface_hub==1.11.0
|
||||
humanize==4.15.0
|
||||
hyperframe==6.1.0
|
||||
hyperopt==0.2.7
|
||||
ibis-framework==9.5.0
|
||||
idna==3.13
|
||||
ImageIO==2.37.3
|
||||
imageio-ffmpeg==0.6.0
|
||||
imagesize==2.0.0
|
||||
imbalanced-learn==0.14.1
|
||||
immutabledict==4.3.1
|
||||
importlib_metadata==8.7.1
|
||||
importlib_resources==7.1.0
|
||||
imutils==0.5.4
|
||||
inequality==1.1.2
|
||||
inflect==7.5.0
|
||||
iniconfig==2.3.0
|
||||
intel-cmplr-lib-ur==2025.3.3
|
||||
intel-openmp==2025.3.3
|
||||
ipyevents==2.0.4
|
||||
ipyfilechooser==0.6.0
|
||||
ipykernel==6.17.1
|
||||
ipyleaflet==0.20.0
|
||||
ipyparallel==8.8.0
|
||||
ipython==7.34.0
|
||||
ipython-genutils==0.2.0
|
||||
ipython-sql==0.5.0
|
||||
ipywidgets==7.7.1
|
||||
isoduration==20.11.0
|
||||
itsdangerous==2.2.0
|
||||
jaraco.classes==3.4.0
|
||||
jaraco.context==6.1.2
|
||||
jaraco.functools==4.4.0
|
||||
jax==0.7.2
|
||||
jax-cuda12-pjrt==0.7.2
|
||||
jax-cuda12-plugin==0.7.2
|
||||
jaxlib==0.7.2
|
||||
jeepney==0.9.0
|
||||
jieba==0.42.1
|
||||
Jinja2==3.1.6
|
||||
jiter==0.14.0
|
||||
joblib==1.5.3
|
||||
jsonpatch==1.33
|
||||
jsonpickle==4.1.1
|
||||
jsonpointer==3.1.1
|
||||
jsonschema==4.26.0
|
||||
jsonschema-specifications==2025.9.1
|
||||
jupyter-console==6.6.3
|
||||
jupyter-events==0.12.1
|
||||
jupyter-leaflet==0.20.0
|
||||
jupyter_client==7.4.9
|
||||
jupyter_core==5.9.1
|
||||
jupyter_kernel_gateway @ git+https://github.com/googlecolab/kernel_gateway@b134e9945df25c2dcb98ade9129399be10788671
|
||||
jupyter_server==2.14.0
|
||||
jupyter_server_terminals==0.5.4
|
||||
jupyterlab_pygments==0.3.0
|
||||
jupyterlab_widgets==3.0.16
|
||||
jupytext==1.19.1
|
||||
kaggle==2.0.2
|
||||
kagglehub==1.0.0
|
||||
kagglesdk==0.1.20
|
||||
keras==3.13.2
|
||||
keras-hub==0.26.0
|
||||
keras-nlp==0.26.0
|
||||
keyring==25.7.0
|
||||
keyrings.google-artifactregistry-auth==1.1.2
|
||||
kiwisolver==1.5.0
|
||||
langchain==1.2.15
|
||||
langchain-core==1.3.1
|
||||
langgraph==1.1.9
|
||||
langgraph-checkpoint==4.0.2
|
||||
langgraph-prebuilt==1.0.10
|
||||
langgraph-sdk==0.3.13
|
||||
langsmith==0.7.34
|
||||
lark==1.3.1
|
||||
launchpadlib==1.10.16
|
||||
lazr.restfulclient==0.14.4
|
||||
lazr.uri==1.0.6
|
||||
lazy-loader==0.5
|
||||
libclang==18.1.1
|
||||
libcudf-cu12==26.2.1
|
||||
libcugraph-cu12==26.2.0
|
||||
libcuml-cu12==26.2.0
|
||||
libcuvs-cu12==26.2.0
|
||||
libkvikio-cu12==26.2.0
|
||||
libpysal==4.14.1
|
||||
libraft-cu12==26.2.0
|
||||
librmm-cu12==26.2.0
|
||||
librosa==0.11.0
|
||||
libucx-cu12==1.19.0
|
||||
libucxx-cu12==0.48.0
|
||||
lightgbm==4.6.0
|
||||
linkify-it-py==2.1.0
|
||||
llvmlite==0.43.0
|
||||
locket==1.0.0
|
||||
logical-unification==0.4.7
|
||||
lxml==6.1.0
|
||||
Mako==1.3.11
|
||||
mapclassify==2.10.0
|
||||
Markdown==3.10.2
|
||||
markdown-it-py==4.0.0
|
||||
MarkupSafe==3.0.3
|
||||
matplotlib==3.10.0
|
||||
matplotlib-inline==0.2.1
|
||||
matplotlib-venn==1.1.2
|
||||
mcp==1.27.0
|
||||
mdit-py-plugins==0.5.0
|
||||
mdurl==0.1.2
|
||||
mgwr==2.2.1
|
||||
miniKanren==1.0.5
|
||||
missingno==0.5.2
|
||||
mistune==3.2.0
|
||||
mizani==0.13.5
|
||||
mkl==2025.3.1
|
||||
ml_dtypes==0.5.4
|
||||
mlxtend==0.23.4
|
||||
mmh3==5.2.1
|
||||
momepy==0.11.0
|
||||
more-itertools==10.8.0
|
||||
moviepy==1.0.3
|
||||
mpmath==1.3.0
|
||||
msgpack==1.1.2
|
||||
multidict==6.7.1
|
||||
multipledispatch==1.0.0
|
||||
multiprocess==0.70.16
|
||||
multitasking==0.0.13
|
||||
murmurhash==1.0.15
|
||||
music21==9.9.1
|
||||
namex==0.1.0
|
||||
narwhals==2.20.0
|
||||
natsort==8.4.0
|
||||
nbclassic==1.3.3
|
||||
nbclient==0.10.4
|
||||
nbconvert==7.17.1
|
||||
nbformat==5.10.4
|
||||
ndindex==1.10.1
|
||||
nest-asyncio==1.6.0
|
||||
networkx==3.6.1
|
||||
nibabel==5.4.2
|
||||
nltk==3.9.1
|
||||
notebook==6.5.7
|
||||
notebook_shim==0.2.4
|
||||
numba==0.60.0
|
||||
numba-cuda==0.22.2
|
||||
numexpr==2.14.1
|
||||
numpy==2.0.2
|
||||
nvidia-cublas-cu12==12.8.4.1
|
||||
nvidia-cuda-cccl-cu12==12.9.27
|
||||
nvidia-cuda-cupti-cu12==12.8.90
|
||||
nvidia-cuda-nvcc-cu12==12.8.93
|
||||
nvidia-cuda-nvrtc-cu12==12.8.93
|
||||
nvidia-cuda-runtime-cu12==12.8.90
|
||||
nvidia-cudnn-cu12==9.10.2.21
|
||||
nvidia-cufft-cu12==11.3.3.83
|
||||
nvidia-cufile-cu12==1.13.1.3
|
||||
nvidia-curand-cu12==10.3.9.90
|
||||
nvidia-cusolver-cu12==11.7.3.90
|
||||
nvidia-cusparse-cu12==12.5.8.93
|
||||
nvidia-cusparselt-cu12==0.7.1
|
||||
nvidia-libnvcomp-cu12==5.1.0.21
|
||||
nvidia-ml-py==13.595.45
|
||||
nvidia-nccl-cu12==2.27.5
|
||||
nvidia-nvimgcodec-cu12==0.7.0.11
|
||||
nvidia-nvjitlink-cu12==12.8.93
|
||||
nvidia-nvshmem-cu12==3.4.5
|
||||
nvidia-nvtx-cu12==12.8.90
|
||||
nvtx==0.2.15
|
||||
nx-cugraph-cu12 @ https://pypi.nvidia.com/nx-cugraph-cu12/nx_cugraph_cu12-26.2.0-py3-none-any.whl
|
||||
oauth2client==4.1.3
|
||||
oauthlib==3.3.1
|
||||
omegaconf==2.3.0
|
||||
onemkl-license==2025.3.1
|
||||
openai==2.32.0
|
||||
opencv-contrib-python==4.13.0.92
|
||||
opencv-python==4.13.0.92
|
||||
opencv-python-headless==4.13.0.92
|
||||
openpyxl==3.1.5
|
||||
opentelemetry-api==1.38.0
|
||||
opentelemetry-exporter-gcp-logging==1.11.0a0
|
||||
opentelemetry-exporter-gcp-monitoring==1.11.0a0
|
||||
opentelemetry-exporter-gcp-trace==1.11.0
|
||||
opentelemetry-exporter-otlp-proto-common==1.38.0
|
||||
opentelemetry-exporter-otlp-proto-http==1.38.0
|
||||
opentelemetry-proto==1.38.0
|
||||
opentelemetry-resourcedetector-gcp==1.11.0a0
|
||||
opentelemetry-sdk==1.38.0
|
||||
opentelemetry-semantic-conventions==0.59b0
|
||||
opt_einsum==3.4.0
|
||||
optax==0.2.8
|
||||
optree==0.19.0
|
||||
orbax-checkpoint==0.11.36
|
||||
orjson==3.11.8
|
||||
ormsgpack==1.12.2
|
||||
osqp==1.1.1
|
||||
overrides==7.7.0
|
||||
packaging==26.1
|
||||
pandas==2.2.2
|
||||
pandas-datareader==0.10.0
|
||||
pandas-gbq==0.30.0
|
||||
pandas-stubs==2.2.2.240909
|
||||
pandocfilters==1.5.1
|
||||
panel==1.8.10
|
||||
param==2.3.3
|
||||
parso==0.8.6
|
||||
parsy==2.2
|
||||
partd==1.4.2
|
||||
patsy==1.0.2
|
||||
peewee==4.0.5
|
||||
peft==0.19.1
|
||||
pexpect==4.9.0
|
||||
pickleshare==0.7.5
|
||||
pillow==11.3.0
|
||||
pip==24.1.2
|
||||
platformdirs==4.9.6
|
||||
plotly==5.24.1
|
||||
plotnine==0.14.5
|
||||
pluggy==1.6.0
|
||||
plum-dispatch==2.8.0
|
||||
pointpats==2.5.5
|
||||
polars==1.35.2
|
||||
polars-runtime-32==1.35.2
|
||||
pooch==1.9.0
|
||||
portpicker==1.5.2
|
||||
preshed==3.0.13
|
||||
prettytable==3.17.0
|
||||
proglog==0.1.12
|
||||
progressbar2==4.5.0
|
||||
prometheus_client==0.25.0
|
||||
promise==2.3
|
||||
prompt_toolkit==3.0.52
|
||||
propcache==0.4.1
|
||||
prophet==1.3.0
|
||||
proto-plus==1.27.2
|
||||
protobuf==5.29.6
|
||||
psutil==5.9.5
|
||||
psycopg2==2.9.12
|
||||
psygnal==0.15.1
|
||||
ptyprocess==0.7.0
|
||||
PuLP==3.3.0
|
||||
py-cpuinfo==9.0.0
|
||||
py4j==0.10.9.9
|
||||
pyarrow==18.1.0
|
||||
pyasn1==0.6.3
|
||||
pyasn1_modules==0.4.2
|
||||
pycairo==1.29.0
|
||||
pycocotools==2.0.11
|
||||
pycparser==3.0
|
||||
pycryptodomex==3.23.0
|
||||
pydantic==2.12.3
|
||||
pydantic-settings==2.14.0
|
||||
pydantic_core==2.41.4
|
||||
pydata-google-auth==1.9.1
|
||||
pydot==4.0.1
|
||||
pydotplus==2.0.2
|
||||
PyDrive2==1.21.3
|
||||
pydub==0.25.1
|
||||
pyerfa==2.0.1.5
|
||||
pygame==2.6.1
|
||||
pygit2==1.19.2
|
||||
Pygments==2.20.0
|
||||
PyGObject==3.48.2
|
||||
pyiceberg==0.11.1
|
||||
PyJWT==2.12.1
|
||||
pylibcudf-cu12==26.2.1
|
||||
pylibcugraph-cu12==26.2.0
|
||||
pylibraft-cu12==26.2.0
|
||||
pymc==5.28.4
|
||||
pynndescent==0.6.0
|
||||
pyogrio==0.12.1
|
||||
pyomo==6.10.0
|
||||
PyOpenGL==3.1.10
|
||||
pyOpenSSL==24.2.1
|
||||
pyparsing==3.3.2
|
||||
pyperclip==1.11.0
|
||||
pyproj==3.7.2
|
||||
pyroaring==1.0.4
|
||||
pysal==25.7
|
||||
pyshp==3.0.3
|
||||
PySocks==1.7.1
|
||||
pyspark==4.0.2
|
||||
pytensor==2.38.2
|
||||
pytest==8.4.2
|
||||
python-apt==0.0.0
|
||||
python-box==7.4.1
|
||||
python-dateutil==2.9.0.post0
|
||||
python-dotenv==1.2.2
|
||||
python-fasthtml==0.12.50
|
||||
python-json-logger==4.1.0
|
||||
python-louvain==0.16
|
||||
python-multipart==0.0.26
|
||||
python-slugify==8.0.4
|
||||
python-snappy==0.7.3
|
||||
python-utils==3.9.1
|
||||
pytz==2025.2
|
||||
pyviz_comms==3.0.6
|
||||
PyWavelets==1.9.0
|
||||
PyYAML==6.0.3
|
||||
pyzmq==26.2.1
|
||||
quantecon==0.11.2
|
||||
raft-dask-cu12==26.2.0
|
||||
rapids-dask-dependency==26.2.0
|
||||
rapids-logger==0.2.3
|
||||
rasterio==1.5.0
|
||||
rasterstats==0.20.0
|
||||
ratelim==0.1.6
|
||||
referencing==0.37.0
|
||||
regex==2025.11.3
|
||||
requests==2.32.4
|
||||
requests-oauthlib==2.0.0
|
||||
requests-toolbelt==1.0.0
|
||||
requirements-parser==0.9.0
|
||||
rfc3339-validator==0.1.4
|
||||
rfc3986-validator==0.1.1
|
||||
rfc3987-syntax==1.1.0
|
||||
rich==13.9.4
|
||||
rmm-cu12==26.2.0
|
||||
roman-numerals==4.1.0
|
||||
roman-numerals-py==4.1.0
|
||||
rpds-py==0.30.0
|
||||
rpy2==3.5.17
|
||||
rsa==4.9.1
|
||||
rtree==1.4.1
|
||||
ruff==0.15.11
|
||||
safehttpx==0.1.7
|
||||
safetensors==0.7.0
|
||||
scikit-image==0.25.2
|
||||
scikit-learn==1.6.1
|
||||
scipy==1.16.3
|
||||
scooby==0.11.2
|
||||
scs==3.2.11
|
||||
seaborn==0.13.2
|
||||
SecretStorage==3.5.0
|
||||
segregation==2.5.4
|
||||
semantic-version==2.10.0
|
||||
Send2Trash==2.1.0
|
||||
sentence-transformers==5.4.1
|
||||
sentencepiece==0.2.1
|
||||
sentry-sdk==2.58.0
|
||||
setuptools==75.2.0
|
||||
shap==0.51.0
|
||||
shapely==2.1.2
|
||||
shellingham==1.5.4
|
||||
simple-parsing==0.1.8
|
||||
simplejson==4.1.0
|
||||
simsimd==6.5.16
|
||||
six==1.17.0
|
||||
sklearn-compat==0.1.5
|
||||
sklearn-pandas==2.2.0
|
||||
slicer==0.0.8
|
||||
smart_open==7.6.0
|
||||
smmap==5.0.3
|
||||
sniffio==1.3.1
|
||||
snowballstemmer==3.0.1
|
||||
sortedcontainers==2.4.0
|
||||
soundfile==0.13.1
|
||||
soupsieve==2.8.3
|
||||
soxr==1.0.0
|
||||
spacy==3.8.14
|
||||
spacy-legacy==3.0.12
|
||||
spacy-loggers==1.0.5
|
||||
spaghetti==1.7.6
|
||||
spanner-graph-notebook==1.1.10
|
||||
spglm==1.1.0
|
||||
Sphinx==8.2.3
|
||||
sphinxcontrib-applehelp==2.0.0
|
||||
sphinxcontrib-devhelp==2.0.0
|
||||
sphinxcontrib-htmlhelp==2.1.0
|
||||
sphinxcontrib-jsmath==1.0.1
|
||||
sphinxcontrib-qthelp==2.0.0
|
||||
sphinxcontrib-serializinghtml==2.0.0
|
||||
spint==1.0.7
|
||||
splot==1.1.7
|
||||
spopt==0.7.0
|
||||
spreg==1.9.0
|
||||
SQLAlchemy==2.0.49
|
||||
sqlalchemy-spanner==1.17.3
|
||||
sqlglot==25.20.2
|
||||
sqlparse==0.5.5
|
||||
srsly==2.5.3
|
||||
sse-starlette==3.3.4
|
||||
stanio==0.5.1
|
||||
starlette==0.52.1
|
||||
statsmodels==0.14.6
|
||||
strictyaml==1.7.3
|
||||
stringzilla==4.6.0
|
||||
stumpy==1.13.0
|
||||
sympy==1.14.0
|
||||
tables==3.10.2
|
||||
tabulate==0.9.0
|
||||
tbb==2022.3.1
|
||||
tblib==3.2.2
|
||||
tcmlib==1.4.1
|
||||
tenacity==9.1.4
|
||||
tensorboard==2.20.0
|
||||
tensorboard-data-server==0.7.2
|
||||
tensorflow==2.20.0
|
||||
tensorflow-datasets==4.9.9
|
||||
tensorflow-hub==0.16.1
|
||||
tensorflow-metadata==1.17.3
|
||||
tensorflow-probability==0.25.0
|
||||
tensorflow-text==2.20.1
|
||||
tensorstore==0.1.82
|
||||
termcolor==3.3.0
|
||||
terminado==0.18.1
|
||||
text-unidecode==1.3
|
||||
textblob==0.19.0
|
||||
tf-slim==1.1.0
|
||||
tf_keras==2.20.0
|
||||
thinc==8.3.13
|
||||
threadpoolctl==3.6.0
|
||||
tifffile==2026.4.11
|
||||
tiktoken==0.12.0
|
||||
timm==1.0.26
|
||||
tinycss2==1.4.0
|
||||
tobler==0.14.0
|
||||
tokenizers==0.22.2
|
||||
toml==0.10.2
|
||||
tomlkit==0.13.3
|
||||
toolz==0.12.1
|
||||
torch==2.10.0+cu128
|
||||
torchao==0.10.0
|
||||
torchaudio==2.10.0+cu128
|
||||
torchcodec==0.10.0+cu128
|
||||
torchdata==0.11.0
|
||||
torchsummary==1.5.1
|
||||
torchtune==0.6.1
|
||||
torchvision==0.25.0+cu128
|
||||
tornado==6.5.1
|
||||
tqdm==4.67.3
|
||||
traitlets==5.7.1
|
||||
traittypes==0.2.3
|
||||
transformers==5.0.0
|
||||
treelite==4.7.0
|
||||
treescope==0.1.10
|
||||
triton==3.6.0
|
||||
tsfresh==0.21.1
|
||||
tweepy==4.16.0
|
||||
typeguard==4.5.1
|
||||
typer==0.24.2
|
||||
typer-slim==0.24.0
|
||||
types-pytz==2026.1.1.20260408
|
||||
types-setuptools==82.0.0.20260408
|
||||
typing-inspection==0.4.2
|
||||
typing_extensions==4.15.0
|
||||
tzdata==2026.1
|
||||
tzlocal==5.3.1
|
||||
uc-micro-py==2.0.0
|
||||
ucxx-cu12==0.48.0
|
||||
umap-learn==0.5.12
|
||||
umf==1.0.3
|
||||
uri-template==1.3.0
|
||||
uritemplate==4.2.0
|
||||
urllib3==2.5.0
|
||||
uuid_utils==0.14.1
|
||||
uvicorn==0.46.0
|
||||
uvloop==0.22.1
|
||||
vega-datasets==0.9.0
|
||||
wadllib==1.3.6
|
||||
wandb==0.26.1
|
||||
wasabi==1.1.3
|
||||
watchdog==6.0.0
|
||||
watchfiles==1.1.1
|
||||
wcwidth==0.6.0
|
||||
weasel==1.0.0
|
||||
webcolors==25.10.0
|
||||
webencodings==0.5.1
|
||||
websocket-client==1.9.0
|
||||
websockets==15.0.1
|
||||
Werkzeug==3.1.8
|
||||
wheel==0.47.0
|
||||
widgetsnbextension==3.6.10
|
||||
wordcloud==1.9.6
|
||||
wrapt==2.1.2
|
||||
xarray==2025.12.0
|
||||
xarray-einstats==0.10.0
|
||||
xgboost==3.2.0
|
||||
xlrd==2.0.2
|
||||
xxhash==3.6.0
|
||||
xyzservices==2026.3.0
|
||||
yarl==1.23.0
|
||||
ydf==0.15.0
|
||||
ydf_tf==2.20.0
|
||||
yellowbrick==1.5
|
||||
yfinance==0.2.66
|
||||
zict==3.0.0
|
||||
zipp==3.23.1
|
||||
zstandard==0.25.0
|
||||
36
scripts/data/colab_to_cpu_pin.json
Normal file
36
scripts/data/colab_to_cpu_pin.json
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
{
|
||||
"_comment": "Maps Colab GPU runtime pinned wheels to CPU equivalents for ubuntu-latest CI smoke jobs. The Colab GPU image ships +cu128 builds that won't install on a CPU-only runner; this map either rewrites the spec to a CPU wheel from https://download.pytorch.org/whl/cpu or falls back to module-spoof for packages with no CPU build.",
|
||||
"rewrite": {
|
||||
"torch": {
|
||||
"from_local_version": "+cu128",
|
||||
"to_index_url": "https://download.pytorch.org/whl/cpu"
|
||||
},
|
||||
"torchvision": {
|
||||
"from_local_version": "+cu128",
|
||||
"to_index_url": "https://download.pytorch.org/whl/cpu"
|
||||
},
|
||||
"torchaudio": {
|
||||
"from_local_version": "+cu128",
|
||||
"to_index_url": "https://download.pytorch.org/whl/cpu"
|
||||
}
|
||||
},
|
||||
"module_spoof": {
|
||||
"torchcodec": "no CPU wheel published; smoke job sys.modules-stubs torchcodec before importing unsloth"
|
||||
},
|
||||
"skip": [
|
||||
"nvidia-cublas-cu12",
|
||||
"nvidia-cuda-cupti-cu12",
|
||||
"nvidia-cuda-nvrtc-cu12",
|
||||
"nvidia-cuda-runtime-cu12",
|
||||
"nvidia-cudnn-cu12",
|
||||
"nvidia-cufft-cu12",
|
||||
"nvidia-curand-cu12",
|
||||
"nvidia-cusolver-cu12",
|
||||
"nvidia-cusparse-cu12",
|
||||
"nvidia-cusparselt-cu12",
|
||||
"nvidia-nccl-cu12",
|
||||
"nvidia-nvjitlink-cu12",
|
||||
"nvidia-nvtx-cu12",
|
||||
"triton"
|
||||
]
|
||||
}
|
||||
|
|
@ -6,12 +6,38 @@ from __future__ import annotations
|
|||
import ast
|
||||
import argparse
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import tokenize
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _atomic_write_text(path: Path, data: str, encoding: str) -> None:
|
||||
"""Write ``data`` to ``path`` atomically.
|
||||
|
||||
Stages a tmp file in the same directory (so it's on the same
|
||||
filesystem as the destination), fsyncs, then `os.replace`s into
|
||||
place. A crash mid-write therefore leaves either the previous
|
||||
content or the fully new content -- never a truncated source file.
|
||||
"""
|
||||
dirpath = str(path.parent) or "."
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".kwargs_fix.", dir=dirpath)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding=encoding) as handle:
|
||||
handle.write(data)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def enforce_spacing(text: str) -> tuple[str, bool]:
|
||||
"""Return updated text with keyword '=' padded by spaces, plus change flag."""
|
||||
lines = text.splitlines(keepends=True)
|
||||
|
|
@ -146,7 +172,7 @@ def process_file(path: Path) -> bool:
|
|||
updated, changed = enforce_spacing(original)
|
||||
updated, removed = remove_redundant_passes(updated)
|
||||
if changed or removed:
|
||||
path.write_text(updated, encoding=encoding)
|
||||
_atomic_write_text(path, updated, encoding)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,15 @@
|
|||
#!/bin/bash
|
||||
set -e
|
||||
set -euo pipefail
|
||||
|
||||
# ============================================================
|
||||
# Gemma 4 MLX — One-command setup + inference
|
||||
#
|
||||
# Supply-chain hardening: the uv installer payload is pinned by
|
||||
# SHA-256. Rotate by running:
|
||||
# curl -sSLf https://astral.sh/uv/install.sh | shasum -a 256
|
||||
# and updating _UV_INSTALLER_SHA256 below.
|
||||
# ============================================================
|
||||
#
|
||||
# Usage:
|
||||
# bash install_gemma4_mlx.sh [--venv-dir DIR]
|
||||
#
|
||||
|
|
@ -104,10 +110,17 @@ else
|
|||
fi
|
||||
|
||||
# ── Install uv ───────────────────────────────────────────────
|
||||
_UV_INSTALLER_SHA256="48cd5aca5d5671a3b3d5f61538cc8622e4434af63319115159990d8b0dd02416"
|
||||
|
||||
if ! command -v uv >/dev/null 2>&1; then
|
||||
step "uv" "installing uv package manager..."
|
||||
_uv_tmp=$(mktemp)
|
||||
curl -LsSf "https://astral.sh/uv/install.sh" -o "$_uv_tmp"
|
||||
_uv_actual=$(shasum -a 256 "$_uv_tmp" | awk '{print $1}')
|
||||
if [ "$_uv_actual" != "$_UV_INSTALLER_SHA256" ]; then
|
||||
rm -f "$_uv_tmp"
|
||||
fail "uv installer SHA-256 mismatch: got $_uv_actual expected $_UV_INSTALLER_SHA256 (refusing to execute)"
|
||||
fi
|
||||
sh "$_uv_tmp" </dev/null >/dev/null 2>&1
|
||||
rm -f "$_uv_tmp"
|
||||
if [ -f "$HOME/.local/bin/env" ]; then
|
||||
|
|
|
|||
|
|
@ -1,9 +1,18 @@
|
|||
#!/bin/bash
|
||||
set -e
|
||||
set -euo pipefail
|
||||
|
||||
# ============================================================
|
||||
# Qwen3.6 MLX — One-command setup + inference
|
||||
#
|
||||
# Supply-chain hardening:
|
||||
# - All third-party downloads (uv installer, mlx_vlm qwen3_5
|
||||
# patches) are pinned to an immutable git commit SHA and verified
|
||||
# against a hardcoded SHA-256. Any mismatch aborts the install
|
||||
# before the bytes are copied into site-packages.
|
||||
# - To rotate any pin, fetch the new file with `curl`, run
|
||||
# `shasum -a 256`, and update the corresponding constant below.
|
||||
# ============================================================
|
||||
#
|
||||
# Usage:
|
||||
# bash install_qwen3_6_mlx.sh [--venv-dir DIR]
|
||||
#
|
||||
|
|
@ -104,10 +113,21 @@ else
|
|||
fi
|
||||
|
||||
# ── Install uv ───────────────────────────────────────────────
|
||||
# Pin the uv installer payload by SHA-256. Rotate by running:
|
||||
# curl -sSLf https://astral.sh/uv/install.sh | shasum -a 256
|
||||
# and updating the constant below. We fetch into a temp file, verify
|
||||
# the digest, and only then execute. Mismatch aborts.
|
||||
_UV_INSTALLER_SHA256="48cd5aca5d5671a3b3d5f61538cc8622e4434af63319115159990d8b0dd02416"
|
||||
|
||||
if ! command -v uv >/dev/null 2>&1; then
|
||||
step "uv" "installing uv package manager..."
|
||||
_uv_tmp=$(mktemp)
|
||||
curl -LsSf "https://astral.sh/uv/install.sh" -o "$_uv_tmp"
|
||||
_uv_actual=$(shasum -a 256 "$_uv_tmp" | awk '{print $1}')
|
||||
if [ "$_uv_actual" != "$_UV_INSTALLER_SHA256" ]; then
|
||||
rm -f "$_uv_tmp"
|
||||
fail "uv installer SHA-256 mismatch: got $_uv_actual expected $_UV_INSTALLER_SHA256 (refusing to execute)"
|
||||
fi
|
||||
sh "$_uv_tmp" </dev/null
|
||||
rm -f "$_uv_tmp"
|
||||
if [ -f "$HOME/.local/bin/env" ]; then
|
||||
|
|
@ -150,21 +170,55 @@ else
|
|||
fi
|
||||
|
||||
# ── Apply patches for multi-turn image chat ──────────────────
|
||||
_PATCH_BASE="https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/fix/ui-fix/unsloth/models/patches/mlx_vlm_qwen3_5"
|
||||
#
|
||||
# Pin every patch to an immutable commit SHA and verify the body
|
||||
# against a hardcoded SHA-256. The mlx_vlm_qwen3_5 patch tree
|
||||
# currently only exists on the upstream `fix/ui-fix` branch; we pin
|
||||
# to the branch HEAD commit, NOT the floating ref, so a forced push
|
||||
# on `fix/ui-fix` cannot swap the bytes under us.
|
||||
#
|
||||
# Rotate by:
|
||||
# _PATCH_COMMIT=<new SHA>
|
||||
# curl -sSLf "https://raw.githubusercontent.com/unslothai/unsloth/$_PATCH_COMMIT/unsloth/models/patches/mlx_vlm_qwen3_5/qwen3_5.py" | shasum -a 256
|
||||
# curl -sSLf "https://raw.githubusercontent.com/unslothai/unsloth/$_PATCH_COMMIT/unsloth/models/patches/mlx_vlm_qwen3_5/generate.py" | shasum -a 256
|
||||
_PATCH_COMMIT="013c99e51bbb8c4b83d88f3b150a1e53251a19d2"
|
||||
_PATCH_BASE="https://raw.githubusercontent.com/unslothai/unsloth/${_PATCH_COMMIT}/unsloth/models/patches/mlx_vlm_qwen3_5"
|
||||
_PATCH_SHA_QWEN35="4b6fbbcc59b1d6b935e7204351aae1476836d25542a11c7885402b672d2efa64"
|
||||
_PATCH_SHA_GENERATE="50c4cbb8c3d94c0c74a4d209db6d2b23b102944c147c6421f2eded427b8edaf7"
|
||||
|
||||
_SITE_PKGS=$("$_VENV_PY" -c "import site; print(site.getsitepackages()[0])")
|
||||
|
||||
step "patch" "fixing multi-turn image chat..."
|
||||
|
||||
if curl -sSLf "${_PATCH_BASE}/qwen3_5.py" -o "${_SITE_PKGS}/mlx_vlm/models/qwen3_5/qwen3_5.py"; then
|
||||
# Stage all downloads in an isolated tmpdir; we only copy into
|
||||
# site-packages after every checksum has matched.
|
||||
_PATCH_TMP=$(mktemp -d)
|
||||
trap 'rm -rf "$_PATCH_TMP"' EXIT
|
||||
|
||||
apply_pinned_patch() {
|
||||
# apply_pinned_patch <remote_basename> <expected_sha256> <dest_abspath>
|
||||
_name="$1"; _expected="$2"; _dest="$3"
|
||||
_staged="$_PATCH_TMP/$_name"
|
||||
if ! curl -sSLf "${_PATCH_BASE}/${_name}" -o "$_staged"; then
|
||||
step "warning" "failed to download ${_name} patch — multi-turn image chat may not work" "$C_WARN"
|
||||
return 1
|
||||
fi
|
||||
_actual=$(shasum -a 256 "$_staged" | awk '{print $1}')
|
||||
if [ "$_actual" != "$_expected" ]; then
|
||||
step "warning" "${_name} SHA-256 mismatch (got $_actual expected $_expected) — refusing to install patch" "$C_WARN"
|
||||
return 1
|
||||
fi
|
||||
mkdir -p "$(dirname "$_dest")"
|
||||
cp "$_staged" "$_dest"
|
||||
return 0
|
||||
}
|
||||
|
||||
if apply_pinned_patch "qwen3_5.py" "$_PATCH_SHA_QWEN35" "${_SITE_PKGS}/mlx_vlm/models/qwen3_5/qwen3_5.py"; then
|
||||
substep "patched qwen3_5.py (MRoPE position reset)"
|
||||
else
|
||||
step "warning" "failed to download qwen3_5.py patch — multi-turn image chat may not work" "$C_WARN"
|
||||
fi
|
||||
|
||||
if curl -sSLf "${_PATCH_BASE}/generate.py" -o "${_SITE_PKGS}/mlx_vlm/generate.py"; then
|
||||
if apply_pinned_patch "generate.py" "$_PATCH_SHA_GENERATE" "${_SITE_PKGS}/mlx_vlm/generate.py"; then
|
||||
substep "patched generate.py (mask trim on cache reuse)"
|
||||
else
|
||||
step "warning" "failed to download generate.py patch — multi-turn image chat may not work" "$C_WARN"
|
||||
fi
|
||||
|
||||
# Clear pycache so patches take effect
|
||||
|
|
|
|||
172
scripts/lint_workflow_triggers.py
Normal file
172
scripts/lint_workflow_triggers.py
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""Refuse dangerous GitHub Actions trigger patterns at PR time.
|
||||
|
||||
Two patterns are banned outright, both of which powered the TanStack
|
||||
GHSA-g7cv-rxg3-hmpx supply-chain compromise:
|
||||
|
||||
1. `pull_request_target` -- runs a fork's workflow YAML against the
|
||||
BASE repository's secrets and permissions. The fork can inject
|
||||
arbitrary code into the base context. The TanStack worm used this
|
||||
to land base-context execution from a fork PR. There is essentially
|
||||
no safe use of this trigger for a public open-source project;
|
||||
`pull_request` is the safe alternative.
|
||||
|
||||
2. `workflow_run` chained to a PR-triggered workflow -- carries the
|
||||
same trust boundary problem one hop later. If a PR-triggered
|
||||
workflow can poison artifacts/caches and a `workflow_run` trigger
|
||||
fires off the result with elevated permissions, the attacker still
|
||||
reaches the trusted context.
|
||||
|
||||
3. Shared cache keys between PR-triggered workflows and publish /
|
||||
release / push-triggered workflows. The TanStack worm poisoned the
|
||||
Actions cache from a fork PR and the legitimate release workflow
|
||||
then restored the poisoned cache. Cache keys must be partitioned
|
||||
so that nothing a PR can write is ever read by a workflow that
|
||||
holds secrets.
|
||||
|
||||
Exit codes
|
||||
==========
|
||||
|
||||
0 no findings
|
||||
1 one or more findings; stderr lists each with file path
|
||||
|
||||
Run from repo root:
|
||||
python3 scripts/lint_workflow_triggers.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import yaml
|
||||
except ImportError:
|
||||
print(
|
||||
"ERROR: PyYAML is required. Install with 'pip install pyyaml'", file = sys.stderr
|
||||
)
|
||||
sys.exit(2)
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
DEFAULT_WORKFLOWS_DIR = REPO_ROOT / ".github" / "workflows"
|
||||
|
||||
BANNED_TRIGGERS: tuple[str, ...] = ("pull_request_target",)
|
||||
RESTRICTED_TRIGGERS: tuple[str, ...] = ("workflow_run",)
|
||||
PUBLISH_WORKFLOW_NAMES: tuple[str, ...] = ("release-desktop.yml",)
|
||||
|
||||
|
||||
def _normalise_on(on_field):
|
||||
if isinstance(on_field, str):
|
||||
return {on_field}
|
||||
if isinstance(on_field, list):
|
||||
return set(on_field)
|
||||
if isinstance(on_field, dict):
|
||||
return set(on_field.keys())
|
||||
return set()
|
||||
|
||||
|
||||
def _load_workflow(path: Path):
|
||||
try:
|
||||
return yaml.safe_load(path.read_text())
|
||||
except Exception as exc:
|
||||
print(f"ERROR: failed to parse {path}: {exc}", file = sys.stderr)
|
||||
sys.exit(2)
|
||||
|
||||
|
||||
def _extract_cache_keys(path: Path) -> list[str]:
|
||||
text = path.read_text()
|
||||
keys: list[str] = []
|
||||
for m in re.finditer(r"(?:^|\n)\s*key:\s*([^\n]+)", text):
|
||||
keys.append(m.group(1).strip())
|
||||
return keys
|
||||
|
||||
|
||||
def _trigger_set(yaml_doc) -> set[str]:
|
||||
on = yaml_doc.get(True)
|
||||
if on is None:
|
||||
on = yaml_doc.get("on")
|
||||
return _normalise_on(on)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description = __doc__)
|
||||
parser.add_argument(
|
||||
"--workflows-dir",
|
||||
type = Path,
|
||||
default = DEFAULT_WORKFLOWS_DIR,
|
||||
help = "Override the workflows directory (used by tests).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
workflows_dir = args.workflows_dir
|
||||
|
||||
findings: list[str] = []
|
||||
workflows = sorted(workflows_dir.glob("*.yml"))
|
||||
pr_triggered: list[tuple[Path, list[str]]] = []
|
||||
publish_triggered: list[tuple[Path, list[str]]] = []
|
||||
|
||||
for path in workflows:
|
||||
doc = _load_workflow(path)
|
||||
triggers = _trigger_set(doc)
|
||||
|
||||
for t in BANNED_TRIGGERS:
|
||||
if t in triggers:
|
||||
findings.append(
|
||||
f"{path.name}: BANNED trigger '{t}' (GHSA-g7cv-rxg3-hmpx "
|
||||
"pattern: fork PRs run in base-repo context). Switch to "
|
||||
"'pull_request' and use a deploy-on-merge workflow for "
|
||||
"any privileged step."
|
||||
)
|
||||
|
||||
for t in RESTRICTED_TRIGGERS:
|
||||
if t in triggers:
|
||||
text = path.read_text()
|
||||
if "lint:workflow_triggers-allow-workflow_run" not in text:
|
||||
findings.append(
|
||||
f"{path.name}: RESTRICTED trigger '{t}' requires an "
|
||||
"explicit `# lint:workflow_triggers-allow-workflow_run` "
|
||||
"comment somewhere in the file, with a justification."
|
||||
)
|
||||
|
||||
if "pull_request" in triggers:
|
||||
pr_triggered.append((path, _extract_cache_keys(path)))
|
||||
is_dispatch_only = "workflow_dispatch" in triggers and not (
|
||||
"push" in triggers or "pull_request" in triggers
|
||||
)
|
||||
if path.name in PUBLISH_WORKFLOW_NAMES or is_dispatch_only:
|
||||
publish_triggered.append((path, _extract_cache_keys(path)))
|
||||
|
||||
pr_keys = {key for _, keys in pr_triggered for key in keys}
|
||||
for pub_path, pub_keys in publish_triggered:
|
||||
for k in pub_keys:
|
||||
if k in pr_keys:
|
||||
findings.append(
|
||||
f"{pub_path.name}: cache key {k!r} is also declared in a "
|
||||
"PR-triggered workflow. A fork PR could poison this cache "
|
||||
"and the publish workflow would restore it on next run. "
|
||||
"Add a unique suffix (e.g. '-publish-only') to partition "
|
||||
"the namespaces."
|
||||
)
|
||||
|
||||
if findings:
|
||||
print(
|
||||
"Workflow trigger lint failed with the following issues:", file = sys.stderr
|
||||
)
|
||||
for f in findings:
|
||||
print(f" - {f}", file = sys.stderr)
|
||||
return 1
|
||||
|
||||
print(
|
||||
f"OK: scanned {len(workflows)} workflow file(s); "
|
||||
f"no pull_request_target, no unjustified workflow_run, "
|
||||
f"no PR/publish cache-key collision."
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
754
scripts/lockfile_supply_chain_audit.py
Normal file
754
scripts/lockfile_supply_chain_audit.py
Normal file
|
|
@ -0,0 +1,754 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""Lockfile supply-chain audit for the Studio frontend and Tauri shell.
|
||||
|
||||
Runs BEFORE `npm ci` / `cargo fetch` in CI. Refuses to proceed when a
|
||||
lockfile contains patterns that indicate the kind of supply-chain
|
||||
injection seen in the npm Shai-Hulud waves and the cargo
|
||||
crates.io brand-squat attempts.
|
||||
|
||||
What it checks
|
||||
==============
|
||||
|
||||
studio/frontend/package-lock.json (lockfileVersion 2 or 3):
|
||||
|
||||
1. `resolved` URL origin. Every entry must resolve through
|
||||
`https://registry.npmjs.org/`. Direct GitHub-hosted dependencies
|
||||
(`git+ssh://`, `git+https://`, `github:owner/repo#sha`,
|
||||
`file:`, `http://`) are refused -- npm's TanStack incident used
|
||||
exactly this vector to land an unaudited GitHub commit hash as
|
||||
an optional dependency.
|
||||
|
||||
2. `integrity` field presence. Every non-workspace entry must carry
|
||||
an `integrity` SHA. A missing integrity means the registry can
|
||||
swap the tarball after lockfile generation and CI will not
|
||||
notice.
|
||||
|
||||
3. Known IOC strings. A hardcoded set of indicator-of-compromise
|
||||
substrings is grepped across the entire lockfile body (file
|
||||
names, dependency keys, URLs). The list is updated as new
|
||||
campaigns surface. Catching one means the local install was
|
||||
about to pull a publicly-known malicious release.
|
||||
|
||||
studio/src-tauri/Cargo.lock:
|
||||
|
||||
4. `source` field origin. Every entry with a `source` must point at
|
||||
`registry+https://github.com/rust-lang/crates.io-index`. Direct
|
||||
git sources (`git+https://...`) and `path+...` for cross-crate
|
||||
paths warrant manual review and are flagged.
|
||||
|
||||
5. Known cargo IOC strings. Same idea as (3), separate list.
|
||||
|
||||
Exit codes
|
||||
==========
|
||||
|
||||
0 no findings, or an opt-out env var (UNSLOTH_LOCKFILE_AUDIT_SKIP)
|
||||
is set to a justification string (>=5 chars, not '1'/'true'/etc).
|
||||
A value like '1' or 'true' is now REJECTED loudly and the audit
|
||||
runs normally
|
||||
1 one or more findings; stderr lists them with file path and line
|
||||
number where derivable
|
||||
2 internal error (missing dependency, malformed JSON, etc.)
|
||||
|
||||
Operational stance
|
||||
==================
|
||||
|
||||
This scanner only PARSES the lockfiles -- it never executes anything
|
||||
in them, never resolves anything against the network. Safe to run
|
||||
ahead of every `npm ci`. The IOC list is short by design; this
|
||||
complements (not replaces) `npm audit`, OSV-Scanner, and the
|
||||
advisory-DB pipeline in `.github/workflows/security-audit.yml`. The
|
||||
shape of the catch is "we refuse to proceed because the lockfile
|
||||
itself is shaped wrong", which fires before any third-party install
|
||||
script gets a chance to run on the runner.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Known IOC strings (case-sensitive substring match).
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
#
|
||||
# Keep these short and FACTUAL. Each entry is tied to a public advisory
|
||||
# and is the literal string an attacker would have to embed for the
|
||||
# attack to work. Adding speculative or generic patterns here would
|
||||
# generate false positives on dependency upgrades.
|
||||
NPM_IOC_STRINGS: tuple[str, ...] = (
|
||||
# Shai-Hulud TanStack wave -- May 11, 2026 (GHSA-g7cv-rxg3-hmpx).
|
||||
"router_init.js",
|
||||
"tanstack_runner.js",
|
||||
"router_runtime.js",
|
||||
"@tanstack/setup",
|
||||
"github:tanstack/router#79ac49eedf774dd4b0cfa308722bc463cfe5885c",
|
||||
# Exfiltration endpoints observed across both Shai-Hulud waves.
|
||||
"filev2.getsession.org",
|
||||
"getsession.org/file/",
|
||||
# Campaign markers; the worm tarballs print this to stdout on run.
|
||||
"A Mini Shai-Hulud has Appeared",
|
||||
# Mini Shai-Hulud May-12 2026 wave.
|
||||
"git-tanstack.com",
|
||||
"transformers.pyz",
|
||||
"/tmp/transformers.pyz",
|
||||
"With Love TeamPCP",
|
||||
# Aikido (May-12 wave): payload SHA-256 hashes + Bun marker.
|
||||
"ab4fcadaec49c03278063dd269ea5eef82d24f2124a8e15d7b90f2fa8601266c",
|
||||
"2ec78d556d696e208927cc503d48e4b5eb56b31abc2870c2ed2e98d6be27fc96",
|
||||
"bun run tanstack_runner.js",
|
||||
"We've been online over 2 hours",
|
||||
)
|
||||
|
||||
# Hard pin-blocks for publicly confirmed malicious versions.
|
||||
# keep in sync with scripts/scan_npm_packages.py
|
||||
BLOCKED_NPM_VERSIONS: dict[str, set[str]] = {
|
||||
# GHSA-g7cv-rxg3-hmpx -- TanStack May-11 2026 (84 versions).
|
||||
"@tanstack/arktype-adapter": {"1.166.12", "1.166.15"},
|
||||
"@tanstack/eslint-plugin-router": {"1.161.9", "1.161.12"},
|
||||
"@tanstack/eslint-plugin-start": {"0.0.4", "0.0.7"},
|
||||
"@tanstack/history": {"1.161.9", "1.161.12"},
|
||||
"@tanstack/nitro-v2-vite-plugin": {"1.154.12", "1.154.15"},
|
||||
"@tanstack/react-router": {"1.169.5", "1.169.8"},
|
||||
"@tanstack/react-router-devtools": {"1.166.16", "1.166.19"},
|
||||
"@tanstack/react-router-ssr-query": {"1.166.15", "1.166.18"},
|
||||
"@tanstack/react-start": {"1.167.68", "1.167.71"},
|
||||
"@tanstack/react-start-client": {"1.166.51", "1.166.54"},
|
||||
"@tanstack/react-start-rsc": {"0.0.47", "0.0.50"},
|
||||
"@tanstack/react-start-server": {"1.166.55", "1.166.58"},
|
||||
"@tanstack/router-cli": {"1.166.46", "1.166.49"},
|
||||
"@tanstack/router-core": {"1.169.5", "1.169.8"},
|
||||
"@tanstack/router-devtools": {"1.166.16", "1.166.19"},
|
||||
"@tanstack/router-devtools-core": {"1.167.6", "1.167.9"},
|
||||
"@tanstack/router-generator": {"1.166.45", "1.166.48"},
|
||||
"@tanstack/router-plugin": {"1.167.38", "1.167.41"},
|
||||
"@tanstack/router-ssr-query-core": {"1.168.3", "1.168.6"},
|
||||
"@tanstack/router-utils": {"1.161.11", "1.161.14"},
|
||||
"@tanstack/router-vite-plugin": {"1.166.53", "1.166.56"},
|
||||
"@tanstack/solid-router": {"1.169.5", "1.169.8"},
|
||||
"@tanstack/solid-router-devtools": {"1.166.16", "1.166.19"},
|
||||
"@tanstack/solid-router-ssr-query": {"1.166.15", "1.166.18"},
|
||||
"@tanstack/solid-start": {"1.167.65", "1.167.68"},
|
||||
"@tanstack/solid-start-client": {"1.166.50", "1.166.53"},
|
||||
"@tanstack/solid-start-server": {"1.166.54", "1.166.57"},
|
||||
"@tanstack/start-client-core": {"1.168.5", "1.168.8"},
|
||||
"@tanstack/start-fn-stubs": {"1.161.9", "1.161.12"},
|
||||
"@tanstack/start-plugin-core": {"1.169.23", "1.169.26"},
|
||||
"@tanstack/start-server-core": {"1.167.33", "1.167.36"},
|
||||
"@tanstack/start-static-server-functions": {"1.166.44", "1.166.47"},
|
||||
"@tanstack/start-storage-context": {"1.166.38", "1.166.41"},
|
||||
"@tanstack/valibot-adapter": {"1.166.12", "1.166.15"},
|
||||
"@tanstack/virtual-file-routes": {"1.161.10", "1.161.13"},
|
||||
"@tanstack/vue-router": {"1.169.5", "1.169.8"},
|
||||
"@tanstack/vue-router-devtools": {"1.166.16", "1.166.19"},
|
||||
"@tanstack/vue-router-ssr-query": {"1.166.15", "1.166.18"},
|
||||
"@tanstack/vue-start": {"1.167.61", "1.167.64"},
|
||||
"@tanstack/vue-start-client": {"1.166.46", "1.166.49"},
|
||||
"@tanstack/vue-start-server": {"1.166.50", "1.166.53"},
|
||||
"@tanstack/zod-adapter": {"1.166.12", "1.166.15"},
|
||||
# Mini Shai-Hulud May-12 wave: OpenSearch JS client.
|
||||
"@opensearch-project/opensearch": {"3.5.3", "3.6.2", "3.7.0", "3.8.0"},
|
||||
# Mini Shai-Hulud May-12 wave: @squawk/* (22 packages, 5 versions each;
|
||||
# https://safedep.io/mass-npm-supply-chain-attack-tanstack-mistral/).
|
||||
"@squawk/airport-data": {"0.7.4", "0.7.5", "0.7.6", "0.7.7", "0.7.8"},
|
||||
"@squawk/airports": {"0.6.2", "0.6.3", "0.6.4", "0.6.5", "0.6.6"},
|
||||
"@squawk/airspace": {"0.8.1", "0.8.2", "0.8.3", "0.8.4", "0.8.5"},
|
||||
"@squawk/airspace-data": {"0.5.3", "0.5.4", "0.5.5", "0.5.6", "0.5.7"},
|
||||
"@squawk/airway-data": {"0.5.4", "0.5.5", "0.5.6", "0.5.7", "0.5.8"},
|
||||
"@squawk/airways": {"0.4.2", "0.4.3", "0.4.4", "0.4.5", "0.4.6"},
|
||||
"@squawk/fix-data": {"0.6.4", "0.6.5", "0.6.6", "0.6.7", "0.6.8"},
|
||||
"@squawk/fixes": {"0.3.2", "0.3.3", "0.3.4", "0.3.5", "0.3.6"},
|
||||
"@squawk/flight-math": {"0.5.4", "0.5.5", "0.5.6", "0.5.7", "0.5.8"},
|
||||
"@squawk/flightplan": {"0.5.2", "0.5.3", "0.5.4", "0.5.5", "0.5.6"},
|
||||
"@squawk/geo": {"0.4.4", "0.4.5", "0.4.6", "0.4.7", "0.4.8"},
|
||||
"@squawk/icao-registry": {"0.5.2", "0.5.3", "0.5.4", "0.5.5", "0.5.6"},
|
||||
"@squawk/icao-registry-data": {"0.8.4", "0.8.5", "0.8.6", "0.8.7", "0.8.8"},
|
||||
"@squawk/mcp": {"0.9.1", "0.9.2", "0.9.3", "0.9.4", "0.9.5"},
|
||||
"@squawk/navaid-data": {"0.6.4", "0.6.5", "0.6.6", "0.6.7", "0.6.8"},
|
||||
"@squawk/navaids": {"0.4.2", "0.4.3", "0.4.4", "0.4.5", "0.4.6"},
|
||||
"@squawk/notams": {"0.3.6", "0.3.7", "0.3.8", "0.3.9", "0.3.10"},
|
||||
"@squawk/procedure-data": {"0.7.3", "0.7.4", "0.7.5", "0.7.6", "0.7.7"},
|
||||
"@squawk/procedures": {"0.5.2", "0.5.3", "0.5.4", "0.5.5", "0.5.6"},
|
||||
"@squawk/types": {"0.8.1", "0.8.2", "0.8.3", "0.8.4", "0.8.5"},
|
||||
"@squawk/units": {"0.4.3", "0.4.4", "0.4.5", "0.4.6", "0.4.7"},
|
||||
"@squawk/weather": {"0.5.6", "0.5.7", "0.5.8", "0.5.9", "0.5.10"},
|
||||
# Mini Shai-Hulud May-12 wave: @uipath/* (64 packages, single version each;
|
||||
# https://www.aikido.dev/blog/mini-shai-hulud-is-back-tanstack-compromised).
|
||||
"@uipath/apollo-react": {"4.24.5"},
|
||||
"@uipath/apollo-wind": {"2.16.2"},
|
||||
"@uipath/cli": {"1.0.1"},
|
||||
"@uipath/rpa-tool": {"0.9.5"},
|
||||
"@uipath/apollo-core": {"5.9.2"},
|
||||
"@uipath/filesystem": {"1.0.1"},
|
||||
"@uipath/solutionpackager-tool-core": {"0.0.34"},
|
||||
"@uipath/solution-tool": {"1.0.1"},
|
||||
"@uipath/maestro-tool": {"1.0.1"},
|
||||
"@uipath/codedapp-tool": {"1.0.1"},
|
||||
"@uipath/agent-tool": {"1.0.1"},
|
||||
"@uipath/orchestrator-tool": {"1.0.1"},
|
||||
"@uipath/integrationservice-tool": {"1.0.2"},
|
||||
"@uipath/rpa-legacy-tool": {"1.0.1"},
|
||||
"@uipath/vertical-solutions-tool": {"1.0.1"},
|
||||
"@uipath/flow-tool": {"1.0.2"},
|
||||
"@uipath/codedagent-tool": {"1.0.1"},
|
||||
"@uipath/common": {"1.0.1"},
|
||||
"@uipath/resource-tool": {"1.0.1"},
|
||||
"@uipath/auth": {"1.0.1"},
|
||||
"@uipath/docsai-tool": {"1.0.1"},
|
||||
"@uipath/case-tool": {"1.0.1"},
|
||||
"@uipath/api-workflow-tool": {"1.0.1"},
|
||||
"@uipath/test-manager-tool": {"1.0.2"},
|
||||
"@uipath/robot": {"1.3.4"},
|
||||
"@uipath/traces-tool": {"1.0.1"},
|
||||
"@uipath/agent-sdk": {"1.0.2"},
|
||||
"@uipath/integrationservice-sdk": {"1.0.2"},
|
||||
"@uipath/maestro-sdk": {"1.0.1"},
|
||||
"@uipath/data-fabric-tool": {"1.0.2"},
|
||||
"@uipath/tasks-tool": {"1.0.1"},
|
||||
"@uipath/insights-tool": {"1.0.1"},
|
||||
"@uipath/insights-sdk": {"1.0.1"},
|
||||
"@uipath/uipath-python-bridge": {"1.0.1"},
|
||||
"@uipath/ap-chat": {"1.5.7"},
|
||||
"@uipath/project-packager": {"1.1.16"},
|
||||
"@uipath/packager-tool-case": {"0.0.9"},
|
||||
"@uipath/packager-tool-workflowcompiler-browser": {"0.0.34"},
|
||||
"@uipath/packager-tool-connector": {"0.0.19"},
|
||||
"@uipath/packager-tool-workflowcompiler": {"0.0.16"},
|
||||
"@uipath/packager-tool-webapp": {"1.0.6"},
|
||||
"@uipath/packager-tool-apiworkflow": {"0.0.19"},
|
||||
"@uipath/packager-tool-functions": {"0.1.1"},
|
||||
"@uipath/widget.sdk": {"1.2.3"},
|
||||
"@uipath/resources-tool": {"0.1.11"},
|
||||
"@uipath/agent.sdk": {"0.0.18"},
|
||||
"@uipath/codedagents-tool": {"0.1.12"},
|
||||
"@uipath/aops-policy-tool": {"0.3.1"},
|
||||
"@uipath/solution-packager": {"0.0.35"},
|
||||
"@uipath/packager-tool-bpmn": {"0.0.9"},
|
||||
"@uipath/packager-tool-flow": {"0.0.19"},
|
||||
"@uipath/telemetry": {"0.0.7"},
|
||||
"@uipath/tool-workflowcompiler": {"0.0.12"},
|
||||
"@uipath/vss": {"0.1.6"},
|
||||
"@uipath/solutionpackager-sdk": {"1.0.11"},
|
||||
"@uipath/ui-widgets-multi-file-upload": {"1.0.1"},
|
||||
"@uipath/access-policy-tool": {"0.3.1"},
|
||||
"@uipath/context-grounding-tool": {"0.1.1"},
|
||||
"@uipath/gov-tool": {"0.3.1"},
|
||||
"@uipath/admin-tool": {"0.1.1"},
|
||||
"@uipath/identity-tool": {"0.1.1"},
|
||||
"@uipath/llmgw-tool": {"1.0.1"},
|
||||
"@uipath/resourcecatalog-tool": {"0.1.1"},
|
||||
"@uipath/functions-tool": {"1.0.1"},
|
||||
"@uipath/access-policy-sdk": {"0.3.1"},
|
||||
"@uipath/platform-tool": {"1.0.1"},
|
||||
# Mini Shai-Hulud May-12 wave: @mistralai/* (npm) — separate from PyPI mistralai
|
||||
# (https://www.aikido.dev/blog/mini-shai-hulud-is-back-tanstack-compromised).
|
||||
"@mistralai/mistralai": {"2.2.2", "2.2.3", "2.2.4"},
|
||||
"@mistralai/mistralai-gcp": {"1.7.1", "1.7.2", "1.7.3"},
|
||||
"@mistralai/mistralai-azure": {"1.7.1", "1.7.2", "1.7.3"},
|
||||
# Mini Shai-Hulud May-12 wave: @tallyui/* (30 entries, 10 packages)
|
||||
# (Aikido enumeration).
|
||||
"@tallyui/components": {"1.0.1", "1.0.2", "1.0.3"},
|
||||
"@tallyui/connector-medusa": {"1.0.1", "1.0.2", "1.0.3"},
|
||||
"@tallyui/connector-shopify": {"1.0.1", "1.0.2", "1.0.3"},
|
||||
"@tallyui/connector-vendure": {"1.0.1", "1.0.2", "1.0.3"},
|
||||
"@tallyui/connector-woocommerce": {"1.0.1", "1.0.2", "1.0.3"},
|
||||
"@tallyui/core": {"0.2.1", "0.2.2", "0.2.3"},
|
||||
"@tallyui/database": {"1.0.1", "1.0.2", "1.0.3"},
|
||||
"@tallyui/pos": {"0.1.1", "0.1.2", "0.1.3"},
|
||||
"@tallyui/storage-sqlite": {"0.2.1", "0.2.2", "0.2.3"},
|
||||
"@tallyui/theme": {"0.2.1", "0.2.2", "0.2.3"},
|
||||
# Mini Shai-Hulud May-12 wave: @beproduct/nestjs-auth (18 versions)
|
||||
# (Aikido enumeration).
|
||||
"@beproduct/nestjs-auth": {
|
||||
"0.1.2",
|
||||
"0.1.3",
|
||||
"0.1.4",
|
||||
"0.1.5",
|
||||
"0.1.6",
|
||||
"0.1.7",
|
||||
"0.1.8",
|
||||
"0.1.9",
|
||||
"0.1.10",
|
||||
"0.1.11",
|
||||
"0.1.12",
|
||||
"0.1.13",
|
||||
"0.1.14",
|
||||
"0.1.15",
|
||||
"0.1.16",
|
||||
"0.1.17",
|
||||
"0.1.18",
|
||||
"0.1.19",
|
||||
},
|
||||
# Mini Shai-Hulud May-12 wave: @draftlab/* + @draftauth/*
|
||||
# (Aikido enumeration).
|
||||
"@draftauth/client": {"0.2.1", "0.2.2"},
|
||||
"@draftauth/core": {"0.13.1", "0.13.2"},
|
||||
"@draftlab/auth": {"0.24.1", "0.24.2"},
|
||||
"@draftlab/auth-router": {"0.5.1", "0.5.2"},
|
||||
"@draftlab/db": {"0.16.1"},
|
||||
# Mini Shai-Hulud May-12 wave: @taskflow-corp/cli + @tolka/cli
|
||||
# (Aikido enumeration).
|
||||
"@taskflow-corp/cli": {"0.1.24", "0.1.25", "0.1.26", "0.1.27", "0.1.28", "0.1.29"},
|
||||
"@tolka/cli": {"1.0.2", "1.0.3", "1.0.4", "1.0.5", "1.0.6"},
|
||||
# Mini Shai-Hulud May-12 wave: @ml-toolkit-ts/* + @mesadev/* + @dirigible-ai/sdk + @supersurkhet/*
|
||||
# (Aikido enumeration).
|
||||
"@dirigible-ai/sdk": {"0.6.2", "0.6.3"},
|
||||
"@mesadev/rest": {"0.28.3"},
|
||||
"@mesadev/saguaro": {"0.4.22"},
|
||||
"@mesadev/sdk": {"0.28.3"},
|
||||
"@ml-toolkit-ts/preprocessing": {"1.0.2", "1.0.3"},
|
||||
"@ml-toolkit-ts/xgboost": {"1.0.3", "1.0.4"},
|
||||
"@supersurkhet/cli": {"0.0.2", "0.0.3", "0.0.4", "0.0.5", "0.0.6", "0.0.7"},
|
||||
"@supersurkhet/sdk": {"0.0.2", "0.0.3", "0.0.4", "0.0.5", "0.0.6", "0.0.7"},
|
||||
# Mini Shai-Hulud May-12 wave: Unscoped packages (10 entries)
|
||||
# (Aikido enumeration).
|
||||
"safe-action": {"0.8.3", "0.8.4"},
|
||||
"ts-dna": {"3.0.1", "3.0.2", "3.0.3", "3.0.4"},
|
||||
"cross-stitch": {"1.1.3", "1.1.4", "1.1.5", "1.1.6"},
|
||||
"cmux-agent-mcp": {"0.1.3", "0.1.4", "0.1.5", "0.1.6", "0.1.7", "0.1.8"},
|
||||
"agentwork-cli": {"0.1.4", "0.1.5"},
|
||||
"git-branch-selector": {"1.3.3", "1.3.4", "1.3.5", "1.3.6", "1.3.7"},
|
||||
"wot-api": {"0.8.1", "0.8.2", "0.8.3", "0.8.4"},
|
||||
"git-git-git": {"1.0.8", "1.0.9", "1.0.10", "1.0.11", "1.0.12"},
|
||||
"nextmove-mcp": {"0.1.3", "0.1.4", "0.1.5", "0.1.6", "0.1.7"},
|
||||
"ml-toolkit-ts": {"1.0.4", "1.0.5"},
|
||||
# Cross-ecosystem Mini Shai-Hulud (Apr-30 wave): npm counterpart of
|
||||
# PyPI lightning 2.6.2/2.6.3. Same threat actor (TeamPCP) per Semgrep,
|
||||
# Aikido, OX Security, Resecurity. Safe version: 7.0.3 and earlier.
|
||||
"intercom-client": {"7.0.4"},
|
||||
}
|
||||
|
||||
CARGO_IOC_STRINGS: tuple[str, ...] = (
|
||||
# Reserved for future cargo-side incidents. Empty by default --
|
||||
# `source` origin check below catches the structural pattern.
|
||||
)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Allowed lockfile origins.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
NPM_REGISTRY_PREFIX = "https://registry.npmjs.org/"
|
||||
|
||||
# Tarballs are also fetched from this mirror on some GH Actions cached
|
||||
# runs (npm rewrites the resolved URL on cache hit). Allow either.
|
||||
NPM_REGISTRY_PREFIXES_ALLOWED: tuple[str, ...] = (NPM_REGISTRY_PREFIX,)
|
||||
|
||||
CARGO_REGISTRY_SOURCE = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Cargo non-registry source allowlist.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
#
|
||||
# Each entry is `(crate_name, exact_source_string)`. The crate must
|
||||
# match by name AND the source must match the full pinned-SHA string
|
||||
# verbatim. Bumping the commit SHA forces a re-review here: the
|
||||
# scanner fires until the new SHA is appended.
|
||||
#
|
||||
# Studio's Tauri shell pulls `fix-path-env` directly from
|
||||
# tauri-apps/fix-path-env-rs because the crate is not published to
|
||||
# crates.io. The pinned commit (c4c45d5) was reviewed at the time it
|
||||
# landed; future bumps need explicit approval.
|
||||
CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = (
|
||||
(
|
||||
"fix-path-env",
|
||||
"git+https://github.com/tauri-apps/fix-path-env-rs#"
|
||||
"c4c45d503ea115a839aae718d02f79e7c7f0f673",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Finding container.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class Finding:
|
||||
__slots__ = ("path", "package", "kind", "detail")
|
||||
|
||||
def __init__(self, path: str, package: str, kind: str, detail: str) -> None:
|
||||
self.path = path
|
||||
self.package = package
|
||||
self.kind = kind
|
||||
self.detail = detail
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
f" [{self.kind}] {self.path}\n"
|
||||
f" package: {self.package}\n"
|
||||
f" detail: {self.detail}"
|
||||
)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# package-lock.json audit.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def audit_npm_lockfile(path: Path) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
if not path.exists():
|
||||
return findings
|
||||
|
||||
raw = path.read_text(encoding = "utf-8")
|
||||
try:
|
||||
lock = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = "<root>",
|
||||
kind = "malformed-lockfile",
|
||||
detail = f"could not parse as JSON: {exc}",
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
lockfile_version = lock.get("lockfileVersion")
|
||||
if lockfile_version not in (2, 3):
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = "<root>",
|
||||
kind = "unsupported-lockfile-version",
|
||||
detail = (f"only lockfileVersion 2 or 3 audited; got {lockfile_version}"),
|
||||
)
|
||||
)
|
||||
|
||||
packages = lock.get("packages") or {}
|
||||
for key, entry in packages.items():
|
||||
# The empty key "" is the project root; workspace entries use
|
||||
# keys like "node_modules/foo" or "studio/frontend/sub-pkg".
|
||||
# Skip the project root (it has no `resolved`).
|
||||
if key == "":
|
||||
continue
|
||||
if entry.get("link"):
|
||||
# Workspace symlink; no tarball to resolve.
|
||||
continue
|
||||
|
||||
resolved = entry.get("resolved")
|
||||
# Entries living inside another package's `node_modules/`
|
||||
# tree are bundled fold-ins -- the parent's tarball ships
|
||||
# their source verbatim and the parent's `integrity` covers
|
||||
# the whole subtree. npm represents them in lockfileVersion 3
|
||||
# as nested entries with no `resolved` and no `integrity` of
|
||||
# their own. Treat them as transparent to this audit.
|
||||
nested = key.count("/node_modules/") >= 1
|
||||
|
||||
# 1. resolved-URL origin.
|
||||
if resolved is None:
|
||||
if nested or entry.get("bundled"):
|
||||
# Bundled / fold-in entry; covered by parent integrity.
|
||||
pass
|
||||
elif entry.get("version"):
|
||||
# Top-level entry without a resolved URL is suspicious.
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = key,
|
||||
kind = "missing-resolved-url",
|
||||
detail = (
|
||||
f"version={entry['version']!r} but no `resolved` "
|
||||
"field; lockfile is incomplete"
|
||||
),
|
||||
)
|
||||
)
|
||||
else:
|
||||
if not any(resolved.startswith(p) for p in NPM_REGISTRY_PREFIXES_ALLOWED):
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = key,
|
||||
kind = "non-registry-resolved-url",
|
||||
detail = (
|
||||
f"resolved={resolved!r}; only "
|
||||
f"{NPM_REGISTRY_PREFIX} is permitted. Direct "
|
||||
"GitHub / git / file references are the "
|
||||
"Shai-Hulud injection vector."
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# 2. integrity-hash presence.
|
||||
if resolved is not None and not entry.get("integrity"):
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = key,
|
||||
kind = "missing-integrity-hash",
|
||||
detail = (
|
||||
"no `integrity` field; npm cannot verify the "
|
||||
"tarball SHA against the registry-published hash"
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# 3. Blocked malicious version list.
|
||||
nm_prefix = "node_modules/"
|
||||
pkg_name = key[len(nm_prefix) :] if key.startswith(nm_prefix) else key
|
||||
version = entry.get("version")
|
||||
blocked = BLOCKED_NPM_VERSIONS.get(pkg_name, set())
|
||||
if version and version in blocked:
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = key,
|
||||
kind = "blocked-known-malicious",
|
||||
detail = (
|
||||
f"{pkg_name}@{version} is on the " "BLOCKED_NPM_VERSIONS list"
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# 4. Known IOC strings: scan the raw file body so we hit fields the
|
||||
# structural pass above doesn't enumerate (scripts, optional
|
||||
# dependencies, etc.). Cheap and complete.
|
||||
for ioc in NPM_IOC_STRINGS:
|
||||
if ioc in raw:
|
||||
# Best-effort line number lookup.
|
||||
line_no = _first_line_containing(raw, ioc)
|
||||
findings.append(
|
||||
Finding(
|
||||
path = f"{path}:{line_no}" if line_no else str(path),
|
||||
package = "<ioc-match>",
|
||||
kind = "known-ioc-string",
|
||||
detail = (
|
||||
f"matched known IOC substring {ioc!r}; this is "
|
||||
"a public indicator of a recent supply-chain "
|
||||
"compromise. Refuse to install."
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return findings
|
||||
|
||||
|
||||
def _first_line_containing(text: str, needle: str) -> int | None:
|
||||
for i, line in enumerate(text.splitlines(), start = 1):
|
||||
if needle in line:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# Cargo.lock audit.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
# Cargo.lock is TOML; parse with stdlib tomllib (Python 3.11+). The
|
||||
# studio's Tauri shell already requires a modern toolchain so this is
|
||||
# always available where CI runs.
|
||||
_PACKAGE_HEADER = re.compile(r"^\[\[package\]\]\s*$")
|
||||
|
||||
|
||||
def audit_cargo_lockfile(path: Path) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
if not path.exists():
|
||||
return findings
|
||||
|
||||
raw = path.read_text(encoding = "utf-8")
|
||||
try:
|
||||
import tomllib # type: ignore[import-not-found]
|
||||
except ImportError:
|
||||
# Python <3.11; fall back to a tomli shim if importable.
|
||||
try:
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
except ImportError:
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = "<root>",
|
||||
kind = "missing-toml-parser",
|
||||
detail = (
|
||||
"Python 3.11+ tomllib or tomli is required to "
|
||||
"parse Cargo.lock; install tomli or upgrade "
|
||||
"Python before re-running this audit"
|
||||
),
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
try:
|
||||
lock = tomllib.loads(raw)
|
||||
except Exception as exc:
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = "<root>",
|
||||
kind = "malformed-lockfile",
|
||||
detail = f"could not parse as TOML: {exc}",
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
for entry in lock.get("package", []):
|
||||
name = entry.get("name") or "<unnamed>"
|
||||
version = entry.get("version") or "<unversioned>"
|
||||
source = entry.get("source")
|
||||
# Workspace-local crates have no `source` field; skip them.
|
||||
if source is None:
|
||||
continue
|
||||
if source != CARGO_REGISTRY_SOURCE:
|
||||
if (name, source) in CARGO_SOURCE_ALLOWLIST:
|
||||
# Pre-approved non-registry source pinned by SHA.
|
||||
pass
|
||||
else:
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = f"{name}@{version}",
|
||||
kind = "non-registry-cargo-source",
|
||||
detail = (
|
||||
f"source={source!r}; only "
|
||||
f"{CARGO_REGISTRY_SOURCE!r} is permitted "
|
||||
"by default, and no allowlist entry covers "
|
||||
"this crate. If the source is legitimate, "
|
||||
"add `(name, source)` to "
|
||||
"CARGO_SOURCE_ALLOWLIST after reviewing the "
|
||||
"pinned commit."
|
||||
),
|
||||
)
|
||||
)
|
||||
if not entry.get("checksum") and source == CARGO_REGISTRY_SOURCE:
|
||||
findings.append(
|
||||
Finding(
|
||||
path = str(path),
|
||||
package = f"{name}@{version}",
|
||||
kind = "missing-cargo-checksum",
|
||||
detail = (
|
||||
"registry crate without checksum; cargo cannot "
|
||||
"verify the downloaded source against the "
|
||||
"registry-published SHA"
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
for ioc in CARGO_IOC_STRINGS:
|
||||
if ioc in raw:
|
||||
line_no = _first_line_containing(raw, ioc)
|
||||
findings.append(
|
||||
Finding(
|
||||
path = f"{path}:{line_no}" if line_no else str(path),
|
||||
package = "<ioc-match>",
|
||||
kind = "known-ioc-string",
|
||||
detail = f"matched known IOC substring {ioc!r}",
|
||||
)
|
||||
)
|
||||
|
||||
return findings
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# CLI.
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
DEFAULT_NPM_LOCKFILES = ("studio/frontend/package-lock.json",)
|
||||
DEFAULT_CARGO_LOCKFILES = ("studio/src-tauri/Cargo.lock",)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description = "Pre-install lockfile supply-chain audit.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--root",
|
||||
default = str(REPO_ROOT),
|
||||
help = "Repo root (default: parent of this script).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--npm-lockfile",
|
||||
action = "append",
|
||||
default = None,
|
||||
help = (
|
||||
"Path to a package-lock.json (repeatable). "
|
||||
"Default: studio/frontend/package-lock.json."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cargo-lockfile",
|
||||
action = "append",
|
||||
default = None,
|
||||
help = (
|
||||
"Path to a Cargo.lock (repeatable). "
|
||||
"Default: studio/src-tauri/Cargo.lock."
|
||||
),
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
# SF4: require a real justification (e.g. JIRA ticket id) for the
|
||||
# skip env var. Treat the trivially-set values ("1", "true", "yes",
|
||||
# "on", empty) as INVALID -- they look like accidental flips and
|
||||
# silently bypassed the supply-chain audit. A valid value is a
|
||||
# non-empty string >=5 chars after stripping that does not match
|
||||
# any of the boolean-shaped tokens above. An invalid value emits a
|
||||
# loud GitHub Actions warning to stderr and FALLS THROUGH to run
|
||||
# the audit normally (fail-safe). A valid value emits a warning
|
||||
# naming the reason and skips with rc=0 (compat).
|
||||
_skip_raw = os.environ.get("UNSLOTH_LOCKFILE_AUDIT_SKIP")
|
||||
if _skip_raw is not None:
|
||||
_skip = _skip_raw.strip()
|
||||
_invalid_tokens = {"", "1", "0", "true", "false", "yes", "no", "on", "off"}
|
||||
if _skip.lower() in _invalid_tokens or len(_skip) < 5:
|
||||
print(
|
||||
"::warning::Lockfile audit skip REQUIRES a justification "
|
||||
f"value (>=5 chars, not '{_skip_raw}'). Proceeding with "
|
||||
"audit. Use e.g. UNSLOTH_LOCKFILE_AUDIT_SKIP=ticket-1234.",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"::warning::Lockfile audit skipped: reason='{_skip}'",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
return 0
|
||||
|
||||
root = Path(args.root).resolve()
|
||||
npm_paths = [root / p for p in (args.npm_lockfile or DEFAULT_NPM_LOCKFILES)]
|
||||
cargo_paths = [root / p for p in (args.cargo_lockfile or DEFAULT_CARGO_LOCKFILES)]
|
||||
|
||||
all_findings: list[Finding] = []
|
||||
for p in npm_paths:
|
||||
print(f"[lockfile-audit] npm: {p}", flush = True)
|
||||
all_findings.extend(audit_npm_lockfile(p))
|
||||
for p in cargo_paths:
|
||||
print(f"[lockfile-audit] cargo: {p}", flush = True)
|
||||
all_findings.extend(audit_cargo_lockfile(p))
|
||||
|
||||
if not all_findings:
|
||||
print(
|
||||
f"[lockfile-audit] OK: 0 findings across "
|
||||
f"{len(npm_paths)} npm + {len(cargo_paths)} cargo lockfile(s)",
|
||||
flush = True,
|
||||
)
|
||||
return 0
|
||||
|
||||
print(
|
||||
f"\n[lockfile-audit] FAIL: {len(all_findings)} finding(s):\n",
|
||||
file = sys.stderr,
|
||||
)
|
||||
for f in all_findings:
|
||||
print(str(f), file = sys.stderr)
|
||||
print(file = sys.stderr)
|
||||
print(
|
||||
"[lockfile-audit] Refusing to proceed. Each finding above is "
|
||||
"either a structural lockfile anomaly or a public indicator-of-"
|
||||
"compromise. Investigate before running `npm ci` or `cargo fetch`.",
|
||||
file = sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
415
scripts/notebook_to_python.py
Normal file
415
scripts/notebook_to_python.py
Normal file
|
|
@ -0,0 +1,415 @@
|
|||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
"""
|
||||
Convert Jupyter notebooks (.ipynb) to executable Python scripts (.py).
|
||||
|
||||
Converts IPython magics to plain Python:
|
||||
!command -> subprocess.run('command', shell=True)
|
||||
%cd path -> os.chdir('path')
|
||||
%env VAR=value -> os.environ['VAR'] = 'value'
|
||||
%%file filename -> with open('filename', 'w') as f: f.write(...)
|
||||
%%capture -> (skipped)
|
||||
/content/... -> _WORKING_DIR + /...
|
||||
"""
|
||||
|
||||
import nbformat
|
||||
import re
|
||||
import shlex
|
||||
import sys
|
||||
import os
|
||||
import urllib.request
|
||||
import urllib.parse
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
# Hosts we are willing to fetch raw notebook JSON from. Anything else
|
||||
# is rejected before `urlopen` so a typoed / hostile URL cannot pull
|
||||
# code from arbitrary infrastructure.
|
||||
_ALLOWED_NOTEBOOK_HOSTS = {
|
||||
"raw.githubusercontent.com",
|
||||
"gist.githubusercontent.com",
|
||||
}
|
||||
|
||||
|
||||
# Shell metacharacters that imply the cell's `!cmd` line cannot be
|
||||
# parsed as a flat argv. If any of these appears, `shlex.split` would
|
||||
# either fail or, worse, silently strip the operator -- so we keep
|
||||
# `shell=True` for that command and emit a review marker.
|
||||
_SHELL_METACHARS_RE = re.compile(r"\$\(|`|\|\||\||&&|>>?|<<?|\*|\?|;")
|
||||
|
||||
|
||||
def needs_fstring(cmd: str) -> bool:
|
||||
"""Check if command has Python variable interpolation like {var_name}."""
|
||||
pattern = r"(?<!\$)\{([a-zA-Z_][a-zA-Z0-9_]*)\}"
|
||||
return bool(re.search(pattern, cmd))
|
||||
|
||||
|
||||
def github_blob_to_raw(url: str) -> str:
|
||||
"""Convert GitHub blob URL to raw URL."""
|
||||
# https://github.com/user/repo/blob/branch/path
|
||||
# -> https://raw.githubusercontent.com/user/repo/branch/path
|
||||
# Compare the parsed host exactly (not as a substring) so a URL
|
||||
# like https://attacker.example.com/github.com/blob/... does NOT
|
||||
# get rewritten to a github raw URL. Closes CodeQL alert
|
||||
# py/incomplete-url-substring-sanitization.
|
||||
parsed = urllib.parse.urlparse(url)
|
||||
if parsed.netloc != "github.com" or "/blob/" not in parsed.path:
|
||||
return url
|
||||
new_path = parsed.path.replace("/blob/", "/", 1)
|
||||
return urllib.parse.urlunparse(
|
||||
parsed._replace(netloc = "raw.githubusercontent.com", path = new_path)
|
||||
)
|
||||
|
||||
|
||||
def download_notebook(url: str) -> tuple[str, str]:
|
||||
"""Download notebook from URL. Returns (content, filename)."""
|
||||
# Convert blob URL to raw if needed
|
||||
raw_url = github_blob_to_raw(url)
|
||||
|
||||
# Extract filename from URL
|
||||
parsed = urllib.parse.urlparse(raw_url)
|
||||
filename = os.path.basename(urllib.parse.unquote(parsed.path))
|
||||
|
||||
# Host allowlist. Refuse to fetch from anywhere the campaign IOC
|
||||
# tables flag (or just anywhere we don't recognise). The blob->raw
|
||||
# conversion above only emits `raw.githubusercontent.com`, so a
|
||||
# rejection here means the caller hand-typed a URL pointing
|
||||
# somewhere we don't trust.
|
||||
host = parsed.hostname
|
||||
if host not in _ALLOWED_NOTEBOOK_HOSTS:
|
||||
raise ValueError(
|
||||
f"Refused notebook fetch from {host!r}: not in allowlist "
|
||||
f"{sorted(_ALLOWED_NOTEBOOK_HOSTS)}"
|
||||
)
|
||||
|
||||
# Download
|
||||
print(f"Downloading {url}...")
|
||||
with urllib.request.urlopen(raw_url, timeout = 60) as response:
|
||||
content = response.read().decode("utf-8")
|
||||
|
||||
return content, filename
|
||||
|
||||
|
||||
def is_url(path: str) -> bool:
|
||||
"""Check if path is a URL."""
|
||||
return path.startswith("http://") or path.startswith("https://")
|
||||
|
||||
|
||||
def replace_colab_paths(source: str) -> str:
|
||||
"""Replace Colab-specific /content/ paths with current working directory."""
|
||||
# Replace /content/ with f-string using _WORKING_DIR
|
||||
source = source.replace('"/content/', 'f"{_WORKING_DIR}/')
|
||||
source = source.replace("'/content/", "f'{_WORKING_DIR}/")
|
||||
return source
|
||||
|
||||
|
||||
def _emit_shell_command(indent: str, full_cmd: str, *, allow_shell: bool) -> list[str]:
|
||||
"""Render a `!cmd` notebook line as one or more Python statements.
|
||||
|
||||
When the command body is f-string-interpolated, contains shell
|
||||
metacharacters, or spans multiple lines, falling back to
|
||||
`shell=True` is the only correct option -- `shlex.split` would
|
||||
either drop operators or fail outright. We surface that with a
|
||||
`# WARNING: shell=True; reviewed for hostile input` comment so a
|
||||
reviewer cannot miss it.
|
||||
|
||||
Otherwise we emit `subprocess.run(shlex.split(cmd), shell=False)`
|
||||
so the converted script is not a re-injection vector if the
|
||||
notebook ever interpolates user-controlled data.
|
||||
|
||||
`allow_shell` defaults to True at the CLI for backwards
|
||||
compatibility. Setting it to False makes `shell=True` emission a
|
||||
hard error (no surprise behaviour).
|
||||
"""
|
||||
needs_f = needs_fstring(full_cmd)
|
||||
has_meta = bool(_SHELL_METACHARS_RE.search(full_cmd))
|
||||
multiline = "\n" in full_cmd
|
||||
|
||||
must_use_shell = needs_f or has_meta or multiline
|
||||
|
||||
if must_use_shell:
|
||||
if not allow_shell:
|
||||
raise ValueError(
|
||||
"Cell uses shell metacharacters / interpolation but "
|
||||
"--no-allow-shell was set; refusing to emit shell=True"
|
||||
)
|
||||
warn = f"{indent}# WARNING: shell=True; reviewed for hostile input"
|
||||
f_prefix = "f" if needs_f else ""
|
||||
if multiline:
|
||||
escaped_cmd = full_cmd.replace('"""', r"\"\"\"")
|
||||
if escaped_cmd.rstrip().endswith('"'):
|
||||
escaped_cmd = escaped_cmd.rstrip() + " "
|
||||
stmt = f'{indent}subprocess.run({f_prefix}"""{escaped_cmd}""", shell=True)'
|
||||
else:
|
||||
stmt = f"{indent}subprocess.run({f_prefix}{full_cmd!r}, shell=True)"
|
||||
return [warn, stmt]
|
||||
|
||||
# Shell-safe argv form.
|
||||
return [f"{indent}subprocess.run(shlex.split({full_cmd!r}), shell=False)"]
|
||||
|
||||
|
||||
def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str:
|
||||
"""Convert a cell's IPython magics to plain Python."""
|
||||
lines = source.split("\n")
|
||||
result = []
|
||||
i = 0
|
||||
|
||||
while i < len(lines):
|
||||
line = lines[i]
|
||||
stripped = line.strip()
|
||||
indent = line[: len(line) - len(line.lstrip())]
|
||||
|
||||
# Skip %%capture
|
||||
if stripped.startswith("%%capture"):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Handle %%file magic
|
||||
if stripped.startswith("%%file "):
|
||||
filename = stripped[7:].strip()
|
||||
file_lines = []
|
||||
i += 1
|
||||
while i < len(lines):
|
||||
file_lines.append(lines[i])
|
||||
i += 1
|
||||
file_content = "\n".join(file_lines)
|
||||
file_content = file_content.replace('"""', r"\"\"\"")
|
||||
result.append(f'{indent}with open({filename!r}, "w") as _f:')
|
||||
result.append(f'{indent} _f.write("""{file_content}""")')
|
||||
continue
|
||||
|
||||
# Handle ! shell commands
|
||||
if stripped.startswith("!"):
|
||||
cmd_lines = [stripped[1:]]
|
||||
while cmd_lines[-1].rstrip().endswith("\\") and i + 1 < len(lines):
|
||||
i += 1
|
||||
cmd_lines.append(lines[i].strip())
|
||||
full_cmd = "\n".join(cmd_lines)
|
||||
|
||||
result.extend(
|
||||
_emit_shell_command(indent, full_cmd, allow_shell = allow_shell)
|
||||
)
|
||||
|
||||
# %cd path -> os.chdir(path)
|
||||
elif stripped.startswith("%cd "):
|
||||
path = stripped[4:].strip()
|
||||
result.append(f"{indent}os.chdir({path!r})")
|
||||
|
||||
# %env VAR=value
|
||||
elif stripped.startswith("%env ") and "=" in stripped:
|
||||
match = re.match(r"%env\s+(\w+)=(.+)", stripped)
|
||||
if match:
|
||||
var, val = match.groups()
|
||||
result.append(f"{indent}os.environ[{var!r}] = {val!r}")
|
||||
|
||||
# %env VAR
|
||||
elif stripped.startswith("%env "):
|
||||
var = stripped[5:].strip()
|
||||
result.append(f"{indent}os.environ.get({var!r})")
|
||||
|
||||
# %pwd
|
||||
elif stripped == "%pwd":
|
||||
result.append(f"{indent}os.getcwd()")
|
||||
|
||||
else:
|
||||
result.append(line)
|
||||
|
||||
i += 1
|
||||
|
||||
return "\n".join(result)
|
||||
|
||||
|
||||
def convert_notebook(
|
||||
notebook_content: str,
|
||||
source_name: str = "notebook",
|
||||
*,
|
||||
allow_shell: bool = True,
|
||||
) -> str:
|
||||
"""Convert notebook JSON content to Python script."""
|
||||
# Parse notebook
|
||||
if isinstance(notebook_content, str):
|
||||
notebook = nbformat.reads(notebook_content, as_version = 4)
|
||||
else:
|
||||
notebook = notebook_content
|
||||
|
||||
lines = [
|
||||
"#!/usr/bin/env python",
|
||||
"# coding: utf-8",
|
||||
f"# Converted from: {source_name}",
|
||||
"",
|
||||
"import shlex",
|
||||
"import subprocess",
|
||||
"import os",
|
||||
"import sys",
|
||||
"import re",
|
||||
"",
|
||||
"# Capture original packages before any installs",
|
||||
"_original_packages = subprocess.run(",
|
||||
" [sys.executable, '-m', 'pip', 'freeze'],",
|
||||
" capture_output=True, text=True",
|
||||
").stdout",
|
||||
"",
|
||||
"# Working directory (replaces Colab's /content/)",
|
||||
"_WORKING_DIR = os.getcwd()",
|
||||
"",
|
||||
]
|
||||
|
||||
for cell in notebook.cells:
|
||||
source = cell.source.strip()
|
||||
if not source:
|
||||
continue
|
||||
|
||||
if cell.cell_type == "code":
|
||||
converted = convert_cell_to_python(source, allow_shell = allow_shell)
|
||||
converted = replace_colab_paths(converted)
|
||||
lines.append(converted)
|
||||
lines.append("")
|
||||
|
||||
elif cell.cell_type == "markdown":
|
||||
for line in source.split("\n"):
|
||||
lines.append(f"# {line}")
|
||||
lines.append("")
|
||||
|
||||
# Add package restoration at the end
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
"# Restore original packages (install one by one, skip failures)",
|
||||
"for _pkg in _original_packages.strip().split('\\n'):",
|
||||
" if _pkg:",
|
||||
" subprocess.run([sys.executable, '-m', 'pip', 'install', _pkg, '-q'],",
|
||||
" stderr=subprocess.DEVNULL)",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def convert_notebook_to_script(
|
||||
source: str,
|
||||
output_dir: str | None = None,
|
||||
*,
|
||||
allow_shell: bool = True,
|
||||
):
|
||||
"""
|
||||
Convert a notebook to Python script.
|
||||
|
||||
Args:
|
||||
source: Local file path or URL to notebook
|
||||
output_dir: Output directory (optional, defaults to current directory)
|
||||
allow_shell: When False, refuse to emit `shell=True` for any
|
||||
`!cmd` cell that uses metacharacters / interpolation.
|
||||
"""
|
||||
if is_url(source):
|
||||
content, filename = download_notebook(source)
|
||||
source_name = source
|
||||
else:
|
||||
filename = os.path.basename(source)
|
||||
with open(source, "r", encoding = "utf-8") as f:
|
||||
content = f.read()
|
||||
source_name = source
|
||||
|
||||
# Generate output filename
|
||||
output_filename = filename.replace(".ipynb", ".py")
|
||||
# Clean up filename
|
||||
output_filename = (
|
||||
output_filename.replace("(", "").replace(")", "").replace("-", "_")
|
||||
)
|
||||
|
||||
# Add output directory if specified
|
||||
if output_dir:
|
||||
output_path = os.path.join(output_dir, output_filename)
|
||||
else:
|
||||
output_path = output_filename
|
||||
|
||||
# Convert
|
||||
script = convert_notebook(content, source_name, allow_shell = allow_shell)
|
||||
|
||||
# Write output
|
||||
with open(output_path, "w", encoding = "utf-8") as f:
|
||||
f.write(script)
|
||||
|
||||
print(f"Converted {source} -> {output_path}")
|
||||
return output_path
|
||||
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
|
||||
class Formatter(
|
||||
argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter
|
||||
):
|
||||
pass
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description = __doc__,
|
||||
formatter_class = Formatter,
|
||||
epilog = """
|
||||
Examples:
|
||||
python notebook_to_python.py notebook.ipynb
|
||||
python notebook_to_python.py -o scripts/ notebook1.ipynb notebook2.ipynb
|
||||
python notebook_to_python.py --output ./converted https://github.com/user/repo/blob/main/notebook.ipynb
|
||||
python notebook_to_python.py https://github.com/unslothai/notebooks/blob/main/nb/Oute_TTS_(1B).ipynb
|
||||
""",
|
||||
)
|
||||
parser.add_argument(
|
||||
"notebooks", nargs = "+", help = "Notebook files or URLs to convert."
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o", "--output", dest = "output_dir", default = ".", help = "Output directory."
|
||||
)
|
||||
# Default True for backwards compatibility: existing Colab notebooks
|
||||
# routinely use pipes / redirection / interpolation in `!cmd` lines
|
||||
# and the converted script needs to keep working. Operators who
|
||||
# convert untrusted notebooks should pass --no-allow-shell to force
|
||||
# a hard error on every metacharacter-bearing cell.
|
||||
parser.add_argument(
|
||||
"--allow-shell",
|
||||
dest = "allow_shell",
|
||||
action = "store_true",
|
||||
default = True,
|
||||
help = "Allow emitting subprocess.run(..., shell=True) for cells "
|
||||
"that use shell metacharacters or interpolation (default).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-allow-shell",
|
||||
dest = "allow_shell",
|
||||
action = "store_false",
|
||||
help = "Refuse to emit shell=True; cells with metacharacters error out.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Create output directory if needed
|
||||
os.makedirs(args.output_dir, exist_ok = True)
|
||||
|
||||
# SF2: track per-notebook failures so a CI invocation that converts
|
||||
# 10 notebooks but silently fails on 3 is no longer reported as
|
||||
# success. Each failure is collected and the loop continues so the
|
||||
# caller sees the full set; final exit status is 1 if anything
|
||||
# failed.
|
||||
failures: list[tuple[str, str]] = []
|
||||
ok = 0
|
||||
total = len(args.notebooks)
|
||||
for source in args.notebooks:
|
||||
try:
|
||||
convert_notebook_to_script(
|
||||
source,
|
||||
output_dir = args.output_dir if args.output_dir != "." else None,
|
||||
allow_shell = args.allow_shell,
|
||||
)
|
||||
ok += 1
|
||||
except Exception as e:
|
||||
print(f"ERROR converting {source}: {e}")
|
||||
failures.append((source, f"{type(e).__name__}: {e}"))
|
||||
|
||||
print(
|
||||
f"converted {ok}/{total}, {len(failures)} failed",
|
||||
file = sys.stderr if failures else sys.stdout,
|
||||
)
|
||||
sys.exit(1 if failures else 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
1349
scripts/notebook_validator.py
Normal file
1349
scripts/notebook_validator.py
Normal file
File diff suppressed because it is too large
Load diff
1457
scripts/scan_npm_packages.py
Normal file
1457
scripts/scan_npm_packages.py
Normal file
File diff suppressed because it is too large
Load diff
2226
scripts/scan_packages.py
Normal file
2226
scripts/scan_packages.py
Normal file
File diff suppressed because it is too large
Load diff
282
scripts/stamp_studio_release.py
Normal file
282
scripts/stamp_studio_release.py
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
#!/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
|
||||
|
||||
"""Stamp and verify display-only Studio release metadata for builds."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tarfile
|
||||
import tempfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _atomic_write_text(path: Path, data: str, encoding: str = "utf-8") -> None:
|
||||
"""Atomic version of ``Path.write_text``.
|
||||
|
||||
A crash or signal mid-write leaves the prior file intact; the
|
||||
Studio build never reads a partial ``_studio_release_build.py``.
|
||||
"""
|
||||
dirpath = str(path.parent) or "."
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix = ".stamp_studio.", dir = dirpath)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding = encoding) as handle:
|
||||
handle.write(data)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
BUILD_INFO_PATH = (
|
||||
REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py"
|
||||
)
|
||||
BUILD_INFO_SUFFIX = "studio/backend/utils/_studio_release_build.py"
|
||||
VERSION_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$")
|
||||
GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$")
|
||||
MAX_VERSION_LENGTH = 64
|
||||
PLACEHOLDER = """# 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-stamped Studio release metadata.
|
||||
|
||||
Release builds may rewrite this module in the build workspace before creating
|
||||
Python artifacts. Keep the committed value neutral so source checkouts do not
|
||||
accidentally report a stale release tag.
|
||||
\"\"\"
|
||||
|
||||
STUDIO_RELEASE_VERSION = None
|
||||
"""
|
||||
|
||||
|
||||
def is_valid_version(value: object) -> bool:
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
version = value.strip()
|
||||
if not version or len(version) > MAX_VERSION_LENGTH:
|
||||
return False
|
||||
if version.endswith("-dirty") or GIT_DESCRIBE_SUFFIX_RE.search(version):
|
||||
return False
|
||||
return VERSION_RE.fullmatch(version) is not None
|
||||
|
||||
|
||||
def _exact_git_tag() -> str | None:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"git",
|
||||
"describe",
|
||||
"--tags",
|
||||
"--exact-match",
|
||||
"--match",
|
||||
"v[0-9]*",
|
||||
"HEAD",
|
||||
],
|
||||
cwd = REPO_ROOT,
|
||||
check = False,
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.DEVNULL,
|
||||
text = True,
|
||||
timeout = 2.0,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
return None
|
||||
if result.returncode != 0:
|
||||
return None
|
||||
tag = result.stdout.strip()
|
||||
return tag if is_valid_version(tag) else None
|
||||
|
||||
|
||||
def _git_worktree_is_dirty() -> bool:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "status", "--porcelain"],
|
||||
cwd = REPO_ROOT,
|
||||
check = False,
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.DEVNULL,
|
||||
text = True,
|
||||
timeout = 2.0,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
return True
|
||||
if result.returncode != 0:
|
||||
return True
|
||||
return bool(result.stdout.strip())
|
||||
|
||||
|
||||
def _github_tag() -> str | None:
|
||||
if os.environ.get("GITHUB_REF_TYPE") != "tag":
|
||||
return None
|
||||
github_ref = os.environ.get("GITHUB_REF_NAME", "").strip()
|
||||
return github_ref or None
|
||||
|
||||
|
||||
def resolve_version() -> tuple[str | None, str]:
|
||||
env_version = os.environ.get("UNSLOTH_STUDIO_RELEASE_VERSION", "").strip()
|
||||
if env_version:
|
||||
return (env_version, "UNSLOTH_STUDIO_RELEASE_VERSION")
|
||||
|
||||
github_ref = _github_tag()
|
||||
if github_ref:
|
||||
return (github_ref, "GITHUB_REF_NAME")
|
||||
|
||||
git_tag = _exact_git_tag()
|
||||
if git_tag:
|
||||
return (git_tag, "exact git tag")
|
||||
|
||||
return (None, "none")
|
||||
|
||||
|
||||
def build_info_source(version: str | None) -> str:
|
||||
literal = repr(version) if version is not None else "None"
|
||||
return f'''# 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-stamped Studio release metadata."""
|
||||
|
||||
STUDIO_RELEASE_VERSION = {literal}
|
||||
'''
|
||||
|
||||
|
||||
def _env_version_conflicts(version: str) -> list[tuple[str, str]]:
|
||||
conflicts: list[tuple[str, str]] = []
|
||||
github_ref = _github_tag()
|
||||
if github_ref and is_valid_version(github_ref) and github_ref != version:
|
||||
conflicts.append(("GITHUB_REF_NAME", github_ref))
|
||||
|
||||
git_tag = _exact_git_tag()
|
||||
if git_tag and git_tag != version:
|
||||
conflicts.append(("exact git tag", git_tag))
|
||||
|
||||
return conflicts
|
||||
|
||||
|
||||
def stamp(require_release: bool) -> int:
|
||||
version, source = resolve_version()
|
||||
if version is not None and not is_valid_version(version):
|
||||
print(
|
||||
f"Invalid Studio release version from {source}: {version!r}",
|
||||
file = sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
if version is not None and source == "UNSLOTH_STUDIO_RELEASE_VERSION":
|
||||
conflicts = _env_version_conflicts(version)
|
||||
if conflicts:
|
||||
details = ", ".join(f"{name}={value!r}" for name, value in conflicts)
|
||||
print(
|
||||
"UNSLOTH_STUDIO_RELEASE_VERSION does not match available "
|
||||
f"release tag metadata: {details}",
|
||||
file = sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
if require_release and source == "exact git tag" and _git_worktree_is_dirty():
|
||||
print(
|
||||
"Refusing to publish from a dirty exact-tag checkout. Set "
|
||||
"UNSLOTH_STUDIO_RELEASE_VERSION explicitly from release automation "
|
||||
"or publish from a clean tag checkout.",
|
||||
file = sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
if version is None:
|
||||
if require_release:
|
||||
print(
|
||||
"No Studio release version available. Set "
|
||||
"UNSLOTH_STUDIO_RELEASE_VERSION, build from a GitHub tag, "
|
||||
"or run from an exact local Studio release tag.",
|
||||
file = sys.stderr,
|
||||
)
|
||||
return 2
|
||||
_atomic_write_text(BUILD_INFO_PATH, PLACEHOLDER, encoding = "utf-8")
|
||||
print("dev")
|
||||
return 0
|
||||
|
||||
_atomic_write_text(BUILD_INFO_PATH, build_info_source(version), encoding = "utf-8")
|
||||
print(f"Stamping Studio release version {version} from {source}", file = sys.stderr)
|
||||
print(version)
|
||||
return 0
|
||||
|
||||
|
||||
def _read_wheel_member(path: Path) -> str | None:
|
||||
with zipfile.ZipFile(path) as archive:
|
||||
for name in archive.namelist():
|
||||
if name.endswith(BUILD_INFO_SUFFIX):
|
||||
return archive.read(name).decode("utf-8")
|
||||
return None
|
||||
|
||||
|
||||
def _read_sdist_member(path: Path) -> str | None:
|
||||
with tarfile.open(path) as archive:
|
||||
for member in archive.getmembers():
|
||||
if member.name.endswith(BUILD_INFO_SUFFIX):
|
||||
extracted = archive.extractfile(member)
|
||||
if extracted is None:
|
||||
return None
|
||||
return extracted.read().decode("utf-8")
|
||||
return None
|
||||
|
||||
|
||||
def verify_dist(expected: str, dist_dir: Path) -> int:
|
||||
if not is_valid_version(expected):
|
||||
print(f"Invalid expected Studio release version: {expected!r}", file = sys.stderr)
|
||||
return 2
|
||||
|
||||
artifacts = list(dist_dir.glob("*.whl")) + list(dist_dir.glob("*.tar.gz"))
|
||||
if not artifacts:
|
||||
print(f"No wheel or sdist artifacts found in {dist_dir}", file = sys.stderr)
|
||||
return 2
|
||||
|
||||
expected_line = f"STUDIO_RELEASE_VERSION = {expected!r}"
|
||||
failures: list[str] = []
|
||||
for artifact in artifacts:
|
||||
if artifact.suffix == ".whl":
|
||||
content = _read_wheel_member(artifact)
|
||||
else:
|
||||
content = _read_sdist_member(artifact)
|
||||
if content is None:
|
||||
failures.append(f"{artifact.name}: missing {BUILD_INFO_SUFFIX}")
|
||||
elif expected_line not in content:
|
||||
failures.append(f"{artifact.name}: Studio release version mismatch")
|
||||
|
||||
if failures:
|
||||
for failure in failures:
|
||||
print(failure, file = sys.stderr)
|
||||
return 2
|
||||
|
||||
print(f"Verified Studio release version {expected} in {len(artifacts)} artifact(s)")
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description = __doc__)
|
||||
parser.add_argument("--require-release", action = "store_true")
|
||||
parser.add_argument("--verify-dist", type = Path)
|
||||
parser.add_argument("--expected")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.verify_dist is not None:
|
||||
if not args.expected:
|
||||
parser.error("--verify-dist requires --expected")
|
||||
return verify_dist(args.expected, args.verify_dist)
|
||||
|
||||
return stamp(require_release = args.require_release)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
247
scripts/verify_comment_only_diff.py
Normal file
247
scripts/verify_comment_only_diff.py
Normal file
|
|
@ -0,0 +1,247 @@
|
|||
# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning
|
||||
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Lesser General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Lesser General Public License for more details.
|
||||
|
||||
"""Deterministic comment / docstring-only verifier.
|
||||
|
||||
Compares a list of changed files between two git refs and reports whether
|
||||
each diff is strictly comments / docstrings (Python) or comments
|
||||
(YAML / GitHub Actions). Useful for gating a "comment trim" /
|
||||
"docstring refactor" PR against accidental code drift.
|
||||
|
||||
Per .py file: parse both revs into AST, strip module / class / function
|
||||
docstrings, then compare ast.unparse output. Pure Python comments are
|
||||
discarded by the parser by construction, so any post-strip diff is real
|
||||
code. Per .yml file: yaml.safe_load both sides and compare the parsed
|
||||
Python object; if scalar values differ, also strip shell comments inside
|
||||
``run: |`` block bodies before comparing. Exit code 0 = all OK, 1 = at
|
||||
least one file has a real (non-comment) diff or an error.
|
||||
|
||||
Usage:
|
||||
python scripts/verify_comment_only_diff.py [--base REF] [--head REF] path ...
|
||||
|
||||
Defaults: --base origin/main, --head HEAD. Paths are repo-relative.
|
||||
|
||||
Example:
|
||||
git diff --name-only origin/main..HEAD \\
|
||||
| xargs python scripts/verify_comment_only_diff.py --base origin/main
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import difflib
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
def _git_show(rev: str, path: str) -> str:
|
||||
return subprocess.check_output(
|
||||
["git", "show", f"{rev}:{path}"],
|
||||
text = True,
|
||||
stderr = subprocess.DEVNULL,
|
||||
)
|
||||
|
||||
|
||||
def _strip_docstrings(tree: ast.AST) -> ast.AST:
|
||||
"""Remove every string-literal docstring (Module / FunctionDef /
|
||||
AsyncFunctionDef / ClassDef). Empty body becomes ``pass`` so
|
||||
ast.unparse stays valid."""
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(
|
||||
node,
|
||||
(ast.Module, ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef),
|
||||
):
|
||||
body = getattr(node, "body", None)
|
||||
if not body:
|
||||
continue
|
||||
first = body[0]
|
||||
if (
|
||||
isinstance(first, ast.Expr)
|
||||
and isinstance(first.value, ast.Constant)
|
||||
and isinstance(first.value.value, str)
|
||||
):
|
||||
node.body = body[1:]
|
||||
if not node.body:
|
||||
node.body = [ast.Pass()]
|
||||
return tree
|
||||
|
||||
|
||||
def _normalize_py(src: str) -> str:
|
||||
tree = ast.parse(src)
|
||||
tree = _strip_docstrings(tree)
|
||||
return ast.unparse(tree)
|
||||
|
||||
|
||||
def _strip_shell_comments(s: str) -> str:
|
||||
"""Strip pure-comment lines and inline trailing comments from a shell
|
||||
snippet, then collapse runs of blank lines. Heuristic only: leaves a
|
||||
line untouched if it has an odd quote count (open string)."""
|
||||
out = []
|
||||
for line in s.splitlines():
|
||||
stripped = line.lstrip()
|
||||
if stripped.startswith("#"):
|
||||
continue
|
||||
has_single = line.count("'") % 2 == 0
|
||||
has_double = line.count('"') % 2 == 0
|
||||
if has_single and has_double:
|
||||
idx = line.find(" #")
|
||||
if idx >= 0:
|
||||
line = line[:idx].rstrip()
|
||||
out.append(line)
|
||||
norm = []
|
||||
prev_blank = False
|
||||
for line in out:
|
||||
if line.strip() == "":
|
||||
if prev_blank:
|
||||
continue
|
||||
prev_blank = True
|
||||
else:
|
||||
prev_blank = False
|
||||
norm.append(line)
|
||||
return "\n".join(norm).strip()
|
||||
|
||||
|
||||
def _normalize_yaml_run_strings(obj: Any) -> Any:
|
||||
"""Walk the parsed YAML object; for any multi-line string (i.e. a
|
||||
``run: |`` script body), strip shell comments. Returns a normalised
|
||||
copy."""
|
||||
if isinstance(obj, dict):
|
||||
return {k: _normalize_yaml_run_strings(v) for k, v in obj.items()}
|
||||
if isinstance(obj, list):
|
||||
return [_normalize_yaml_run_strings(x) for x in obj]
|
||||
if isinstance(obj, str) and "\n" in obj:
|
||||
return _strip_shell_comments(obj)
|
||||
return obj
|
||||
|
||||
|
||||
def _walk_yaml_diff(b: Any, a: Any, prefix: str = "") -> None:
|
||||
"""Print a path-keyed summary of the first structural / scalar diff."""
|
||||
if type(b) is not type(a):
|
||||
print(
|
||||
f" type-diff at {prefix or '/'}: "
|
||||
f"{type(b).__name__} -> {type(a).__name__}",
|
||||
)
|
||||
return
|
||||
if isinstance(b, dict):
|
||||
keys = sorted((set(b.keys()) | set(a.keys())), key = lambda x: str(x))
|
||||
for k in keys:
|
||||
if k not in b:
|
||||
print(f" added key {prefix}/{k}")
|
||||
elif k not in a:
|
||||
print(f" removed key {prefix}/{k}")
|
||||
else:
|
||||
_walk_yaml_diff(b[k], a[k], f"{prefix}/{k}")
|
||||
elif isinstance(b, list):
|
||||
if len(b) != len(a):
|
||||
print(
|
||||
f" list len at {prefix or '/'}: " f"{len(b)} -> {len(a)}",
|
||||
)
|
||||
for i, (bi, ai) in enumerate(zip(b, a)):
|
||||
_walk_yaml_diff(bi, ai, f"{prefix}[{i}]")
|
||||
elif b != a:
|
||||
bs = repr(b)[:300]
|
||||
as_ = repr(a)[:300]
|
||||
print(f" scalar at {prefix or '/'}:")
|
||||
print(f" before: {bs}")
|
||||
print(f" after: {as_}")
|
||||
|
||||
|
||||
def _verify_python(path: str, before: str, after: str) -> bool:
|
||||
try:
|
||||
norm_before = _normalize_py(before)
|
||||
norm_after = _normalize_py(after)
|
||||
except SyntaxError as exc:
|
||||
print(f"FAIL {path}: SyntaxError parsing -- {exc}")
|
||||
return False
|
||||
if norm_before == norm_after:
|
||||
print(f"OK {path} (AST identical after docstring strip)")
|
||||
return True
|
||||
diff = list(
|
||||
difflib.unified_diff(
|
||||
norm_before.splitlines(),
|
||||
norm_after.splitlines(),
|
||||
fromfile = f"{path}@before",
|
||||
tofile = f"{path}@after",
|
||||
n = 2,
|
||||
)
|
||||
)
|
||||
print(f"FAIL {path}: AST differs after docstring strip:")
|
||||
for line in diff[:40]:
|
||||
print(f" {line}")
|
||||
return False
|
||||
|
||||
|
||||
def _verify_yaml(path: str, before: str, after: str) -> bool:
|
||||
try:
|
||||
raw_before = yaml.safe_load(before)
|
||||
raw_after = yaml.safe_load(after)
|
||||
except yaml.YAMLError as exc:
|
||||
print(f"FAIL {path}: YAML parse error -- {exc}")
|
||||
return False
|
||||
if raw_before == raw_after:
|
||||
print(f"OK {path} (YAML parsed object identical)")
|
||||
return True
|
||||
norm_before = _normalize_yaml_run_strings(raw_before)
|
||||
norm_after = _normalize_yaml_run_strings(raw_after)
|
||||
if norm_before == norm_after:
|
||||
print(
|
||||
f"OK {path} (YAML parsed object identical after "
|
||||
f"stripping shell comments from run: bodies)",
|
||||
)
|
||||
return True
|
||||
print(
|
||||
f"FAIL {path}: YAML parsed objects still differ after stripping "
|
||||
f"shell comments from `run:` bodies.",
|
||||
)
|
||||
_walk_yaml_diff(norm_before, norm_after)
|
||||
return False
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description = "Verify each path's diff between BASE and HEAD is "
|
||||
"strictly comments / docstrings.",
|
||||
)
|
||||
parser.add_argument("--base", default = "origin/main", help = "base git ref")
|
||||
parser.add_argument("--head", default = "HEAD", help = "head git ref")
|
||||
parser.add_argument("paths", nargs = "+", help = "repo-relative paths")
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
rc = 0
|
||||
print(f"Comparing {len(args.paths)} files: {args.base} vs {args.head}\n")
|
||||
for path in args.paths:
|
||||
try:
|
||||
before = _git_show(args.base, path)
|
||||
after = _git_show(args.head, path)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
print(f"SKIP {path}: {exc}")
|
||||
continue
|
||||
|
||||
if path.endswith(".py"):
|
||||
if not _verify_python(path, before, after):
|
||||
rc = 1
|
||||
elif path.endswith((".yml", ".yaml")):
|
||||
if not _verify_yaml(path, before, after):
|
||||
rc = 1
|
||||
else:
|
||||
print(f"NOTE {path}: not .py or .yaml -- skipped automated check.")
|
||||
|
||||
return rc
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -480,6 +480,37 @@ def save_refresh_token(
|
|||
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.
|
||||
|
|
|
|||
|
|
@ -9,16 +9,14 @@ Export backend - handles model exporting in various formats
|
|||
import glob
|
||||
import json
|
||||
import structlog
|
||||
import tempfile
|
||||
from loggers import get_logger
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple, List
|
||||
from peft import PeftModel, PeftModelForCausalLM
|
||||
from unsloth import FastLanguageModel, FastVisionModel
|
||||
from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX
|
||||
from huggingface_hub import HfApi, ModelCard
|
||||
from transformers.modeling_utils import PushToHubMixin
|
||||
import torch
|
||||
from utils.hardware import clear_gpu_cache
|
||||
|
||||
from utils.models import is_vision_model, get_base_model_from_lora
|
||||
|
|
@ -26,6 +24,12 @@ from utils.models.model_config import detect_audio_type
|
|||
from utils.paths import ensure_dir, outputs_root, resolve_export_dir, resolve_output_dir
|
||||
from core.inference import get_inference_backend
|
||||
|
||||
# GPU-only imports — guarded for Apple Silicon where these aren't needed
|
||||
if not _IS_MLX:
|
||||
from peft import PeftModel, PeftModelForCausalLM
|
||||
from transformers.modeling_utils import PushToHubMixin
|
||||
import torch
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
|
||||
|
|
@ -225,7 +229,7 @@ class ExportBackend:
|
|||
model, tokenizer = FastModel.from_pretrained(
|
||||
model_name = checkpoint_path,
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = torch.float32,
|
||||
dtype = None if _IS_MLX else torch.float32,
|
||||
load_in_4bit = False,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -262,8 +266,12 @@ class ExportBackend:
|
|||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
||||
# Check if PEFT model
|
||||
self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM))
|
||||
# Check if PEFT / LoRA model
|
||||
if _IS_MLX:
|
||||
# MLX doesn't use PeftModel — detect LoRA via adapter_config.json
|
||||
self.is_peft = adapter_config.exists()
|
||||
else:
|
||||
self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM))
|
||||
|
||||
# Store loaded model
|
||||
self.current_model = model
|
||||
|
|
@ -325,9 +333,7 @@ class ExportBackend:
|
|||
private: Whether to make the repo private
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message, output_path). output_path is the
|
||||
resolved absolute on-disk directory of the saved model when
|
||||
``save_directory`` was set, else None.
|
||||
Tuple of (success: bool, message: str, output_path: Optional[str])
|
||||
"""
|
||||
if not self.current_model or not self.current_tokenizer:
|
||||
return False, "No model loaded. Please select a checkpoint first.", None
|
||||
|
|
@ -341,14 +347,17 @@ class ExportBackend:
|
|||
|
||||
output_path: Optional[str] = None
|
||||
try:
|
||||
# Determine save method
|
||||
if format_type == "4-bit (FP4)":
|
||||
save_method = "merged_4bit_forced"
|
||||
elif self._audio_type == "whisper":
|
||||
# Whisper uses save_method=None for local 16-bit merged save
|
||||
save_method = None
|
||||
else: # 16-bit (FP16)
|
||||
save_method = "merged_16bit"
|
||||
if _IS_MLX:
|
||||
mlx_save_method = (
|
||||
"merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
|
||||
)
|
||||
else:
|
||||
if format_type == "4-bit (FP4)":
|
||||
save_method = "merged_4bit_forced"
|
||||
elif self._audio_type == "whisper":
|
||||
save_method = None
|
||||
else:
|
||||
save_method = "merged_16bit"
|
||||
|
||||
# Save locally if requested
|
||||
if save_directory:
|
||||
|
|
@ -356,11 +365,17 @@ class ExportBackend:
|
|||
logger.info(f"Saving merged model locally to: {save_directory}")
|
||||
ensure_dir(Path(save_directory))
|
||||
|
||||
self.current_model.save_pretrained_merged(
|
||||
save_directory, self.current_tokenizer, save_method = save_method
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
# 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())
|
||||
|
|
@ -376,17 +391,40 @@ class ExportBackend:
|
|||
|
||||
logger.info(f"Pushing merged model to Hub: {repo_id}")
|
||||
|
||||
# Whisper uses save_method=None for local but "merged_16bit" for hub push
|
||||
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,
|
||||
)
|
||||
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,
|
||||
)
|
||||
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
|
||||
|
|
@ -411,9 +449,7 @@ class ExportBackend:
|
|||
Export base model (for non-PEFT models).
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message, output_path). output_path is the
|
||||
resolved absolute on-disk directory of the saved model when
|
||||
``save_directory`` was set, else None.
|
||||
Tuple of (success: bool, message: str, output_path: Optional[str])
|
||||
"""
|
||||
if not self.current_model or not self.current_tokenizer:
|
||||
return False, "No model loaded. Please select a checkpoint first.", None
|
||||
|
|
@ -433,8 +469,16 @@ class ExportBackend:
|
|||
logger.info(f"Saving base model locally to: {save_directory}")
|
||||
ensure_dir(Path(save_directory))
|
||||
|
||||
self.current_model.save_pretrained(save_directory)
|
||||
self.current_tokenizer.save_pretrained(save_directory)
|
||||
if _IS_MLX:
|
||||
# MLX: save_pretrained_merged handles non-LoRA models too
|
||||
# (fuse() is a no-op when there are no LoRA layers)
|
||||
self.current_model.save_pretrained_merged(
|
||||
save_directory,
|
||||
self.current_tokenizer,
|
||||
)
|
||||
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)
|
||||
|
|
@ -452,44 +496,73 @@ class ExportBackend:
|
|||
|
||||
logger.info(f"Pushing base model to Hub: {repo_id}")
|
||||
|
||||
# Get base model name from request or model config
|
||||
base_model = (
|
||||
base_model_id
|
||||
or self.current_model.config._name_or_path
|
||||
or "unknown"
|
||||
)
|
||||
|
||||
# Create repo
|
||||
hf_api = HfApi(token = hf_token)
|
||||
repo_id = PushToHubMixin._create_repo(
|
||||
PushToHubMixin,
|
||||
repo_id = repo_id,
|
||||
private = private,
|
||||
token = hf_token,
|
||||
)
|
||||
username = repo_id.split("/")[0]
|
||||
|
||||
# Create and push model card
|
||||
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"
|
||||
)
|
||||
|
||||
# Upload model files
|
||||
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}")
|
||||
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,
|
||||
)
|
||||
self.current_model.push_to_hub_merged(
|
||||
repo_id,
|
||||
self.current_tokenizer,
|
||||
save_directory = tmp_dir,
|
||||
token = hf_token,
|
||||
private = private,
|
||||
)
|
||||
else:
|
||||
return False, "Local save directory required for Hub upload", None
|
||||
# Get base model name from request or model config
|
||||
base_model = (
|
||||
base_model_id
|
||||
or self.current_model.config._name_or_path
|
||||
or "unknown"
|
||||
)
|
||||
|
||||
# Create repo
|
||||
hf_api = HfApi(token = hf_token)
|
||||
repo_id = PushToHubMixin._create_repo(
|
||||
PushToHubMixin,
|
||||
repo_id = repo_id,
|
||||
private = private,
|
||||
token = hf_token,
|
||||
)
|
||||
username = repo_id.split("/")[0]
|
||||
|
||||
# Create and push model card
|
||||
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"
|
||||
)
|
||||
|
||||
# Upload model files
|
||||
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
|
||||
|
||||
|
|
@ -519,9 +592,7 @@ class ExportBackend:
|
|||
hf_token: Hugging Face token
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message, output_path). output_path is the
|
||||
resolved absolute on-disk directory containing the .gguf
|
||||
files when ``save_directory`` was set, else None.
|
||||
Tuple of (success: bool, message: str, output_path: Optional[str])
|
||||
"""
|
||||
if not self.current_model or not self.current_tokenizer:
|
||||
return False, "No model loaded. Please select a checkpoint first.", None
|
||||
|
|
@ -692,9 +763,7 @@ class ExportBackend:
|
|||
Export LoRA adapter only (not merged).
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message, output_path). output_path is the
|
||||
resolved absolute on-disk directory of the saved adapter
|
||||
when ``save_directory`` was set, else None.
|
||||
Tuple of (success: bool, message: str, output_path: Optional[str])
|
||||
"""
|
||||
if not self.current_model or not self.current_tokenizer:
|
||||
return False, "No model loaded. Please select a checkpoint first.", None
|
||||
|
|
@ -710,8 +779,13 @@ class ExportBackend:
|
|||
logger.info(f"Saving LoRA adapter locally to: {save_directory}")
|
||||
ensure_dir(Path(save_directory))
|
||||
|
||||
self.current_model.save_pretrained(save_directory)
|
||||
self.current_tokenizer.save_pretrained(save_directory)
|
||||
if _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())
|
||||
|
||||
|
|
@ -726,10 +800,24 @@ class ExportBackend:
|
|||
|
||||
logger.info(f"Pushing LoRA adapter to Hub: {repo_id}")
|
||||
|
||||
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
|
||||
)
|
||||
if _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
|
||||
|
|
|
|||
3102
studio/backend/core/inference/external_provider.py
Normal file
3102
studio/backend/core/inference/external_provider.py
Normal file
File diff suppressed because it is too large
Load diff
127
studio/backend/core/inference/key_exchange.py
Normal file
127
studio/backend/core/inference/key_exchange.py
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""
|
||||
RSA key pair for encrypting API keys in transit.
|
||||
|
||||
The frontend encrypts API keys with the server's public key before
|
||||
including them in requests. The backend decrypts with its private key
|
||||
before forwarding to external providers.
|
||||
|
||||
The key pair is generated at server startup and lives only in memory —
|
||||
it is regenerated on each restart. The frontend fetches the public key
|
||||
via GET /api/providers/public-key on load.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import logging
|
||||
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa, padding
|
||||
from cryptography.hazmat.primitives import serialization, hashes
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_private_key: rsa.RSAPrivateKey | None = None
|
||||
_public_key_pem: str | None = None
|
||||
_public_key_fingerprint: str | None = None
|
||||
|
||||
|
||||
def _compute_fingerprint(pem: str) -> str:
|
||||
"""SHA256 of the PEM bytes, truncated for log compactness."""
|
||||
return hashlib.sha256(pem.encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
def init_key_pair() -> None:
|
||||
"""Generate an RSA-2048 key pair. Called once at server startup."""
|
||||
global _private_key, _public_key_pem, _public_key_fingerprint
|
||||
if _private_key is not None:
|
||||
# Re-entry is suspicious — every fresh keypair invalidates all
|
||||
# in-flight ciphertext encrypted against the previous public key.
|
||||
# Log loudly so a regression that calls init twice is visible.
|
||||
logger.warning(
|
||||
"init_key_pair called again — replacing existing RSA keypair "
|
||||
"(previous fingerprint=%s). Any frontend that cached the old "
|
||||
"public key will start hitting decryption failures.",
|
||||
_public_key_fingerprint,
|
||||
)
|
||||
_private_key = rsa.generate_private_key(
|
||||
public_exponent = 65537,
|
||||
key_size = 2048,
|
||||
)
|
||||
_public_key_pem = (
|
||||
_private_key.public_key()
|
||||
.public_bytes(
|
||||
serialization.Encoding.PEM,
|
||||
serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
.decode("utf-8")
|
||||
)
|
||||
_public_key_fingerprint = _compute_fingerprint(_public_key_pem)
|
||||
logger.info(
|
||||
"RSA key pair generated for API key encryption (fingerprint=%s)",
|
||||
_public_key_fingerprint,
|
||||
)
|
||||
|
||||
|
||||
def get_public_key_fingerprint() -> str | None:
|
||||
"""Short SHA256 of the current public key PEM; None before init."""
|
||||
return _public_key_fingerprint
|
||||
|
||||
|
||||
def get_public_key_pem() -> str:
|
||||
"""Return the PEM-encoded public key for the frontend."""
|
||||
if _public_key_pem is None:
|
||||
raise RuntimeError("Key pair not initialized. Call init_key_pair() first.")
|
||||
return _public_key_pem
|
||||
|
||||
|
||||
def decrypt_api_key(encrypted_b64: str) -> str:
|
||||
"""
|
||||
Decrypt an API key that was encrypted with the public key.
|
||||
|
||||
Args:
|
||||
encrypted_b64: Base64-encoded RSA-OAEP ciphertext.
|
||||
|
||||
Returns:
|
||||
The plaintext API key string.
|
||||
"""
|
||||
if _private_key is None:
|
||||
raise RuntimeError("Key pair not initialized. Call init_key_pair() first.")
|
||||
|
||||
try:
|
||||
ciphertext = base64.b64decode(encrypted_b64)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"decrypt_api_key: base64 decode failed (input_len=%d, fingerprint=%s): %s: %s",
|
||||
len(encrypted_b64),
|
||||
_public_key_fingerprint,
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
raise
|
||||
|
||||
try:
|
||||
plaintext = _private_key.decrypt(
|
||||
ciphertext,
|
||||
padding.OAEP(
|
||||
mgf = padding.MGF1(algorithm = hashes.SHA256()),
|
||||
algorithm = hashes.SHA256(),
|
||||
label = None,
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
# Surface enough state to distinguish key mismatch (wrong public key
|
||||
# used on encrypt) from a padding/algo mismatch or corrupted bytes.
|
||||
# Expected ciphertext length for RSA-2048 is exactly 256 bytes.
|
||||
logger.warning(
|
||||
"decrypt_api_key: RSA decrypt failed (ciphertext_len=%d, expected=256, "
|
||||
"fingerprint=%s, exc=%s): %s",
|
||||
len(ciphertext),
|
||||
_public_key_fingerprint,
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
raise
|
||||
|
||||
return plaintext.decode("utf-8")
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -69,7 +69,15 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
|
|||
# Single-model server -- Studio runs one model per llama-server
|
||||
# process and serves its own UI. Enabling multi-model loading or
|
||||
# llama-server's built-in web UI changes the surface clients see.
|
||||
# ``--webui``/``--no-webui`` are the legacy spelling; current
|
||||
# upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions.
|
||||
# Keep both so the denylist matches old and new llama-server
|
||||
# binaries (Studio's prebuilt vs system-llama.cpp).
|
||||
frozenset({"--webui", "--no-webui"}),
|
||||
frozenset({"--ui", "--no-ui"}),
|
||||
frozenset({"--ui-config"}),
|
||||
frozenset({"--ui-config-file"}),
|
||||
frozenset({"--ui-mcp-proxy", "--no-ui-mcp-proxy"}),
|
||||
frozenset({"--models-dir"}),
|
||||
frozenset({"--models-preset"}),
|
||||
frozenset({"--models-max"}),
|
||||
|
|
@ -118,3 +126,101 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
|
|||
def is_managed_flag(flag: str) -> bool:
|
||||
"""True if ``flag`` is a Studio-managed llama-server flag."""
|
||||
return flag in _DENYLIST
|
||||
|
||||
|
||||
# Pass-through flags that shadow first-class ``LoadRequest`` fields
|
||||
# (max_seq_length, cache_type_kv, speculative_type,
|
||||
# chat_template_override). Stripped from inherited extras so they
|
||||
# can't last-wins-override an Apply that re-sets the same first-class
|
||||
# field.
|
||||
_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
|
||||
_CACHE_FLAGS: frozenset[str] = frozenset(
|
||||
{"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}
|
||||
)
|
||||
_SPEC_FLAGS: frozenset[str] = frozenset(
|
||||
{
|
||||
"--spec-default",
|
||||
"--spec-type",
|
||||
"--spec-ngram-size-n",
|
||||
"--spec-ngram-size",
|
||||
"--draft-min",
|
||||
"--draft-max",
|
||||
# MTP path (llama.cpp #22673).
|
||||
"--spec-draft-n-max",
|
||||
"--spec-draft-n-min",
|
||||
"--spec-ngram-mod-n-match",
|
||||
"--spec-ngram-mod-n-min",
|
||||
"--spec-ngram-mod-n-max",
|
||||
}
|
||||
)
|
||||
_TEMPLATE_FLAGS: frozenset[str] = frozenset(
|
||||
{
|
||||
"--chat-template",
|
||||
"--chat-template-file",
|
||||
"--chat-template-kwargs",
|
||||
"--jinja",
|
||||
"--no-jinja",
|
||||
}
|
||||
)
|
||||
|
||||
_SHADOWING_FLAGS: frozenset[str] = (
|
||||
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
|
||||
)
|
||||
|
||||
# Boolean flags inside _SHADOWING_FLAGS that take no value. The
|
||||
# value-consuming heuristic in strip_shadowing_flags must skip just the
|
||||
# flag for these, never the following token.
|
||||
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
|
||||
{"--spec-default", "--jinja", "--no-jinja"}
|
||||
)
|
||||
|
||||
|
||||
def strip_shadowing_flags(
|
||||
args: Iterable[str],
|
||||
*,
|
||||
strip_context: bool = True,
|
||||
strip_cache: bool = True,
|
||||
strip_spec: bool = True,
|
||||
strip_template: bool = True,
|
||||
) -> list[str]:
|
||||
"""Strip flags that shadow first-class Studio settings.
|
||||
|
||||
Used when the route inherits a previous load's ``llama_extra_args``
|
||||
so that an inherited ``-c 4096`` cannot override the current
|
||||
request's ``max_seq_length`` (and equivalents for cache /
|
||||
speculative / chat template). Each ``strip_*`` flag controls one
|
||||
group; the route only strips groups whose corresponding first-class
|
||||
field was actually supplied by the caller, so an inherited
|
||||
``--chat-template-file`` survives an Apply that omits both
|
||||
``llama_extra_args`` and ``chat_template_override``.
|
||||
"""
|
||||
shadowing: set[str] = set()
|
||||
if strip_context:
|
||||
shadowing |= _CONTEXT_FLAGS
|
||||
if strip_cache:
|
||||
shadowing |= _CACHE_FLAGS
|
||||
if strip_spec:
|
||||
shadowing |= _SPEC_FLAGS
|
||||
if strip_template:
|
||||
shadowing |= _TEMPLATE_FLAGS
|
||||
|
||||
tokens = [str(a) for a in (args or [])]
|
||||
out: list[str] = []
|
||||
i, n = 0, len(tokens)
|
||||
while i < n:
|
||||
tok = tokens[i]
|
||||
flag = _flag_name(tok)
|
||||
if flag is None or flag not in shadowing:
|
||||
out.append(tok)
|
||||
i += 1
|
||||
continue
|
||||
# Drop this token. Boolean shadowing flags never carry a value;
|
||||
# other shadowing flags consume the next token when it isn't a
|
||||
# flag and the value isn't already packed as ``--key=value``.
|
||||
if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok:
|
||||
i += 1
|
||||
elif i + 1 < n and _flag_name(tokens[i + 1]) is None:
|
||||
i += 2
|
||||
else:
|
||||
i += 1
|
||||
return out
|
||||
|
|
|
|||
417
studio/backend/core/inference/mlx_inference.py
Normal file
417
studio/backend/core/inference/mlx_inference.py
Normal file
|
|
@ -0,0 +1,417 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
"""MLX inference backend for Apple Silicon.
|
||||
|
||||
Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm
|
||||
instead of torch/transformers for model loading and generation.
|
||||
"""
|
||||
|
||||
import threading
|
||||
from typing import Optional, Generator
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class MLXInferenceBackend:
|
||||
def __init__(self):
|
||||
self.models = {}
|
||||
self.active_model_name = None
|
||||
self.loading_models = set()
|
||||
self.loaded_local_models = []
|
||||
self.device = "mlx"
|
||||
self._generation_lock = threading.Lock()
|
||||
|
||||
# MLX state
|
||||
self._model = None
|
||||
self._tokenizer = None
|
||||
self._processor = None
|
||||
self._is_vlm = False
|
||||
self._config = {}
|
||||
|
||||
# Recorded for unload to release pinned memory back to the OS.
|
||||
self._memory_limits_applied = {}
|
||||
|
||||
def _configure_memory_limits(self):
|
||||
"""Apply Metal memory caps before loading a model.
|
||||
|
||||
Mirrors MLXTrainer._configure_memory_limits's defaults:
|
||||
memory_limit = 85% of recommended working-set,
|
||||
wired_limit = min(recommended, memory_limit). Recorded so unload
|
||||
can lower wired_limit back to release pinned RAM.
|
||||
"""
|
||||
import mlx.core as mx
|
||||
|
||||
if not mx.metal.is_available():
|
||||
return
|
||||
info = mx.device_info()
|
||||
rec_bytes = info.get("max_recommended_working_set_size")
|
||||
if not rec_bytes or rec_bytes <= 0:
|
||||
return
|
||||
rec_gb = rec_bytes / 1e9
|
||||
memory_limit_gb = rec_gb * 0.85
|
||||
wired_limit_gb = min(rec_gb, memory_limit_gb)
|
||||
mx.set_memory_limit(int(memory_limit_gb * 1e9))
|
||||
mx.set_wired_limit(int(wired_limit_gb * 1e9))
|
||||
self._memory_limits_applied = {
|
||||
"memory_limit_gb": memory_limit_gb,
|
||||
"wired_limit_gb": wired_limit_gb,
|
||||
"recommended_gb": rec_gb,
|
||||
}
|
||||
logger.info(
|
||||
"MLX memory caps: memory_limit=%.2f GB, wired_limit=%.2f GB",
|
||||
memory_limit_gb,
|
||||
wired_limit_gb,
|
||||
)
|
||||
|
||||
def load_model(
|
||||
self,
|
||||
config,
|
||||
max_seq_length = 2048,
|
||||
load_in_4bit = True,
|
||||
hf_token = None,
|
||||
trust_remote_code = False,
|
||||
gpu_ids = None,
|
||||
dtype = None,
|
||||
) -> bool:
|
||||
import mlx.core as mx
|
||||
|
||||
model_name = config.identifier if hasattr(config, "identifier") else str(config)
|
||||
is_vision = getattr(config, "is_vision", False)
|
||||
|
||||
# GGUF guard. GGUF models are served via llama-server in the
|
||||
# parent process, NOT via mlx-lm in this MLX subprocess. The
|
||||
# route at studio/backend/routes/inference.py:592 (`if config.
|
||||
# is_gguf:`) is responsible for sending GGUF traffic to the
|
||||
# llama-server backend before reaching the MLX orchestrator.
|
||||
# If we end up here with is_gguf=True, the route's
|
||||
# `detect_gguf_model_remote` returned None on its first call
|
||||
# (transient HF Hub flake) but the subprocess re-detection
|
||||
# succeeded. The subprocess cannot reach into the parent's
|
||||
# llama-server, so all we can do is raise loudly so the caller
|
||||
# gets a clear error instead of a cryptic
|
||||
# "config.json does not exist" from mlx_lm.utils.load_model.
|
||||
if getattr(config, "is_gguf", False):
|
||||
raise RuntimeError(
|
||||
f"MLXInferenceBackend cannot load GGUF model '{model_name}': "
|
||||
f"GGUF models must be served by llama-server in the parent "
|
||||
f"process. The /api/inference/load route should have "
|
||||
f"detected this repo as GGUF before dispatching to the MLX "
|
||||
f"orchestrator -- this fallback indicates a transient HF "
|
||||
f"Hub failure during initial detection. Retry the request."
|
||||
)
|
||||
|
||||
if hf_token:
|
||||
import os
|
||||
|
||||
os.environ["HF_TOKEN"] = hf_token
|
||||
self._configure_memory_limits()
|
||||
|
||||
is_lora = getattr(config, "is_lora", False)
|
||||
|
||||
logger.info(
|
||||
"Loading %s via %s (is_lora=%s)",
|
||||
model_name,
|
||||
"mlx-vlm" if is_vision else "mlx-lm",
|
||||
is_lora,
|
||||
)
|
||||
|
||||
try:
|
||||
from unsloth_zoo.mlx.loader import FastMLXModel
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Unsloth: MLX inference requires unsloth-zoo with the MLX modules "
|
||||
"(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon."
|
||||
) from e
|
||||
|
||||
model, tokenizer_or_processor = FastMLXModel.from_pretrained(
|
||||
model_name,
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
load_in_4bit = load_in_4bit,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
text_only = False if is_vision else True,
|
||||
)
|
||||
|
||||
if is_vision:
|
||||
processor = tokenizer_or_processor
|
||||
self._model = model
|
||||
self._processor = processor
|
||||
self._tokenizer = getattr(processor, "tokenizer", processor)
|
||||
self._is_vlm = True
|
||||
else:
|
||||
tokenizer = tokenizer_or_processor
|
||||
self._model = model
|
||||
self._tokenizer = tokenizer
|
||||
self._processor = None
|
||||
self._is_vlm = False
|
||||
|
||||
self.active_model_name = model_name
|
||||
self.models[model_name] = {
|
||||
"model": self._model,
|
||||
"tokenizer": self._tokenizer,
|
||||
"processor": self._processor,
|
||||
"is_vision": is_vision,
|
||||
"is_lora": getattr(config, "is_lora", False),
|
||||
"is_audio": False,
|
||||
"audio_type": None,
|
||||
"has_audio_input": False,
|
||||
}
|
||||
|
||||
logger.info("Model %s loaded successfully", model_name)
|
||||
return True
|
||||
|
||||
def unload_model(self, model_name: str) -> bool:
|
||||
import mlx.core as mx
|
||||
import gc
|
||||
|
||||
if model_name in self.models:
|
||||
del self.models[model_name]
|
||||
self._model = None
|
||||
self._tokenizer = None
|
||||
self._processor = None
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
if mx.metal.is_available() and self._memory_limits_applied and not self.models:
|
||||
try:
|
||||
mx.set_wired_limit(0)
|
||||
logger.info("MLX wired_limit released back to OS on unload")
|
||||
except Exception as e:
|
||||
logger.warning("Failed to release wired_limit: %s", e)
|
||||
self._memory_limits_applied = {}
|
||||
logger.info("Model %s unloaded", model_name)
|
||||
return True
|
||||
|
||||
def generate_chat_response(
|
||||
self,
|
||||
messages,
|
||||
system_prompt = "",
|
||||
image = None,
|
||||
temperature = 0.7,
|
||||
top_p = 0.9,
|
||||
top_k = 40,
|
||||
min_p = 0.0,
|
||||
max_new_tokens = 256,
|
||||
repetition_penalty = 1.0,
|
||||
cancel_event = None,
|
||||
) -> Generator[str, None, None]:
|
||||
if self._model is None:
|
||||
raise RuntimeError("No model loaded")
|
||||
|
||||
# Build messages with system prompt
|
||||
full_messages = []
|
||||
if system_prompt:
|
||||
full_messages.append({"role": "system", "content": system_prompt})
|
||||
full_messages.extend(messages)
|
||||
|
||||
# Inject image into the last user message for VLM
|
||||
if self._is_vlm and image is not None:
|
||||
for msg in reversed(full_messages):
|
||||
if msg.get("role") == "user":
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str):
|
||||
msg["content"] = [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": content},
|
||||
]
|
||||
elif isinstance(content, list):
|
||||
# Prepend image if not already there
|
||||
has_image = any(
|
||||
p.get("type") == "image"
|
||||
for p in content
|
||||
if isinstance(p, dict)
|
||||
)
|
||||
if not has_image:
|
||||
content.insert(0, {"type": "image"})
|
||||
break
|
||||
|
||||
if self._is_vlm:
|
||||
yield from self._generate_vlm(
|
||||
full_messages,
|
||||
image,
|
||||
temperature,
|
||||
top_p,
|
||||
top_k,
|
||||
min_p,
|
||||
max_new_tokens,
|
||||
repetition_penalty,
|
||||
cancel_event,
|
||||
)
|
||||
else:
|
||||
yield from self._generate_text(
|
||||
full_messages,
|
||||
temperature,
|
||||
top_p,
|
||||
top_k,
|
||||
min_p,
|
||||
max_new_tokens,
|
||||
repetition_penalty,
|
||||
cancel_event,
|
||||
)
|
||||
|
||||
def _generate_text(
|
||||
self,
|
||||
messages,
|
||||
temperature,
|
||||
top_p,
|
||||
top_k,
|
||||
min_p,
|
||||
max_new_tokens,
|
||||
repetition_penalty,
|
||||
cancel_event,
|
||||
):
|
||||
from mlx_lm import stream_generate
|
||||
from mlx_lm.sample_utils import make_sampler, make_logits_processors
|
||||
|
||||
prompt = self._tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
)
|
||||
if prompt is None:
|
||||
raise RuntimeError(
|
||||
"apply_chat_template returned None — tokenizer may be incompatible"
|
||||
)
|
||||
|
||||
sampler = make_sampler(
|
||||
temp = temperature,
|
||||
top_p = top_p,
|
||||
top_k = int(top_k or 0),
|
||||
min_p = float(min_p or 0.0),
|
||||
min_tokens_to_keep = 1,
|
||||
)
|
||||
# Only build a logits processor when we actually have a non-trivial
|
||||
# repetition penalty (1.0 is the no-op value).
|
||||
logits_processors = None
|
||||
if repetition_penalty is not None and float(repetition_penalty) not in (
|
||||
0.0,
|
||||
1.0,
|
||||
):
|
||||
logits_processors = make_logits_processors(
|
||||
repetition_penalty = float(repetition_penalty),
|
||||
)
|
||||
|
||||
token_ids = []
|
||||
logger.info(
|
||||
"Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s",
|
||||
len(prompt),
|
||||
max_new_tokens,
|
||||
type(self._model).__name__,
|
||||
type(self._tokenizer).__name__,
|
||||
)
|
||||
with self._generation_lock:
|
||||
try:
|
||||
gen_kwargs = dict(
|
||||
prompt = prompt,
|
||||
max_tokens = max_new_tokens,
|
||||
sampler = sampler,
|
||||
)
|
||||
if logits_processors is not None:
|
||||
gen_kwargs["logits_processors"] = logits_processors
|
||||
for response in stream_generate(
|
||||
self._model,
|
||||
self._tokenizer,
|
||||
**gen_kwargs,
|
||||
):
|
||||
token_ids.append(response.token)
|
||||
# Decode full sequence with skip_special_tokens — same as GPU
|
||||
cumulative = self._tokenizer.decode(
|
||||
token_ids,
|
||||
skip_special_tokens = True,
|
||||
)
|
||||
yield cumulative
|
||||
|
||||
if cancel_event and cancel_event.is_set():
|
||||
break
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
logger.error("stream_generate failed:\n%s", traceback.format_exc())
|
||||
raise
|
||||
|
||||
def _generate_vlm(
|
||||
self,
|
||||
messages,
|
||||
image,
|
||||
temperature,
|
||||
top_p,
|
||||
top_k,
|
||||
min_p,
|
||||
max_new_tokens,
|
||||
repetition_penalty,
|
||||
cancel_event,
|
||||
):
|
||||
from mlx_vlm import stream_generate as vlm_stream
|
||||
|
||||
# Apply chat template
|
||||
chat_fn = getattr(self._processor, "apply_chat_template", None)
|
||||
if (
|
||||
chat_fn is None
|
||||
or not hasattr(self._processor, "chat_template")
|
||||
or self._processor.chat_template is None
|
||||
):
|
||||
tok = getattr(self._processor, "tokenizer", self._processor)
|
||||
chat_fn = tok.apply_chat_template
|
||||
|
||||
prompt = chat_fn(messages, tokenize = False, add_generation_prompt = True)
|
||||
|
||||
# For VLM: always use mlx_vlm's stream_generate which handles
|
||||
# pixel_values properly (passes None for text-only, image for VLM)
|
||||
images = [image] if image is not None else None
|
||||
|
||||
cumulative = ""
|
||||
logger.info(
|
||||
"VLM generating: prompt_len=%d, has_image=%s",
|
||||
len(prompt),
|
||||
image is not None,
|
||||
)
|
||||
# mlx_vlm.stream_generate forwards **kwargs into generate_step, which
|
||||
# accepts temp/top_p/top_k/repetition_penalty (and builds the sampler
|
||||
# + logits_processors internally). Pass them through.
|
||||
# NOTE: mlx_vlm.generate_step expects ``temperature=`` (long form) —
|
||||
# passing ``temp=`` silently falls into **kwargs and is ignored,
|
||||
# leaving generation stuck at the default 0.0 (greedy).
|
||||
vlm_kwargs = dict(
|
||||
max_tokens = max_new_tokens,
|
||||
temperature = temperature,
|
||||
top_p = top_p,
|
||||
top_k = int(top_k or 0),
|
||||
min_p = float(min_p or 0.0),
|
||||
)
|
||||
if repetition_penalty is not None and float(repetition_penalty) not in (
|
||||
0.0,
|
||||
1.0,
|
||||
):
|
||||
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
|
||||
|
||||
with self._generation_lock:
|
||||
for response in vlm_stream(
|
||||
self._model,
|
||||
self._processor,
|
||||
prompt,
|
||||
images,
|
||||
**vlm_kwargs,
|
||||
):
|
||||
token_text = (
|
||||
response.text if hasattr(response, "text") else str(response)
|
||||
)
|
||||
cumulative += token_text
|
||||
yield cumulative
|
||||
if cancel_event and cancel_event.is_set():
|
||||
break
|
||||
|
||||
def generate_with_adapter_control(
|
||||
self, use_adapter = None, cancel_event = None, **gen_kwargs
|
||||
) -> Generator[str, None, None]:
|
||||
# MLX LoRA adapter toggling not yet supported — generate normally
|
||||
yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs)
|
||||
|
||||
def reset_generation_state(self):
|
||||
import mlx.core as mx
|
||||
import gc
|
||||
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
317
studio/backend/core/inference/providers.py
Normal file
317
studio/backend/core/inference/providers.py
Normal file
|
|
@ -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
|
||||
|
||||
"""
|
||||
Static registry of supported external LLM providers.
|
||||
|
||||
All providers expose OpenAI-compatible /v1/chat/completions endpoints
|
||||
with Bearer token authentication and SSE streaming support.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
|
||||
"openai": {
|
||||
"display_name": "OpenAI",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"default_models": [
|
||||
"gpt-5.5",
|
||||
"gpt-5.4",
|
||||
"gpt-5.4-mini",
|
||||
"o3",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
# Keep the model picker scoped to the current generation. The remote
|
||||
# /v1/models listing returns dozens of historical snapshots, fine-tunes
|
||||
# and non-chat models (embeddings, TTS, image, moderation) that we
|
||||
# never want to surface in the chat UI. Filtering here so backend
|
||||
# is the single source of truth.
|
||||
"model_id_allowlist": re.compile(r"^(gpt-5\.[345]|gpt-4\.5|o3)(?:[-.]|$)"),
|
||||
# Hide dated snapshots and the retired plain gpt-5.3 id.
|
||||
"model_id_denylist": re.compile(r"^(gpt-5\.3)$|-\d{4}-\d{2}-\d{2}$"),
|
||||
},
|
||||
"anthropic": {
|
||||
"display_name": "Anthropic",
|
||||
"base_url": "https://api.anthropic.com/v1",
|
||||
"default_models": [
|
||||
"claude-opus-4-7",
|
||||
"claude-opus-4-6",
|
||||
"claude-sonnet-4-6",
|
||||
"claude-opus-4-5",
|
||||
"claude-sonnet-4-5",
|
||||
"claude-haiku-4-5",
|
||||
],
|
||||
# Anthropic /v1/models returns dated snapshot ids alongside the
|
||||
# canonical names (e.g. claude-3-5-sonnet-20241022). Hide the
|
||||
# YYYYMMDD-suffixed variants from the picker — same intent as the
|
||||
# OpenAI denylist, just a different date format (no dashes between
|
||||
# year/month/day).
|
||||
"model_id_denylist": re.compile(r"-\d{8}$"),
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": False,
|
||||
"auth_header": "x-api-key",
|
||||
"auth_prefix": "",
|
||||
"extra_headers": {
|
||||
"anthropic-version": "2023-06-01",
|
||||
},
|
||||
"openai_compatible": False,
|
||||
"notes": "Native Anthropic Messages API. Uses x-api-key header and /v1/messages endpoint with SSE translation.",
|
||||
},
|
||||
"gemini": {
|
||||
"display_name": "Google Gemini",
|
||||
"base_url": "https://generativelanguage.googleapis.com/v1beta/openai",
|
||||
# Curated lineup — Google's /v1beta/openai/models returns dozens
|
||||
# of historical / experimental / embedding ids. Cap to the current
|
||||
# 3.x family plus the rolling `*-latest` aliases.
|
||||
"default_models": [
|
||||
"gemini-3.1-pro-preview",
|
||||
"gemini-3.1-flash-lite",
|
||||
"gemini-3-flash-preview",
|
||||
"gemini-pro-latest",
|
||||
"gemini-flash-latest",
|
||||
"gemini-flash-lite-latest",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
"notes": "OpenAI-compatible endpoint. API key from https://aistudio.google.com/apikey.",
|
||||
"model_id_allowlist": re.compile(
|
||||
r"^(gemini-3\.1-flash-lite|gemini-3-flash-preview|"
|
||||
r"gemini-3\.1-pro-preview|gemini-pro-latest|"
|
||||
r"gemini-flash-latest|gemini-flash-lite-latest)$"
|
||||
),
|
||||
},
|
||||
"deepseek": {
|
||||
"display_name": "DeepSeek",
|
||||
"base_url": "https://api.deepseek.com/v1",
|
||||
"default_models": [
|
||||
"deepseek-chat",
|
||||
"deepseek-reasoner",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": False,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
"notes": "OpenAI-compatible API. deepseek-chat = V3, deepseek-reasoner = R1 thinking mode.",
|
||||
},
|
||||
"mistral": {
|
||||
"display_name": "Mistral AI",
|
||||
"base_url": "https://api.mistral.ai/v1",
|
||||
"default_models": [
|
||||
"codestral-latest",
|
||||
"devstral-latest",
|
||||
"devstral-medium-latest",
|
||||
"magistral-medium-latest",
|
||||
"ministral-14b-latest",
|
||||
"ministral-3b-latest",
|
||||
"ministral-8b-latest",
|
||||
"mistral-large-latest",
|
||||
"mistral-medium-latest",
|
||||
"mistral-small-latest",
|
||||
"mistral-tiny-latest",
|
||||
"mistral-vibe-cli-latest",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
"model_id_allowlist": re.compile(
|
||||
r"^(codestral-latest|devstral-latest|devstral-medium-latest|"
|
||||
r"magistral-medium-latest|ministral-(?:14b|3b|8b)-latest|"
|
||||
r"mistral-(?:large|medium|small|tiny)-latest|"
|
||||
r"mistral-vibe-cli-latest)$"
|
||||
),
|
||||
},
|
||||
"kimi": {
|
||||
"display_name": "Kimi",
|
||||
"base_url": "https://api.moonshot.ai/v1",
|
||||
# Current Kimi model lineup per the official docs:
|
||||
# https://platform.kimi.ai/docs/models
|
||||
# Listing/overview endpoints used to enumerate them:
|
||||
# https://platform.kimi.ai/docs/api/list-models
|
||||
# https://platform.kimi.ai/docs/api/overview
|
||||
# kimi-k2.6 and kimi-k2.5 are the two SoTA multimodal models we
|
||||
# surface in the picker; everything else (moonshot-v1-*, dated
|
||||
# k2 previews) is filtered out by model_id_allowlist below.
|
||||
"default_models": [
|
||||
"kimi-k2.6",
|
||||
"kimi-k2.5",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
"notes": "Moonshot API key. China: use base URL https://api.moonshot.cn/v1",
|
||||
"model_id_allowlist": re.compile(r"^kimi-k2\.[56]$"),
|
||||
# Both k2.6 and k2.5 are reasoning-class. The API rejects custom
|
||||
# sampling: "invalid temperature: only 1 is allowed for this model"
|
||||
# (and the same shape for top_p). Strip both fields from the
|
||||
# outbound body so the server falls back to its required defaults.
|
||||
"body_omit": ("temperature", "top_p"),
|
||||
},
|
||||
"qwen": {
|
||||
"display_name": "Qwen",
|
||||
"base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
|
||||
"default_models": [
|
||||
"qwen-plus",
|
||||
"qwen-turbo",
|
||||
"qwen-max",
|
||||
"qwen2.5-72b-instruct",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
"notes": "DashScope API key. China mainland: override base URL to https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
},
|
||||
"huggingface": {
|
||||
"display_name": "Hugging Face",
|
||||
"base_url": "https://router.huggingface.co/v1",
|
||||
# Seed the picker with a few popular ids so something is selectable
|
||||
# before the live /v1/models call resolves. The remote listing is
|
||||
# the source of truth — see model_list_mode below.
|
||||
"default_models": [
|
||||
"openai/gpt-oss-120b",
|
||||
"deepseek-ai/DeepSeek-V3",
|
||||
"meta-llama/Llama-3.3-70B-Instruct",
|
||||
"Qwen/Qwen2.5-72B-Instruct",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
"notes": (
|
||||
"HF token from huggingface.co/settings/tokens. Uses the "
|
||||
"OpenAI-compatible router at /v1/chat/completions; /v1/models "
|
||||
"returns the cross-provider chat catalog. See "
|
||||
"https://huggingface.co/docs/inference-providers/index."
|
||||
),
|
||||
# /v1/models works on the HF router and returns the full chat-model
|
||||
# catalog (state.org/model[:policy] ids). Switch to remote so users
|
||||
# see live availability — the picker has a search box, and
|
||||
# loadModels() merges defaults so default_models entries remain
|
||||
# visible if the remote call fails.
|
||||
"model_list_mode": "remote",
|
||||
# Scope the catalog to first-party org repos we trust as primary
|
||||
# sources. The HF /v1/models response is otherwise hundreds of
|
||||
# ids long (community fine-tunes, mirrors, fp8 variants, etc.).
|
||||
"model_id_allowlist": re.compile(
|
||||
r"^(openai|deepseek-ai|google|meta-llama|Qwen|moonshotai|"
|
||||
r"mistralai|zai-org)/"
|
||||
),
|
||||
# Cap the post-filter list. /v1/models has no server-side limit
|
||||
# or popularity sort, so this is just "first N matches" — pair it
|
||||
# with the default_models seed so the most useful flagship ids
|
||||
# are always among the top regardless of the API's order.
|
||||
"model_id_limit": 15,
|
||||
},
|
||||
"vllm": {
|
||||
"display_name": "vLLM",
|
||||
# User-supplied via provider_base_url; the route layer already falls
|
||||
# back to the payload's base_url when the registry entry has none.
|
||||
"base_url": "",
|
||||
"default_models": [],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
# Force /v1/chat/completions in stream_chat_completion — vLLM's
|
||||
# /v1/responses rebuilds messages and runs them through the loaded
|
||||
# model's chat template, which 400s on strict-alternation templates
|
||||
# (Gemma 3 raises "Conversation roles must alternate user/assistant
|
||||
# /user/assistant/..."). The chat-completions path takes messages
|
||||
# verbatim and avoids that template gauntlet.
|
||||
"notes": "Self-hosted vLLM server. Always routed to /v1/chat/completions.",
|
||||
# Surfaced through the frontend's CUSTOM_PROVIDER_PRESETS, not the
|
||||
# /api/providers/registry dropdown — see list_available_providers.
|
||||
"hidden": True,
|
||||
},
|
||||
"openrouter": {
|
||||
"display_name": "OpenRouter",
|
||||
"base_url": "https://openrouter.ai/api/v1",
|
||||
# Curated list for Studio's picker (explicitly locked, not live /models).
|
||||
"default_models": [
|
||||
"openrouter/free",
|
||||
"openai/gpt-4o",
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
"google/gemini-2.5-flash",
|
||||
"mistralai/mistral-large-2411",
|
||||
"deepseek/deepseek-r1",
|
||||
"mistralai/mistral-small-3.1-24b-instruct",
|
||||
"perceptron/perceptron-mk1",
|
||||
"inclusionai/ring-2.6-1t:free",
|
||||
"google/gemini-3.1-flash-lite",
|
||||
"baidu/cobuddy:free",
|
||||
"openai/gpt-chat-latest",
|
||||
"x-ai/grok-4.3",
|
||||
"ibm-granite/granite-4.1-8b",
|
||||
"openrouter/owl-alpha",
|
||||
"poolside/laguna-xs.2:free",
|
||||
"~google/gemini-pro-latest",
|
||||
"~moonshotai/kimi-latest",
|
||||
],
|
||||
"supports_streaming": True,
|
||||
"supports_vision": True,
|
||||
"supports_tool_calling": True,
|
||||
"auth_header": "Authorization",
|
||||
"auth_prefix": "Bearer ",
|
||||
"extra_headers": {
|
||||
"HTTP-Referer": "https://unsloth.ai",
|
||||
"X-Title": "Unsloth Studio",
|
||||
},
|
||||
"notes": "Unified gateway to 300+ models across all major providers. HTTP-Referer and X-Title headers sent for attribution.",
|
||||
"model_list_mode": "curated",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_provider_info(provider_type: str) -> dict[str, Any] | None:
|
||||
"""Return the registry entry for a provider type, or None if unknown."""
|
||||
return PROVIDER_REGISTRY.get(provider_type)
|
||||
|
||||
|
||||
def get_base_url(provider_type: str) -> str | None:
|
||||
"""Return the default base URL for a provider type."""
|
||||
info = PROVIDER_REGISTRY.get(provider_type)
|
||||
return info["base_url"] if info else None
|
||||
|
||||
|
||||
def list_available_providers() -> list[dict[str, Any]]:
|
||||
"""Return all registered providers (for the /registry endpoint).
|
||||
|
||||
Hidden entries (``"hidden": True``) are filtered out — they exist in the
|
||||
registry only for backend lookups (e.g. ``supports_vision`` for vLLM) and
|
||||
are surfaced in the frontend via ``CUSTOM_PROVIDER_PRESETS`` instead of
|
||||
the cloud-provider dropdown.
|
||||
"""
|
||||
result = []
|
||||
for provider_type, info in PROVIDER_REGISTRY.items():
|
||||
if info.get("hidden"):
|
||||
continue
|
||||
result.append(
|
||||
{
|
||||
"provider_type": provider_type,
|
||||
"display_name": info["display_name"],
|
||||
"base_url": info["base_url"],
|
||||
"default_models": info["default_models"],
|
||||
"supports_streaming": info["supports_streaming"],
|
||||
"supports_vision": info.get("supports_vision", False),
|
||||
"supports_tool_calling": info.get("supports_tool_calling", False),
|
||||
"model_list_mode": info.get("model_list_mode", "remote"),
|
||||
}
|
||||
)
|
||||
return result
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -648,6 +648,36 @@ def run_inference_process(
|
|||
os.environ["HF_HUB_DISABLE_XET"] = "1"
|
||||
logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)")
|
||||
|
||||
# Offline auto-detect: skip 25s of hf_hub_download retries per file
|
||||
# if DNS is dead; cached files resolve instantly under HF_HUB_OFFLINE=1.
|
||||
# Scope is this subprocess only -- orchestrator spawns a fresh worker
|
||||
# per load (see core/inference/orchestrator.py), so the env cannot
|
||||
# persist across loads.
|
||||
if "HF_HUB_OFFLINE" not in os.environ:
|
||||
import socket as _socket
|
||||
import threading as _threading
|
||||
|
||||
# Probe on a daemon thread so concurrent sockets in the parent
|
||||
# interpreter are not affected by socket.setdefaulttimeout.
|
||||
_result: list = [None]
|
||||
|
||||
def _probe() -> None:
|
||||
try:
|
||||
_socket.gethostbyname("huggingface.co")
|
||||
_result[0] = False
|
||||
except Exception:
|
||||
_result[0] = True
|
||||
|
||||
_t = _threading.Thread(target = _probe, daemon = True)
|
||||
_t.start()
|
||||
_t.join(2.0)
|
||||
if _result[0] is None or _result[0] is True:
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
||||
logger.warning(
|
||||
"huggingface.co unreachable; HF_HUB_OFFLINE=1 set for this worker."
|
||||
)
|
||||
|
||||
import warnings
|
||||
from loggers.config import LogConfig
|
||||
|
||||
|
|
@ -663,6 +693,98 @@ def run_inference_process(
|
|||
|
||||
model_name = config["model_name"]
|
||||
|
||||
# ── 0. MLX fast-path — skip torch/transformers entirely ──
|
||||
backend_path = str(Path(__file__).resolve().parent.parent.parent)
|
||||
if backend_path not in sys.path:
|
||||
sys.path.insert(0, backend_path)
|
||||
|
||||
from utils.hardware import hardware as _hw
|
||||
|
||||
_hw.detect_hardware()
|
||||
if _hw.DEVICE == _hw.DeviceType.MLX:
|
||||
try:
|
||||
_activate_transformers_version(model_name)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
from core.inference.mlx_inference import MLXInferenceBackend
|
||||
|
||||
backend = MLXInferenceBackend()
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{"type": "status", "message": "Loading model...", "ts": time.time()},
|
||||
)
|
||||
_handle_load(backend, config, resp_queue)
|
||||
except Exception as exc:
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "error",
|
||||
"error": f"MLX inference init failed: {exc}",
|
||||
"stack": traceback.format_exc(limit = 20),
|
||||
"ts": time.time(),
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
# Enter same command loop as GPU path
|
||||
logger.info("MLX inference subprocess ready, entering command loop")
|
||||
while True:
|
||||
try:
|
||||
cmd = cmd_queue.get(timeout = 1.0)
|
||||
except _queue.Empty:
|
||||
continue
|
||||
except (EOFError, OSError):
|
||||
return
|
||||
if cmd is None:
|
||||
continue
|
||||
cmd_type = cmd.get("type", "")
|
||||
try:
|
||||
if cmd_type == "generate":
|
||||
cancel_event.clear()
|
||||
_handle_generate(backend, cmd, resp_queue, cancel_event)
|
||||
elif cmd_type == "load":
|
||||
if backend.active_model_name:
|
||||
backend.unload_model(backend.active_model_name)
|
||||
_handle_load(backend, cmd, resp_queue)
|
||||
elif cmd_type == "unload":
|
||||
_handle_unload(backend, cmd, resp_queue)
|
||||
elif cmd_type == "cancel":
|
||||
cancel_event.set()
|
||||
elif cmd_type == "reset":
|
||||
cancel_event.set()
|
||||
backend.reset_generation_state()
|
||||
_send_response(resp_queue, {"type": "reset_ack", "ts": time.time()})
|
||||
elif cmd_type == "status":
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "status_response",
|
||||
"active_model": backend.active_model_name,
|
||||
"models": {
|
||||
k: {kk: vv for kk, vv in v.items() if kk != "model"}
|
||||
for k, v in backend.models.items()
|
||||
},
|
||||
"loading": list(backend.loading_models),
|
||||
"ts": time.time(),
|
||||
},
|
||||
)
|
||||
elif cmd_type == "shutdown":
|
||||
return
|
||||
except Exception as exc:
|
||||
logger.error("MLX command error (%s): %s", cmd_type, exc)
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "gen_error" if cmd_type == "generate" else "error",
|
||||
"request_id": cmd.get("request_id"),
|
||||
"error": str(exc),
|
||||
"stack": traceback.format_exc(limit = 20),
|
||||
"ts": time.time(),
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
# ── 1. Activate correct transformers version BEFORE any ML imports ──
|
||||
try:
|
||||
_activate_transformers_version(model_name)
|
||||
|
|
|
|||
|
|
@ -59,9 +59,12 @@ from dataclasses import dataclass
|
|||
import pandas as pd
|
||||
from datasets import Dataset, load_dataset
|
||||
|
||||
from core.inference.llama_cpp import _hf_offline_if_dns_dead
|
||||
from utils.models import is_vision_model, detect_audio_type
|
||||
from utils.models.model_config import _env_offline
|
||||
from utils.datasets import format_and_template_dataset
|
||||
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
|
||||
from utils.datasets.raw_text import prepare_raw_text_dataset
|
||||
from utils.paths import (
|
||||
ensure_dir,
|
||||
resolve_dataset_path,
|
||||
|
|
@ -125,6 +128,7 @@ class UnslothTrainer:
|
|||
self.load_in_4bit = True # Track quantization mode for metadata
|
||||
|
||||
# Model state tracking
|
||||
self.is_cpt = False # Set to True for Continued Pretraining
|
||||
self.is_vlm = False
|
||||
self.is_audio = False
|
||||
self.is_audio_vlm = (
|
||||
|
|
@ -615,7 +619,8 @@ class UnslothTrainer:
|
|||
|
||||
# Proactive gated-model check: verify access BEFORE from_pretrained.
|
||||
# Catches ALL gated/private models (text, vision, audio) globally.
|
||||
if "/" in model_name: # Only check HF repo IDs, not local paths
|
||||
# Skip when offline -- from_pretrained will use the cache.
|
||||
if "/" in model_name and not _env_offline():
|
||||
try:
|
||||
from huggingface_hub import model_info as hf_model_info
|
||||
|
||||
|
|
@ -925,6 +930,7 @@ class UnslothTrainer:
|
|||
use_gradient_checkpointing: str = "unsloth",
|
||||
use_rslora: bool = False,
|
||||
use_loftq: bool = False,
|
||||
modules_to_save: list = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Prepare model for training (with optional LoRA).
|
||||
|
|
@ -1121,11 +1127,14 @@ class UnslothTrainer:
|
|||
loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
|
||||
if use_loftq
|
||||
else None,
|
||||
modules_to_save = modules_to_save,
|
||||
)
|
||||
else:
|
||||
# Text model LoRA
|
||||
logger.info(f"Text model LoRA configuration:")
|
||||
logger.info(f" - Target modules: {target_modules}\n")
|
||||
if modules_to_save:
|
||||
logger.info(f" - Modules to save: {modules_to_save}\n")
|
||||
|
||||
self.model = FastLanguageModel.get_peft_model(
|
||||
self.model,
|
||||
|
|
@ -1140,6 +1149,7 @@ class UnslothTrainer:
|
|||
loftq_config = {"loftq_bits": 4, "loftq_iter": 1}
|
||||
if use_loftq
|
||||
else None,
|
||||
modules_to_save = modules_to_save,
|
||||
)
|
||||
|
||||
# Check if stopped during LoRA preparation
|
||||
|
|
@ -2342,6 +2352,7 @@ class UnslothTrainer:
|
|||
eval_steps: float = 0.00,
|
||||
dataset_slice_start: int = None,
|
||||
dataset_slice_end: int = None,
|
||||
is_cpt: bool = False,
|
||||
) -> Optional[tuple]:
|
||||
"""
|
||||
Load and prepare dataset for training.
|
||||
|
|
@ -2360,6 +2371,35 @@ class UnslothTrainer:
|
|||
False # True if eval comes from a separate HF split
|
||||
)
|
||||
eval_enabled = eval_steps is not None and eval_steps > 0
|
||||
raw_text_mode = is_cpt or format_type == "raw"
|
||||
|
||||
def _raw_mode_label() -> str:
|
||||
return "CPT" if is_cpt else "raw text"
|
||||
|
||||
def _apply_raw_text_prep(ds: Dataset, split_name: str) -> Dataset:
|
||||
try:
|
||||
result = prepare_raw_text_dataset(
|
||||
ds,
|
||||
mode_label = _raw_mode_label(),
|
||||
split_name = split_name,
|
||||
eos_token = getattr(self.tokenizer, "eos_token", None),
|
||||
append_eos = True,
|
||||
)
|
||||
except ValueError as exc:
|
||||
error_msg = str(exc)
|
||||
logger.error(error_msg)
|
||||
self._update_progress(error = error_msg)
|
||||
raise
|
||||
|
||||
for notice in result.notices:
|
||||
if notice.level == "warning":
|
||||
logger.warning(notice.message)
|
||||
if notice.update_status:
|
||||
self._update_progress(status_message = notice.message)
|
||||
else:
|
||||
logger.info(f"{notice.message}\n")
|
||||
|
||||
return result.dataset
|
||||
|
||||
if local_datasets:
|
||||
# Load local datasets using load_dataset() so the result is
|
||||
|
|
@ -2534,6 +2574,48 @@ class UnslothTrainer:
|
|||
processed = self._preprocess_dac_dataset(dataset, custom_format_mapping)
|
||||
return ({"dataset": processed, "final_format": "audio_dac"}, None)
|
||||
|
||||
# ========== RAW TEXT BYPASS ==========
|
||||
if raw_text_mode:
|
||||
logger.info(
|
||||
f"{_raw_mode_label().capitalize()} mode: bypassing chat template, "
|
||||
"using raw text\n"
|
||||
)
|
||||
dataset = _apply_raw_text_prep(dataset, "train")
|
||||
if has_separate_eval_source and eval_dataset is not None:
|
||||
eval_dataset = _apply_raw_text_prep(eval_dataset, "eval")
|
||||
|
||||
dataset_info = {
|
||||
"dataset": dataset,
|
||||
"detected_format": "raw_text",
|
||||
"final_format": "raw_text",
|
||||
"success": True,
|
||||
}
|
||||
|
||||
if has_separate_eval_source and eval_dataset is not None:
|
||||
logger.info(
|
||||
f"{_raw_mode_label().capitalize()}: eval dataset "
|
||||
f"({len(eval_dataset)} rows) kept as raw text\n"
|
||||
)
|
||||
elif eval_enabled and not has_separate_eval_source:
|
||||
split_result = self._resolve_eval_split_from_dataset(dataset)
|
||||
if split_result is not None:
|
||||
train_portion, eval_dataset = split_result
|
||||
dataset_info["dataset"] = train_portion
|
||||
|
||||
train_dataset = dataset_info["dataset"]
|
||||
n = len(train_dataset) if hasattr(train_dataset, "__len__") else None
|
||||
n_display = f"{n:,}" if isinstance(n, int) else "streaming"
|
||||
self._update_progress(
|
||||
status_message = f"Dataset ready ({n_display} samples, raw text)"
|
||||
)
|
||||
logger.info(f"Raw-text dataset ready ({n_display} samples)\n")
|
||||
|
||||
if "text" not in train_dataset.column_names:
|
||||
raise ValueError(
|
||||
f"Raw-text dataset missing 'text' column: {train_dataset.column_names}"
|
||||
)
|
||||
return (dataset_info, eval_dataset)
|
||||
|
||||
elif self.is_audio_vlm:
|
||||
formatted = self._format_audio_vlm_dataset(
|
||||
dataset, custom_format_mapping
|
||||
|
|
@ -2676,6 +2758,7 @@ class UnslothTrainer:
|
|||
output_dir: str | None = None,
|
||||
num_epochs: int = 3,
|
||||
learning_rate: float = 2e-4,
|
||||
embedding_learning_rate: float | None = None,
|
||||
batch_size: int = 2,
|
||||
gradient_accumulation_steps: int = 4,
|
||||
warmup_steps: int = None,
|
||||
|
|
@ -2728,6 +2811,7 @@ class UnslothTrainer:
|
|||
"output_dir": output_dir,
|
||||
"num_epochs": num_epochs,
|
||||
"learning_rate": learning_rate,
|
||||
"embedding_learning_rate": embedding_learning_rate,
|
||||
"batch_size": batch_size,
|
||||
"gradient_accumulation_steps": gradient_accumulation_steps,
|
||||
"warmup_steps": warmup_steps,
|
||||
|
|
@ -2945,6 +3029,13 @@ class UnslothTrainer:
|
|||
|
||||
logger.info("Configuring data collator...\n")
|
||||
|
||||
dataset_final_format = (
|
||||
str(dataset.get("final_format", "")).lower()
|
||||
if isinstance(dataset, dict)
|
||||
else ""
|
||||
)
|
||||
raw_text_mode = dataset_final_format == "raw_text"
|
||||
|
||||
data_collator = None # Default to built-in data collator
|
||||
if is_deepseek_ocr:
|
||||
# Special DeepSeek OCR collator - auto-install if needed
|
||||
|
|
@ -2984,7 +3075,7 @@ class UnslothTrainer:
|
|||
self._update_progress(error = error_msg, is_training = False)
|
||||
return
|
||||
|
||||
elif self.is_audio_vlm:
|
||||
elif self.is_audio_vlm and not raw_text_mode:
|
||||
# Audio VLM collator (e.g. Gemma 3N with audio data)
|
||||
# Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook
|
||||
logger.info("Configuring audio VLM data collator...\n")
|
||||
|
|
@ -3026,7 +3117,7 @@ class UnslothTrainer:
|
|||
data_collator = audio_vlm_collate_fn
|
||||
logger.info("Audio VLM data collator configured\n")
|
||||
|
||||
elif self.is_vlm:
|
||||
elif self.is_vlm and not raw_text_mode:
|
||||
# Standard VLM collator (images)
|
||||
logger.info("Using UnslothVisionDataCollator for vision model\n")
|
||||
from unsloth.trainer import UnslothVisionDataCollator
|
||||
|
|
@ -3120,6 +3211,9 @@ class UnslothTrainer:
|
|||
if eval_steps_val > 0:
|
||||
config_args["eval_strategy"] = "steps"
|
||||
config_args["eval_steps"] = eval_steps_val
|
||||
config_args["per_device_eval_batch_size"] = config_args[
|
||||
"per_device_train_batch_size"
|
||||
]
|
||||
logger.info(
|
||||
f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n"
|
||||
)
|
||||
|
|
@ -3137,8 +3231,9 @@ class UnslothTrainer:
|
|||
optim_value = training_args.get("optim", "adamw_8bit")
|
||||
lr_scheduler_type_value = training_args.get("lr_scheduler_type", "linear")
|
||||
|
||||
if self.is_vlm or self.is_audio_vlm:
|
||||
if (self.is_vlm or self.is_audio_vlm) and not raw_text_mode:
|
||||
# Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns)
|
||||
# Raw-text runs on VLM-capable models are routed to the text path below.
|
||||
label = "audio VLM" if self.is_audio_vlm else "vision"
|
||||
logger.info(f"Configuring {label} model training parameters\n")
|
||||
# Use provided values or defaults for vision models
|
||||
|
|
@ -3160,7 +3255,14 @@ class UnslothTrainer:
|
|||
}
|
||||
)
|
||||
else:
|
||||
logger.info("Configuring text model training parameters\n")
|
||||
is_cpt = training_args.get("is_cpt", False)
|
||||
self.is_cpt = is_cpt
|
||||
if is_cpt:
|
||||
logger.info("Configuring Continued Pretraining (CPT) parameters\n")
|
||||
elif raw_text_mode:
|
||||
logger.info("Configuring raw-text training parameters\n")
|
||||
else:
|
||||
logger.info("Configuring text model training parameters\n")
|
||||
config_args.update(
|
||||
{
|
||||
"optim": optim_value,
|
||||
|
|
@ -3189,9 +3291,10 @@ class UnslothTrainer:
|
|||
|
||||
logger.info("Training configuration prepared\n")
|
||||
# ========== TRAINER INITIALIZATION ==========
|
||||
if self.is_audio_vlm:
|
||||
if self.is_audio_vlm and not raw_text_mode:
|
||||
# Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset
|
||||
# Notebook uses processing_class=processor.tokenizer (text tokenizer only)
|
||||
# Raw-text runs are routed to the text path below.
|
||||
train_dataset = (
|
||||
dataset if isinstance(dataset, Dataset) else dataset["dataset"]
|
||||
)
|
||||
|
|
@ -3210,8 +3313,9 @@ class UnslothTrainer:
|
|||
if eval_dataset is not None:
|
||||
trainer_kwargs["eval_dataset"] = eval_dataset
|
||||
self.trainer = SFTTrainer(**trainer_kwargs)
|
||||
elif self.is_vlm:
|
||||
elif self.is_vlm and not raw_text_mode:
|
||||
# Image VLM: dataset is dict wrapper from format_and_template_dataset
|
||||
# Raw-text runs are routed to the text path below.
|
||||
train_dataset = (
|
||||
dataset["dataset"] if isinstance(dataset, dict) else dataset
|
||||
)
|
||||
|
|
@ -3242,16 +3346,48 @@ class UnslothTrainer:
|
|||
)
|
||||
sft_tokenizer = self.tokenizer.tokenizer
|
||||
|
||||
trainer_kwargs = {
|
||||
"model": self.model,
|
||||
"tokenizer": sft_tokenizer,
|
||||
"train_dataset": dataset["dataset"],
|
||||
"data_collator": data_collator,
|
||||
"args": SFTConfig(**config_args),
|
||||
}
|
||||
if eval_dataset is not None:
|
||||
trainer_kwargs["eval_dataset"] = eval_dataset
|
||||
self.trainer = SFTTrainer(**trainer_kwargs)
|
||||
if is_cpt:
|
||||
try:
|
||||
from unsloth import (
|
||||
UnslothTrainer as _UnslothCPTTrainer,
|
||||
UnslothTrainingArguments as _UnslothTrainingArguments,
|
||||
)
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"CPT requires a newer Unsloth install that exports "
|
||||
"`UnslothTrainer` and `UnslothTrainingArguments` "
|
||||
"(for embedding_learning_rate support). "
|
||||
"Upgrade with: `pip install -U unsloth unsloth_zoo`."
|
||||
) from exc
|
||||
|
||||
embedding_lr = training_args.get("embedding_learning_rate")
|
||||
logger.info(
|
||||
f"CPT: using UnslothTrainer with embedding_learning_rate={embedding_lr}\n"
|
||||
)
|
||||
trainer_kwargs = {
|
||||
"model": self.model,
|
||||
"tokenizer": sft_tokenizer,
|
||||
"train_dataset": dataset["dataset"],
|
||||
"data_collator": data_collator,
|
||||
"args": _UnslothTrainingArguments(
|
||||
embedding_learning_rate = embedding_lr,
|
||||
**config_args,
|
||||
),
|
||||
}
|
||||
if eval_dataset is not None:
|
||||
trainer_kwargs["eval_dataset"] = eval_dataset
|
||||
self.trainer = _UnslothCPTTrainer(**trainer_kwargs)
|
||||
else:
|
||||
trainer_kwargs = {
|
||||
"model": self.model,
|
||||
"tokenizer": sft_tokenizer,
|
||||
"train_dataset": dataset["dataset"],
|
||||
"data_collator": data_collator,
|
||||
"args": SFTConfig(**config_args),
|
||||
}
|
||||
if eval_dataset is not None:
|
||||
trainer_kwargs["eval_dataset"] = eval_dataset
|
||||
self.trainer = SFTTrainer(**trainer_kwargs)
|
||||
# Restore the full processor as processing_class so checkpoint
|
||||
# saves include preprocessor_config.json (needed for GGUF export).
|
||||
if sft_tokenizer is not self.tokenizer:
|
||||
|
|
@ -3260,19 +3396,32 @@ class UnslothTrainer:
|
|||
|
||||
# ========== TRAIN ON RESPONSES ONLY ==========
|
||||
# Determine if we should train on responses only
|
||||
# Raw-text datasets always train on all tokens.
|
||||
instruction_part = None
|
||||
response_part = None
|
||||
train_on_responses_enabled = training_args.get(
|
||||
"train_on_completions", False
|
||||
is_cpt = training_args.get("is_cpt", False)
|
||||
train_on_responses_enabled = (
|
||||
False
|
||||
if (is_cpt or raw_text_mode)
|
||||
else training_args.get("train_on_completions", False)
|
||||
)
|
||||
|
||||
if is_cpt:
|
||||
logger.info(
|
||||
"CPT mode: skipping train_on_responses_only — training on all tokens\n"
|
||||
)
|
||||
elif raw_text_mode:
|
||||
logger.info(
|
||||
"Raw-text mode: skipping train_on_responses_only — training on all tokens\n"
|
||||
)
|
||||
|
||||
# DeepSeek OCR handles this internally in its collator, so skip
|
||||
# Audio VLM handles label masking in its collator, so skip
|
||||
if (
|
||||
train_on_responses_enabled
|
||||
and not self.is_audio_vlm
|
||||
and not self.is_audio
|
||||
and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca")
|
||||
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
|
||||
):
|
||||
try:
|
||||
logger.info("Configuring train on responses only...\n")
|
||||
|
|
@ -3318,7 +3467,7 @@ class UnslothTrainer:
|
|||
and response_part
|
||||
and not self.is_audio_vlm
|
||||
and not self.is_audio
|
||||
and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca")
|
||||
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
|
||||
):
|
||||
try:
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
|
|
@ -3451,7 +3600,9 @@ class UnslothTrainer:
|
|||
config = json.load(f)
|
||||
|
||||
# Determine the training method
|
||||
if self.load_in_4bit:
|
||||
if self.is_cpt:
|
||||
method = "CPT"
|
||||
elif self.load_in_4bit:
|
||||
method = "qlora"
|
||||
else:
|
||||
method = "lora"
|
||||
|
|
|
|||
|
|
@ -17,7 +17,10 @@ Pattern follows core/data_recipe/jobs/manager.py.
|
|||
import json as _json
|
||||
import math
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
import structlog
|
||||
|
|
@ -33,9 +36,56 @@ from utils.native_path_leases import (
|
|||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
from utils.paths import outputs_root
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
_HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$")
|
||||
|
||||
|
||||
def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
|
||||
"""Remove only HF Trainer ``tmp-checkpoint-<step>/`` partials after a cancel.
|
||||
|
||||
Completed ``checkpoint-<int>/`` dirs and any non-numeric-suffix tmp dir
|
||||
are user-owned and survive. Symlinked output_dir / children are skipped
|
||||
so containment cannot be bypassed.
|
||||
"""
|
||||
out = Path(output_dir)
|
||||
if not out.exists() or not out.is_dir() or out.is_symlink():
|
||||
return
|
||||
try:
|
||||
out_real = out.resolve()
|
||||
out_root_real = Path(outputs_root()).resolve()
|
||||
except OSError:
|
||||
return
|
||||
try:
|
||||
out_real.relative_to(out_root_real)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"Skipping checkpoint cleanup - %s is not under outputs_root %s",
|
||||
out_real,
|
||||
out_root_real,
|
||||
)
|
||||
return
|
||||
removed = 0
|
||||
for entry in out.iterdir():
|
||||
if not entry.is_dir() or entry.is_symlink():
|
||||
continue
|
||||
if not _HF_TMP_CHECKPOINT_RE.match(entry.name):
|
||||
continue
|
||||
try:
|
||||
shutil.rmtree(entry, ignore_errors = False)
|
||||
removed += 1
|
||||
except OSError as exc:
|
||||
logger.warning("Could not remove %s: %s", entry, exc)
|
||||
logger.info(
|
||||
"Cancelled-run cleanup removed %d in-flight tmp-checkpoint dir(s) under %s",
|
||||
removed,
|
||||
out,
|
||||
)
|
||||
|
||||
|
||||
_CTX = mp.get_context("spawn")
|
||||
|
||||
# Plot styling constants
|
||||
|
|
@ -62,6 +112,7 @@ class TrainingProgress:
|
|||
grad_norm: Optional[float] = None
|
||||
num_tokens: Optional[int] = None
|
||||
eval_loss: Optional[float] = None
|
||||
peak_memory_gb: Optional[float] = None
|
||||
|
||||
|
||||
class TrainingBackend:
|
||||
|
|
@ -158,6 +209,7 @@ class TrainingBackend:
|
|||
"is_embedding": kwargs.get("is_embedding", False),
|
||||
"num_epochs": kwargs.get("num_epochs", 3),
|
||||
"learning_rate": kwargs.get("learning_rate", "2e-4"),
|
||||
"embedding_learning_rate": kwargs.get("embedding_learning_rate"),
|
||||
"batch_size": kwargs.get("batch_size", 2),
|
||||
"gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4),
|
||||
"warmup_steps": kwargs.get("warmup_steps"),
|
||||
|
|
@ -165,6 +217,7 @@ class TrainingBackend:
|
|||
"max_steps": kwargs.get("max_steps", 0),
|
||||
"save_steps": kwargs.get("save_steps", 0),
|
||||
"weight_decay": kwargs.get("weight_decay", 0.001),
|
||||
"max_grad_norm": kwargs.get("max_grad_norm", 0.0),
|
||||
"random_seed": kwargs.get("random_seed", 3407),
|
||||
"packing": kwargs.get("packing", False),
|
||||
"optim": kwargs.get("optim", "adamw_8bit"),
|
||||
|
|
@ -194,26 +247,33 @@ class TrainingBackend:
|
|||
"gpu_ids": kwargs.get("gpu_ids"),
|
||||
}
|
||||
|
||||
# Derive load_in_4bit from training_type
|
||||
if config["training_type"] != "LoRA/QLoRA":
|
||||
# Full finetuning always runs in 16-bit. LoRA/QLoRA and CPT preserve the
|
||||
# explicit request so 4-bit adapter/raw-text runs remain possible.
|
||||
if config["training_type"] == "Full Finetuning":
|
||||
config["load_in_4bit"] = False
|
||||
|
||||
# Spawn subprocess — use locals so state is untouched on failure
|
||||
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
|
||||
kwargs.get("gpu_ids"),
|
||||
model_name = config["model_name"],
|
||||
hf_token = config["hf_token"] or None,
|
||||
training_type = config["training_type"],
|
||||
load_in_4bit = config["load_in_4bit"],
|
||||
batch_size = config.get("batch_size", 4),
|
||||
max_seq_length = config.get("max_seq_length", 2048),
|
||||
lora_rank = config.get("lora_r", 16),
|
||||
target_modules = config.get("target_modules"),
|
||||
gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
|
||||
optimizer = config.get("optim", "adamw_8bit"),
|
||||
)
|
||||
config["resolved_gpu_ids"] = resolved_gpu_ids
|
||||
config["gpu_selection"] = gpu_selection
|
||||
from utils.hardware import hardware as _hw
|
||||
|
||||
if _hw.DEVICE == _hw.DeviceType.MLX:
|
||||
config["resolved_gpu_ids"] = None
|
||||
config["gpu_selection"] = None
|
||||
else:
|
||||
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
|
||||
kwargs.get("gpu_ids"),
|
||||
model_name = config["model_name"],
|
||||
hf_token = config["hf_token"] or None,
|
||||
training_type = config["training_type"],
|
||||
load_in_4bit = config["load_in_4bit"],
|
||||
batch_size = config.get("batch_size", 4),
|
||||
max_seq_length = config.get("max_seq_length", 2048),
|
||||
lora_rank = config.get("lora_r", 16),
|
||||
target_modules = config.get("target_modules"),
|
||||
gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
|
||||
optimizer = config.get("optim", "adamw_8bit"),
|
||||
)
|
||||
config["resolved_gpu_ids"] = resolved_gpu_ids
|
||||
config["gpu_selection"] = gpu_selection
|
||||
|
||||
from .worker import run_training_process
|
||||
|
||||
|
|
@ -307,6 +367,8 @@ class TrainingBackend:
|
|||
)
|
||||
self._proc.terminate()
|
||||
proc = self._proc
|
||||
cancelled = self._cancel_requested
|
||||
output_dir = self._output_dir
|
||||
|
||||
if proc is not None:
|
||||
proc.join(timeout = 5.0)
|
||||
|
|
@ -319,6 +381,15 @@ class TrainingBackend:
|
|||
if self._pump_thread is not None and self._pump_thread.is_alive():
|
||||
self._pump_thread.join(timeout = 8.0)
|
||||
|
||||
if cancelled and output_dir:
|
||||
try:
|
||||
_cleanup_cancelled_checkpoints(output_dir)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to clean up cancelled-run checkpoints under %s",
|
||||
output_dir,
|
||||
)
|
||||
|
||||
def is_training_active(self) -> bool:
|
||||
"""Check if training is currently active."""
|
||||
with self._lock:
|
||||
|
|
@ -512,6 +583,12 @@ class TrainingBackend:
|
|||
self._progress.grad_norm = event.get("grad_norm")
|
||||
self._progress.num_tokens = event.get("num_tokens")
|
||||
self._progress.eval_loss = event.get("eval_loss")
|
||||
_peak = event.get("peak_memory_gb")
|
||||
if _peak is not None:
|
||||
try:
|
||||
self._progress.peak_memory_gb = float(_peak)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
self._progress.is_training = True
|
||||
status = event.get("status_message", "")
|
||||
if status:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -78,7 +78,7 @@ class LoggingMiddleware(BaseHTTPMiddleware):
|
|||
|
||||
|
||||
def filter_sensitive_data(logger, method_name, event_dict):
|
||||
"""Structlog processor to filter out base64 data from logs."""
|
||||
"""Structlog processor to redact native path leases from logs."""
|
||||
|
||||
def filter_value(value):
|
||||
if isinstance(value, str):
|
||||
|
|
@ -87,13 +87,7 @@ def filter_sensitive_data(logger, method_name, event_dict):
|
|||
except Exception:
|
||||
pass
|
||||
value = _NATIVE_PATH_LEASE_RE.sub(r"\1<redacted native path lease>", value)
|
||||
if (
|
||||
isinstance(value, str)
|
||||
and len(value) > 100
|
||||
and ("," in value or "/" in value)
|
||||
):
|
||||
# Likely base64 data, truncate it
|
||||
return value[:20] + "..."
|
||||
return value
|
||||
elif isinstance(value, dict):
|
||||
return {
|
||||
k: "<redacted native path lease>"
|
||||
|
|
|
|||
|
|
@ -23,12 +23,68 @@ if _backend_dir not in sys.path:
|
|||
# See: https://github.com/python/cpython/issues/102396
|
||||
import _platform_compat # noqa: F401
|
||||
|
||||
# Direct `uvicorn main:app` launches bypass run.py, so re-export here too
|
||||
# (mirrors run.py). Required BEFORE the unsloth-zoo import below, since
|
||||
# its LLAMA_CPP_DEFAULT_DIR binding is import-time.
|
||||
from utils.paths.storage_roots import studio_root as _studio_root
|
||||
|
||||
try:
|
||||
_LEGACY_STUDIO_ROOT = (_Path.home() / ".unsloth" / "studio").resolve()
|
||||
except (OSError, ValueError):
|
||||
_LEGACY_STUDIO_ROOT = _Path.home() / ".unsloth" / "studio"
|
||||
try:
|
||||
_STUDIO_ROOT_RESOLVED = _studio_root().resolve()
|
||||
except (OSError, ValueError):
|
||||
_STUDIO_ROOT_RESOLVED = _studio_root()
|
||||
if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
|
||||
if not os.environ.get("UNSLOTH_STUDIO_HOME"):
|
||||
os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED)
|
||||
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
|
||||
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp")
|
||||
|
||||
import hashlib
|
||||
import mimetypes
|
||||
import re as _re
|
||||
import shutil
|
||||
import warnings
|
||||
from contextlib import asynccontextmanager
|
||||
from importlib.metadata import PackageNotFoundError, version as package_version
|
||||
|
||||
|
||||
_STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$")
|
||||
|
||||
|
||||
def _read_studio_install_id() -> str:
|
||||
"""Per-install opaque id written by install.sh / install.ps1 at
|
||||
$STUDIO_HOME/share/studio_install_id. Returns "" when the file is
|
||||
absent (pre-PR install, fresh tree never run through the installer)
|
||||
or contains anything other than a 64-char lowercase-hex token --
|
||||
in which case /api/health emits "" and the launcher's _check_health
|
||||
falls back to the existing "no baked id, accept any healthy
|
||||
Unsloth backend" path. This intentionally replaces a previous
|
||||
sha256(resolved_install_path) so the field carries no install-path
|
||||
information for callers reaching /api/health (relevant when Studio
|
||||
is run with -H 0.0.0.0)."""
|
||||
try:
|
||||
token = (
|
||||
(_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
|
||||
)
|
||||
except (OSError, ValueError):
|
||||
return ""
|
||||
return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else ""
|
||||
|
||||
|
||||
_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()
|
||||
|
||||
|
||||
def _studio_root_id() -> str:
|
||||
"""Same-install discriminator for /api/health: a per-install opaque
|
||||
token written once by the installer and read once at module import.
|
||||
Empty when no installer-written token is present; the launcher
|
||||
contract treats "" as "no baked id, accept any healthy backend"."""
|
||||
return _STUDIO_ROOT_ID_CACHE
|
||||
|
||||
|
||||
# Fix broken Windows registry MIME types. Some Windows installs map .js to
|
||||
# "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes
|
||||
# module reads from the registry, and FastAPI/Starlette's StaticFiles uses
|
||||
|
|
@ -48,7 +104,7 @@ if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
|
|||
# warnings.filterwarnings("ignore", category=DeprecationWarning)
|
||||
# warnings.filterwarnings("ignore", module="triton.*")
|
||||
|
||||
from fastapi import Depends, FastAPI, Request
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.responses import FileResponse, HTMLResponse, Response
|
||||
|
|
@ -64,6 +120,7 @@ from routes import (
|
|||
inference_router,
|
||||
inference_studio_router,
|
||||
models_router,
|
||||
providers_router,
|
||||
training_history_router,
|
||||
training_router,
|
||||
)
|
||||
|
|
@ -79,6 +136,11 @@ import utils.hardware.hardware as _hw_module
|
|||
|
||||
from utils.cache_cleanup import clear_unsloth_compiled_cache
|
||||
from utils.native_path_leases import native_path_leases_supported
|
||||
from utils.update_status import (
|
||||
get_studio_install_source_status,
|
||||
get_studio_update_status,
|
||||
)
|
||||
from utils.studio_version import get_studio_version
|
||||
|
||||
|
||||
def get_unsloth_version() -> str:
|
||||
|
|
@ -100,6 +162,25 @@ def get_unsloth_version() -> str:
|
|||
|
||||
|
||||
UNSLOTH_VERSION = get_unsloth_version()
|
||||
STUDIO_VERSION = get_studio_version()
|
||||
|
||||
|
||||
def _load_desktop_owner() -> dict[str, str] | None:
|
||||
token = os.environ.pop("UNSLOTH_STUDIO_DESKTOP_OWNER_TOKEN", "")
|
||||
kind = os.environ.pop("UNSLOTH_STUDIO_DESKTOP_OWNER_KIND", "")
|
||||
if kind != "tauri" or not token:
|
||||
return None
|
||||
return {
|
||||
"kind": "tauri",
|
||||
"token_sha256": hashlib.sha256(token.encode("utf-8")).hexdigest(),
|
||||
}
|
||||
|
||||
|
||||
_DESKTOP_OWNER = _load_desktop_owner()
|
||||
|
||||
|
||||
def _desktop_owner() -> dict[str, str] | None:
|
||||
return _DESKTOP_OWNER
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
|
|
@ -117,6 +198,43 @@ async def lifespan(app: FastAPI):
|
|||
# Detect hardware first — sets DEVICE global used everywhere
|
||||
detect_hardware()
|
||||
|
||||
# llama.cpp probes: capability (MTP support) + freshness (release age).
|
||||
# Both cached; freshness has a 24h disk TTL.
|
||||
try:
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
from utils.llama_cpp_freshness import (
|
||||
check_prebuilt_freshness,
|
||||
format_stale_warning,
|
||||
)
|
||||
|
||||
_bin = LlamaCppBackend._find_llama_server_binary()
|
||||
_caps = LlamaCppBackend.probe_server_capabilities(_bin)
|
||||
app.state.llama_cpp_capabilities = _caps
|
||||
_freshness = check_prebuilt_freshness(_bin)
|
||||
app.state.llama_cpp_freshness = _freshness
|
||||
|
||||
import structlog as _structlog
|
||||
|
||||
_log = _structlog.get_logger(__name__)
|
||||
if _caps.get("found") and not _caps.get("supports_mtp"):
|
||||
_msg = (
|
||||
"llama.cpp prebuilt lacks MTP support "
|
||||
"(--spec-type mtp/draft-mtp). Run `unsloth studio update`. "
|
||||
"MTP GGUFs will load without speculative decoding."
|
||||
)
|
||||
_log.warning(_msg)
|
||||
print(f"WARNING: {_msg}", flush = True)
|
||||
if _freshness.get("stale"):
|
||||
_msg = format_stale_warning(_freshness)
|
||||
_log.warning(_msg)
|
||||
print(f"WARNING: {_msg}", flush = True)
|
||||
except Exception as _probe_exc:
|
||||
import structlog as _structlog
|
||||
|
||||
_structlog.get_logger(__name__).debug(
|
||||
"llama.cpp startup probes failed: %s", _probe_exc
|
||||
)
|
||||
|
||||
from storage.studio_db import cleanup_orphaned_runs
|
||||
|
||||
try:
|
||||
|
|
@ -142,6 +260,11 @@ async def lifespan(app: FastAPI):
|
|||
|
||||
threading.Thread(target = _precache, daemon = True).start()
|
||||
|
||||
# Initialize RSA key pair for API key encryption (external providers)
|
||||
from core.inference.key_exchange import init_key_pair
|
||||
|
||||
init_key_pair()
|
||||
|
||||
if storage.ensure_default_admin():
|
||||
bootstrap_pw = storage.get_bootstrap_password()
|
||||
app.state.bootstrap_password = bootstrap_pw
|
||||
|
|
@ -180,6 +303,182 @@ logger = LogConfig.setup_logging(
|
|||
|
||||
app.add_middleware(LoggingMiddleware)
|
||||
|
||||
|
||||
# Citation favicons load from www.google.com/s2/favicons; *.gstatic.com is
|
||||
# kept for legacy web-search faviconV2 paths. Everything else is same-origin.
|
||||
from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402
|
||||
from starlette.requests import Request as _StarletteRequest # noqa: E402
|
||||
|
||||
|
||||
_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
|
||||
|
||||
|
||||
def _build_csp(script_nonce: "str | None" = None) -> str:
|
||||
script_src = "script-src 'self'"
|
||||
if script_nonce:
|
||||
script_src += f" 'nonce-{script_nonce}'"
|
||||
return (
|
||||
"default-src 'self'; "
|
||||
"img-src 'self' data: blob: https://t0.gstatic.com "
|
||||
"https://t1.gstatic.com https://t2.gstatic.com "
|
||||
"https://t3.gstatic.com https://www.google.com; "
|
||||
"connect-src 'self' https://huggingface.co https://datasets-server.huggingface.co; "
|
||||
"style-src 'self' 'unsafe-inline'; "
|
||||
f"{script_src}; "
|
||||
"font-src 'self' data:; "
|
||||
"frame-ancestors 'none'; "
|
||||
"form-action 'self'; "
|
||||
"base-uri 'self'"
|
||||
)
|
||||
|
||||
|
||||
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
|
||||
"""Set baseline security headers; splice per-response inline-script nonces into CSP."""
|
||||
|
||||
async def dispatch(self, request: _StarletteRequest, call_next):
|
||||
response = await call_next(request)
|
||||
# Strip the internal nonce hand-off header so it never reaches the client.
|
||||
nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER)
|
||||
if nonce is not None:
|
||||
del response.headers[_CSP_SCRIPT_NONCE_HEADER]
|
||||
response.headers.setdefault("Content-Security-Policy", _build_csp(nonce))
|
||||
response.headers.setdefault("X-Frame-Options", "DENY")
|
||||
response.headers.setdefault("X-Content-Type-Options", "nosniff")
|
||||
response.headers.setdefault("Referrer-Policy", "no-referrer")
|
||||
response.headers.setdefault(
|
||||
"Permissions-Policy",
|
||||
"camera=(), microphone=(), geolocation=(), interest-cohort=()",
|
||||
)
|
||||
response.headers["server"] = "unsloth-studio"
|
||||
return response
|
||||
|
||||
|
||||
app.add_middleware(SecurityHeadersMiddleware)
|
||||
|
||||
|
||||
# Cap upload body on protected POSTs; default 500 MB, env-tunable.
|
||||
import json as _json_for_413 # noqa: E402
|
||||
|
||||
|
||||
_MAX_BODY_BYTES = int(os.environ.get("UNSLOTH_STUDIO_MAX_BODY_MB", "500")) * 1024 * 1024
|
||||
_BODY_PROTECTED_PREFIXES = (
|
||||
"/v1/chat/completions",
|
||||
"/v1/completions",
|
||||
"/api/inference",
|
||||
"/api/data-recipe",
|
||||
"/api/datasets",
|
||||
"/api/train",
|
||||
"/api/export",
|
||||
)
|
||||
|
||||
|
||||
async def _send_413(send, total_bytes: int) -> None:
|
||||
payload = _json_for_413.dumps(
|
||||
{
|
||||
"detail": (
|
||||
f"Request body too large "
|
||||
f"({total_bytes:,} bytes; max {_MAX_BODY_BYTES:,})."
|
||||
)
|
||||
},
|
||||
).encode("utf-8")
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 413,
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"content-length", str(len(payload)).encode("ascii")),
|
||||
],
|
||||
}
|
||||
)
|
||||
await send({"type": "http.response.body", "body": payload, "more_body": False})
|
||||
|
||||
|
||||
class MaxBodyMiddleware:
|
||||
"""Reject oversized bodies on protected POST/PUT/PATCH; raw ASGI so chunked uploads cannot bypass the cap."""
|
||||
|
||||
def __init__(self, app, max_bytes: int, protected_prefixes: tuple):
|
||||
self.app = app
|
||||
self.max_bytes = max_bytes
|
||||
self.protected_prefixes = protected_prefixes
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
method = scope.get("method", "").upper()
|
||||
path = scope.get("path", "")
|
||||
if method not in ("POST", "PUT", "PATCH") or not any(
|
||||
path.startswith(p) for p in self.protected_prefixes
|
||||
):
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
declared = None
|
||||
for name, value in scope.get("headers", []):
|
||||
if name == b"content-length":
|
||||
try:
|
||||
declared = int(value.decode("latin-1"))
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
declared = None
|
||||
break
|
||||
if declared is not None and declared > self.max_bytes:
|
||||
await _send_413(send, declared)
|
||||
return
|
||||
|
||||
chunks: list = []
|
||||
total = 0
|
||||
while True:
|
||||
msg = await receive()
|
||||
mtype = msg.get("type")
|
||||
if mtype == "http.disconnect":
|
||||
return
|
||||
if mtype != "http.request":
|
||||
# Mid-stream unexpected frame: forwarding would corrupt downstream.
|
||||
return
|
||||
body = msg.get("body", b"") or b""
|
||||
if body:
|
||||
total += len(body)
|
||||
if total > self.max_bytes:
|
||||
await _send_413(send, total)
|
||||
return
|
||||
chunks.append(body)
|
||||
if not msg.get("more_body", False):
|
||||
break
|
||||
|
||||
replayed = {"sent": False}
|
||||
|
||||
async def replay_receive():
|
||||
if not replayed["sent"]:
|
||||
replayed["sent"] = True
|
||||
return {
|
||||
"type": "http.request",
|
||||
"body": b"".join(chunks),
|
||||
"more_body": False,
|
||||
}
|
||||
# After replay, fall through so http.disconnect still propagates.
|
||||
return await receive()
|
||||
|
||||
await self.app(scope, replay_receive, send)
|
||||
|
||||
|
||||
app.add_middleware(
|
||||
MaxBodyMiddleware,
|
||||
max_bytes = _MAX_BODY_BYTES,
|
||||
protected_prefixes = _BODY_PROTECTED_PREFIXES,
|
||||
)
|
||||
|
||||
|
||||
from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402
|
||||
|
||||
|
||||
@app.get("/recipes", include_in_schema = False)
|
||||
@app.get("/recipes/{rest:path}", include_in_schema = False)
|
||||
async def _recipes_redirect(rest: str = ""):
|
||||
target = "/data-recipes" + (("/" + rest) if rest else "")
|
||||
return _RedirectResponse(url = target, status_code = 308)
|
||||
|
||||
|
||||
# CORS middleware
|
||||
_api_only = os.environ.get("UNSLOTH_API_ONLY") == "1"
|
||||
_cors_origins = ["*"]
|
||||
|
|
@ -219,6 +518,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["
|
|||
# so external tools (Open WebUI, SillyTavern, etc.) can use the
|
||||
# standard /v1/chat/completions path.
|
||||
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
|
||||
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
|
||||
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
|
||||
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
|
||||
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
|
||||
|
|
@ -231,22 +531,69 @@ app.include_router(
|
|||
|
||||
|
||||
@app.get("/api/health")
|
||||
async def health_check():
|
||||
"""Health check endpoint"""
|
||||
platform_map = {"darwin": "mac", "win32": "windows", "linux": "linux"}
|
||||
device_type = platform_map.get(sys.platform, sys.platform)
|
||||
async def health_check(request: Request):
|
||||
"""Liveness plus launcher capability bits; install fingerprint gated on a valid bearer.
|
||||
|
||||
return {
|
||||
Unauthenticated callers (Tauri watchdog, frontend bootstrap polls) need
|
||||
``service`` / ``studio_root_id`` / ``chat_only`` / ``desktop_*`` / ``native_path_leases_supported``
|
||||
to (a) re-adopt a sibling backend across restarts and (b) gate UI surfaces
|
||||
before any token is available. None of those leak install path or version.
|
||||
``version`` / ``studio_version`` / ``device_type`` still require a bearer
|
||||
because they fingerprint the host.
|
||||
"""
|
||||
base = {
|
||||
"status": "healthy",
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"service": "Unsloth UI Backend",
|
||||
"version": UNSLOTH_VERSION,
|
||||
"device_type": device_type,
|
||||
"chat_only": _hw_module.CHAT_ONLY,
|
||||
"desktop_protocol_version": 1,
|
||||
"desktop_manageability_version": 1,
|
||||
"supports_desktop_auth": True,
|
||||
"supports_desktop_backend_ownership": True,
|
||||
# Opaque per-install id; launchers reject sibling Studios on the same port.
|
||||
"studio_root_id": _studio_root_id(),
|
||||
"native_path_leases_supported": native_path_leases_supported(),
|
||||
**({"desktop_owner": owner} if (owner := _desktop_owner()) else {}),
|
||||
}
|
||||
auth = request.headers.get("authorization", "")
|
||||
if not auth.lower().startswith("bearer "):
|
||||
return base
|
||||
try:
|
||||
from auth.authentication import get_current_subject as _gcs
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
|
||||
creds = HTTPAuthorizationCredentials(
|
||||
scheme = "Bearer", credentials = auth.split(" ", 1)[1]
|
||||
)
|
||||
# Must await: a bare coroutine is truthy and would skip the auth check.
|
||||
subject = await _gcs(creds)
|
||||
except HTTPException:
|
||||
return base
|
||||
except Exception:
|
||||
return base
|
||||
if not subject:
|
||||
return base
|
||||
|
||||
platform_map = {"darwin": "mac", "win32": "windows", "linux": "linux"}
|
||||
device_type = platform_map.get(sys.platform, sys.platform)
|
||||
return {
|
||||
**base,
|
||||
"version": UNSLOTH_VERSION,
|
||||
"studio_version": STUDIO_VERSION,
|
||||
"device_type": device_type,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/studio/install-source")
|
||||
def studio_install_source(_current_subject: str = Depends(get_current_subject)):
|
||||
"""Return source-aware install metadata without remote update checks."""
|
||||
return get_studio_install_source_status(UNSLOTH_VERSION)
|
||||
|
||||
|
||||
@app.get("/api/studio/update-status")
|
||||
def studio_update_status(_current_subject: str = Depends(get_current_subject)):
|
||||
"""Return source-aware manual update status for browser-served Studio."""
|
||||
return get_studio_update_status(UNSLOTH_VERSION)
|
||||
|
||||
|
||||
@app.post("/api/shutdown")
|
||||
|
|
@ -278,8 +625,17 @@ async def shutdown_server(
|
|||
|
||||
|
||||
@app.get("/api/system")
|
||||
async def get_system_info():
|
||||
"""Get system information"""
|
||||
async def get_system_info(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Get system information.
|
||||
|
||||
Gated behind auth: the response includes platform, Python version,
|
||||
GPU name, memory total, and ML package set -- enough to fingerprint
|
||||
a host. Studio's chat-only-mode design assumes only the local user
|
||||
reaches /api/system; in -H 0.0.0.0 / Colab / Tauri-relayed setups
|
||||
that assumption breaks unless we require a bearer.
|
||||
"""
|
||||
import platform
|
||||
import psutil
|
||||
from utils.hardware import get_device
|
||||
|
|
@ -319,8 +675,14 @@ async def get_gpu_visibility(
|
|||
|
||||
|
||||
@app.get("/api/system/hardware")
|
||||
async def get_hardware_info():
|
||||
"""Return GPU name, total VRAM, and key ML package versions."""
|
||||
async def get_hardware_info(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Return GPU name, total VRAM, and key ML package versions.
|
||||
|
||||
Gated behind auth alongside /api/system -- same fingerprinting
|
||||
concern. /api/system/gpu-visibility is also auth-gated already.
|
||||
"""
|
||||
from utils.hardware import get_gpu_summary, get_package_versions
|
||||
|
||||
return {
|
||||
|
|
@ -348,21 +710,22 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes:
|
|||
return html.encode("utf-8")
|
||||
|
||||
|
||||
def _inject_bootstrap(html_bytes: bytes, app: FastAPI) -> bytes:
|
||||
"""Inject bootstrap credentials into HTML when password change is required.
|
||||
def _inject_bootstrap(html_bytes: bytes, app: FastAPI):
|
||||
"""Inject bootstrap credentials when password change is pending.
|
||||
|
||||
The script tag is only injected while the default admin account still
|
||||
has ``must_change_password=True``. Once the user changes the password
|
||||
the HTML is served clean — no credentials leak.
|
||||
Returns ``(html_bytes, script_nonce_or_None)``. Callers must forward
|
||||
the nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so the inline script is
|
||||
not blocked by CSP.
|
||||
"""
|
||||
import json as _json
|
||||
import secrets as _secrets
|
||||
|
||||
if not storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME):
|
||||
return html_bytes
|
||||
return html_bytes, None
|
||||
|
||||
bootstrap_pw = getattr(app.state, "bootstrap_password", None)
|
||||
if not bootstrap_pw:
|
||||
return html_bytes
|
||||
return html_bytes, None
|
||||
|
||||
payload = _json.dumps(
|
||||
{
|
||||
|
|
@ -370,10 +733,11 @@ def _inject_bootstrap(html_bytes: bytes, app: FastAPI) -> bytes:
|
|||
"password": bootstrap_pw,
|
||||
}
|
||||
)
|
||||
tag = f"<script>window.__UNSLOTH_BOOTSTRAP__={payload}</script>"
|
||||
nonce = _secrets.token_urlsafe(16)
|
||||
tag = f'<script nonce="{nonce}">window.__UNSLOTH_BOOTSTRAP__={payload}</script>'
|
||||
html = html_bytes.decode("utf-8")
|
||||
html = html.replace("</head>", f"{tag}</head>", 1)
|
||||
return html.encode("utf-8")
|
||||
return html.encode("utf-8"), nonce
|
||||
|
||||
|
||||
def setup_frontend(app: FastAPI, build_path: Path):
|
||||
|
|
@ -386,17 +750,23 @@ def setup_frontend(app: FastAPI, build_path: Path):
|
|||
if assets_dir.exists():
|
||||
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
|
||||
|
||||
@app.get("/")
|
||||
async def serve_root():
|
||||
def _build_index_response() -> Response:
|
||||
content = (build_path / "index.html").read_bytes()
|
||||
content = _strip_crossorigin(content)
|
||||
content = _inject_bootstrap(content, app)
|
||||
content, nonce = _inject_bootstrap(content, app)
|
||||
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"}
|
||||
if nonce:
|
||||
headers[_CSP_SCRIPT_NONCE_HEADER] = nonce
|
||||
return Response(
|
||||
content = content,
|
||||
media_type = "text/html",
|
||||
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
|
||||
headers = headers,
|
||||
)
|
||||
|
||||
@app.get("/")
|
||||
async def serve_root():
|
||||
return _build_index_response()
|
||||
|
||||
@app.get("/{full_path:path}")
|
||||
async def serve_frontend(full_path: str):
|
||||
if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")):
|
||||
|
|
@ -412,13 +782,6 @@ def setup_frontend(app: FastAPI, build_path: Path):
|
|||
return FileResponse(file_path)
|
||||
|
||||
# Serve index.html as bytes — avoids Content-Length mismatch
|
||||
content = (build_path / "index.html").read_bytes()
|
||||
content = _strip_crossorigin(content)
|
||||
content = _inject_bootstrap(content, app)
|
||||
return Response(
|
||||
content = content,
|
||||
media_type = "text/html",
|
||||
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"},
|
||||
)
|
||||
return _build_index_response()
|
||||
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from .training import (
|
|||
TrainingRunMetrics,
|
||||
TrainingRunDetailResponse,
|
||||
TrainingRunDeleteResponse,
|
||||
TrainingRunUpdateRequest,
|
||||
)
|
||||
from .models import (
|
||||
CheckpointInfo,
|
||||
|
|
@ -81,6 +82,7 @@ __all__ = [
|
|||
"TrainingRunMetrics",
|
||||
"TrainingRunDetailResponse",
|
||||
"TrainingRunDeleteResponse",
|
||||
"TrainingRunUpdateRequest",
|
||||
# Model management schemas
|
||||
"ModelDetails",
|
||||
"LocalModelInfo",
|
||||
|
|
|
|||
|
|
@ -37,7 +37,10 @@ class AuthStatusResponse(BaseModel):
|
|||
initialized: bool = Field(
|
||||
..., description = "True if the auth database contains a login user"
|
||||
)
|
||||
default_username: str = Field(..., description = "Default seeded admin username")
|
||||
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",
|
||||
|
|
|
|||
|
|
@ -5,10 +5,36 @@
|
|||
Pydantic schemas for Export API.
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
from typing import List, Optional, Literal, Dict, Any
|
||||
|
||||
|
||||
def _validate_save_directory(value: str) -> str:
|
||||
"""Reject save_directory values that escape the export root."""
|
||||
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")
|
||||
if len(raw) > 255:
|
||||
raise ValueError("save_directory must be <= 255 characters")
|
||||
path = Path(raw).expanduser()
|
||||
if path.is_absolute():
|
||||
raise ValueError(
|
||||
"save_directory must be a name or relative path under the "
|
||||
"export root; absolute paths are rejected"
|
||||
)
|
||||
if ".." in path.parts:
|
||||
raise ValueError("save_directory may not contain '..' segments")
|
||||
return raw
|
||||
|
||||
|
||||
class LoadCheckpointRequest(BaseModel):
|
||||
"""Request for loading a checkpoint into the export backend."""
|
||||
|
||||
|
|
@ -64,6 +90,12 @@ class ExportCommonOptions(BaseModel):
|
|||
...,
|
||||
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",
|
||||
|
|
@ -108,6 +140,12 @@ class ExportGGUFRequest(BaseModel):
|
|||
...,
|
||||
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: str = Field(
|
||||
"Q4_K_M",
|
||||
description = 'GGUF quantization method (e.g. "Q4_K_M")',
|
||||
|
|
|
|||
|
|
@ -11,7 +11,14 @@ import time
|
|||
import uuid
|
||||
from typing import Annotated, Any, Dict, Literal, Optional, List, Union
|
||||
|
||||
from pydantic import BaseModel, Discriminator, Field, Tag, model_validator
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
Discriminator,
|
||||
Field,
|
||||
Tag,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
|
||||
class LoadRequest(BaseModel):
|
||||
|
|
@ -43,6 +50,16 @@ class LoadRequest(BaseModel):
|
|||
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')",
|
||||
|
|
@ -299,10 +316,6 @@ class InferenceStatusResponse(BaseModel):
|
|||
supports_tools: bool = Field(
|
||||
False, description = "Whether the active model supports tool calling"
|
||||
)
|
||||
chat_template: Optional[str] = Field(
|
||||
None,
|
||||
description = "Jinja2 chat template string for the active model",
|
||||
)
|
||||
context_length: Optional[int] = Field(
|
||||
None, description = "Context length of the active model"
|
||||
)
|
||||
|
|
@ -314,10 +327,43 @@ class InferenceStatusResponse(BaseModel):
|
|||
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 = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
|
||||
)
|
||||
llama_cpp_supports_mtp: bool = Field(
|
||||
True,
|
||||
description = (
|
||||
"Whether llama.cpp supports MTP (--spec-type mtp/draft-mtp). "
|
||||
"False -> recommend `unsloth studio update`."
|
||||
),
|
||||
)
|
||||
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.",
|
||||
)
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
|
@ -369,15 +415,12 @@ ContentPart = Annotated[
|
|||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
"""
|
||||
A single message in the conversation.
|
||||
"""Single message in a chat conversation.
|
||||
|
||||
``content`` may be a plain string (text-only) or a list of
|
||||
content parts for multimodal messages (OpenAI vision format).
|
||||
Assistant messages that only contain tool calls may set ``content``
|
||||
to ``None`` with ``tool_calls`` populated. ``role="tool"`` messages
|
||||
carry the result of a client-executed tool call and require
|
||||
``tool_call_id`` per the OpenAI spec.
|
||||
``content`` is a string or a list of multimodal content parts. Assistant
|
||||
messages with only ``tool_calls`` populated may set ``content=None``.
|
||||
Missing ``tool_call_id`` on ``role="tool"`` is resolved at the
|
||||
``ChatCompletionRequest`` layer by walking back to the preceding assistant.
|
||||
"""
|
||||
|
||||
role: Literal["system", "user", "assistant", "tool"] = Field(
|
||||
|
|
@ -401,14 +444,6 @@ class ChatMessage(BaseModel):
|
|||
|
||||
@model_validator(mode = "after")
|
||||
def _validate_role_shape(self) -> "ChatMessage":
|
||||
# Enforce the per-role OpenAI spec shape at the request boundary.
|
||||
# Without this, malformed messages (e.g. user entries with no
|
||||
# content, tool_calls on a user/system role, role="tool" without
|
||||
# tool_call_id) would be silently forwarded to llama-server via
|
||||
# the passthrough path, surfacing as opaque upstream errors or
|
||||
# broken tool-call reconciliation downstream.
|
||||
|
||||
# Tool-call metadata must appear only on the appropriate role.
|
||||
if self.tool_calls is not None and self.role != "assistant":
|
||||
raise ValueError('"tool_calls" is only valid on role="assistant" messages.')
|
||||
if self.tool_call_id is not None and self.role != "tool":
|
||||
|
|
@ -416,23 +451,14 @@ class ChatMessage(BaseModel):
|
|||
if self.name is not None and self.role != "tool":
|
||||
raise ValueError('"name" is only valid on role="tool" messages.')
|
||||
|
||||
# Per-role content requirements. OpenAI-compatible clients may send
|
||||
# ``content=""`` for image-only turns when the image travels in a
|
||||
# companion field such as Studio's ``image_base64`` extension, so treat
|
||||
# empty strings as present content for user/system messages.
|
||||
if self.role == "tool":
|
||||
if not self.tool_call_id:
|
||||
raise ValueError(
|
||||
'role="tool" messages require "tool_call_id" per the OpenAI spec.'
|
||||
)
|
||||
# tool_call_id resolution happens at ChatCompletionRequest scope.
|
||||
if not self.content:
|
||||
raise ValueError('role="tool" messages require non-empty "content".')
|
||||
elif self.role == "assistant":
|
||||
# Assistant messages may omit content when tool_calls is set.
|
||||
if not self.content and not self.tool_calls:
|
||||
raise ValueError(
|
||||
'role="assistant" messages require either "content" or "tool_calls".'
|
||||
)
|
||||
# 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".')
|
||||
|
|
@ -518,9 +544,11 @@ class ChatCompletionRequest(BaseModel):
|
|||
None,
|
||||
description = "[x-unsloth] Enable/disable thinking/reasoning mode for supported models",
|
||||
)
|
||||
reasoning_effort: Optional[Literal["low", "medium", "high"]] = Field(
|
||||
reasoning_effort: Optional[
|
||||
Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"]
|
||||
] = Field(
|
||||
None,
|
||||
description = "[x-unsloth] Reasoning effort level ('low'|'medium'|'high') for Harmony-style reasoning models (e.g. gpt-oss). Overrides enable_thinking when the active model uses reasoning_effort style.",
|
||||
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,
|
||||
|
|
@ -557,6 +585,198 @@ class ChatCompletionRequest(BaseModel):
|
|||
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[bool] = Field(
|
||||
None,
|
||||
description = (
|
||||
"[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, "
|
||||
"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 this flag is informational. "
|
||||
"Ignored for every other provider (mistral, gemini, kimi, openrouter, "
|
||||
"vllm, local, etc.). Treated as enabled when omitted."
|
||||
),
|
||||
)
|
||||
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."
|
||||
),
|
||||
)
|
||||
|
||||
@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 random id only if no candidate exists.
|
||||
Crossing a user turn breaks the lookup.
|
||||
"""
|
||||
# Pre-mark explicit ids first so a sibling missing-id result does not
|
||||
# steal one already claimed by name.
|
||||
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
|
||||
|
||||
|
||||
# ── OpenAI shell-tool container management ─────────────────────
|
||||
|
||||
|
||||
class OpenAIContainerRequest(BaseModel):
|
||||
"""
|
||||
Shared body for the three OpenAI container endpoints (list / create
|
||||
/ delete). Carries the encrypted API key + base URL so the route
|
||||
handler can decrypt it and proxy to the user's OpenAI account.
|
||||
Same pattern as the inference proxy endpoints — keeps the key off
|
||||
persistent storage on the backend.
|
||||
"""
|
||||
|
||||
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 ────────────────────────────────────
|
||||
|
||||
|
|
|
|||
130
studio/backend/models/providers.py
Normal file
130
studio/backend/models/providers.py
Normal file
|
|
@ -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 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)"
|
||||
)
|
||||
|
||||
|
||||
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)"
|
||||
)
|
||||
|
|
@ -5,10 +5,43 @@
|
|||
Pydantic schemas for Training API
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
from typing import Any, Optional, List, Dict, Literal
|
||||
|
||||
|
||||
_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
|
||||
|
||||
|
||||
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"""
|
||||
|
||||
|
|
@ -16,8 +49,11 @@ class TrainingStartRequest(BaseModel):
|
|||
model_name: str = Field(
|
||||
..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
|
||||
)
|
||||
training_type: str = Field(
|
||||
..., description = "Training type: 'LoRA/QLoRA' or 'Full Finetuning'"
|
||||
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")
|
||||
|
|
@ -61,6 +97,150 @@ class TrainingStartRequest(BaseModel):
|
|||
values.setdefault("train_split", values.pop("split"))
|
||||
return values
|
||||
|
||||
@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 meaning "use max_steps instead"; the frontend's
|
||||
# steps-vs-epochs toggle sends it.
|
||||
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("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 = (
|
||||
|
|
@ -82,10 +262,22 @@ class TrainingStartRequest(BaseModel):
|
|||
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.",
|
||||
)
|
||||
random_seed: int = Field(42, description = "Random seed")
|
||||
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)")
|
||||
|
|
@ -137,6 +329,16 @@ class TrainingStartRequest(BaseModel):
|
|||
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.",
|
||||
)
|
||||
|
||||
@model_validator(mode = "after")
|
||||
def _check_steps_or_epochs(self) -> "TrainingStartRequest":
|
||||
# num_epochs and max_steps each accept 0 as a "use the other one"
|
||||
# sentinel. If both resolve to 0 there's nothing to train against.
|
||||
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"""
|
||||
|
|
@ -214,6 +416,7 @@ class TrainingRunSummary(BaseModel):
|
|||
status: Literal["running", "completed", "stopped", "error"]
|
||||
model_name: str
|
||||
dataset_name: str
|
||||
display_name: Optional[str] = None
|
||||
started_at: str
|
||||
ended_at: Optional[str] = None
|
||||
total_steps: Optional[int] = None
|
||||
|
|
@ -227,6 +430,14 @@ class TrainingRunSummary(BaseModel):
|
|||
resumed_later: bool = False
|
||||
|
||||
|
||||
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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -48,13 +48,31 @@ class ScrapeConfig:
|
|||
max_comments_per_item: int
|
||||
|
||||
|
||||
def _resolve_token(token: str) -> str:
|
||||
tok = token or os.environ.get("GH_TOKEN", "") or os.environ.get("GITHUB_TOKEN", "")
|
||||
if not tok:
|
||||
raise ValueError(
|
||||
"GitHub token is required. Set it in the recipe config or the GH_TOKEN / GITHUB_TOKEN env var."
|
||||
@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)",
|
||||
)
|
||||
return tok
|
||||
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):
|
||||
|
|
@ -155,7 +173,7 @@ def _flatten_commit_row(r: dict, repo: str) -> dict:
|
|||
def scrape(cfg: ScrapeConfig, base_dir: Path):
|
||||
token = _resolve_token(cfg.token)
|
||||
GitHubClient, RepoScraper = _load_impl()
|
||||
client = GitHubClient(token = token)
|
||||
client = GitHubClient(token = token.value, token_source = token.source)
|
||||
base_dir.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
# Per-resource trial limits. limit <= 0 means "all": use a very large cap.
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ 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
|
||||
|
|
@ -29,16 +31,46 @@ 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,
|
||||
):
|
||||
token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
|
||||
if not token:
|
||||
raise RuntimeError("GH_TOKEN not set in environment")
|
||||
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}"}
|
||||
|
|
@ -59,6 +91,49 @@ class GitHubClient:
|
|||
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:
|
||||
"""Distinguish auth failures from rate limiting on 401/403 responses.
|
||||
|
||||
- 401: always an auth failure (invalid / expired / wrong-scope token).
|
||||
- 403: an auth failure UNLESS the response carries a clear rate-limit signal
|
||||
(Retry-After header, X-RateLimit-Remaining: 0, or GitHub's secondary /
|
||||
abuse rate-limit response text).
|
||||
"""
|
||||
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
|
||||
|
|
@ -112,13 +187,14 @@ class GitHubClient:
|
|||
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:
|
||||
# Check for secondary/abuse
|
||||
retry_after = r.headers.get("Retry-After")
|
||||
if retry_after:
|
||||
t = int(retry_after)
|
||||
log.warning("Secondary rate limit. Sleep %ds.", t)
|
||||
time.sleep(t + 2)
|
||||
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)
|
||||
|
|
@ -188,12 +264,15 @@ class GitHubClient:
|
|||
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 = r.headers.get("Retry-After")
|
||||
if retry_after:
|
||||
t = int(retry_after)
|
||||
log.warning("Secondary rate limit on REST. Sleep %ds.", t)
|
||||
time.sleep(t + 2)
|
||||
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
|
||||
# Check if primary rate
|
||||
if self.rest_remaining == 0 and self.rest_reset:
|
||||
|
|
|
|||
|
|
@ -8,7 +8,28 @@
|
|||
|
||||
# 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
|
||||
# pydantic 2.x deps. With --no-deps, `import pydantic` blows up
|
||||
# with `ModuleNotFoundError: 'pydantic_core'` (compiled Rust core,
|
||||
# separate wheel), then `'annotated_types'`, then
|
||||
# `'typing_inspection'` (used by pydantic 2.10+ for fields).
|
||||
pydantic-core
|
||||
annotated-types>=0.6
|
||||
typing-inspection>=0.4
|
||||
pyyaml
|
||||
nest-asyncio
|
||||
|
||||
|
|
@ -42,7 +63,9 @@ anyio
|
|||
sniffio
|
||||
h11
|
||||
|
||||
tokenizers
|
||||
# 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
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ typer
|
|||
fastapi
|
||||
uvicorn
|
||||
pydantic
|
||||
packaging
|
||||
matplotlib
|
||||
pandas
|
||||
nest_asyncio
|
||||
|
|
@ -15,3 +16,5 @@ huggingface-hub==0.36.2
|
|||
structlog>=24.1.0
|
||||
diceware
|
||||
ddgs
|
||||
cryptography>=42.0.0
|
||||
httpx>=0.27.0
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ 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.providers import router as providers_router
|
||||
|
||||
__all__ = [
|
||||
"training_router",
|
||||
|
|
@ -25,4 +26,5 @@ __all__ = [
|
|||
"data_recipe_router",
|
||||
"export_router",
|
||||
"training_history_router",
|
||||
"providers_router",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -5,8 +5,13 @@
|
|||
Authentication API routes
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
|
||||
|
||||
import ipaddress
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from models.auth import (
|
||||
|
|
@ -33,14 +38,160 @@ from auth.authentication import (
|
|||
router = APIRouter()
|
||||
|
||||
|
||||
# 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 we prune stale entries; if still full the
|
||||
# failure folds into the per-IP aggregate only.
|
||||
_LOGIN_MAX_BUCKETS = 4096
|
||||
# Unrepresentable as a real username (leading NUL); folds unknown-user attempts
|
||||
# into one slot so attacker cardinality cannot blow the bucket dict.
|
||||
_UNKNOWN_LOGIN_USER = "\x00unknown-user"
|
||||
|
||||
|
||||
def _trust_forwarded_for() -> bool:
|
||||
"""Honour X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is set.
|
||||
|
||||
Off by default so a direct caller cannot spoof the header.
|
||||
"""
|
||||
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 and takes the 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 cannot 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 _record_login_failure(key: tuple[str, str]) -> int:
|
||||
now = time.monotonic()
|
||||
ip, _username = key
|
||||
with _LOGIN_BUCKETS_LOCK:
|
||||
ip_bucket = _LOGIN_IP_BUCKETS.setdefault(ip, deque())
|
||||
_prune_bucket(ip_bucket, now)
|
||||
ip_bucket.append(now)
|
||||
|
||||
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)
|
||||
# Bucket dict is at its cap; per-IP cap still applies via ip_bucket.
|
||||
return len(ip_bucket)
|
||||
|
||||
|
||||
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:
|
||||
return max(
|
||||
_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS),
|
||||
_blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS),
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@router.get("/status", response_model = AuthStatusResponse)
|
||||
async def auth_status() -> AuthStatusResponse:
|
||||
"""
|
||||
Check whether auth has already been initialized.
|
||||
|
||||
- initialized = False -> frontend should wait for the seeded admin bootstrap.
|
||||
- initialized = True -> frontend should show login or force the first password change.
|
||||
"""
|
||||
"""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,
|
||||
|
|
@ -53,12 +204,28 @@ async def auth_status() -> AuthStatusResponse:
|
|||
|
||||
|
||||
@router.post("/login", response_model = Token)
|
||||
async def login(payload: AuthLoginRequest) -> Token:
|
||||
"""
|
||||
Login with username/password and receive access + refresh tokens.
|
||||
"""
|
||||
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 is intentionally not interpolated into the body; behind a
|
||||
# proxy or NAT it is either misleading or an info leak.
|
||||
detail = (
|
||||
f"Too many failed login attempts. "
|
||||
f"Try again in {blocked_for} seconds."
|
||||
),
|
||||
headers = {"Retry-After": str(blocked_for)},
|
||||
)
|
||||
|
||||
record = storage.get_user_and_secret(payload.username)
|
||||
if record is None:
|
||||
# Record under a single sentinel key per IP so attacker-controlled
|
||||
# username cardinality does not allocate buckets without bound.
|
||||
_record_login_failure(unknown_key)
|
||||
raise HTTPException(
|
||||
status_code = status.HTTP_401_UNAUTHORIZED,
|
||||
detail = "Incorrect password. Run 'unsloth studio reset-password' in your terminal to reset it.",
|
||||
|
|
@ -66,11 +233,14 @@ async def login(payload: AuthLoginRequest) -> Token:
|
|||
|
||||
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 = "Incorrect password. Run 'unsloth studio reset-password' in your terminal to reset it.",
|
||||
)
|
||||
|
||||
_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(
|
||||
|
|
@ -81,6 +251,23 @@ async def login(payload: AuthLoginRequest) -> Token:
|
|||
)
|
||||
|
||||
|
||||
@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."""
|
||||
|
|
@ -101,21 +288,20 @@ async def desktop_login(payload: DesktopLoginRequest) -> Token:
|
|||
|
||||
@router.post("/refresh", response_model = Token)
|
||||
async def refresh(payload: RefreshTokenRequest) -> Token:
|
||||
"""
|
||||
Exchange a valid refresh token for a new access token.
|
||||
|
||||
The refresh token itself is reusable until it expires (7 days).
|
||||
"""
|
||||
new_access_token, username, is_desktop = refresh_access_token(payload.refresh_token)
|
||||
if new_access_token is None or username is None:
|
||||
"""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 = payload.refresh_token,
|
||||
refresh_token = new_refresh_token,
|
||||
token_type = "bearer",
|
||||
must_change_password = False
|
||||
if is_desktop
|
||||
|
|
@ -126,6 +312,7 @@ async def refresh(payload: RefreshTokenRequest) -> Token:
|
|||
@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."""
|
||||
|
|
@ -150,6 +337,10 @@ async def change_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(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ Export API routes: checkpoint discovery and model export operations.
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
|
@ -184,14 +185,18 @@ async def get_export_status(
|
|||
|
||||
|
||||
def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]:
|
||||
"""Wrap the resolved on-disk export path into the details dict the
|
||||
frontend reads to populate the Export Complete screen. Returns None
|
||||
when the export had no local component (Hub-only push) so the
|
||||
Pydantic field stays absent rather than ``{"output_path": null}``.
|
||||
"""
|
||||
"""Return the export path relative to exports_root so the install path is not leaked."""
|
||||
if not output_path:
|
||||
return None
|
||||
return {"output_path": output_path}
|
||||
try:
|
||||
from utils.paths.storage_roots import exports_root
|
||||
|
||||
rel = os.path.relpath(output_path, exports_root())
|
||||
if rel.startswith(".."):
|
||||
rel = os.path.basename(output_path)
|
||||
return {"output_path": rel}
|
||||
except Exception:
|
||||
return {"output_path": os.path.basename(output_path)}
|
||||
|
||||
|
||||
@router.post("/export/merged", response_model = ExportOperationResponse)
|
||||
|
|
|
|||
|
|
@ -117,9 +117,13 @@ try:
|
|||
LlamaCppBackend,
|
||||
_DEFAULT_MAX_TOKENS_FLOOR,
|
||||
_DEFAULT_T_MAX_PREDICT_MS,
|
||||
_hf_offline_if_dns_dead,
|
||||
detect_reasoning_flags,
|
||||
)
|
||||
from core.inference.llama_server_args import validate_extra_args
|
||||
from core.inference.llama_server_args import (
|
||||
strip_shadowing_flags,
|
||||
validate_extra_args,
|
||||
)
|
||||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
from utils.models.model_config import load_model_defaults
|
||||
|
|
@ -139,9 +143,13 @@ except ImportError:
|
|||
LlamaCppBackend,
|
||||
_DEFAULT_MAX_TOKENS_FLOOR,
|
||||
_DEFAULT_T_MAX_PREDICT_MS,
|
||||
_hf_offline_if_dns_dead,
|
||||
detect_reasoning_flags,
|
||||
)
|
||||
from core.inference.llama_server_args import validate_extra_args
|
||||
from core.inference.llama_server_args import (
|
||||
strip_shadowing_flags,
|
||||
validate_extra_args,
|
||||
)
|
||||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
from utils.models.model_config import load_model_defaults
|
||||
|
|
@ -194,6 +202,11 @@ from models.inference import (
|
|||
AnthropicResponseTextBlock,
|
||||
AnthropicResponseToolUseBlock,
|
||||
AnthropicUsage,
|
||||
CreateOpenAIContainerBody,
|
||||
DeleteOpenAIContainerBody,
|
||||
ListOpenAIContainersResponse,
|
||||
OpenAIContainerRequest,
|
||||
OpenAIContainerSummary,
|
||||
)
|
||||
from core.inference.anthropic_compat import (
|
||||
anthropic_messages_to_openai,
|
||||
|
|
@ -204,6 +217,11 @@ from core.inference.anthropic_compat import (
|
|||
)
|
||||
from auth.authentication import get_current_subject
|
||||
|
||||
from core.inference.key_exchange import decrypt_api_key
|
||||
from core.inference.providers import get_provider_info, get_base_url
|
||||
from core.inference.external_provider import ExternalProviderClient
|
||||
from storage import providers_db
|
||||
|
||||
import io
|
||||
import wave
|
||||
import base64
|
||||
|
|
@ -396,6 +414,57 @@ def _validate_native_mmproj_companion(
|
|||
) from exc
|
||||
|
||||
|
||||
def _normalise_settings_str(value: Optional[str]) -> Optional[str]:
|
||||
"""Lowercase + strip a settings string, mapping blank/None to None."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
stripped = value.strip().lower()
|
||||
return stripped or None
|
||||
return value
|
||||
|
||||
|
||||
def _request_matches_loaded_settings(
|
||||
request: LoadRequest, llama_backend: LlamaCppBackend
|
||||
) -> bool:
|
||||
"""True iff every runtime setting on the request matches the loaded
|
||||
server. Caller has already checked model+variant+is_loaded. See #5401."""
|
||||
# Compare requested n_ctx (not effective) so VRAM-cap doesn't mask
|
||||
# an Auto-vs-explicit slider flip.
|
||||
if request.max_seq_length != llama_backend.requested_n_ctx:
|
||||
return False
|
||||
if _normalise_settings_str(request.cache_type_kv) != _normalise_settings_str(
|
||||
llama_backend.cache_type_kv
|
||||
):
|
||||
return False
|
||||
# Vision loads silently drop speculative decoding (llama_cpp.py gates
|
||||
# spec on ``not is_vision``), so treat the request as ``off`` against
|
||||
# the backend's ``None`` to avoid forcing a redundant reload.
|
||||
if llama_backend.is_vision:
|
||||
req_spec = "off"
|
||||
else:
|
||||
req_spec = _normalise_settings_str(request.speculative_type) or "off"
|
||||
backend_spec = _normalise_settings_str(llama_backend.speculative_type) or "off"
|
||||
if req_spec != backend_spec:
|
||||
return False
|
||||
if (request.chat_template_override or None) != (
|
||||
llama_backend.chat_template_override or None
|
||||
):
|
||||
return False
|
||||
# llama_extra_args=None means "inherit"; only an explicit list that
|
||||
# differs forces a reload. On the inherit path, refuse to match if
|
||||
# stored extras contain any shadow flag, so the reload path can
|
||||
# strip them instead of leaving a stale override in effect.
|
||||
backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else []
|
||||
if request.llama_extra_args is None:
|
||||
if backend_extra and strip_shadowing_flags(backend_extra) != backend_extra:
|
||||
return False
|
||||
else:
|
||||
if list(request.llama_extra_args) != backend_extra:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _resolve_model_identifier_for_request(
|
||||
request: LoadRequest | ValidateModelRequest,
|
||||
*,
|
||||
|
|
@ -451,6 +520,11 @@ async def load_model(
|
|||
extra_llama_args = validate_extra_args(request.llama_extra_args)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code = 400, detail = str(exc))
|
||||
# Re-narrow []-from-None back to None so the inheritance path
|
||||
# below can tell "caller omitted" from "caller explicit []".
|
||||
extra_llama_args: Optional[list[str]] = (
|
||||
None if request.llama_extra_args is None else extra_llama_args
|
||||
)
|
||||
|
||||
model_identifier, model_log_label, native_grant_backed = (
|
||||
_resolve_model_identifier_for_request(request, operation = "load-model")
|
||||
|
|
@ -469,12 +543,14 @@ async def load_model(
|
|||
and llama_backend.hf_variant.lower() == request.gguf_variant.lower()
|
||||
and llama_backend.model_identifier
|
||||
and llama_backend.model_identifier.lower() == model_identifier.lower()
|
||||
# Also require runtime settings to match so Apply changes
|
||||
# aren't silently dropped (#5401).
|
||||
and _request_matches_loaded_settings(request, llama_backend)
|
||||
):
|
||||
logger.info(
|
||||
f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload"
|
||||
)
|
||||
inference_config = load_inference_config(llama_backend.model_identifier)
|
||||
from utils.models import is_audio_input_type
|
||||
|
||||
_gguf_audio = (
|
||||
llama_backend._audio_type
|
||||
|
|
@ -495,9 +571,7 @@ async def load_model(
|
|||
is_gguf = True,
|
||||
is_audio = _gguf_is_audio,
|
||||
audio_type = _gguf_audio,
|
||||
has_audio_input = is_audio_input_type(_gguf_audio)
|
||||
if _gguf_audio
|
||||
else False,
|
||||
has_audio_input = False,
|
||||
inference = inference_config,
|
||||
requires_trust_remote_code = bool(
|
||||
inference_config.get("trust_remote_code", False)
|
||||
|
|
@ -571,13 +645,15 @@ async def load_model(
|
|||
chat_template = _chat_template,
|
||||
)
|
||||
|
||||
# Create config using clean factory method
|
||||
# is_lora is auto-detected from adapter_config.json on disk/HF
|
||||
config = ModelConfig.from_identifier(
|
||||
model_id = model_identifier,
|
||||
hf_token = request.hf_token,
|
||||
gguf_variant = request.gguf_variant,
|
||||
)
|
||||
# is_lora auto-detected from adapter_config.json on disk/HF.
|
||||
# DNS-probe wrap so offline loads skip 30-60s of soft-failed
|
||||
# network checks before the worker starts.
|
||||
with _hf_offline_if_dns_dead():
|
||||
config = ModelConfig.from_identifier(
|
||||
model_id = model_identifier,
|
||||
hf_token = request.hf_token,
|
||||
gguf_variant = request.gguf_variant,
|
||||
)
|
||||
|
||||
if not config:
|
||||
raise HTTPException(
|
||||
|
|
@ -606,6 +682,70 @@ async def load_model(
|
|||
)
|
||||
unsloth_backend.unload_model(unsloth_backend.active_model_name)
|
||||
|
||||
# Inherit llama_extra_args from the previous load when the
|
||||
# request omits the field (the chat-settings Apply path
|
||||
# does not round-trip them; explicit [] still clears).
|
||||
# Inheritance is gated on (model_identifier, hf_variant)
|
||||
# to refuse cross-model pickup, and shadowing flags are
|
||||
# stripped so an inherited override can't win the last-wins
|
||||
# CLI parse against a freshly-supplied first-class field.
|
||||
if request.llama_extra_args is None and llama_backend.extra_args:
|
||||
source = llama_backend.extra_args_source
|
||||
# Compare against the resolved variant, not the request
|
||||
# field: callers commonly omit gguf_variant for local
|
||||
# ``.gguf`` paths and HF auto-pick flows. ``config.gguf_
|
||||
# variant`` is the variant load_model was actually
|
||||
# invoked with (see the HF / local branches below), so
|
||||
# both sides of the comparison key off the same string.
|
||||
resolved_variant = config.gguf_variant
|
||||
same_source = bool(
|
||||
source
|
||||
and source[0]
|
||||
and source[0].lower() == model_identifier.lower()
|
||||
and (source[1] or "").lower() == (resolved_variant or "").lower()
|
||||
)
|
||||
if not same_source:
|
||||
logger.info(
|
||||
"Not inheriting llama_extra_args: stored args came "
|
||||
"from %s, loading %s",
|
||||
source,
|
||||
(model_identifier, resolved_variant),
|
||||
)
|
||||
# Cross-model: clear explicitly so the backend
|
||||
# doesn't inherit via "no opinion" semantics.
|
||||
extra_llama_args = []
|
||||
else:
|
||||
# Strip only the groups whose first-class field
|
||||
# was actually set by the caller, so an inherited
|
||||
# --chat-template-file survives an Apply that omits
|
||||
# chat_template_override.
|
||||
fields_set = getattr(request, "model_fields_set", set())
|
||||
stripped = strip_shadowing_flags(
|
||||
llama_backend.extra_args,
|
||||
strip_context = "max_seq_length" in fields_set,
|
||||
strip_cache = "cache_type_kv" in fields_set,
|
||||
strip_spec = "speculative_type" in fields_set,
|
||||
strip_template = "chat_template_override" in fields_set,
|
||||
)
|
||||
try:
|
||||
extra_llama_args = validate_extra_args(stripped)
|
||||
except ValueError:
|
||||
# Should not happen on already-validated args; degrade
|
||||
# to no-extras rather than 400 if managed flags changed.
|
||||
logger.warning(
|
||||
"Stored llama_extra_args failed revalidation; "
|
||||
"loading without them: %s",
|
||||
stripped,
|
||||
)
|
||||
extra_llama_args = []
|
||||
else:
|
||||
if extra_llama_args:
|
||||
logger.info(
|
||||
"Inheriting llama_extra_args from previous "
|
||||
"load (same model, shadow-stripped): %s",
|
||||
extra_llama_args,
|
||||
)
|
||||
|
||||
# Route to HF mode or local mode based on config
|
||||
# Run in a thread so the event loop stays free for progress
|
||||
# polling and other requests during the (potentially long)
|
||||
|
|
@ -638,6 +778,10 @@ async def load_model(
|
|||
llama_backend.load_model,
|
||||
gguf_path = config.gguf_file,
|
||||
mmproj_path = config.gguf_mmproj_file,
|
||||
# Pass the resolved variant so _extra_args_source
|
||||
# is keyed off the same string the inheritance
|
||||
# check at the top of /load uses (#5401 followup).
|
||||
hf_variant = config.gguf_variant,
|
||||
model_identifier = config.identifier,
|
||||
is_vision = config.is_vision,
|
||||
n_ctx = request.max_seq_length,
|
||||
|
|
@ -658,9 +802,10 @@ async def load_model(
|
|||
f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}"
|
||||
)
|
||||
|
||||
# Detect TTS audio by probing the loaded model's vocabulary
|
||||
from utils.models import is_audio_input_type
|
||||
|
||||
# Detect TTS/audio marker tokens by probing the loaded model's vocabulary.
|
||||
# GGUF audio input is not wired through the chat path yet, so do not
|
||||
# advertise has_audio_input for GGUF models until uploaded audio is
|
||||
# actually forwarded to llama-server.
|
||||
_gguf_audio = llama_backend.detect_audio_type()
|
||||
_gguf_is_audio = _gguf_audio in ("snac", "bicodec", "dac")
|
||||
llama_backend._is_audio = _gguf_is_audio
|
||||
|
|
@ -681,12 +826,12 @@ async def load_model(
|
|||
display_name = model_log_label
|
||||
if native_grant_backed
|
||||
else config.display_name,
|
||||
is_vision = config.is_vision,
|
||||
is_vision = llama_backend.is_vision,
|
||||
is_lora = False,
|
||||
is_gguf = True,
|
||||
is_audio = _gguf_is_audio,
|
||||
audio_type = _gguf_audio,
|
||||
has_audio_input = is_audio_input_type(_gguf_audio),
|
||||
has_audio_input = False,
|
||||
inference = inference_config,
|
||||
requires_trust_remote_code = bool(
|
||||
inference_config.get("trust_remote_code", False)
|
||||
|
|
@ -1141,6 +1286,24 @@ async def get_status(
|
|||
try:
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
|
||||
# MTP probe + freshness check (both cached). Drive the UI banner.
|
||||
try:
|
||||
_bin = type(llama_backend)._find_llama_server_binary()
|
||||
_caps = type(llama_backend).probe_server_capabilities(_bin)
|
||||
_supports_mtp = bool(_caps.get("supports_mtp", False))
|
||||
except Exception:
|
||||
_bin = None
|
||||
_supports_mtp = True # fail open
|
||||
try:
|
||||
from utils.llama_cpp_freshness import check_prebuilt_freshness
|
||||
|
||||
_freshness = check_prebuilt_freshness(_bin)
|
||||
except Exception:
|
||||
_freshness = {}
|
||||
_stale = bool(_freshness.get("stale"))
|
||||
_installed_tag = _freshness.get("installed_tag")
|
||||
_latest_tag = _freshness.get("latest_tag")
|
||||
|
||||
# If a GGUF model is loaded via llama-server, report that
|
||||
if llama_backend.is_loaded:
|
||||
_model_id = llama_backend.model_identifier
|
||||
|
|
@ -1156,13 +1319,15 @@ async def get_status(
|
|||
):
|
||||
_display_model_id = os.path.basename(_model_id)
|
||||
_inference_cfg = load_inference_config(_model_id) if _model_id else None
|
||||
_audio_type = getattr(llama_backend, "_audio_type", None)
|
||||
return InferenceStatusResponse(
|
||||
active_model = _display_model_id,
|
||||
is_vision = llama_backend.is_vision,
|
||||
is_gguf = True,
|
||||
gguf_variant = llama_backend.hf_variant,
|
||||
is_audio = getattr(llama_backend, "_is_audio", False),
|
||||
audio_type = getattr(llama_backend, "_audio_type", None),
|
||||
audio_type = _audio_type,
|
||||
has_audio_input = False,
|
||||
loading = [],
|
||||
loaded = [_display_model_id] if _display_model_id else [],
|
||||
inference = _inference_cfg,
|
||||
|
|
@ -1178,7 +1343,13 @@ async def get_status(
|
|||
context_length = llama_backend.context_length,
|
||||
max_context_length = llama_backend.max_context_length,
|
||||
native_context_length = llama_backend.native_context_length,
|
||||
cache_type_kv = llama_backend.cache_type_kv,
|
||||
chat_template_override = llama_backend.chat_template_override,
|
||||
speculative_type = llama_backend.speculative_type,
|
||||
llama_cpp_supports_mtp = _supports_mtp,
|
||||
llama_cpp_prebuilt_stale = _stale,
|
||||
llama_cpp_installed_tag = _installed_tag,
|
||||
llama_cpp_latest_tag = _latest_tag,
|
||||
)
|
||||
|
||||
# Otherwise, report Unsloth backend status
|
||||
|
|
@ -1239,6 +1410,10 @@ async def get_status(
|
|||
supports_preserve_thinking = False,
|
||||
supports_tools = False,
|
||||
chat_template = chat_template,
|
||||
llama_cpp_supports_mtp = _supports_mtp,
|
||||
llama_cpp_prebuilt_stale = _stale,
|
||||
llama_cpp_installed_tag = _installed_tag,
|
||||
llama_cpp_latest_tag = _latest_tag,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -1462,6 +1637,346 @@ def _extract_content_parts(
|
|||
return system_prompt, chat_messages, first_image_b64
|
||||
|
||||
|
||||
# ── External provider proxy ──────────────────────────────────────
|
||||
|
||||
|
||||
def _build_external_messages(
|
||||
messages: list,
|
||||
supports_vision: bool,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Convert ChatMessage list to OpenAI-compatible dicts for external providers.
|
||||
|
||||
- Vision providers: preserve multimodal content arrays (image_url parts intact).
|
||||
- Non-vision providers: flatten to text-only (images silently dropped).
|
||||
"""
|
||||
result = []
|
||||
for msg in messages:
|
||||
if isinstance(msg.content, str):
|
||||
# Skip assistant messages with empty content (some providers reject them)
|
||||
if msg.role == "assistant" and not msg.content.strip():
|
||||
continue
|
||||
result.append({"role": msg.role, "content": msg.content})
|
||||
elif isinstance(msg.content, list):
|
||||
if supports_vision:
|
||||
parts = []
|
||||
for part in msg.content:
|
||||
if part.type == "text":
|
||||
parts.append({"type": "text", "text": part.text})
|
||||
elif part.type == "image_url":
|
||||
parts.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": part.image_url.url},
|
||||
}
|
||||
)
|
||||
result.append({"role": msg.role, "content": parts})
|
||||
else:
|
||||
# Non-vision provider — strip images, keep text only
|
||||
text = "\n".join(p.text for p in msg.content if p.type == "text")
|
||||
result.append({"role": msg.role, "content": text})
|
||||
return result
|
||||
|
||||
|
||||
async def _proxy_to_external_provider(
|
||||
payload: ChatCompletionRequest,
|
||||
request: Request,
|
||||
) -> StreamingResponse:
|
||||
"""
|
||||
Proxy a chat completion request to an external LLM provider.
|
||||
|
||||
Resolves provider config (from DB or registry), decrypts the API key,
|
||||
and streams the response back in OpenAI SSE format.
|
||||
"""
|
||||
# Resolve provider type and base URL
|
||||
provider_type = payload.provider_type
|
||||
base_url = payload.provider_base_url
|
||||
|
||||
if payload.provider_id:
|
||||
config = providers_db.get_provider(payload.provider_id)
|
||||
if config is None:
|
||||
raise HTTPException(
|
||||
status_code = 404,
|
||||
detail = f"Provider config not found: {payload.provider_id}",
|
||||
)
|
||||
if not config["is_enabled"]:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Provider '{config['display_name']}' is disabled.",
|
||||
)
|
||||
provider_type = provider_type or config["provider_type"]
|
||||
base_url = base_url or config["base_url"]
|
||||
|
||||
if not provider_type:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Either provider_id or provider_type is required for external provider routing.",
|
||||
)
|
||||
|
||||
# Fall back to registry default base URL
|
||||
if not base_url:
|
||||
base_url = get_base_url(provider_type)
|
||||
if not base_url:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Unknown provider type: {provider_type}",
|
||||
)
|
||||
|
||||
api_key = ""
|
||||
if payload.encrypted_api_key:
|
||||
try:
|
||||
api_key = decrypt_api_key(payload.encrypted_api_key)
|
||||
except Exception as exc:
|
||||
logger.warning("external_provider.decrypt_failed", error = str(exc))
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.",
|
||||
)
|
||||
|
||||
model = payload.external_model or payload.model
|
||||
if model == "default":
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "external_model is required when using an external provider.",
|
||||
)
|
||||
|
||||
# Build messages preserving multimodal content for vision-capable providers
|
||||
from core.inference.providers import get_provider_info as _get_provider_info
|
||||
|
||||
_pinfo = _get_provider_info(provider_type) or {}
|
||||
_supports_vision = _pinfo.get("supports_vision", False)
|
||||
chat_messages = _build_external_messages(payload.messages, _supports_vision)
|
||||
|
||||
client = ExternalProviderClient(
|
||||
provider_type = provider_type,
|
||||
base_url = base_url,
|
||||
api_key = api_key,
|
||||
)
|
||||
|
||||
async def _stream():
|
||||
gen = client.stream_chat_completion(
|
||||
messages = chat_messages,
|
||||
model = model,
|
||||
temperature = payload.temperature,
|
||||
top_p = payload.top_p,
|
||||
max_tokens = payload.max_tokens,
|
||||
presence_penalty = payload.presence_penalty,
|
||||
top_k = payload.top_k,
|
||||
enable_thinking = payload.enable_thinking,
|
||||
reasoning_effort = payload.reasoning_effort,
|
||||
enabled_tools = payload.enabled_tools,
|
||||
enable_prompt_caching = payload.enable_prompt_caching,
|
||||
openai_code_exec_container_id = payload.openai_code_exec_container_id,
|
||||
anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id,
|
||||
stream = payload.stream,
|
||||
)
|
||||
try:
|
||||
sent_done = False
|
||||
async for line in gen:
|
||||
yield f"{line}\n\n"
|
||||
if "[DONE]" in line:
|
||||
sent_done = True
|
||||
if not sent_done:
|
||||
yield "data: [DONE]\n\n"
|
||||
except Exception as exc:
|
||||
logger.error("external_provider.stream_error", error = str(exc))
|
||||
finally:
|
||||
try:
|
||||
await gen.aclose()
|
||||
except RuntimeError:
|
||||
pass # suppress httpcore asyncgen cleanup error (Python 3.13 + httpcore 1.0.x)
|
||||
await client.close()
|
||||
|
||||
return StreamingResponse(
|
||||
_stream(),
|
||||
media_type = "text/event-stream",
|
||||
headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ── OpenAI shell-tool container management ───────────────────────
|
||||
|
||||
|
||||
def _resolve_openai_cloud_client(
|
||||
body: OpenAIContainerRequest,
|
||||
) -> ExternalProviderClient:
|
||||
"""
|
||||
Decrypt the API key + validate the base URL points at OpenAI cloud,
|
||||
then build an ExternalProviderClient for the three container CRUD
|
||||
endpoints below. The shell tool only exists on api.openai.com, so
|
||||
rejecting non-cloud bases up front prevents confusing 404s on
|
||||
ollama / llama.cpp / vLLM / custom presets.
|
||||
"""
|
||||
base_url = body.provider_base_url or get_base_url("openai")
|
||||
if not base_url or "api.openai.com" not in base_url:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = (
|
||||
"OpenAI container management is only available on the "
|
||||
"managed cloud (api.openai.com). The provider's base URL "
|
||||
f"points at {base_url!r}."
|
||||
),
|
||||
)
|
||||
try:
|
||||
api_key = decrypt_api_key(body.encrypted_api_key)
|
||||
except Exception as exc:
|
||||
logger.warning("external_provider.decrypt_failed", error = str(exc))
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.",
|
||||
)
|
||||
return ExternalProviderClient(
|
||||
provider_type = "openai",
|
||||
base_url = base_url,
|
||||
api_key = api_key,
|
||||
)
|
||||
|
||||
|
||||
def _summarize_container(raw: dict) -> OpenAIContainerSummary:
|
||||
expires = raw.get("expires_after")
|
||||
expires_minutes: Optional[int] = None
|
||||
if isinstance(expires, dict):
|
||||
minutes = expires.get("minutes")
|
||||
if isinstance(minutes, int):
|
||||
expires_minutes = minutes
|
||||
return OpenAIContainerSummary(
|
||||
id = str(raw.get("id") or ""),
|
||||
name = raw.get("name"),
|
||||
created_at = raw.get("created_at")
|
||||
if isinstance(raw.get("created_at"), int)
|
||||
else None,
|
||||
last_active_at = raw.get("last_active_at")
|
||||
if isinstance(raw.get("last_active_at"), int)
|
||||
else None,
|
||||
expires_after_minutes = expires_minutes,
|
||||
status = raw.get("status") if isinstance(raw.get("status"), str) else None,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/external/openai/containers/list",
|
||||
response_model = ListOpenAIContainersResponse,
|
||||
)
|
||||
async def list_openai_containers(
|
||||
body: OpenAIContainerRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> ListOpenAIContainersResponse:
|
||||
"""List the user's OpenAI shell-tool containers."""
|
||||
client = _resolve_openai_cloud_client(body)
|
||||
try:
|
||||
try:
|
||||
raw = await client.list_openai_containers()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
detail = exc.response.text[:500] if exc.response is not None else str(exc)
|
||||
raise HTTPException(
|
||||
status_code = exc.response.status_code if exc.response else 502,
|
||||
detail = f"OpenAI rejected /containers list: {detail}",
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"Failed to reach OpenAI: {exc}",
|
||||
)
|
||||
# OpenAI keeps expired containers in /v1/containers indefinitely
|
||||
# with status="expired" — they're effectively dead but still
|
||||
# listed. Hide them so the picker only shows usable containers.
|
||||
return ListOpenAIContainersResponse(
|
||||
containers = [
|
||||
_summarize_container(c)
|
||||
for c in raw
|
||||
if isinstance(c, dict) and c.get("status") != "expired"
|
||||
],
|
||||
)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
@router.post(
|
||||
"/external/openai/containers/create",
|
||||
response_model = OpenAIContainerSummary,
|
||||
)
|
||||
async def create_openai_container(
|
||||
body: CreateOpenAIContainerBody,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> OpenAIContainerSummary:
|
||||
"""Create a named container with the user-chosen idle TTL."""
|
||||
client = _resolve_openai_cloud_client(body)
|
||||
try:
|
||||
try:
|
||||
raw = await client.create_openai_container(
|
||||
name = body.name,
|
||||
ttl_minutes = body.ttl_minutes,
|
||||
)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
detail = exc.response.text[:500] if exc.response is not None else str(exc)
|
||||
raise HTTPException(
|
||||
status_code = exc.response.status_code if exc.response else 502,
|
||||
detail = f"OpenAI rejected /containers create: {detail}",
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"Failed to reach OpenAI: {exc}",
|
||||
)
|
||||
if not isinstance(raw, dict):
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = "OpenAI returned an unexpected container payload.",
|
||||
)
|
||||
return _summarize_container(raw)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
@router.post("/external/openai/containers/delete", status_code = 204)
|
||||
async def delete_openai_container(
|
||||
body: DeleteOpenAIContainerBody,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> None:
|
||||
"""Delete a named container by id."""
|
||||
logger.info(
|
||||
"openai_container_delete.request subject=%s container_id=%s base_url=%s",
|
||||
current_subject,
|
||||
body.container_id,
|
||||
body.provider_base_url,
|
||||
)
|
||||
client = _resolve_openai_cloud_client(body)
|
||||
try:
|
||||
try:
|
||||
await client.delete_openai_container(body.container_id)
|
||||
logger.info(
|
||||
"openai_container_delete.success container_id=%s",
|
||||
body.container_id,
|
||||
)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
detail = exc.response.text[:500] if exc.response is not None else str(exc)
|
||||
logger.warning(
|
||||
"openai_container_delete.openai_rejected container_id=%s status=%s body=%s",
|
||||
body.container_id,
|
||||
exc.response.status_code if exc.response else None,
|
||||
detail,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = exc.response.status_code if exc.response else 502,
|
||||
detail = f"OpenAI rejected /containers delete: {detail}",
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
logger.warning(
|
||||
"openai_container_delete.transport_error container_id=%s error=%s",
|
||||
body.container_id,
|
||||
exc,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"Failed to reach OpenAI: {exc}",
|
||||
)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
@router.post("/chat/completions")
|
||||
async def openai_chat_completions(
|
||||
payload: ChatCompletionRequest,
|
||||
|
|
@ -1474,13 +1989,21 @@ async def openai_chat_completions(
|
|||
Supports multimodal messages: ``content`` may be a plain string or a
|
||||
list of content parts (``text`` / ``image_url``).
|
||||
|
||||
Streaming (default): returns SSE chunks matching OpenAI's format.
|
||||
Non-streaming: returns a single ChatCompletion JSON object.
|
||||
Non-streaming (default): returns a single ChatCompletion JSON object.
|
||||
Streaming: returns SSE chunks matching OpenAI's format.
|
||||
|
||||
``stream`` defaults to ``false`` to match OpenAI's spec; clients opt
|
||||
into SSE by sending ``stream: true``.
|
||||
|
||||
Automatically routes to the correct backend:
|
||||
- GGUF models → llama-server via LlamaCppBackend
|
||||
- Other models → Unsloth/transformers via InferenceBackend
|
||||
"""
|
||||
# ── External provider routing ────────────────────────────────
|
||||
# encrypted_api_key is optional — local providers (llama.cpp / vLLM / Ollama) may run without auth.
|
||||
if payload.provider_id or payload.provider_type:
|
||||
return await _proxy_to_external_provider(payload, request)
|
||||
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
using_gguf = llama_backend.is_loaded
|
||||
|
||||
|
|
@ -1669,6 +2192,12 @@ async def openai_chat_completions(
|
|||
and not _effective_enable_tools(payload)
|
||||
and (_tools_passthrough or _has_response_format)
|
||||
):
|
||||
if payload.audio_base64:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Audio input is not supported for GGUF chat models yet.",
|
||||
)
|
||||
|
||||
# Preserve the vision guard that would otherwise run in the
|
||||
# non-passthrough path below: text-only tool-capable GGUFs
|
||||
# should return a clear 400 here rather than forwarding the
|
||||
|
|
@ -1688,6 +2217,9 @@ async def openai_chat_completions(
|
|||
|
||||
cancel_event = threading.Event()
|
||||
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
|
||||
# `stream` defaults to False on ChatCompletionRequest (OpenAI spec
|
||||
# parity). Naive curl / .NET / System.Text.Json clients omitting
|
||||
# the field used to get SSE here and choke on deserialization (#5047).
|
||||
if payload.stream:
|
||||
return await _openai_passthrough_stream(
|
||||
request,
|
||||
|
|
@ -1716,6 +2248,12 @@ async def openai_chat_completions(
|
|||
|
||||
# ── GGUF path: proxy to llama-server /v1/chat/completions ──
|
||||
if using_gguf:
|
||||
if payload.audio_base64:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Audio input is not supported for GGUF chat models yet.",
|
||||
)
|
||||
|
||||
# Reject images if this GGUF model doesn't support vision
|
||||
image_b64 = extracted_image_b64 or payload.image_base64
|
||||
if image_b64 and not llama_backend.is_vision:
|
||||
|
|
@ -1729,7 +2267,7 @@ async def openai_chat_completions(
|
|||
try:
|
||||
import base64 as _b64
|
||||
from io import BytesIO as _BytesIO
|
||||
from PIL import Image as _Image
|
||||
from PIL import Image as _Image, UnidentifiedImageError as _UIE
|
||||
|
||||
raw = _b64.b64decode(image_b64)
|
||||
# Normalize to RGB so PNG encoding succeeds regardless of
|
||||
|
|
@ -1740,9 +2278,15 @@ async def openai_chat_completions(
|
|||
buf = _BytesIO()
|
||||
img.save(buf, format = "PNG")
|
||||
image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii")
|
||||
except Exception as e:
|
||||
except _UIE:
|
||||
raise HTTPException(
|
||||
status_code = 400, detail = f"Failed to process image: {e}"
|
||||
status_code = 400,
|
||||
detail = "Unsupported or corrupt image format.",
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Failed to process image.",
|
||||
)
|
||||
|
||||
# Build message list with system prompt prepended
|
||||
|
|
@ -3031,6 +3575,17 @@ async def _responses_stream(
|
|||
),
|
||||
)
|
||||
|
||||
# Direct pass-through bypasses the openai_chat_completions image gate.
|
||||
if not llama_backend.is_vision and any(
|
||||
isinstance(m.content, list)
|
||||
and any(isinstance(p, ImageContentPart) for p in m.content)
|
||||
for m in messages
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Image provided but current GGUF model does not support vision.",
|
||||
)
|
||||
|
||||
body = _build_openai_passthrough_body(
|
||||
chat_req, backend_ctx = llama_backend.context_length
|
||||
)
|
||||
|
|
@ -3412,10 +3967,10 @@ def _normalize_anthropic_openai_images(
|
|||
buf = io.BytesIO()
|
||||
img.save(buf, format = "PNG")
|
||||
png_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Failed to process image: {e}",
|
||||
detail = "Failed to process image.",
|
||||
)
|
||||
part["image_url"] = {"url": f"data:image/png;base64,{png_b64}"}
|
||||
|
||||
|
|
@ -3451,6 +4006,7 @@ async def anthropic_messages(
|
|||
[m.model_dump() for m in payload.messages],
|
||||
payload.system,
|
||||
)
|
||||
openai_messages = _drop_empty_assistant_sentinels(openai_messages)
|
||||
|
||||
# Enforce vision guard + re-encode embedded images to PNG so the
|
||||
# Anthropic endpoint matches the behavior of /v1/chat/completions.
|
||||
|
|
@ -4176,6 +4732,19 @@ async def _anthropic_passthrough_non_streaming(
|
|||
# =====================================================================
|
||||
|
||||
|
||||
def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]:
|
||||
"""Drop bare ``{"role":"assistant"}`` Stop-button sentinels; passthrough backends reject them."""
|
||||
out: list[dict] = []
|
||||
for m in messages:
|
||||
if m.get("role") == "assistant":
|
||||
has_content = bool(m.get("content"))
|
||||
has_tool_calls = bool(m.get("tool_calls"))
|
||||
if not has_content and not has_tool_calls:
|
||||
continue
|
||||
out.append(m)
|
||||
return out
|
||||
|
||||
|
||||
def _openai_messages_for_passthrough(payload) -> list[dict]:
|
||||
"""Build OpenAI-format message dicts for the /v1/chat/completions
|
||||
passthrough path.
|
||||
|
|
@ -4192,7 +4761,9 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
|
|||
``image_url`` content part so vision + function-calling requests work
|
||||
transparently.
|
||||
"""
|
||||
messages = [m.model_dump(exclude_none = True) for m in payload.messages]
|
||||
messages = _drop_empty_assistant_sentinels(
|
||||
[m.model_dump(exclude_none = True) for m in payload.messages]
|
||||
)
|
||||
|
||||
if not payload.image_base64:
|
||||
return messages
|
||||
|
|
@ -4207,10 +4778,10 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
|
|||
buf = _BytesIO()
|
||||
img.save(buf, format = "PNG")
|
||||
png_b64 = _b64.b64encode(buf.getvalue()).decode("ascii")
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Failed to process image: {e}",
|
||||
detail = "Failed to process image.",
|
||||
)
|
||||
|
||||
data_url = f"data:image/png;base64,{png_b64}"
|
||||
|
|
|
|||
|
|
@ -26,6 +26,22 @@ def _is_valid_repo_id(repo_id: str) -> bool:
|
|||
return bool(_VALID_REPO_ID.fullmatch(repo_id))
|
||||
|
||||
|
||||
def _safe_is_dir(path) -> bool:
|
||||
"""``Path.is_dir()`` that returns ``False`` instead of raising.
|
||||
|
||||
On Python >= 3.12 ``is_dir()``'s ``os.stat`` only suppresses
|
||||
"not found"-class errors and now propagates ``PermissionError``
|
||||
(EACCES); on Python <= 3.11 it returned ``False``. The folder-scan
|
||||
endpoints probe well-known system locations (e.g. a root-owned,
|
||||
mode-700 ``/usr/share/ollama/.ollama/models``) and must treat an
|
||||
un-stat-able path as "not a directory", never 500.
|
||||
"""
|
||||
try:
|
||||
return Path(path).is_dir()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
# Add backend directory to path
|
||||
backend_path = Path(__file__).parent.parent.parent
|
||||
if str(backend_path) not in sys.path:
|
||||
|
|
@ -882,7 +898,7 @@ async def get_recommended_folders(
|
|||
return
|
||||
if resolved in seen:
|
||||
return
|
||||
if Path(resolved).is_dir() and os.access(resolved, os.R_OK | os.X_OK):
|
||||
if _safe_is_dir(resolved) and os.access(resolved, os.R_OK | os.X_OK):
|
||||
seen.add(resolved)
|
||||
folders.append(resolved)
|
||||
|
||||
|
|
@ -1056,7 +1072,7 @@ def _build_browse_allowlist() -> list[Path]:
|
|||
resolved = p.resolve()
|
||||
except OSError:
|
||||
return
|
||||
if resolved.is_dir():
|
||||
if _safe_is_dir(resolved):
|
||||
candidates.append(resolved)
|
||||
|
||||
_add(Path.home())
|
||||
|
|
@ -1389,7 +1405,7 @@ async def browse_folders(
|
|||
return
|
||||
if resolved in seen_sug:
|
||||
return
|
||||
if Path(resolved).is_dir():
|
||||
if _safe_is_dir(resolved):
|
||||
seen_sug.add(resolved)
|
||||
suggestions.append(resolved)
|
||||
|
||||
|
|
|
|||
346
studio/backend/routes/providers.py
Normal file
346
studio/backend/routes/providers.py
Normal file
|
|
@ -0,0 +1,346 @@
|
|||
# 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 for external LLM provider management.
|
||||
|
||||
Provides endpoints for:
|
||||
- Discovering available provider types (registry)
|
||||
- CRUD for saved provider configurations (no API keys stored)
|
||||
- Fetching the RSA public key for API key encryption
|
||||
- Testing provider connectivity
|
||||
- Listing models from a provider
|
||||
"""
|
||||
|
||||
import uuid
|
||||
import structlog
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from core.inference.key_exchange import (
|
||||
decrypt_api_key,
|
||||
get_public_key_fingerprint,
|
||||
get_public_key_pem,
|
||||
)
|
||||
from core.inference.providers import (
|
||||
get_base_url,
|
||||
get_provider_info,
|
||||
list_available_providers,
|
||||
)
|
||||
from core.inference.external_provider import ExternalProviderClient
|
||||
from models.providers import (
|
||||
ProviderCreate,
|
||||
ProviderModelsRequest,
|
||||
ProviderModelInfo,
|
||||
ProviderResponse,
|
||||
ProviderRegistryEntry,
|
||||
ProviderTestRequest,
|
||||
ProviderTestResult,
|
||||
ProviderUpdate,
|
||||
)
|
||||
from storage import providers_db
|
||||
|
||||
logger = structlog.get_logger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
# ── Public key for API key encryption ─────────────────────────────
|
||||
|
||||
|
||||
@router.get("/public-key")
|
||||
async def get_public_key(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Return the RSA public key PEM for client-side API key encryption.
|
||||
|
||||
The ``fingerprint`` field is a short SHA256 of the PEM and is meant
|
||||
purely for diagnostics — a mismatch between what the frontend
|
||||
captured at encrypt time and what the server reports here is a
|
||||
clear signal that the keypair rotated mid-flight (e.g. the server
|
||||
re-ran ``init_key_pair`` for any reason).
|
||||
"""
|
||||
return {
|
||||
"public_key": get_public_key_pem(),
|
||||
"fingerprint": get_public_key_fingerprint(),
|
||||
}
|
||||
|
||||
|
||||
# ── Provider registry (static) ───────────────────────────────────
|
||||
|
||||
|
||||
@router.get("/registry", response_model = list[ProviderRegistryEntry])
|
||||
async def list_registry(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""List all supported provider types with their default configurations."""
|
||||
return list_available_providers()
|
||||
|
||||
|
||||
# ── Provider config CRUD ──────────────────────────────────────────
|
||||
|
||||
|
||||
@router.get("/", response_model = list[ProviderResponse])
|
||||
async def list_provider_configs(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""List all saved provider configurations."""
|
||||
rows = providers_db.list_providers()
|
||||
return [
|
||||
ProviderResponse(
|
||||
id = row["id"],
|
||||
provider_type = row["provider_type"],
|
||||
display_name = row["display_name"],
|
||||
base_url = row["base_url"],
|
||||
is_enabled = bool(row["is_enabled"]),
|
||||
created_at = row["created_at"],
|
||||
updated_at = row["updated_at"],
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
@router.post("/", response_model = ProviderResponse, status_code = 201)
|
||||
async def create_provider_config(
|
||||
payload: ProviderCreate,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Create a new saved provider configuration (no API key stored)."""
|
||||
info = get_provider_info(payload.provider_type)
|
||||
if info is None:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Unknown provider type: {payload.provider_type}. "
|
||||
f"Use GET /api/providers/registry to see available types.",
|
||||
)
|
||||
|
||||
provider_id = uuid.uuid4().hex[:16]
|
||||
base_url = payload.base_url or info["base_url"]
|
||||
|
||||
providers_db.create_provider(
|
||||
id = provider_id,
|
||||
provider_type = payload.provider_type,
|
||||
display_name = payload.display_name,
|
||||
base_url = base_url,
|
||||
)
|
||||
|
||||
row = providers_db.get_provider(provider_id)
|
||||
return ProviderResponse(
|
||||
id = row["id"],
|
||||
provider_type = row["provider_type"],
|
||||
display_name = row["display_name"],
|
||||
base_url = row["base_url"],
|
||||
is_enabled = bool(row["is_enabled"]),
|
||||
created_at = row["created_at"],
|
||||
updated_at = row["updated_at"],
|
||||
)
|
||||
|
||||
|
||||
@router.put("/{provider_id}", response_model = ProviderResponse)
|
||||
async def update_provider_config(
|
||||
provider_id: str,
|
||||
payload: ProviderUpdate,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Update a saved provider configuration."""
|
||||
existing = providers_db.get_provider(provider_id)
|
||||
if not existing:
|
||||
raise HTTPException(status_code = 404, detail = "Provider not found")
|
||||
|
||||
updated = providers_db.update_provider(
|
||||
id = provider_id,
|
||||
display_name = payload.display_name,
|
||||
base_url = payload.base_url,
|
||||
is_enabled = payload.is_enabled,
|
||||
)
|
||||
if not updated:
|
||||
raise HTTPException(status_code = 400, detail = "No fields to update")
|
||||
|
||||
row = providers_db.get_provider(provider_id)
|
||||
return ProviderResponse(
|
||||
id = row["id"],
|
||||
provider_type = row["provider_type"],
|
||||
display_name = row["display_name"],
|
||||
base_url = row["base_url"],
|
||||
is_enabled = bool(row["is_enabled"]),
|
||||
created_at = row["created_at"],
|
||||
updated_at = row["updated_at"],
|
||||
)
|
||||
|
||||
|
||||
@router.delete("/{provider_id}", status_code = 204)
|
||||
async def delete_provider_config(
|
||||
provider_id: str,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Delete a saved provider configuration."""
|
||||
deleted = providers_db.delete_provider(provider_id)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code = 404, detail = "Provider not found")
|
||||
|
||||
|
||||
# ── Test connectivity ─────────────────────────────────────────────
|
||||
|
||||
|
||||
@router.post("/test", response_model = ProviderTestResult)
|
||||
async def test_provider(
|
||||
payload: ProviderTestRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
Test connectivity to an external provider.
|
||||
|
||||
Makes a lightweight GET /models call to verify the API key works.
|
||||
The encrypted_api_key is decrypted server-side and never stored.
|
||||
"""
|
||||
info = get_provider_info(payload.provider_type)
|
||||
if info is None:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Unknown provider type: {payload.provider_type}",
|
||||
)
|
||||
|
||||
api_key = ""
|
||||
if payload.encrypted_api_key:
|
||||
try:
|
||||
api_key = decrypt_api_key(payload.encrypted_api_key)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to decrypt API key (%s): %s", type(exc).__name__, exc
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
|
||||
)
|
||||
|
||||
base_url = payload.base_url or info["base_url"]
|
||||
client = ExternalProviderClient(
|
||||
provider_type = payload.provider_type,
|
||||
base_url = base_url,
|
||||
api_key = api_key,
|
||||
timeout = 15.0,
|
||||
)
|
||||
|
||||
try:
|
||||
if info.get("model_list_mode") == "curated":
|
||||
await client.verify_models_endpoint_lightweight()
|
||||
return ProviderTestResult(
|
||||
success = True,
|
||||
message = (
|
||||
"Connected successfully. Full model list is not fetched for this provider — "
|
||||
"use suggestions and manual model IDs in the dialog."
|
||||
),
|
||||
models_count = None,
|
||||
)
|
||||
models = await client.list_models()
|
||||
return ProviderTestResult(
|
||||
success = True,
|
||||
message = f"Connected successfully. Found {len(models)} model(s).",
|
||||
models_count = len(models),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Provider test failed for %s: %s", payload.provider_type, exc)
|
||||
return ProviderTestResult(
|
||||
success = False,
|
||||
message = f"Connection failed: {exc}",
|
||||
models_count = None,
|
||||
)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
# ── List models from provider ─────────────────────────────────────
|
||||
|
||||
|
||||
@router.post("/models", response_model = list[ProviderModelInfo])
|
||||
async def list_provider_models(
|
||||
payload: ProviderModelsRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""
|
||||
List models available from an external provider.
|
||||
|
||||
The encrypted_api_key is decrypted server-side and never stored.
|
||||
"""
|
||||
info = get_provider_info(payload.provider_type)
|
||||
if info is None:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Unknown provider type: {payload.provider_type}",
|
||||
)
|
||||
|
||||
api_key = ""
|
||||
if payload.encrypted_api_key:
|
||||
try:
|
||||
api_key = decrypt_api_key(payload.encrypted_api_key)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to decrypt API key (%s): %s", type(exc).__name__, exc
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
|
||||
)
|
||||
|
||||
if info.get("model_list_mode") == "curated":
|
||||
return [
|
||||
ProviderModelInfo(
|
||||
id = m,
|
||||
display_name = m,
|
||||
context_length = None,
|
||||
owned_by = None,
|
||||
)
|
||||
for m in info.get("default_models", [])
|
||||
]
|
||||
|
||||
base_url = payload.base_url or info["base_url"]
|
||||
client = ExternalProviderClient(
|
||||
provider_type = payload.provider_type,
|
||||
base_url = base_url,
|
||||
api_key = api_key,
|
||||
timeout = 15.0,
|
||||
)
|
||||
|
||||
try:
|
||||
models = await client.list_models()
|
||||
allow_prefixes = info.get("model_id_allow_prefixes")
|
||||
if allow_prefixes is not None:
|
||||
prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p))
|
||||
if prefix_tuple:
|
||||
models = [m for m in models if m.get("id", "").startswith(prefix_tuple)]
|
||||
allowlist = info.get("model_id_allowlist")
|
||||
if allowlist is not None:
|
||||
models = [m for m in models if allowlist.match(m.get("id", ""))]
|
||||
deny_exact = info.get("model_id_deny_exact")
|
||||
if deny_exact is not None:
|
||||
deny_ids = {str(m) for m in deny_exact if str(m)}
|
||||
if deny_ids:
|
||||
models = [m for m in models if m.get("id", "") not in deny_ids]
|
||||
denylist = info.get("model_id_denylist")
|
||||
if denylist is not None:
|
||||
models = [m for m in models if not denylist.search(m.get("id", ""))]
|
||||
# Apply an optional cap after filtering so registry entries with a
|
||||
# large remote catalog (e.g. HF Inference Providers) can stay
|
||||
# picker-sized. No popularity sort happens server-side, so this is
|
||||
# "first N matches" — pair with default_models for any must-have
|
||||
# flagship ids.
|
||||
limit = info.get("model_id_limit")
|
||||
if isinstance(limit, int) and limit > 0:
|
||||
models = models[:limit]
|
||||
return [
|
||||
ProviderModelInfo(
|
||||
id = m.get("id", ""),
|
||||
display_name = m.get("id", ""),
|
||||
context_length = m.get("context_length") or m.get("context_window"),
|
||||
owned_by = m.get("owned_by"),
|
||||
)
|
||||
for m in models
|
||||
]
|
||||
except Exception as exc:
|
||||
logger.error("Failed to list models from %s: %s", payload.provider_type, exc)
|
||||
raise HTTPException(
|
||||
status_code = 502,
|
||||
detail = f"Failed to list models from {payload.provider_type}: {exc}",
|
||||
)
|
||||
finally:
|
||||
await client.close()
|
||||
|
|
@ -207,6 +207,7 @@ async def start_training(
|
|||
"custom_format_mapping": request.custom_format_mapping,
|
||||
"num_epochs": request.num_epochs,
|
||||
"learning_rate": request.learning_rate,
|
||||
"embedding_learning_rate": request.embedding_learning_rate,
|
||||
"batch_size": request.batch_size,
|
||||
"gradient_accumulation_steps": request.gradient_accumulation_steps,
|
||||
"warmup_steps": request.warmup_steps,
|
||||
|
|
@ -214,6 +215,7 @@ async def start_training(
|
|||
"max_steps": request.max_steps,
|
||||
"save_steps": request.save_steps,
|
||||
"weight_decay": request.weight_decay,
|
||||
"max_grad_norm": request.max_grad_norm,
|
||||
"random_seed": request.random_seed,
|
||||
"packing": request.packing,
|
||||
"optim": request.optim,
|
||||
|
|
|
|||
|
|
@ -18,8 +18,15 @@ from models import (
|
|||
TrainingRunListResponse,
|
||||
TrainingRunMetrics,
|
||||
TrainingRunSummary,
|
||||
TrainingRunUpdateRequest,
|
||||
)
|
||||
from storage.studio_db import (
|
||||
delete_run,
|
||||
get_run,
|
||||
get_run_metrics,
|
||||
list_runs,
|
||||
update_run_display_name,
|
||||
)
|
||||
from storage.studio_db import delete_run, get_run, get_run_metrics, list_runs
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -73,6 +80,34 @@ async def get_training_run_detail(
|
|||
)
|
||||
|
||||
|
||||
@router.patch("/runs/{run_id}", response_model = TrainingRunSummary)
|
||||
async def update_training_run(
|
||||
run_id: str,
|
||||
payload: TrainingRunUpdateRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Update mutable fields on a training run (currently only display_name)."""
|
||||
run = get_run(run_id)
|
||||
if run is None:
|
||||
raise HTTPException(status_code = 404, detail = f"Run {run_id} not found")
|
||||
|
||||
if "display_name" in payload.model_fields_set:
|
||||
next_display = payload.display_name
|
||||
if next_display is not None:
|
||||
next_display = next_display.strip() or None
|
||||
update_run_display_name(run_id, next_display)
|
||||
|
||||
refreshed = get_run(run_id)
|
||||
if refreshed is None:
|
||||
raise HTTPException(status_code = 404, detail = f"Run {run_id} not found")
|
||||
return TrainingRunSummary(
|
||||
**{
|
||||
**{k: v for k, v in refreshed.items() if k != "config_json"},
|
||||
"can_resume": can_resume_run(refreshed),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@router.delete("/runs/{run_id}", response_model = TrainingRunDeleteResponse)
|
||||
async def delete_training_run(
|
||||
run_id: str,
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ if str(backend_dir) not in sys.path:
|
|||
import _platform_compat # noqa: F401
|
||||
|
||||
from loggers import get_logger
|
||||
from startup_banner import print_studio_access_banner
|
||||
from startup_banner import print_studio_access_banner, print_studio_stop_hint
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -74,6 +74,255 @@ def _resolve_external_ip() -> str:
|
|||
return "0.0.0.0"
|
||||
|
||||
|
||||
def _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> None:
|
||||
"""Rewrite Uvicorn's startup log line: swap wildcard bind for the
|
||||
externally-reachable address, replace the CTRL+C suffix with our Mac-aware
|
||||
stop hint, and rename the prefix to "Unsloth Studio running on"."""
|
||||
import logging
|
||||
import re
|
||||
|
||||
rewrite_host = (
|
||||
bind_host in ("0.0.0.0", "::")
|
||||
and bool(display_host)
|
||||
and display_host != bind_host
|
||||
)
|
||||
new_suffix = "(To stop: press Ctrl+C -- on macOS, Control+C not Command+C)"
|
||||
old_suffix_re = re.compile(r"\(Press CTRL\+C to quit\)")
|
||||
old_prefix = "Uvicorn running on "
|
||||
new_prefix = "Unsloth Studio running on "
|
||||
|
||||
def _rewrite(text: str) -> str:
|
||||
if text.startswith(old_prefix):
|
||||
text = new_prefix + text[len(old_prefix) :]
|
||||
return old_suffix_re.sub(new_suffix, text)
|
||||
|
||||
class _UvicornStartupRewrite(logging.Filter):
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
try:
|
||||
msg = record.msg if isinstance(record.msg, str) else ""
|
||||
if (
|
||||
msg.startswith(old_prefix)
|
||||
and isinstance(record.args, tuple)
|
||||
and len(record.args) >= 3
|
||||
):
|
||||
if rewrite_host and record.args[1] == bind_host:
|
||||
record.args = (
|
||||
record.args[0],
|
||||
display_host,
|
||||
record.args[2],
|
||||
*record.args[3:],
|
||||
)
|
||||
record.msg = _rewrite(msg)
|
||||
cmsg = getattr(record, "color_message", None)
|
||||
if isinstance(cmsg, str):
|
||||
record.color_message = _rewrite(cmsg)
|
||||
except Exception:
|
||||
pass
|
||||
return True
|
||||
|
||||
f = _UvicornStartupRewrite()
|
||||
for name in ("uvicorn", "uvicorn.error"):
|
||||
logging.getLogger(name).addFilter(f)
|
||||
|
||||
|
||||
def _local_port_open(host: str, port: int, timeout: float = 1.0) -> bool:
|
||||
"""Return True iff a TCP connection to (host, port) succeeds within timeout."""
|
||||
import socket
|
||||
|
||||
try:
|
||||
with socket.create_connection((host, port), timeout = timeout):
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _working_local_url(port: int) -> "str | None":
|
||||
"""Return a working loopback URL on this machine, or None if neither
|
||||
127.0.0.1 nor ::1 responds. Used as a fallback when external reachability fails."""
|
||||
if _local_port_open("127.0.0.1", port):
|
||||
return f"http://127.0.0.1:{port}"
|
||||
if _local_port_open("::1", port):
|
||||
return f"http://[::1]:{port}"
|
||||
return None
|
||||
|
||||
|
||||
def _stdout_color_ok() -> bool:
|
||||
"""Whether to emit ANSI color codes on stdout. Mirrors startup_banner."""
|
||||
if os.environ.get("NO_COLOR", "").strip():
|
||||
return False
|
||||
if os.environ.get("FORCE_COLOR", "").strip():
|
||||
return True
|
||||
try:
|
||||
return sys.stdout.isatty()
|
||||
except (AttributeError, OSError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def _verify_global_reachability(display_host: str, port: int) -> None:
|
||||
"""Probe check-host.net to confirm display_host:port is reachable from the
|
||||
public internet. Synchronous so the caller can render output between the
|
||||
banner URL section and the trailing stop hint. Bounded at ~15s; failures
|
||||
are swallowed (the verifier failing is not Studio failing). Only meaningful
|
||||
when bound to a wildcard host."""
|
||||
import ipaddress
|
||||
import json
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
|
||||
if not display_host or display_host in ("0.0.0.0", "::"):
|
||||
return
|
||||
|
||||
use_color = _stdout_color_ok()
|
||||
dim = "\033[38;5;245m" if use_color else ""
|
||||
ok_c = "\033[38;5;120;1m" if use_color else ""
|
||||
err_c = "\033[38;5;203;1m" if use_color else ""
|
||||
warn_c = "\033[38;5;215;1m" if use_color else ""
|
||||
local_url_c = "\033[38;5;108;1m" if use_color else "" # matches banner's URL color
|
||||
reset = "\033[0m" if use_color else ""
|
||||
|
||||
url = f"http://{display_host}:{port}"
|
||||
|
||||
# Private / loopback / link-local addresses are not globally routable.
|
||||
try:
|
||||
addr = ipaddress.ip_address(display_host)
|
||||
if addr.is_loopback or addr.is_private or addr.is_link_local:
|
||||
print(
|
||||
f"{dim} Note: {display_host} is a private/LAN address -- "
|
||||
f"reachable on this network only, not from the public internet."
|
||||
f"{reset}",
|
||||
flush = True,
|
||||
)
|
||||
return
|
||||
except ValueError:
|
||||
# Not an IP literal; probe by hostname.
|
||||
pass
|
||||
|
||||
try:
|
||||
qs = urllib.parse.urlencode({"host": f"{display_host}:{port}", "max_nodes": 3})
|
||||
req = urllib.request.Request(
|
||||
f"https://check-host.net/check-tcp?{qs}",
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "unsloth-studio-reachability/1",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout = 5) as resp:
|
||||
init = json.loads(resp.read().decode("utf-8", errors = "replace"))
|
||||
req_id = init.get("request_id")
|
||||
if not req_id:
|
||||
return
|
||||
|
||||
results = {}
|
||||
deadline = time.monotonic() + 15.0
|
||||
poll_req = urllib.request.Request(
|
||||
f"https://check-host.net/check-result/{req_id}",
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "unsloth-studio-reachability/1",
|
||||
},
|
||||
)
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(1.5)
|
||||
try:
|
||||
with urllib.request.urlopen(poll_req, timeout = 5) as resp:
|
||||
results = json.loads(resp.read().decode("utf-8", errors = "replace"))
|
||||
except Exception:
|
||||
continue
|
||||
if results and all(v is not None for v in results.values()):
|
||||
break
|
||||
# Two decisive nodes is enough; stop polling early.
|
||||
decisive = [
|
||||
v
|
||||
for v in results.values()
|
||||
if isinstance(v, list)
|
||||
and v
|
||||
and isinstance(v[0], dict)
|
||||
and ("time" in v[0] or "error" in v[0])
|
||||
]
|
||||
if len(decisive) >= 2:
|
||||
break
|
||||
|
||||
ok_nodes = err_nodes = 0
|
||||
for v in results.values():
|
||||
if not isinstance(v, list) or not v or not isinstance(v[0], dict):
|
||||
continue
|
||||
if "time" in v[0]:
|
||||
ok_nodes += 1
|
||||
elif "error" in v[0]:
|
||||
err_nodes += 1
|
||||
total = ok_nodes + err_nodes
|
||||
|
||||
print("", flush = True)
|
||||
if ok_nodes:
|
||||
print(
|
||||
f"{ok_c} Reachability check: {url}/ is reachable from the "
|
||||
f"public internet ({ok_nodes}/{total} probe nodes connected).{reset}",
|
||||
flush = True,
|
||||
)
|
||||
elif err_nodes:
|
||||
print(
|
||||
f"{err_c} Reachability check: {url}/ is NOT reachable from "
|
||||
f"the public internet ({err_nodes}/{total} probe nodes failed).{reset}",
|
||||
flush = True,
|
||||
)
|
||||
print(f"{dim} Common causes:{reset}", flush = True)
|
||||
print(
|
||||
f"{dim} * AWS -- the instance's Security Group doesn't "
|
||||
f"allow inbound TCP {port}.{reset}",
|
||||
flush = True,
|
||||
)
|
||||
print(
|
||||
f"{dim} * GCP -- no firewall rule allowing TCP {port} "
|
||||
f"for the instance's network tag.{reset}",
|
||||
flush = True,
|
||||
)
|
||||
print(
|
||||
f"{dim} * Azure / other clouds -- equivalent NSG / "
|
||||
f"firewall rule missing.{reset}",
|
||||
flush = True,
|
||||
)
|
||||
print(
|
||||
f"{dim} * Home -- your router isn't port-forwarding "
|
||||
f"{port} to this machine.{reset}",
|
||||
flush = True,
|
||||
)
|
||||
print(
|
||||
f"{dim} Workaround that needs no firewall changes -- "
|
||||
f"SSH local-forward from your laptop:{reset}",
|
||||
flush = True,
|
||||
)
|
||||
print(
|
||||
f"{dim} ssh -L {port}:localhost:{port} "
|
||||
f"<user>@{display_host}{reset}",
|
||||
flush = True,
|
||||
)
|
||||
print(
|
||||
f"{dim} then open http://localhost:{port}/ in your browser.{reset}",
|
||||
flush = True,
|
||||
)
|
||||
# Only offer the local URL if loopback actually answers.
|
||||
local_url = _working_local_url(port)
|
||||
if local_url:
|
||||
print(
|
||||
f"{local_url_c} You can access Unsloth Studio locally "
|
||||
f"in the meantime: {local_url}{reset}",
|
||||
flush = True,
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"{warn_c} Reachability check: probe nodes did not respond "
|
||||
f"in time -- could not verify {url}/.{reset}",
|
||||
flush = True,
|
||||
)
|
||||
except urllib.error.URLError:
|
||||
# Outbound HTTPS blocked; skip silently.
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
|
||||
"""Return (pid, process_name) of the process listening on *port*, or None.
|
||||
|
||||
|
|
@ -159,7 +408,27 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int:
|
|||
)
|
||||
|
||||
|
||||
_PID_FILE = Path.home() / ".unsloth" / "studio" / "studio.pid"
|
||||
from utils.paths.storage_roots import studio_root as _studio_root
|
||||
|
||||
_PID_FILE = _studio_root() / "studio.pid"
|
||||
|
||||
# Direct backend launches bypass the CLI's env re-export; do it here for
|
||||
# real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR
|
||||
# picks up the custom build. Skip for legacy-default to avoid flipping
|
||||
# default-mode installs into env-override.
|
||||
try:
|
||||
_LEGACY_STUDIO_ROOT = (Path.home() / ".unsloth" / "studio").resolve()
|
||||
except (OSError, ValueError):
|
||||
_LEGACY_STUDIO_ROOT = Path.home() / ".unsloth" / "studio"
|
||||
try:
|
||||
_STUDIO_ROOT_RESOLVED = _studio_root().resolve()
|
||||
except (OSError, ValueError):
|
||||
_STUDIO_ROOT_RESOLVED = _studio_root()
|
||||
if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
|
||||
if not os.environ.get("UNSLOTH_STUDIO_HOME"):
|
||||
os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED)
|
||||
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
|
||||
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp")
|
||||
|
||||
|
||||
def _write_pid_file():
|
||||
|
|
@ -287,7 +556,6 @@ def run_server(
|
|||
|
||||
import asyncio
|
||||
from threading import Thread, Event
|
||||
import time
|
||||
import uvicorn
|
||||
|
||||
from main import app, setup_frontend
|
||||
|
|
@ -316,10 +584,6 @@ def run_server(
|
|||
print("=" * 50)
|
||||
print("")
|
||||
|
||||
# Output port for Tauri to parse when in api-only mode
|
||||
if api_only:
|
||||
print(f"TAURI_PORT={port}", flush = True)
|
||||
|
||||
# Setup frontend if path provided (skip in api-only mode)
|
||||
if frontend_path and not api_only:
|
||||
if setup_frontend(app, frontend_path):
|
||||
|
|
@ -329,11 +593,30 @@ def run_server(
|
|||
if not silent:
|
||||
print(f"[WARNING] Frontend not found at {frontend_path}")
|
||||
|
||||
# Create the uvicorn server and expose it for signal handlers
|
||||
# Resolve once; shared by the log rewrite and the banner.
|
||||
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
||||
_install_uvicorn_startup_log_rewrite(host, display_host)
|
||||
|
||||
ready_event = Event()
|
||||
startup_failed = Event()
|
||||
startup_errors = []
|
||||
|
||||
class _ReadyServer(uvicorn.Server):
|
||||
async def startup(self, *args, **kwargs):
|
||||
await super().startup(*args, **kwargs)
|
||||
if getattr(self, "started", False) and not self.should_exit:
|
||||
ready_event.set()
|
||||
|
||||
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
|
||||
config = uvicorn.Config(
|
||||
app, host = host, port = port, log_level = "info", access_log = False
|
||||
app,
|
||||
host = host,
|
||||
port = port,
|
||||
log_level = "info",
|
||||
access_log = False,
|
||||
server_header = False,
|
||||
)
|
||||
_server = uvicorn.Server(config)
|
||||
_server = _ReadyServer(config)
|
||||
_shutdown_event = Event()
|
||||
|
||||
# Expose the actual bound port so request-handling code can build
|
||||
|
|
@ -345,21 +628,8 @@ def run_server(
|
|||
app.state.server_port = port if port and port > 0 else None
|
||||
app.state.llama_parallel_slots = llama_parallel_slots
|
||||
|
||||
# Run server in a daemon thread
|
||||
def _run():
|
||||
asyncio.run(_server.serve())
|
||||
|
||||
thread = Thread(target = _run, daemon = True)
|
||||
thread.start()
|
||||
time.sleep(3)
|
||||
|
||||
_write_pid_file()
|
||||
import atexit
|
||||
|
||||
atexit.register(_remove_pid_file)
|
||||
|
||||
# Expose a shutdown callable via app.state so the /api/shutdown endpoint
|
||||
# can trigger graceful shutdown without circular imports.
|
||||
# Expose a shutdown callable via app.state before the server can accept
|
||||
# requests so /api/shutdown is available as soon as readiness is published.
|
||||
def _trigger_shutdown():
|
||||
_graceful_shutdown(_server)
|
||||
if _shutdown_event is not None:
|
||||
|
|
@ -367,13 +637,60 @@ def run_server(
|
|||
|
||||
app.state.trigger_shutdown = _trigger_shutdown
|
||||
|
||||
# Run server in a daemon thread
|
||||
def _run():
|
||||
try:
|
||||
asyncio.run(_server.serve())
|
||||
except BaseException as exc:
|
||||
startup_errors.append(exc)
|
||||
startup_failed.set()
|
||||
finally:
|
||||
if not ready_event.is_set():
|
||||
startup_failed.set()
|
||||
|
||||
thread = Thread(target = _run, daemon = True)
|
||||
thread.start()
|
||||
|
||||
# Wait until uvicorn has completed lifespan startup and bound sockets, or
|
||||
# until the server exits/fails before startup. This intentionally has no
|
||||
# correctness deadline: a slow but live startup should remain in progress.
|
||||
try:
|
||||
while not ready_event.is_set():
|
||||
if startup_failed.is_set() or not thread.is_alive():
|
||||
if startup_errors:
|
||||
raise RuntimeError(
|
||||
"Uvicorn server failed before startup completed"
|
||||
) from startup_errors[0]
|
||||
raise RuntimeError("Uvicorn server exited before startup completed")
|
||||
ready_event.wait(timeout = 0.1)
|
||||
except KeyboardInterrupt:
|
||||
_graceful_shutdown(_server)
|
||||
_shutdown_event.set()
|
||||
raise
|
||||
|
||||
_write_pid_file()
|
||||
import atexit
|
||||
|
||||
atexit.register(_remove_pid_file)
|
||||
|
||||
# Output port for Tauri to parse when in api-only mode. Emit only after
|
||||
# uvicorn sockets are bound and FastAPI lifespan/startup has completed.
|
||||
if api_only:
|
||||
print(f"TAURI_PORT={port}", flush = True)
|
||||
|
||||
if not silent:
|
||||
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
||||
wildcard_bind = host in ("0.0.0.0", "::")
|
||||
# For wildcard binds, run the reachability check between the URL
|
||||
# section and the stop hint so the stop hint stays last on screen.
|
||||
print_studio_access_banner(
|
||||
port = port,
|
||||
bind_host = host,
|
||||
display_host = display_host,
|
||||
include_stop_hint = not wildcard_bind,
|
||||
)
|
||||
if wildcard_bind:
|
||||
_verify_global_reachability(display_host, port)
|
||||
print_studio_stop_hint()
|
||||
|
||||
return app
|
||||
|
||||
|
|
|
|||
|
|
@ -33,18 +33,49 @@ def print_port_in_use_notice(original_port: int, new_port: int) -> None:
|
|||
print(msg)
|
||||
|
||||
|
||||
def print_studio_stop_hint() -> None:
|
||||
"""Print the trailing stop hint + closing divider. Separate from the main
|
||||
banner so callers can interleave content (e.g. a reachability check)."""
|
||||
use_color = stdout_supports_color()
|
||||
dim = "\033[38;5;245m"
|
||||
stop_hint_style = "\033[38;5;215;1m"
|
||||
reset = "\033[0m"
|
||||
|
||||
def style(text: str, code: str) -> str:
|
||||
return f"{code}{text}{reset}" if use_color else text
|
||||
|
||||
print(
|
||||
"\n".join(
|
||||
[
|
||||
"",
|
||||
style(
|
||||
" To stop Unsloth Studio: press Ctrl+C in this terminal.",
|
||||
stop_hint_style,
|
||||
),
|
||||
style(" (On macOS this is Control+C, not Command+C.)", dim),
|
||||
style("─" * 52, dim),
|
||||
"",
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def print_studio_access_banner(
|
||||
*,
|
||||
port: int,
|
||||
bind_host: str,
|
||||
display_host: str,
|
||||
include_stop_hint: bool = True,
|
||||
) -> None:
|
||||
"""Pretty-print URLs after the server is listening (beginner-friendly)."""
|
||||
"""Pretty-print URLs after the server is listening. Set
|
||||
``include_stop_hint=False`` to omit the trailing stop block; pair with
|
||||
:func:`print_studio_stop_hint` after inserting your own content."""
|
||||
use_color = stdout_supports_color()
|
||||
dim = "\033[38;5;245m"
|
||||
title = "\033[38;5;150m"
|
||||
local_url_style = "\033[38;5;108;1m"
|
||||
secondary = "\033[38;5;109m"
|
||||
stop_hint_style = "\033[38;5;215;1m"
|
||||
reset = "\033[0m"
|
||||
|
||||
def style(text: str, code: str) -> str:
|
||||
|
|
@ -116,8 +147,48 @@ def print_studio_access_banner(
|
|||
f" Tip: if you are on this computer, open {tip_url}/ in your browser.",
|
||||
dim,
|
||||
),
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
if loopback_bind and not listen_all:
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
style(
|
||||
" Studio is only reachable on this machine (bound to 127.0.0.1).",
|
||||
secondary,
|
||||
),
|
||||
style(
|
||||
" To deploy and access globally:",
|
||||
secondary,
|
||||
),
|
||||
style(
|
||||
" 1. press Ctrl+C to stop Studio",
|
||||
secondary,
|
||||
),
|
||||
style(
|
||||
f" 2. relaunch with: unsloth studio -H 0.0.0.0 -p {port}",
|
||||
secondary,
|
||||
),
|
||||
style(
|
||||
" Only do this on trusted networks -- it exposes the API on every interface.",
|
||||
secondary,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
if include_stop_hint:
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
style(
|
||||
" To stop Unsloth Studio: press Ctrl+C in this terminal.",
|
||||
stop_hint_style,
|
||||
),
|
||||
style(" (On macOS this is Control+C, not Command+C.)", dim),
|
||||
style("─" * 52, dim),
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
print("\n".join(lines))
|
||||
|
|
|
|||
153
studio/backend/storage/providers_db.py
Normal file
153
studio/backend/storage/providers_db.py
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
# 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 external LLM provider configurations.
|
||||
|
||||
Follows the same pattern as studio_db.py — module-level functions,
|
||||
raw sqlite3, WAL mode, per-function connections.
|
||||
|
||||
NOTE: API keys are NOT stored here. They live only in the browser
|
||||
(localStorage) and are sent encrypted per-request.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sqlite3
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from utils.paths import studio_db_path, ensure_dir
|
||||
|
||||
_schema_lock = threading.Lock()
|
||||
_schema_ready = False
|
||||
|
||||
|
||||
def _ensure_schema(conn: sqlite3.Connection) -> None:
|
||||
"""Create the llm_providers table if it doesn't exist. Called once per process."""
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS llm_providers (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
provider_type TEXT NOT NULL,
|
||||
display_name TEXT NOT NULL,
|
||||
base_url TEXT NOT NULL,
|
||||
is_enabled INTEGER NOT NULL DEFAULT 1,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def get_connection() -> sqlite3.Connection:
|
||||
"""Open studio.db with WAL mode, create table once per process."""
|
||||
global _schema_ready
|
||||
db_path = studio_db_path()
|
||||
ensure_dir(db_path.parent)
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
conn.row_factory = sqlite3.Row
|
||||
if not _schema_ready:
|
||||
with _schema_lock:
|
||||
if not _schema_ready:
|
||||
try:
|
||||
_ensure_schema(conn)
|
||||
_schema_ready = True
|
||||
except Exception:
|
||||
conn.close()
|
||||
raise
|
||||
return conn
|
||||
|
||||
|
||||
def create_provider(
|
||||
id: str,
|
||||
provider_type: str,
|
||||
display_name: str,
|
||||
base_url: str,
|
||||
) -> None:
|
||||
"""Insert a new provider configuration."""
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
conn = get_connection()
|
||||
try:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO llm_providers (id, provider_type, display_name, base_url, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(id, provider_type, display_name, base_url, now, now),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def update_provider(
|
||||
id: str,
|
||||
display_name: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
is_enabled: Optional[bool] = None,
|
||||
) -> bool:
|
||||
"""Update fields on an existing provider. Returns True if a row was updated."""
|
||||
updates = []
|
||||
params = []
|
||||
if display_name is not None:
|
||||
updates.append("display_name = ?")
|
||||
params.append(display_name)
|
||||
if base_url is not None:
|
||||
updates.append("base_url = ?")
|
||||
params.append(base_url)
|
||||
if is_enabled is not None:
|
||||
updates.append("is_enabled = ?")
|
||||
params.append(1 if is_enabled else 0)
|
||||
if not updates:
|
||||
return False
|
||||
updates.append("updated_at = ?")
|
||||
params.append(datetime.now(timezone.utc).isoformat())
|
||||
params.append(id)
|
||||
|
||||
conn = get_connection()
|
||||
try:
|
||||
cursor = conn.execute(
|
||||
f"UPDATE llm_providers SET {', '.join(updates)} WHERE id = ?",
|
||||
params,
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def delete_provider(id: str) -> bool:
|
||||
"""Delete a provider by ID. Returns True if a row was deleted."""
|
||||
conn = get_connection()
|
||||
try:
|
||||
cursor = conn.execute("DELETE FROM llm_providers WHERE id = ?", (id,))
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def get_provider(id: str) -> Optional[dict]:
|
||||
"""Fetch a single provider by ID."""
|
||||
conn = get_connection()
|
||||
try:
|
||||
row = conn.execute("SELECT * FROM llm_providers WHERE id = ?", (id,)).fetchone()
|
||||
return dict(row) if row else None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def list_providers() -> list[dict]:
|
||||
"""List all provider configurations, ordered by creation time."""
|
||||
conn = get_connection()
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM llm_providers ORDER BY created_at"
|
||||
).fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
finally:
|
||||
conn.close()
|
||||
|
|
@ -75,10 +75,16 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
output_dir TEXT,
|
||||
error_message TEXT,
|
||||
duration_seconds REAL,
|
||||
loss_sparkline TEXT
|
||||
loss_sparkline TEXT,
|
||||
display_name TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
existing_cols = {
|
||||
row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()
|
||||
}
|
||||
if "display_name" not in existing_cols:
|
||||
conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT")
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS training_metrics (
|
||||
|
|
@ -261,6 +267,18 @@ def insert_metrics_batch(run_id: str, metrics: list[dict]) -> None:
|
|||
conn.close()
|
||||
|
||||
|
||||
def update_run_display_name(id: str, display_name: Optional[str]) -> None:
|
||||
conn = get_connection()
|
||||
try:
|
||||
conn.execute(
|
||||
"UPDATE training_runs SET display_name = ? WHERE id = ?",
|
||||
(display_name, id),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def list_runs(limit: int = 50, offset: int = 0) -> dict:
|
||||
conn = get_connection()
|
||||
try:
|
||||
|
|
@ -270,7 +288,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
|
|||
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
|
||||
r.ended_at, r.total_steps, r.final_step, r.final_loss,
|
||||
r.output_dir, r.duration_seconds, r.error_message,
|
||||
r.loss_sparkline,
|
||||
r.loss_sparkline, r.display_name,
|
||||
CASE
|
||||
WHEN r.status = 'stopped'
|
||||
AND r.output_dir IS NOT NULL
|
||||
|
|
|
|||
419
studio/backend/tests/test_anthropic_code_execution.py
Normal file
419
studio/backend/tests/test_anthropic_code_execution.py
Normal file
|
|
@ -0,0 +1,419 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""
|
||||
Unit tests for Anthropic's server-side `code_execution_20250825` tool
|
||||
translation in `_stream_anthropic`.
|
||||
|
||||
Covers:
|
||||
- Request body: when ``enabled_tools=["code_execution"]``, the outbound
|
||||
``tools`` array carries ``{"type": "code_execution_20250825", "name":
|
||||
"code_execution"}`` and the ``anthropic-beta`` header includes
|
||||
``code-execution-2025-08-25``.
|
||||
- Combined request: ``enabled_tools=["web_search", "code_execution"]``
|
||||
sends both tool entries; the beta header still merges the code-exec
|
||||
flag onto whatever the registry contributed.
|
||||
- SSE translation: a `bash_code_execution` server_tool_use +
|
||||
`bash_code_execution_tool_result` pair emits one tool_start and one
|
||||
tool_end ``_toolEvent`` chunk with the expected arguments and result.
|
||||
- SSE translation: a `text_editor_code_execution` create + result emits
|
||||
a tool_start with ``kind="text_editor"`` + parsed args, and tool_end
|
||||
with ``"Created"`` (or ``"Updated"``) based on the ``is_file_update``
|
||||
flag.
|
||||
- Error path: a ``bash_code_execution_tool_result_error`` with
|
||||
``error_code="container_expired"`` renders as ``"Error:
|
||||
container_expired"`` in the tool_end ``result``.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
from core.inference import external_provider as ep_mod
|
||||
from core.inference.external_provider import ExternalProviderClient
|
||||
|
||||
|
||||
def _drive(coro):
|
||||
return asyncio.new_event_loop().run_until_complete(coro)
|
||||
|
||||
|
||||
async def _collect(agen):
|
||||
out = []
|
||||
async for line in agen:
|
||||
out.append(line)
|
||||
return out
|
||||
|
||||
|
||||
def _mock_http_client(monkeypatch, handler):
|
||||
transport = httpx.MockTransport(handler)
|
||||
monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport))
|
||||
|
||||
|
||||
def _make_client() -> ExternalProviderClient:
|
||||
return ExternalProviderClient(
|
||||
provider_type = "anthropic",
|
||||
base_url = "https://api.anthropic.com/v1",
|
||||
api_key = "sk-ant-test",
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_sse(events: list[dict]) -> bytes:
|
||||
chunks: list[str] = []
|
||||
for event in events:
|
||||
chunks.append(f"event: {event['type']}")
|
||||
chunks.append(f"data: {json.dumps(event)}")
|
||||
chunks.append("")
|
||||
return ("\n".join(chunks) + "\n").encode("utf-8")
|
||||
|
||||
|
||||
def _tool_events(lines: list[str]) -> list[dict]:
|
||||
"""Extract `_toolEvent` payloads from emitted SSE data lines."""
|
||||
out: list[dict] = []
|
||||
for line in lines:
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
raw = line[len("data:") :].strip()
|
||||
if not raw or raw == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(parsed, dict) and "_toolEvent" in parsed:
|
||||
out.append(parsed["_toolEvent"])
|
||||
return out
|
||||
|
||||
|
||||
def test_code_execution_tool_appended_to_request_body(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
captured["headers"] = dict(request.headers)
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "compute 2 + 2"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
enabled_tools = ["code_execution"],
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
body = captured["body"]
|
||||
tools = body.get("tools") or []
|
||||
assert {
|
||||
"type": "code_execution_20250825",
|
||||
"name": "code_execution",
|
||||
} in tools
|
||||
# No web_search entry when only code_execution is enabled.
|
||||
assert all(t.get("type") != "web_search_20250305" for t in tools)
|
||||
# Beta header carries the documented flag.
|
||||
beta_header = captured["headers"].get("anthropic-beta", "")
|
||||
assert "code-execution-2025-08-25" in beta_header
|
||||
|
||||
|
||||
def test_code_execution_with_web_search_sends_both_tools(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
captured["headers"] = dict(request.headers)
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "look it up and chart it"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
enabled_tools = ["web_search", "code_execution"],
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
tools = captured["body"].get("tools") or []
|
||||
tool_types = {t.get("type") for t in tools if isinstance(t, dict)}
|
||||
assert "web_search_20250305" in tool_types
|
||||
assert "code_execution_20250825" in tool_types
|
||||
assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "")
|
||||
|
||||
|
||||
def test_no_code_execution_tool_when_pill_off(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
captured["headers"] = dict(request.headers)
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
tools = captured["body"].get("tools") or []
|
||||
assert all(t.get("type") != "code_execution_20250825" for t in tools)
|
||||
# Beta header must NOT mention code-execution when the tool isn't on
|
||||
# — that flag is opt-in only.
|
||||
assert "code-execution-2025-08-25" not in captured["headers"].get(
|
||||
"anthropic-beta", ""
|
||||
)
|
||||
|
||||
|
||||
def test_bash_code_execution_emits_tool_start_and_end(monkeypatch):
|
||||
sse_events = [
|
||||
{"type": "message_start", "message": {"usage": {}}},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_1",
|
||||
"name": "bash_code_execution",
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": '{"command": "ls -la"}',
|
||||
},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 1,
|
||||
"content_block": {
|
||||
"type": "bash_code_execution_tool_result",
|
||||
"tool_use_id": "srvtoolu_1",
|
||||
"content": {
|
||||
"type": "bash_code_execution_result",
|
||||
"stdout": "total 24\ndrwxr-xr-x 2 user user 4096 Jan 1 12:00 .",
|
||||
"stderr": "",
|
||||
"return_code": 0,
|
||||
},
|
||||
},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
{"type": "message_stop"},
|
||||
]
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse(sse_events),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
return await _collect(
|
||||
client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "list files"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
enabled_tools = ["code_execution"],
|
||||
)
|
||||
)
|
||||
|
||||
lines = _drive(run())
|
||||
events = _tool_events(lines)
|
||||
|
||||
assert len(events) == 2
|
||||
start, end = events
|
||||
assert start["type"] == "tool_start"
|
||||
assert start["tool_name"] == "code_execution"
|
||||
assert start["tool_call_id"] == "srvtoolu_1"
|
||||
assert start["arguments"] == {"kind": "bash", "command": "ls -la"}
|
||||
|
||||
assert end["type"] == "tool_end"
|
||||
assert end["tool_call_id"] == "srvtoolu_1"
|
||||
assert "total 24" in end["result"]
|
||||
# Non-zero return_code not present, so no return_code line.
|
||||
assert "return_code:" not in end["result"]
|
||||
|
||||
|
||||
def test_text_editor_create_emits_kind_and_status(monkeypatch):
|
||||
sse_events = [
|
||||
{"type": "message_start", "message": {"usage": {}}},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_2",
|
||||
"name": "text_editor_code_execution",
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": (
|
||||
'{"command": "create", "path": "new_file.txt", '
|
||||
'"file_text": "hi"}'
|
||||
),
|
||||
},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 1,
|
||||
"content_block": {
|
||||
"type": "text_editor_code_execution_tool_result",
|
||||
"tool_use_id": "srvtoolu_2",
|
||||
"content": {
|
||||
"type": "text_editor_code_execution_result",
|
||||
"is_file_update": False,
|
||||
},
|
||||
},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
{"type": "message_stop"},
|
||||
]
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse(sse_events),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
return await _collect(
|
||||
client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "write a file"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
enabled_tools = ["code_execution"],
|
||||
)
|
||||
)
|
||||
|
||||
lines = _drive(run())
|
||||
events = _tool_events(lines)
|
||||
|
||||
assert len(events) == 2
|
||||
start, end = events
|
||||
assert start["arguments"]["kind"] == "text_editor"
|
||||
assert start["arguments"]["command"] == "create"
|
||||
assert start["arguments"]["path"] == "new_file.txt"
|
||||
assert end["result"] == "Created"
|
||||
|
||||
|
||||
def test_code_execution_error_renders_error_code(monkeypatch):
|
||||
sse_events = [
|
||||
{"type": "message_start", "message": {"usage": {}}},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_3",
|
||||
"name": "bash_code_execution",
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": '{"command": "echo broken"}',
|
||||
},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 1,
|
||||
"content_block": {
|
||||
"type": "bash_code_execution_tool_result",
|
||||
"tool_use_id": "srvtoolu_3",
|
||||
"content": {
|
||||
"type": "bash_code_execution_tool_result_error",
|
||||
"error_code": "container_expired",
|
||||
},
|
||||
},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
{"type": "message_stop"},
|
||||
]
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse(sse_events),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
return await _collect(
|
||||
client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "run it"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
enabled_tools = ["code_execution"],
|
||||
)
|
||||
)
|
||||
|
||||
lines = _drive(run())
|
||||
events = _tool_events(lines)
|
||||
|
||||
assert len(events) == 2
|
||||
end = events[1]
|
||||
assert end["type"] == "tool_end"
|
||||
assert end["result"] == "Error: container_expired"
|
||||
404
studio/backend/tests/test_anthropic_thinking_translation.py
Normal file
404
studio/backend/tests/test_anthropic_thinking_translation.py
Normal file
|
|
@ -0,0 +1,404 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""
|
||||
Unit tests for the Anthropic extended-thinking translation in
|
||||
external_provider.
|
||||
|
||||
Covers:
|
||||
- Adaptive-mode request body nests effort under
|
||||
``output_config: {effort: "<level>"}`` per the Messages API
|
||||
reference (a top-level ``effort`` field 400s with
|
||||
"effort: Extra inputs are not permitted").
|
||||
- Streaming SSE: ``content_block_delta`` with
|
||||
``delta.type == "thinking_delta"`` is translated into inline
|
||||
``<think>...</think>`` chat-completion chunks so the frontend's
|
||||
reasoning-panel pipeline lifts it correctly.
|
||||
- The ``<think>`` tag closes when the first ``text_delta`` arrives,
|
||||
on ``content_block_stop``, on ``message_delta``, or on
|
||||
``message_stop``.
|
||||
- Thinking is paired with ``temperature=1`` and no ``top_p`` /
|
||||
``top_k`` on the wire (Anthropic extended-thinking contract).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
from core.inference import external_provider as ep_mod
|
||||
from core.inference.external_provider import ExternalProviderClient
|
||||
|
||||
|
||||
def _drive(coro):
|
||||
return asyncio.new_event_loop().run_until_complete(coro)
|
||||
|
||||
|
||||
async def _collect(agen):
|
||||
out = []
|
||||
async for line in agen:
|
||||
out.append(line)
|
||||
return out
|
||||
|
||||
|
||||
def _mock_http_client(monkeypatch, handler):
|
||||
transport = httpx.MockTransport(handler)
|
||||
monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport))
|
||||
|
||||
|
||||
def _make_client() -> ExternalProviderClient:
|
||||
return ExternalProviderClient(
|
||||
provider_type = "anthropic",
|
||||
base_url = "https://api.anthropic.com/v1",
|
||||
api_key = "sk-ant-test",
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_sse(events: list[dict]) -> bytes:
|
||||
"""Serialize a list of Messages-API event dicts as an SSE byte stream."""
|
||||
chunks: list[str] = []
|
||||
for event in events:
|
||||
chunks.append(f"event: {event['type']}")
|
||||
chunks.append(f"data: {json.dumps(event)}")
|
||||
chunks.append("")
|
||||
return ("\n".join(chunks) + "\n").encode("utf-8")
|
||||
|
||||
|
||||
def _payloads_from_lines(lines: list[str]) -> list:
|
||||
out = []
|
||||
for line in lines:
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
raw = line[len("data:") :].strip()
|
||||
if not raw:
|
||||
continue
|
||||
if raw == "[DONE]":
|
||||
out.append("[DONE]")
|
||||
else:
|
||||
out.append(json.loads(raw))
|
||||
return out
|
||||
|
||||
|
||||
def test_adaptive_thinking_body_uses_output_config_effort_shape(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-opus-4-6",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
top_k = None,
|
||||
enable_thinking = None,
|
||||
reasoning_effort = "medium",
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
body = captured["body"]
|
||||
# display=summarized is set explicitly so Opus 4.7 (which defaults to
|
||||
# "omitted") still emits thinking_delta events for the reasoning panel.
|
||||
assert body["thinking"] == {"type": "adaptive", "display": "summarized"}
|
||||
# Documented shape: effort is nested under output_config.
|
||||
# A top-level `effort` field produces a 400:
|
||||
# "effort: Extra inputs are not permitted".
|
||||
assert body["output_config"] == {"effort": "medium"}
|
||||
assert "effort" not in body
|
||||
# Extended-thinking contract: temperature=1, no top_p / top_k.
|
||||
assert body["temperature"] == 1
|
||||
assert "top_p" not in body
|
||||
assert "top_k" not in body
|
||||
|
||||
|
||||
def test_adaptive_thinking_maps_xhigh_to_max_on_claude_4_6(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-sonnet-4-6",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
top_k = None,
|
||||
enable_thinking = None,
|
||||
reasoning_effort = "xhigh",
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
assert captured["body"]["output_config"] == {"effort": "max"}
|
||||
|
||||
|
||||
def test_adaptive_thinking_keeps_max_on_claude_4_6(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-opus-4-6",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
top_k = None,
|
||||
enable_thinking = None,
|
||||
reasoning_effort = "max",
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
assert captured["body"]["output_config"] == {"effort": "max"}
|
||||
|
||||
|
||||
def test_adaptive_thinking_keeps_xhigh_on_claude_4_7(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
top_k = None,
|
||||
enable_thinking = None,
|
||||
reasoning_effort = "xhigh",
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
body = captured["body"]
|
||||
assert body["output_config"] == {"effort": "xhigh"}
|
||||
assert "effort" not in body
|
||||
|
||||
|
||||
def test_manual_thinking_body_uses_budget_tokens_on_4_5(monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content.decode("utf-8"))
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse([{"type": "message_stop"}]),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
async for _ in client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-opus-4-5",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 1024,
|
||||
top_k = None,
|
||||
enable_thinking = None,
|
||||
reasoning_effort = "high",
|
||||
):
|
||||
pass
|
||||
await client.close()
|
||||
|
||||
_drive(run())
|
||||
|
||||
body = captured["body"]
|
||||
assert body["thinking"] == {"type": "enabled", "budget_tokens": 4096}
|
||||
# max_tokens must be strictly greater than budget_tokens; we shipped 1024
|
||||
# and budget is 4096, so the wrapper should bump max_tokens.
|
||||
assert body["max_tokens"] > body["thinking"]["budget_tokens"]
|
||||
# Manual-thinking path does not use output_config / effort — those are
|
||||
# the adaptive-mode controls (Claude 4.6 / 4.7).
|
||||
assert "effort" not in body
|
||||
assert "output_config" not in body
|
||||
|
||||
|
||||
def test_thinking_delta_wrapped_in_think_tags(monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
events = [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "thinking", "thinking": "", "signature": ""},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "thinking_delta", "thinking": "First "},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "thinking_delta", "thinking": "I plan."},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "signature_delta", "signature": "abc123"},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 1,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 1,
|
||||
"delta": {"type": "text_delta", "text": "Answer."},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}},
|
||||
{"type": "message_stop"},
|
||||
]
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse(events),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
lines = await _collect(
|
||||
client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-opus-4-6",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
top_k = None,
|
||||
enable_thinking = True,
|
||||
reasoning_effort = None,
|
||||
)
|
||||
)
|
||||
await client.close()
|
||||
return lines
|
||||
|
||||
lines = _drive(run())
|
||||
payloads = _payloads_from_lines(lines)
|
||||
|
||||
combined = "".join(
|
||||
p["choices"][0]["delta"].get("content", "")
|
||||
for p in payloads
|
||||
if isinstance(p, dict) and p["choices"][0]["delta"]
|
||||
)
|
||||
|
||||
# Reasoning text should be wrapped in <think>...</think>, followed by the
|
||||
# answer text, and the stream should terminate with [DONE].
|
||||
assert "<think>First I plan.</think>" in combined
|
||||
assert combined.endswith("Answer.")
|
||||
# signature_delta is intentionally dropped — no leaked signature text.
|
||||
assert "abc123" not in combined
|
||||
assert "[DONE]" in payloads
|
||||
|
||||
|
||||
def test_thinking_only_turn_closes_tag_without_text_delta(monkeypatch):
|
||||
"""display=omitted on Claude 4.7 emits a signature_delta and no text.
|
||||
|
||||
The <think> open is still triggered by the (synthetic) thinking_delta;
|
||||
we want content_block_stop to close it cleanly so the tag never leaks
|
||||
into the next chunk."""
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
events = [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "thinking", "thinking": "", "signature": ""},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "thinking_delta", "thinking": "internal"},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}},
|
||||
{"type": "message_stop"},
|
||||
]
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = _anthropic_sse(events),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
_mock_http_client(monkeypatch, handler)
|
||||
|
||||
async def run():
|
||||
client = _make_client()
|
||||
lines = await _collect(
|
||||
client._stream_anthropic(
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
model = "claude-opus-4-7",
|
||||
temperature = 0.7,
|
||||
top_p = 0.95,
|
||||
max_tokens = 4096,
|
||||
top_k = None,
|
||||
enable_thinking = True,
|
||||
reasoning_effort = None,
|
||||
)
|
||||
)
|
||||
await client.close()
|
||||
return lines
|
||||
|
||||
payloads = _payloads_from_lines(_drive(run()))
|
||||
combined = "".join(
|
||||
p["choices"][0]["delta"].get("content", "")
|
||||
for p in payloads
|
||||
if isinstance(p, dict) and p["choices"][0]["delta"]
|
||||
)
|
||||
assert combined == "<think>internal</think>"
|
||||
180
studio/backend/tests/test_cleanup_cancelled_checkpoints.py
Normal file
180
studio/backend/tests/test_cleanup_cancelled_checkpoints.py
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""Tests for core/training/training.py:_cleanup_cancelled_checkpoints."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(_BACKEND_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND_ROOT))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def outputs_setup(tmp_path, monkeypatch):
|
||||
"""Point outputs_root() at a temp dir so cleanup is allowed to run on it.
|
||||
|
||||
The training module binds ``outputs_root`` at import time
|
||||
(``from utils.paths import outputs_root``), so we have to patch
|
||||
the symbol on the importer module, not on storage_roots.
|
||||
"""
|
||||
from core.training import training as training_mod
|
||||
|
||||
monkeypatch.setattr(training_mod, "outputs_root", lambda: tmp_path)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _mk_dir(parent: Path, name: str) -> Path:
|
||||
p = parent / name
|
||||
p.mkdir()
|
||||
(p / "marker.txt").write_text(name)
|
||||
return p
|
||||
|
||||
|
||||
def test_completed_checkpoints_are_preserved(outputs_setup):
|
||||
"""The big regression: prior to this fix, every completed
|
||||
checkpoint-N/ was rmtree'd on Cancel, destroying resume points."""
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
out = outputs_setup / "run-1"
|
||||
out.mkdir()
|
||||
ckpts = [_mk_dir(out, f"checkpoint-{n}") for n in (200, 400, 600)]
|
||||
tmp = _mk_dir(out, "tmp-checkpoint-800")
|
||||
|
||||
_cleanup_cancelled_checkpoints(out)
|
||||
|
||||
for c in ckpts:
|
||||
assert c.exists(), f"completed {c.name} was destroyed"
|
||||
assert (c / "marker.txt").exists()
|
||||
assert not tmp.exists(), "in-flight tmp-checkpoint-800 should be removed"
|
||||
|
||||
|
||||
def test_in_flight_tmp_checkpoints_removed(outputs_setup):
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
out = outputs_setup / "run-2"
|
||||
out.mkdir()
|
||||
_mk_dir(out, "tmp-checkpoint-100")
|
||||
_mk_dir(out, "tmp-checkpoint-200")
|
||||
_mk_dir(out, "checkpoint-50") # completed, kept
|
||||
|
||||
_cleanup_cancelled_checkpoints(out)
|
||||
|
||||
assert not (out / "tmp-checkpoint-100").exists()
|
||||
assert not (out / "tmp-checkpoint-200").exists()
|
||||
assert (out / "checkpoint-50").exists()
|
||||
|
||||
|
||||
def test_non_checkpoint_dirs_left_alone(outputs_setup):
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
out = outputs_setup / "run-3"
|
||||
out.mkdir()
|
||||
_mk_dir(out, "logs")
|
||||
_mk_dir(out, "tensorboard")
|
||||
_mk_dir(out, "checkpoint-final") # non-int suffix, kept
|
||||
_mk_dir(out, "checkpoint-best")
|
||||
_mk_dir(out, "tmp-checkpoint-99")
|
||||
|
||||
_cleanup_cancelled_checkpoints(out)
|
||||
|
||||
for n in ("logs", "tensorboard", "checkpoint-final", "checkpoint-best"):
|
||||
assert (out / n).exists(), f"{n} should be preserved"
|
||||
assert not (out / "tmp-checkpoint-99").exists()
|
||||
|
||||
|
||||
def test_output_dir_outside_outputs_root_is_refused(tmp_path, monkeypatch):
|
||||
"""Containment check: even if a bug passed an output_dir outside
|
||||
outputs_root, the cleanup must refuse to touch it."""
|
||||
from core.training import training as training_mod
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
inside = tmp_path / "inside"
|
||||
inside.mkdir()
|
||||
monkeypatch.setattr(training_mod, "outputs_root", lambda: inside)
|
||||
|
||||
outside = tmp_path / "outside"
|
||||
outside.mkdir()
|
||||
_mk_dir(outside, "tmp-checkpoint-1")
|
||||
|
||||
_cleanup_cancelled_checkpoints(outside)
|
||||
|
||||
assert (
|
||||
outside / "tmp-checkpoint-1"
|
||||
).exists(), "must not rmtree under a path outside outputs_root"
|
||||
|
||||
|
||||
def test_symlinked_output_dir_skipped(outputs_setup):
|
||||
"""A symlinked output_dir is skipped so the realpath check can't be
|
||||
leveraged to delete content via a symlink trick."""
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
real = outputs_setup / "real-run"
|
||||
real.mkdir()
|
||||
_mk_dir(real, "tmp-checkpoint-1")
|
||||
|
||||
link = outputs_setup / "link-run"
|
||||
try:
|
||||
link.symlink_to(real, target_is_directory = True)
|
||||
except (OSError, NotImplementedError):
|
||||
pytest.skip("symlinks not supported on this filesystem / platform")
|
||||
|
||||
_cleanup_cancelled_checkpoints(link)
|
||||
|
||||
assert (real / "tmp-checkpoint-1").exists(), "symlinked output_dir must be skipped"
|
||||
|
||||
|
||||
def test_missing_output_dir_is_noop(outputs_setup):
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
_cleanup_cancelled_checkpoints(outputs_setup / "does-not-exist")
|
||||
# Should not raise; nothing to assert beyond non-failure.
|
||||
|
||||
|
||||
def test_symlinked_child_skipped(outputs_setup):
|
||||
"""A symlinked tmp-checkpoint-* child must not be deleted, so the
|
||||
realpath bypass cannot redirect rmtree to arbitrary content."""
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
out = outputs_setup / "run-symchild"
|
||||
out.mkdir()
|
||||
target = outputs_setup / "external"
|
||||
target.mkdir()
|
||||
(target / "important.txt").write_text("keep me")
|
||||
|
||||
link = out / "tmp-checkpoint-99"
|
||||
try:
|
||||
link.symlink_to(target, target_is_directory = True)
|
||||
except (OSError, NotImplementedError):
|
||||
pytest.skip("symlinks not supported on this filesystem / platform")
|
||||
|
||||
_cleanup_cancelled_checkpoints(out)
|
||||
|
||||
assert (
|
||||
target / "important.txt"
|
||||
).exists(), "symlink target outside outputs_root must not be rmtree'd"
|
||||
|
||||
|
||||
def test_non_numeric_tmp_checkpoint_suffix_preserved(outputs_setup):
|
||||
"""HF Trainer's partials are tmp-checkpoint-<step>. A user-named
|
||||
tmp-checkpoint-final / tmp-checkpoint-backup / tmp-checkpoint-notes
|
||||
must NOT be deleted by the cancel cleanup."""
|
||||
from core.training.training import _cleanup_cancelled_checkpoints
|
||||
|
||||
out = outputs_setup / "run-non-numeric"
|
||||
out.mkdir()
|
||||
numeric = _mk_dir(out, "tmp-checkpoint-100")
|
||||
user_final = _mk_dir(out, "tmp-checkpoint-final")
|
||||
user_backup = _mk_dir(out, "tmp-checkpoint-backup")
|
||||
user_notes = _mk_dir(out, "tmp-checkpoint-user-notes")
|
||||
|
||||
_cleanup_cancelled_checkpoints(out)
|
||||
|
||||
assert not numeric.exists(), "in-flight tmp-checkpoint-100 should be removed"
|
||||
assert user_final.exists(), "user dir tmp-checkpoint-final must be preserved"
|
||||
assert user_backup.exists(), "user dir tmp-checkpoint-backup must be preserved"
|
||||
assert user_notes.exists(), "user dir tmp-checkpoint-user-notes must be preserved"
|
||||
|
|
@ -227,6 +227,60 @@ def test_desktop_refresh_preserves_desktop_marker():
|
|||
assert payload["desktop"] is True
|
||||
|
||||
|
||||
def test_consume_refresh_token_second_call_returns_none():
|
||||
"""Single-use rotation rejects the same token on a second consume."""
|
||||
seed_user()
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
raw = secrets.token_urlsafe(48)
|
||||
expires = (datetime.now(timezone.utc) + timedelta(days = 30)).isoformat()
|
||||
storage.save_refresh_token(raw, storage.DEFAULT_ADMIN_USERNAME, expires)
|
||||
|
||||
first = storage.consume_refresh_token(raw)
|
||||
assert first == (storage.DEFAULT_ADMIN_USERNAME, False)
|
||||
second = storage.consume_refresh_token(raw)
|
||||
assert second is None
|
||||
|
||||
|
||||
def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatch):
|
||||
"""64-thread pile-up against one token; DELETE RETURNING permits one winner."""
|
||||
seed_user()
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
raw = secrets.token_urlsafe(48)
|
||||
expires = (datetime.now(timezone.utc) + timedelta(days = 30)).isoformat()
|
||||
storage.save_refresh_token(raw, storage.DEFAULT_ADMIN_USERNAME, expires)
|
||||
|
||||
workers = 64
|
||||
|
||||
def attempt(_idx: int):
|
||||
try:
|
||||
return storage.consume_refresh_token(raw)
|
||||
except sqlite3.OperationalError:
|
||||
# "database is locked" under heavy contention; treat as losing the race.
|
||||
return None
|
||||
|
||||
with ThreadPoolExecutor(max_workers = workers) as pool:
|
||||
results = list(pool.map(attempt, range(workers)))
|
||||
|
||||
successes = [r for r in results if r is not None]
|
||||
assert (
|
||||
len(successes) == 1
|
||||
), f"expected exactly one consumer to win, got {len(successes)}"
|
||||
assert successes[0] == (storage.DEFAULT_ADMIN_USERNAME, False)
|
||||
|
||||
|
||||
def test_consume_refresh_token_expired_returns_none():
|
||||
seed_user()
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
raw = secrets.token_urlsafe(48)
|
||||
expires = (datetime.now(timezone.utc) - timedelta(hours = 1)).isoformat()
|
||||
storage.save_refresh_token(raw, storage.DEFAULT_ADMIN_USERNAME, expires)
|
||||
assert storage.consume_refresh_token(raw) is None
|
||||
|
||||
|
||||
def test_desktop_session_uses_real_admin_identity_for_api_keys():
|
||||
seed_user(must_change_password = True)
|
||||
raw = storage.create_desktop_secret()
|
||||
|
|
@ -383,6 +437,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
|
|||
inference_router = APIRouter(),
|
||||
inference_studio_router = APIRouter(),
|
||||
models_router = APIRouter(),
|
||||
providers_router = APIRouter(),
|
||||
training_history_router = APIRouter(),
|
||||
training_router = APIRouter(),
|
||||
)
|
||||
|
|
@ -392,7 +447,21 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
|
|||
|
||||
monkeypatch.setattr(backend_main._hw_module, "CHAT_ONLY", False)
|
||||
|
||||
body = asyncio.run(backend_main.health_check())
|
||||
seed_user()
|
||||
from auth.authentication import create_access_token
|
||||
|
||||
token = create_access_token(storage.DEFAULT_ADMIN_USERNAME)
|
||||
|
||||
app = FastAPI()
|
||||
app.add_api_route("/api/health", backend_main.health_check, methods = ["GET"])
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.get(
|
||||
"/api/health",
|
||||
headers = {"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
|
||||
assert body["desktop_protocol_version"] == 1
|
||||
assert body["supports_desktop_auth"] is True
|
||||
|
|
|
|||
326
studio/backend/tests/test_detect_mmproj_file.py
Normal file
326
studio/backend/tests/test_detect_mmproj_file.py
Normal file
|
|
@ -0,0 +1,326 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for :func:`utils.models.model_config.detect_mmproj_file` (#5347)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import struct
|
||||
|
||||
from utils.models.model_config import (
|
||||
_detect_family_token,
|
||||
detect_mmproj_file,
|
||||
mmproj_matches_model_family,
|
||||
)
|
||||
|
||||
|
||||
_GGUF_MAGIC = 0x46554747
|
||||
|
||||
|
||||
def _gguf_with_general(path: Path, fields: dict) -> Path:
|
||||
"""Write a minimal GGUF with only ``general.*`` string KVs."""
|
||||
body = b""
|
||||
for k, v in fields.items():
|
||||
kb = k.encode("utf-8")
|
||||
vb = v.encode("utf-8")
|
||||
body += struct.pack("<Q", len(kb)) + kb
|
||||
body += struct.pack("<I", 8) # STRING vtype
|
||||
body += struct.pack("<Q", len(vb)) + vb
|
||||
header = struct.pack("<IIQQ", _GGUF_MAGIC, 3, 0, len(fields))
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_bytes(header + body)
|
||||
return path
|
||||
|
||||
|
||||
def _touch(path: Path) -> Path:
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_bytes(b"")
|
||||
return path
|
||||
|
||||
|
||||
def test_returns_none_when_no_mmproj(tmp_path: Path):
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
assert detect_mmproj_file(str(model)) is None
|
||||
|
||||
|
||||
def test_single_matching_family_mmproj_picked(tmp_path: Path):
|
||||
"""Single same-family projector: returned (historical behaviour)."""
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
mmproj = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
|
||||
|
||||
|
||||
def test_hf_style_unprefixed_mmproj_still_works(tmp_path: Path):
|
||||
"""HF convention: weight + ``mmproj-F16.gguf`` sibling."""
|
||||
model = _touch(tmp_path / "model.gguf")
|
||||
mmproj = _touch(tmp_path / "mmproj-F16.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
|
||||
|
||||
|
||||
def test_blocks_single_cross_family_projector(tmp_path: Path):
|
||||
"""#5347 core: Qwen weight + lone Gemma mmproj returns None."""
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
assert detect_mmproj_file(str(model)) is None
|
||||
|
||||
|
||||
def test_picks_matching_family_among_mixed_candidates(tmp_path: Path):
|
||||
"""Mixed Qwen + Gemma projectors: pick Qwen, drop Gemma."""
|
||||
model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf")
|
||||
qwen_mm = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(qwen_mm.resolve())
|
||||
|
||||
|
||||
def test_prefers_longest_prefix_within_same_family(tmp_path: Path):
|
||||
"""Same family, different sizes: longest shared stem prefix wins."""
|
||||
model = _touch(tmp_path / "Qwen3.5-35B-A3B-UD-Q4_K_L.gguf")
|
||||
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
big_mm = _touch(tmp_path / "Qwen3.5-35B-A3B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(big_mm.resolve())
|
||||
|
||||
|
||||
def test_unrecognised_family_does_not_break_detection(tmp_path: Path):
|
||||
"""Unknown model family must not return None on a sole candidate."""
|
||||
model = _touch(tmp_path / "MyCustomBrand-7B-Q4_K_M.gguf")
|
||||
mmproj = _touch(tmp_path / "MyCustomBrand-7B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(mmproj.resolve())
|
||||
|
||||
|
||||
def test_directory_path_returns_first_candidate(tmp_path: Path):
|
||||
"""Directory path: no model stem to compare; legacy first-candidate."""
|
||||
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
result = detect_mmproj_file(str(tmp_path))
|
||||
assert result is not None
|
||||
assert "mmproj" in Path(result).name.lower()
|
||||
|
||||
|
||||
def test_search_root_walk_still_works(tmp_path: Path):
|
||||
"""Snapshot layout: weight in quant subdir, mmproj at snapshot root."""
|
||||
snapshot = tmp_path / "snapshot"
|
||||
weight = _touch(snapshot / "BF16" / "Qwen3.5-9B-BF16.gguf")
|
||||
mmproj = _touch(snapshot / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
result = detect_mmproj_file(str(weight), search_root = str(snapshot))
|
||||
assert result == str(mmproj.resolve())
|
||||
|
||||
|
||||
# -- Family token detection: word-bounded matching ----------------------
|
||||
|
||||
|
||||
def test_family_token_phi_does_not_match_sapphire():
|
||||
"""``phi`` substring inside ``sapphire`` must not tag Phi."""
|
||||
assert _detect_family_token("sapphire-7b-q4_k_m.gguf") is None
|
||||
|
||||
|
||||
def test_family_token_yi_does_not_match_tinyish_names():
|
||||
"""``yi`` must not cross letter boundaries (``yip``)."""
|
||||
assert _detect_family_token("yip-7b.gguf") is None
|
||||
assert _detect_family_token("yi-vl-6b.gguf") == "yi"
|
||||
|
||||
|
||||
def test_family_token_mimo_does_not_match_mimosa():
|
||||
"""``mimo`` must not tag ``mimosa``."""
|
||||
assert _detect_family_token("mimosa-rosa-7b.gguf") is None
|
||||
assert _detect_family_token("MiMo-VL-7B-RL-BF16.gguf") == "mimo"
|
||||
|
||||
|
||||
def test_family_token_mistral_does_not_match_ministral():
|
||||
"""Pin Mistral-derivative tagging."""
|
||||
assert _detect_family_token("Ministral-3-8B-Instruct-2512-BF16.gguf") == "ministral"
|
||||
assert _detect_family_token("Mistral-7B-Instruct-v0.3.gguf") == "mistral"
|
||||
assert _detect_family_token("Magistral-Small-2506-BF16.gguf") == "magistral"
|
||||
assert (
|
||||
_detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf")
|
||||
== "devstral"
|
||||
)
|
||||
|
||||
|
||||
def test_family_token_picks_leftmost_when_multiple_present():
|
||||
"""Leftmost family token wins, not tuple order."""
|
||||
assert _detect_family_token("llama-phi-merge.gguf") == "llama"
|
||||
assert _detect_family_token("phi-llama-merge.gguf") == "phi"
|
||||
assert _detect_family_token("llama3-3b-instruct.gguf") == "llama"
|
||||
|
||||
|
||||
def test_family_token_new_families_recognised():
|
||||
"""Catalogue-audit additions tag correctly."""
|
||||
assert _detect_family_token("NVIDIA-Nemotron-3-Nano-Omni-30B.gguf") == "nemotron"
|
||||
assert _detect_family_token("Kimi-K2.6-BF16.gguf") == "kimi"
|
||||
assert _detect_family_token("Nanonets-OCR-s-BF16.gguf") == "nanonets"
|
||||
assert _detect_family_token("Cosmos-Reason1-7B-BF16.gguf") == "cosmos"
|
||||
assert _detect_family_token("Apriel-1.5-15b-Thinker-BF16.gguf") == "apriel"
|
||||
assert _detect_family_token("LFM2.5-VL-1.6B-BF16.gguf") == "lfm"
|
||||
|
||||
|
||||
# -- Cross-family rejection with the expanded token list ----------------
|
||||
|
||||
|
||||
def test_blocks_cross_family_for_new_token_pair(tmp_path: Path):
|
||||
"""Nemotron weight + lone Gemma projector returns None."""
|
||||
model = _touch(
|
||||
tmp_path / "NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-MXFP4_MOE.gguf"
|
||||
)
|
||||
_touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf")
|
||||
assert detect_mmproj_file(str(model)) is None
|
||||
|
||||
|
||||
def test_picks_devstral_mmproj_in_mixed_dir(tmp_path: Path):
|
||||
"""Devstral weight + Devstral mmproj + a Qwen mmproj: pick Devstral."""
|
||||
model = _touch(tmp_path / "Devstral-Small-2-24B-Instruct-2512-BF16.gguf")
|
||||
dev_mm = _touch(tmp_path / "Devstral-Small-2-mmproj-bf16.gguf")
|
||||
_touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf")
|
||||
assert detect_mmproj_file(str(model)) == str(dev_mm.resolve())
|
||||
|
||||
|
||||
# -- Launcher-level family guard ----------------------------------------
|
||||
|
||||
|
||||
def test_mmproj_family_guard_blocks_cross_family():
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Qwen3.5-9B-Q4_K_M.gguf",
|
||||
"/models/gemma-4-26B-A4B-it.mmproj-q8_0.gguf",
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_mmproj_family_guard_allows_same_family():
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Qwen3.5-9B-Q4_K_M.gguf",
|
||||
"/models/Qwen3.5-9B-BF16-mmproj.gguf",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_mmproj_family_guard_allows_generic_hf_mmproj():
|
||||
"""No family token on the projector: wildcard."""
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Qwen3.5-9B-Q4_K_M.gguf",
|
||||
"/models/mmproj-F16.gguf",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_mmproj_family_guard_allows_unrecognised_model_family():
|
||||
"""No family token on the model: wildcard."""
|
||||
assert (
|
||||
mmproj_matches_model_family(
|
||||
"/models/Apriel-1.5-15b-Thinker-BF16.gguf",
|
||||
"/models/mmproj-F16.gguf",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
# -- Metadata-primary pairing in detect_mmproj_file ---------------------
|
||||
|
||||
|
||||
def test_metadata_url_match_picked_over_filename_lookalike(tmp_path: Path):
|
||||
"""URL match beats a longer-prefix sibling."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B-Q4_K_M.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
# Closer filename prefix, wrong upstream.
|
||||
_gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B-mmproj-bf16.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-1.5B",
|
||||
},
|
||||
)
|
||||
# Matching upstream.
|
||||
correct = _gguf_with_general(
|
||||
tmp_path / "mmproj-BF16.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) == str(correct.resolve())
|
||||
|
||||
|
||||
def test_metadata_url_mismatch_dropped(tmp_path: Path):
|
||||
"""Filenames match family but metadata disagrees: returns None."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "qwen-9b.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
_gguf_with_general(
|
||||
tmp_path / "qwen-9b-mmproj.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) is None
|
||||
|
||||
|
||||
def test_metadata_identifies_mmproj_without_filename_hint(tmp_path: Path):
|
||||
"""Projector named ``vision-projector.gguf`` discovered via header."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
projector = _gguf_with_general(
|
||||
tmp_path / "vision-projector.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) == str(projector.resolve())
|
||||
|
||||
|
||||
def test_metadata_score_outranks_filename_prefix(tmp_path: Path):
|
||||
"""Score 100 (URL match) beats score 0 (long filename prefix)."""
|
||||
weight = _gguf_with_general(
|
||||
tmp_path / "Qwen3.5-9B-Q4_K_M.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
# Headerless: long shared stem, score 0.
|
||||
_touch(tmp_path / "Qwen3.5-9B-Q4_K_M-mmproj.gguf")
|
||||
# Headered: generic name, score 100.
|
||||
correct = _gguf_with_general(
|
||||
tmp_path / "mmproj-BF16.gguf",
|
||||
{
|
||||
"general.architecture": "clip",
|
||||
"general.type": "mmproj",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
},
|
||||
)
|
||||
assert detect_mmproj_file(str(weight)) == str(correct.resolve())
|
||||
216
studio/backend/tests/test_gguf_metadata.py
Normal file
216
studio/backend/tests/test_gguf_metadata.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for :mod:`utils.models.gguf_metadata`. Synthesise small GGUF
|
||||
headers in tmp dirs so we never depend on real model files."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Mapping
|
||||
|
||||
from utils.models.gguf_metadata import (
|
||||
is_mmproj_by_metadata,
|
||||
pairing_score,
|
||||
read_gguf_general_metadata,
|
||||
)
|
||||
|
||||
|
||||
_GGUF_MAGIC = 0x46554747
|
||||
_VTYPE_STRING = 8
|
||||
_VTYPE_UINT32 = 4
|
||||
_VTYPE_ARRAY = 9
|
||||
|
||||
|
||||
def _enc_string(s: str) -> bytes:
|
||||
b = s.encode("utf-8")
|
||||
return struct.pack("<Q", len(b)) + b
|
||||
|
||||
|
||||
def _enc_kv_string(key: str, value: str) -> bytes:
|
||||
return _enc_string(key) + struct.pack("<I", _VTYPE_STRING) + _enc_string(value)
|
||||
|
||||
|
||||
def _enc_kv_uint32(key: str, value: int) -> bytes:
|
||||
return (
|
||||
_enc_string(key) + struct.pack("<I", _VTYPE_UINT32) + struct.pack("<I", value)
|
||||
)
|
||||
|
||||
|
||||
def _enc_kv_string_array(key: str, values: Iterable[str]) -> bytes:
|
||||
vals = list(values)
|
||||
out = _enc_string(key) + struct.pack("<I", _VTYPE_ARRAY)
|
||||
out += struct.pack("<I", _VTYPE_STRING) + struct.pack("<Q", len(vals))
|
||||
for v in vals:
|
||||
out += _enc_string(v)
|
||||
return out
|
||||
|
||||
|
||||
def _write_synthetic_gguf(
|
||||
path: Path,
|
||||
general_strings: Mapping[str, str],
|
||||
*,
|
||||
extra_uint32: Mapping[str, int] | None = None,
|
||||
extra_string_arrays: Mapping[str, Iterable[str]] | None = None,
|
||||
) -> Path:
|
||||
"""Minimal GGUF: header + KV body, no tensors."""
|
||||
extra_uint32 = extra_uint32 or {}
|
||||
extra_string_arrays = extra_string_arrays or {}
|
||||
kv_count = len(general_strings) + len(extra_uint32) + len(extra_string_arrays)
|
||||
body = b""
|
||||
for k, v in general_strings.items():
|
||||
body += _enc_kv_string(k, v)
|
||||
for k, v in extra_uint32.items():
|
||||
body += _enc_kv_uint32(k, v)
|
||||
for k, v in extra_string_arrays.items():
|
||||
body += _enc_kv_string_array(k, v)
|
||||
header = struct.pack(
|
||||
"<IIQQ",
|
||||
_GGUF_MAGIC,
|
||||
3, # version
|
||||
0, # tensor_count
|
||||
kv_count,
|
||||
)
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_bytes(header + body)
|
||||
return path
|
||||
|
||||
|
||||
# --- read_gguf_general_metadata ----------------------------------------
|
||||
|
||||
|
||||
def test_returns_none_for_missing_file(tmp_path: Path):
|
||||
assert read_gguf_general_metadata(str(tmp_path / "nope.gguf")) is None
|
||||
|
||||
|
||||
def test_returns_none_for_non_gguf(tmp_path: Path):
|
||||
p = tmp_path / "garbage.gguf"
|
||||
p.write_bytes(b"not a gguf file at all, just bytes")
|
||||
assert read_gguf_general_metadata(str(p)) is None
|
||||
|
||||
|
||||
def test_extracts_general_string_fields(tmp_path: Path):
|
||||
p = _write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{
|
||||
"general.architecture": "qwen2vl",
|
||||
"general.type": "model",
|
||||
"general.basename": "Qwen3.5",
|
||||
"general.organization": "Qwen",
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
"general.base_model.0.name": "Qwen3.5 9B",
|
||||
"general.base_model.0.organization": "Qwen",
|
||||
},
|
||||
)
|
||||
meta = read_gguf_general_metadata(str(p))
|
||||
assert meta is not None
|
||||
assert meta["general.architecture"] == "qwen2vl"
|
||||
assert meta["general.basename"] == "Qwen3.5"
|
||||
assert (
|
||||
meta["general.base_model.0.repo_url"]
|
||||
== "https://huggingface.co/Qwen/Qwen3.5-9B"
|
||||
)
|
||||
|
||||
|
||||
def test_skips_unrelated_fields_without_breaking(tmp_path: Path):
|
||||
"""Skip unwanted arrays and uint32s without losing position."""
|
||||
p = _write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{"general.basename": "Foo"},
|
||||
extra_uint32 = {"qwen2vl.context_length": 32768},
|
||||
extra_string_arrays = {"tokenizer.ggml.tokens": ["a", "bc", "def"]},
|
||||
)
|
||||
meta = read_gguf_general_metadata(str(p))
|
||||
assert meta == {"general.basename": "Foo"}
|
||||
|
||||
|
||||
def test_metadata_is_cached(tmp_path: Path):
|
||||
"""Cache invalidates on size change."""
|
||||
p = _write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{"general.basename": "First"},
|
||||
)
|
||||
first = read_gguf_general_metadata(str(p))
|
||||
assert first == {"general.basename": "First"}
|
||||
# Force size change so the (path, mtime, size) key invalidates.
|
||||
_write_synthetic_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
{"general.basename": "Second", "general.organization": "X"},
|
||||
)
|
||||
second = read_gguf_general_metadata(str(p))
|
||||
assert second == {"general.basename": "Second", "general.organization": "X"}
|
||||
|
||||
|
||||
# --- is_mmproj_by_metadata --------------------------------------------
|
||||
|
||||
|
||||
def test_is_mmproj_by_metadata_signals():
|
||||
assert is_mmproj_by_metadata({"general.type": "mmproj"}) is True
|
||||
assert is_mmproj_by_metadata({"general.type": "MMProj"}) is True
|
||||
assert is_mmproj_by_metadata({"general.type": "model"}) is False
|
||||
assert is_mmproj_by_metadata({"general.basename": "foo"}) is None
|
||||
assert is_mmproj_by_metadata({}) is None
|
||||
assert is_mmproj_by_metadata(None) is None
|
||||
|
||||
|
||||
# --- pairing_score -----------------------------------------------------
|
||||
|
||||
|
||||
def test_pairing_score_base_model_url_match():
|
||||
weight = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
mmproj = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == 100
|
||||
|
||||
|
||||
def test_pairing_score_base_model_url_mismatch():
|
||||
weight = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
mmproj = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == -1
|
||||
|
||||
|
||||
def test_pairing_score_base_model_url_trailing_slash_normalised():
|
||||
weight = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B/",
|
||||
}
|
||||
mmproj = {
|
||||
"general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == 100
|
||||
|
||||
|
||||
def test_pairing_score_basename_plus_org_fallback():
|
||||
weight = {
|
||||
"general.basename": "Nanonets-Ocr-S",
|
||||
"general.base_model.0.organization": "Nanonets",
|
||||
}
|
||||
mmproj = {
|
||||
"general.basename": "Nanonets-Ocr-S",
|
||||
"general.base_model.0.organization": "Nanonets",
|
||||
}
|
||||
assert pairing_score(weight, mmproj) == 80
|
||||
|
||||
|
||||
def test_pairing_score_basename_only_fallback():
|
||||
assert (
|
||||
pairing_score(
|
||||
{"general.basename": "Nanonets-Ocr-S"},
|
||||
{"general.basename": "Nanonets-Ocr-S"},
|
||||
)
|
||||
== 60
|
||||
)
|
||||
|
||||
|
||||
def test_pairing_score_no_overlap_returns_zero():
|
||||
"""One side empty: scorer punts to filename fallback."""
|
||||
assert pairing_score({"general.basename": "Foo"}, {}) == 0
|
||||
assert pairing_score({}, {"general.basename": "Foo"}) == 0
|
||||
assert pairing_score(None, {"general.basename": "Foo"}) == 0
|
||||
237
studio/backend/tests/test_gguf_reload_inheritance.py
Normal file
237
studio/backend/tests/test_gguf_reload_inheritance.py
Normal file
|
|
@ -0,0 +1,237 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Backend contract for the GGUF reload duplicate-load guard.
|
||||
|
||||
``LlamaCppBackend._already_in_target_state`` is the in-process
|
||||
short-circuit that prevents a serialised duplicate /load from killing
|
||||
the just-spawned llama-server. These tests pin the local-file
|
||||
identity, the HF-mode hf_variant fallback, and the ``extra_args``
|
||||
None-vs-[] inherit semantics so the guard cannot silently regress.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types as _types
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
_structlog_stub = _types.ModuleType("structlog")
|
||||
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
|
||||
sys.modules.setdefault("structlog", _structlog_stub)
|
||||
|
||||
_httpx_stub = _types.ModuleType("httpx")
|
||||
for _exc in (
|
||||
"ConnectError",
|
||||
"TimeoutException",
|
||||
"ReadTimeout",
|
||||
"ReadError",
|
||||
"RemoteProtocolError",
|
||||
"CloseError",
|
||||
):
|
||||
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
|
||||
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
|
||||
_httpx_stub.Client = type(
|
||||
"C",
|
||||
(),
|
||||
{
|
||||
"__init__": lambda s, **kw: None,
|
||||
"__enter__": lambda s: s,
|
||||
"__exit__": lambda s, *a: None,
|
||||
},
|
||||
)
|
||||
sys.modules.setdefault("httpx", _httpx_stub)
|
||||
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
|
||||
class _FakeProcess:
|
||||
"""Stand-in for subprocess.Popen so atexit cleanup doesn't crash."""
|
||||
|
||||
def terminate(self):
|
||||
pass
|
||||
|
||||
def wait(self, timeout = None):
|
||||
return 0
|
||||
|
||||
def kill(self):
|
||||
pass
|
||||
|
||||
def poll(self):
|
||||
return 0
|
||||
|
||||
|
||||
def _loaded_backend(**overrides):
|
||||
backend = LlamaCppBackend()
|
||||
backend._process = _FakeProcess() # is_loaded only checks "is not None"
|
||||
backend._healthy = True
|
||||
backend._model_identifier = "owner/repo"
|
||||
backend._hf_variant = "Q4_K_M"
|
||||
backend._requested_n_ctx = 8192
|
||||
backend._cache_type_kv = None
|
||||
backend._speculative_type = None
|
||||
backend._chat_template_override = None
|
||||
backend._is_vision = False
|
||||
backend._extra_args = None
|
||||
backend._extra_args_source = None
|
||||
backend._gguf_path = None
|
||||
for key, value in overrides.items():
|
||||
setattr(backend, key, value)
|
||||
return backend
|
||||
|
||||
|
||||
# ── Local-file identity via gguf_path ────────────────────────────────
|
||||
|
||||
|
||||
def test_already_in_target_state_uses_gguf_path_when_present(tmp_path):
|
||||
gguf_file = tmp_path / "model.Q4_K_M.gguf"
|
||||
gguf_file.write_bytes(b"")
|
||||
backend = _loaded_backend(
|
||||
_hf_variant = "Q4_K_M",
|
||||
_gguf_path = str(gguf_file),
|
||||
)
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = str(gguf_file),
|
||||
model_identifier = "owner/repo",
|
||||
hf_variant = None,
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_rejects_different_gguf_path(tmp_path):
|
||||
a = tmp_path / "a.gguf"
|
||||
a.write_bytes(b"")
|
||||
b = tmp_path / "b.gguf"
|
||||
b.write_bytes(b"")
|
||||
backend = _loaded_backend(_gguf_path = str(a))
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = str(b),
|
||||
model_identifier = "owner/repo",
|
||||
hf_variant = None,
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
# ── HF mode falls back to hf_variant comparison ──────────────────────
|
||||
|
||||
|
||||
def test_already_in_target_state_falls_back_to_hf_variant_for_hf_loads():
|
||||
backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None)
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "owner/repo",
|
||||
hf_variant = "Q8_0",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_hf_same_variant_matches():
|
||||
backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None)
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "owner/repo",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
# ── extra_args: None inherits, [] forces reload, list enforces ───────
|
||||
|
||||
|
||||
def test_already_in_target_state_none_extras_inherits_stored():
|
||||
backend = _loaded_backend(_extra_args = ["--top-k", "20"])
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "owner/repo",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_empty_extras_forces_reload_when_stored():
|
||||
backend = _loaded_backend(_extra_args = ["--top-k", "20"])
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "owner/repo",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = [],
|
||||
is_vision = False,
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_explicit_extras_match():
|
||||
backend = _loaded_backend(_extra_args = ["--top-k", "20"])
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "owner/repo",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = ["--top-k", "20"],
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_extra_args_source_default_is_none():
|
||||
backend = LlamaCppBackend()
|
||||
assert backend.extra_args_source is None
|
||||
226
studio/backend/tests/test_inference_model_validation.py
Normal file
226
studio/backend/tests/test_inference_model_validation.py
Normal file
|
|
@ -0,0 +1,226 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
_backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
sys.path.insert(0, _backend)
|
||||
|
||||
from models.inference import LoadRequest
|
||||
|
||||
|
||||
def _base_load_request(**overrides):
|
||||
data = {
|
||||
"model_path": "unsloth/test-model-GGUF",
|
||||
"hf_token": None,
|
||||
"max_seq_length": 4096,
|
||||
"load_in_4bit": True,
|
||||
"is_lora": False,
|
||||
"gguf_variant": "Q4_K_M",
|
||||
}
|
||||
data.update(overrides)
|
||||
return LoadRequest.model_validate(data)
|
||||
|
||||
|
||||
def test_blank_chat_template_override_normalizes_to_none():
|
||||
req = _base_load_request(chat_template_override = " \n\t")
|
||||
|
||||
assert req.chat_template_override is None
|
||||
|
||||
|
||||
def test_nonblank_chat_template_override_is_preserved_verbatim():
|
||||
template = " {{ messages }} "
|
||||
req = _base_load_request(chat_template_override = template)
|
||||
|
||||
assert req.chat_template_override == template
|
||||
|
||||
|
||||
# ---------- ChatCompletionRequest tool_call_id walkback ----------
|
||||
|
||||
from models.inference import ChatCompletionRequest
|
||||
|
||||
|
||||
def _req(messages, **overrides):
|
||||
payload = {"model": "x", "messages": messages, **overrides}
|
||||
return ChatCompletionRequest.model_validate(payload)
|
||||
|
||||
|
||||
def test_tool_message_inherits_id_from_prior_assistant_tool_call():
|
||||
req = _req(
|
||||
[
|
||||
{"role": "user", "content": "what is 2+2"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_real123",
|
||||
"type": "function",
|
||||
"function": {"name": "calc", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "name": "calc", "content": "4"}, # no tool_call_id
|
||||
]
|
||||
)
|
||||
assert req.messages[-1].tool_call_id == "call_real123"
|
||||
|
||||
|
||||
def test_tool_message_with_explicit_id_unchanged():
|
||||
req = _req(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_a",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_user_supplied", "content": "ok"},
|
||||
]
|
||||
)
|
||||
assert req.messages[-1].tool_call_id == "call_user_supplied"
|
||||
|
||||
|
||||
def test_walkback_prefers_function_name_match():
|
||||
req = _req(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_x",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": "{}"},
|
||||
},
|
||||
{
|
||||
"id": "call_y",
|
||||
"type": "function",
|
||||
"function": {"name": "calc", "arguments": "{}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{"role": "tool", "name": "calc", "content": "4"},
|
||||
]
|
||||
)
|
||||
assert req.messages[-1].tool_call_id == "call_y"
|
||||
|
||||
|
||||
def test_walkback_takes_first_unconsumed_when_no_name():
|
||||
req = _req(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_a",
|
||||
"type": "function",
|
||||
"function": {"name": "calc", "arguments": "{}"},
|
||||
},
|
||||
{
|
||||
"id": "call_b",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": "{}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{"role": "tool", "content": "first result"},
|
||||
{"role": "tool", "content": "second result"},
|
||||
]
|
||||
)
|
||||
assert req.messages[-2].tool_call_id == "call_a"
|
||||
assert req.messages[-1].tool_call_id == "call_b"
|
||||
|
||||
|
||||
def test_walkback_falls_back_to_synth_when_no_assistant_turn():
|
||||
req = _req(
|
||||
[
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "tool", "content": "orphan"},
|
||||
]
|
||||
)
|
||||
tcid = req.messages[-1].tool_call_id
|
||||
assert tcid is not None and tcid.startswith("call_") and len(tcid) > 5
|
||||
|
||||
|
||||
def test_walkback_does_not_cross_user_turn():
|
||||
req = _req(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "old_call",
|
||||
"type": "function",
|
||||
"function": {"name": "calc", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "old_call", "content": "4"},
|
||||
{"role": "user", "content": "next turn"},
|
||||
{"role": "tool", "content": "no parent in this turn"},
|
||||
]
|
||||
)
|
||||
last = req.messages[-1].tool_call_id
|
||||
# The walkback must NOT pick old_call because a user turn intervenes;
|
||||
# falls back to synth.
|
||||
assert last is not None
|
||||
assert last != "old_call"
|
||||
assert last.startswith("call_")
|
||||
|
||||
|
||||
def test_walkback_skips_explicitly_consumed_tool_call_id():
|
||||
"""Sibling tool result with an explicit id must reserve its assistant
|
||||
slot so a follow-up missing-id result picks the OTHER tool call."""
|
||||
req = _req(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_a",
|
||||
"type": "function",
|
||||
"function": {"name": "calc", "arguments": "{}"},
|
||||
},
|
||||
{
|
||||
"id": "call_b",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": "{}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_a", "content": "4"},
|
||||
{"role": "tool", "content": "second result"},
|
||||
]
|
||||
)
|
||||
assert [m.tool_call_id for m in req.messages if m.role == "tool"] == [
|
||||
"call_a",
|
||||
"call_b",
|
||||
]
|
||||
|
||||
|
||||
def test_walkback_handles_malformed_function_string():
|
||||
"""A tool_call with ``function`` as a string (provider quirk) must not
|
||||
raise; resolution falls back to fallback id selection."""
|
||||
req = _req(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": "call_a", "type": "function", "function": "calc"},
|
||||
],
|
||||
},
|
||||
{"role": "tool", "name": "calc", "content": "4"},
|
||||
]
|
||||
)
|
||||
assert req.messages[-1].tool_call_id == "call_a"
|
||||
|
|
@ -192,6 +192,7 @@ def _drive(
|
|||
else:
|
||||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||||
matched = False
|
||||
pin_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
|
||||
for n_gpus in range(1, len(ranked) + 1):
|
||||
subset = ranked[:n_gpus]
|
||||
pool_mib = sum(free for _, free in subset)
|
||||
|
|
@ -203,7 +204,7 @@ def _drive(
|
|||
)
|
||||
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
|
||||
total_mib = (model_size + kv) / (1024 * 1024)
|
||||
if total_mib <= pool_mib * 0.90:
|
||||
if total_mib <= pool_mib * pin_fraction:
|
||||
effective_ctx = capped
|
||||
gpu_indices = sorted(idx for idx, _ in subset)
|
||||
use_fit = False
|
||||
|
|
@ -211,6 +212,17 @@ def _drive(
|
|||
break
|
||||
if not matched:
|
||||
effective_ctx = min(FALLBACK_CTX, effective_ctx)
|
||||
# Mirror llama_cpp.py: re-check fit at FALLBACK_CTX.
|
||||
if effective_ctx > 0:
|
||||
for n_gpus in range(1, len(ranked) + 1):
|
||||
subset = ranked[:n_gpus]
|
||||
pool_mib = sum(free for _, free in subset)
|
||||
kv = inst._estimate_kv_cache_bytes(effective_ctx, cache_type_kv)
|
||||
total_mib = (model_size + kv) / (1024 * 1024)
|
||||
if total_mib <= pool_mib * pin_fraction:
|
||||
gpu_indices = sorted(idx for idx, _ in subset)
|
||||
use_fit = False
|
||||
break
|
||||
elif gpus:
|
||||
gpu_indices, use_fit = inst._select_gpus(model_size, gpus)
|
||||
if use_fit and not explicit_ctx:
|
||||
|
|
@ -378,6 +390,52 @@ class TestFittableAutoPickRegressions:
|
|||
assert plan["gpu_indices"] == [0]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# #5106 regression: 91-95% utilization must still pin GPU.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTightFitPinsToGPU:
|
||||
"""Models that fit at 91-95% of free VRAM must use the GPU."""
|
||||
|
||||
def test_rtx_4090_qwen_24gb_class(self):
|
||||
# noahterbest's #5106 log: 20.8 GB model on 22805 MiB free
|
||||
# GPU, ctx=4096 -> ~94% utilization, ~1.4 GiB headroom.
|
||||
plan = _drive(
|
||||
n_ctx = 0,
|
||||
model_gib = 20.8,
|
||||
gpus = [(0, 22_805)],
|
||||
native_ctx = 131072,
|
||||
kv_per_token_bytes = 25_000,
|
||||
)
|
||||
assert plan["use_fit"] is False
|
||||
assert plan["gpu_indices"] == [0]
|
||||
|
||||
def test_explicit_ctx_at_94_pct_pins_to_gpu(self):
|
||||
# Explicit-ctx branch must agree with auto-ctx on headroom.
|
||||
plan = _drive(
|
||||
n_ctx = 4096,
|
||||
model_gib = 20.8,
|
||||
gpus = [(0, 22_805)],
|
||||
native_ctx = 131072,
|
||||
kv_per_token_bytes = 25_000,
|
||||
)
|
||||
assert plan["use_fit"] is False
|
||||
assert plan["gpu_indices"] == [0]
|
||||
|
||||
def test_genuine_overflow_still_uses_fit(self):
|
||||
# Beyond 95% must still defer to --fit on.
|
||||
plan = _drive(
|
||||
n_ctx = 4096,
|
||||
model_gib = 23,
|
||||
gpus = [(0, 22_000)],
|
||||
native_ctx = 131072,
|
||||
kv_per_token_bytes = 25_000,
|
||||
)
|
||||
assert plan["use_fit"] is True
|
||||
assert plan["gpu_indices"] is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Platform-agnostic input shape
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -391,3 +449,81 @@ def test_identical_decision_across_platforms(platform_tag):
|
|||
plan_a = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
|
||||
plan_b = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
|
||||
assert plan_a == plan_b, platform_tag
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _classify_gpu_offload: detect silent CPU fallback (#5106).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClassifyGpuOffload:
|
||||
def _backend(self, stdout_lines):
|
||||
inst = LlamaCppBackend.__new__(LlamaCppBackend)
|
||||
inst._stdout_lines = list(stdout_lines)
|
||||
return inst
|
||||
|
||||
def test_cuda_buffer_present_returns_true(self):
|
||||
inst = self._backend(
|
||||
[
|
||||
"load_tensors: offloaded 33/33 layers to GPU",
|
||||
"load_tensors: CUDA0 model buffer size = 21000.0 MiB",
|
||||
"load_tensors: CPU_Mapped model buffer size = 0.6 MiB",
|
||||
]
|
||||
)
|
||||
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
|
||||
|
||||
def test_cpu_only_buffer_returns_false(self):
|
||||
# llama-server printed buffer lines but only CPU buffers --
|
||||
# this is the silent CPU fallback symptom we want to catch.
|
||||
inst = self._backend(
|
||||
[
|
||||
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
|
||||
"load_tensors: CPU model buffer size = 0.6 MiB",
|
||||
]
|
||||
)
|
||||
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
|
||||
|
||||
def test_no_buffer_lines_returns_none(self):
|
||||
# If we can't see buffer-allocation lines at all, don't guess.
|
||||
inst = self._backend(
|
||||
[
|
||||
"INFO [main] starting server",
|
||||
"load_tensors: file format = GGUF V3",
|
||||
]
|
||||
)
|
||||
assert inst._classify_gpu_offload(True, [(0, 22805)]) is None
|
||||
|
||||
def test_no_gpus_detected_returns_none(self):
|
||||
# CPU-only systems are valid; suppress the warning entirely.
|
||||
inst = self._backend(
|
||||
[
|
||||
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
|
||||
]
|
||||
)
|
||||
assert inst._classify_gpu_offload(False, []) is None
|
||||
|
||||
def test_user_did_not_intend_gpu_returns_none(self):
|
||||
# Studio called start_llama_server without expecting GPU use;
|
||||
# don't warn.
|
||||
inst = self._backend(
|
||||
[
|
||||
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
|
||||
]
|
||||
)
|
||||
assert inst._classify_gpu_offload(False, [(0, 22805)]) is None
|
||||
|
||||
def test_rocm_buffer_marker_returns_true(self):
|
||||
inst = self._backend(
|
||||
[
|
||||
"load_tensors: ROCm0 model buffer size = 21000.0 MiB",
|
||||
]
|
||||
)
|
||||
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
|
||||
|
||||
def test_metal_buffer_marker_returns_true(self):
|
||||
inst = self._backend(
|
||||
[
|
||||
"load_tensors: Metal model buffer size = 8000.0 MiB",
|
||||
]
|
||||
)
|
||||
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
|
||||
|
|
|
|||
328
studio/backend/tests/test_llama_cpp_freshness.py
Normal file
328
studio/backend/tests/test_llama_cpp_freshness.py
Normal file
|
|
@ -0,0 +1,328 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for the llama.cpp prebuilt freshness check.
|
||||
|
||||
Pins the marker parser, the disk+memory cache, the stale decision
|
||||
matrix, and fail-open behaviour on missing data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import types as _types
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
_structlog_stub = _types.ModuleType("structlog")
|
||||
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
|
||||
sys.modules.setdefault("structlog", _structlog_stub)
|
||||
|
||||
import pytest
|
||||
|
||||
from utils import llama_cpp_freshness as fr
|
||||
|
||||
|
||||
# Helpers.
|
||||
|
||||
|
||||
def _write_marker(install_dir: Path, **overrides) -> Path:
|
||||
payload = {
|
||||
"requested_tag": "latest",
|
||||
"tag": "b9190",
|
||||
"release_tag": "b9190",
|
||||
"published_repo": "unslothai/llama.cpp",
|
||||
"asset": "app-b9190-linux-x64-cuda13-newer.tar.gz",
|
||||
"asset_sha256": None,
|
||||
"source": "published",
|
||||
"installed_at_utc": (datetime.now(tz = timezone.utc) - timedelta(days = 1))
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z"),
|
||||
}
|
||||
payload.update(overrides)
|
||||
install_dir.mkdir(parents = True, exist_ok = True)
|
||||
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(payload))
|
||||
return install_dir / "UNSLOTH_PREBUILT_INFO.json"
|
||||
|
||||
|
||||
def _fake_binary(install_dir: Path, *, layout: str = "cmake") -> Path:
|
||||
"""Stub llama-server under one of the supported install layouts."""
|
||||
if layout == "cmake":
|
||||
bin_dir = install_dir / "build" / "bin"
|
||||
bin_name = "llama-server"
|
||||
elif layout == "root":
|
||||
bin_dir = install_dir
|
||||
bin_name = "llama-server"
|
||||
elif layout == "windows":
|
||||
bin_dir = install_dir / "build" / "bin" / "Release"
|
||||
bin_name = "llama-server.exe"
|
||||
else:
|
||||
raise ValueError(f"unknown layout {layout}")
|
||||
bin_dir.mkdir(parents = True, exist_ok = True)
|
||||
bin_path = bin_dir / bin_name
|
||||
bin_path.write_text("stub\n")
|
||||
return bin_path
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _reset(monkeypatch, tmp_path):
|
||||
# Isolate disk cache per-test; never touch the user's real cache.
|
||||
monkeypatch.setattr(fr, "_cache_dir", lambda: tmp_path / ".freshness")
|
||||
fr.reset_caches()
|
||||
yield
|
||||
fr.reset_caches()
|
||||
|
||||
|
||||
# read_install_marker.
|
||||
|
||||
|
||||
def test_read_install_marker_finds_cmake_layout(tmp_path):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(install_dir, tag = "b9190")
|
||||
bin_path = _fake_binary(install_dir, layout = "cmake")
|
||||
marker = fr.read_install_marker(str(bin_path))
|
||||
assert marker is not None
|
||||
assert marker["tag"] == "b9190"
|
||||
assert marker["published_repo"] == "unslothai/llama.cpp"
|
||||
|
||||
|
||||
def test_read_install_marker_finds_root_layout(tmp_path):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(install_dir, tag = "b9999")
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
marker = fr.read_install_marker(str(bin_path))
|
||||
assert marker is not None
|
||||
assert marker["tag"] == "b9999"
|
||||
|
||||
|
||||
def test_read_install_marker_finds_windows_cmake_layout(tmp_path):
|
||||
# Windows cmake puts the .exe under build/bin/Release/, so the
|
||||
# marker is four levels above the binary.
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(install_dir, tag = "b8888")
|
||||
bin_path = _fake_binary(install_dir, layout = "windows")
|
||||
marker = fr.read_install_marker(str(bin_path))
|
||||
assert marker is not None
|
||||
assert marker["tag"] == "b8888"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"])
|
||||
def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo):
|
||||
# The freshness check queries whichever release repo the marker
|
||||
# records, so CUDA Linux (unslothai), CPU Linux x86_64 / macOS
|
||||
# (ggml-org), and ROCm source-build (unslothai upstream label)
|
||||
# all surface the right "latest" tag.
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(install_dir, tag = "b9000", published_repo = repo)
|
||||
bin_path = _fake_binary(install_dir, layout = "cmake")
|
||||
marker = fr.read_install_marker(str(bin_path))
|
||||
assert marker is not None
|
||||
assert marker["published_repo"] == repo
|
||||
|
||||
|
||||
def test_read_install_marker_missing_returns_none(tmp_path):
|
||||
bin_path = _fake_binary(tmp_path / "no_marker", layout = "root")
|
||||
assert fr.read_install_marker(str(bin_path)) is None
|
||||
|
||||
|
||||
def test_read_install_marker_handles_invalid_json(tmp_path):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
install_dir.mkdir(parents = True)
|
||||
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text("not json")
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
assert fr.read_install_marker(str(bin_path)) is None
|
||||
|
||||
|
||||
def test_read_install_marker_handles_none_path():
|
||||
assert fr.read_install_marker(None) is None
|
||||
|
||||
|
||||
# latest_published_release (with monkeypatched fetcher).
|
||||
|
||||
|
||||
def test_latest_published_release_uses_disk_cache(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def _fake_fetch(repo, timeout = 5.0):
|
||||
calls.append(repo)
|
||||
return "b9999"
|
||||
|
||||
monkeypatch.setattr(fr, "_fetch_latest_release_tag", _fake_fetch)
|
||||
first = fr.latest_published_release("unslothai/llama.cpp")
|
||||
second = fr.latest_published_release("unslothai/llama.cpp")
|
||||
assert first == "b9999"
|
||||
assert second == "b9999"
|
||||
# Memo + disk cache -> only one fetch.
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_latest_published_release_returns_none_on_network_failure(monkeypatch):
|
||||
monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
|
||||
assert fr.latest_published_release("unslothai/llama.cpp") is None
|
||||
|
||||
|
||||
def test_latest_published_release_keeps_old_cache_on_transient_failure(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
# Disk entry older than TTL + network fail -> return cached value.
|
||||
cache_dir = tmp_path / ".freshness"
|
||||
cache_dir.mkdir()
|
||||
cache_file = cache_dir / "unslothai__llama.cpp.json"
|
||||
yesterday = time.time() - 25 * 60 * 60 # > 24h
|
||||
cache_file.write_text(json.dumps({"fetched_at": yesterday, "latest_tag": "b9000"}))
|
||||
monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
|
||||
assert fr.latest_published_release("unslothai/llama.cpp") == "b9000"
|
||||
|
||||
|
||||
# check_prebuilt_freshness end-to-end.
|
||||
|
||||
|
||||
def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(
|
||||
install_dir,
|
||||
tag = "b9190",
|
||||
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5))
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z"),
|
||||
)
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
monkeypatch.setattr(
|
||||
fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
|
||||
)
|
||||
info = fr.check_prebuilt_freshness(str(bin_path))
|
||||
assert info["has_marker"] is True
|
||||
assert info["stale"] is True
|
||||
assert info["installed_tag"] == "b9190"
|
||||
assert info["latest_tag"] == "b9300"
|
||||
assert info["age_days"] == 5
|
||||
assert info["published_repo"] == "unslothai/llama.cpp"
|
||||
|
||||
|
||||
def test_check_prebuilt_freshness_not_stale_when_tag_matches(monkeypatch, tmp_path):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(
|
||||
install_dir,
|
||||
tag = "b9300",
|
||||
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 30))
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z"),
|
||||
)
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
monkeypatch.setattr(
|
||||
fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
|
||||
)
|
||||
info = fr.check_prebuilt_freshness(str(bin_path))
|
||||
assert info["stale"] is False
|
||||
assert info["installed_tag"] == "b9300"
|
||||
assert info["latest_tag"] == "b9300"
|
||||
|
||||
|
||||
def test_check_prebuilt_freshness_not_stale_within_threshold(monkeypatch, tmp_path):
|
||||
# Behind by tag but within the 3-day grace window.
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(
|
||||
install_dir,
|
||||
tag = "b9190",
|
||||
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 1))
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z"),
|
||||
)
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
monkeypatch.setattr(
|
||||
fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
|
||||
)
|
||||
info = fr.check_prebuilt_freshness(str(bin_path))
|
||||
assert info["stale"] is False
|
||||
assert info["age_days"] == 1
|
||||
|
||||
|
||||
def test_check_prebuilt_freshness_fails_open_without_marker(tmp_path):
|
||||
bin_path = _fake_binary(tmp_path / "custom_build", layout = "root")
|
||||
info = fr.check_prebuilt_freshness(str(bin_path))
|
||||
assert info["has_marker"] is False
|
||||
assert info["stale"] is False
|
||||
|
||||
|
||||
def test_check_prebuilt_freshness_fails_open_when_github_unreachable(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(
|
||||
install_dir,
|
||||
tag = "b9190",
|
||||
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 10))
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z"),
|
||||
)
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None)
|
||||
info = fr.check_prebuilt_freshness(str(bin_path))
|
||||
assert info["has_marker"] is True
|
||||
assert info["stale"] is False
|
||||
assert info["latest_tag"] is None
|
||||
|
||||
|
||||
def test_check_prebuilt_freshness_handles_unparseable_install_timestamp(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(install_dir, tag = "b9190", installed_at_utc = "not-a-date")
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
monkeypatch.setattr(
|
||||
fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
|
||||
)
|
||||
info = fr.check_prebuilt_freshness(str(bin_path))
|
||||
assert info["stale"] is False
|
||||
assert info["age_days"] is None
|
||||
|
||||
|
||||
def test_check_prebuilt_freshness_respects_custom_threshold(monkeypatch, tmp_path):
|
||||
install_dir = tmp_path / "llama.cpp"
|
||||
_write_marker(
|
||||
install_dir,
|
||||
tag = "b9190",
|
||||
installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 2))
|
||||
.isoformat()
|
||||
.replace("+00:00", "Z"),
|
||||
)
|
||||
bin_path = _fake_binary(install_dir, layout = "root")
|
||||
monkeypatch.setattr(
|
||||
fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
|
||||
)
|
||||
info = fr.check_prebuilt_freshness(str(bin_path), threshold_days = 1)
|
||||
assert info["stale"] is True
|
||||
|
||||
|
||||
# format_stale_warning.
|
||||
|
||||
|
||||
def test_format_stale_warning_contains_actionable_command():
|
||||
msg = fr.format_stale_warning(
|
||||
{"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 5}
|
||||
)
|
||||
assert "b9190" in msg
|
||||
assert "b9300" in msg
|
||||
assert "5 days" in msg
|
||||
assert "unsloth studio update" in msg
|
||||
|
||||
|
||||
def test_format_stale_warning_singular_day():
|
||||
msg = fr.format_stale_warning(
|
||||
{"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1}
|
||||
)
|
||||
assert "1 day" in msg
|
||||
assert "1 days" not in msg
|
||||
496
studio/backend/tests/test_llama_cpp_mtp_detection.py
Normal file
496
studio/backend/tests/test_llama_cpp_mtp_detection.py
Normal file
|
|
@ -0,0 +1,496 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for the MTP auto-detection path (llama.cpp #22673).
|
||||
|
||||
Pins three contracts: name-based detector, user-override detector, and
|
||||
the _already_in_target_state mirror that prevents needless reloads.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
import sys
|
||||
import types as _types
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
_structlog_stub = _types.ModuleType("structlog")
|
||||
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
|
||||
sys.modules.setdefault("structlog", _structlog_stub)
|
||||
|
||||
_httpx_stub = _types.ModuleType("httpx")
|
||||
for _exc in (
|
||||
"ConnectError",
|
||||
"TimeoutException",
|
||||
"ReadTimeout",
|
||||
"ReadError",
|
||||
"RemoteProtocolError",
|
||||
"CloseError",
|
||||
):
|
||||
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
|
||||
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
|
||||
_httpx_stub.Client = type(
|
||||
"C",
|
||||
(),
|
||||
{
|
||||
"__init__": lambda s, **kw: None,
|
||||
"__enter__": lambda s: s,
|
||||
"__exit__": lambda s, *a: None,
|
||||
},
|
||||
)
|
||||
sys.modules.setdefault("httpx", _httpx_stub)
|
||||
|
||||
import pytest
|
||||
|
||||
from core.inference.llama_cpp import (
|
||||
LlamaCppBackend,
|
||||
_extra_args_set_spec_type,
|
||||
_is_mtp_model_name,
|
||||
)
|
||||
|
||||
|
||||
# Synthetic GGUF helper (mirrors test_gguf_metadata.py).
|
||||
|
||||
_GGUF_MAGIC = 0x46554747
|
||||
_VTYPE_STRING = 8
|
||||
_VTYPE_UINT32 = 4
|
||||
|
||||
|
||||
def _enc_string(s: str) -> bytes:
|
||||
b = s.encode("utf-8")
|
||||
return struct.pack("<Q", len(b)) + b
|
||||
|
||||
|
||||
def _enc_kv_string(key: str, value: str) -> bytes:
|
||||
return _enc_string(key) + struct.pack("<I", _VTYPE_STRING) + _enc_string(value)
|
||||
|
||||
|
||||
def _enc_kv_uint32(key: str, value: int) -> bytes:
|
||||
return (
|
||||
_enc_string(key) + struct.pack("<I", _VTYPE_UINT32) + struct.pack("<I", value)
|
||||
)
|
||||
|
||||
|
||||
def _write_minimal_gguf(
|
||||
path: Path,
|
||||
*,
|
||||
arch: str,
|
||||
nextn: int | None,
|
||||
extra_uint32: dict[str, int] | None = None,
|
||||
) -> Path:
|
||||
"""Header-only GGUF with arch + optional nextn_predict_layers."""
|
||||
extra_uint32 = dict(extra_uint32 or {})
|
||||
body = _enc_kv_string("general.architecture", arch)
|
||||
kv_count = 1
|
||||
if nextn is not None:
|
||||
body += _enc_kv_uint32(f"{arch}.nextn_predict_layers", nextn)
|
||||
kv_count += 1
|
||||
for k, v in extra_uint32.items():
|
||||
body += _enc_kv_uint32(k, v)
|
||||
kv_count += 1
|
||||
header = struct.pack("<IIQQ", _GGUF_MAGIC, 3, 0, kv_count)
|
||||
path.write_bytes(header + body)
|
||||
return path
|
||||
|
||||
|
||||
# _is_mtp_model_name helper.
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"identifier",
|
||||
[
|
||||
"unsloth/Qwen3.6-27B-MTP-GGUF",
|
||||
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
|
||||
"unsloth/qwen3.6-27b-mtp-gguf",
|
||||
"unsloth/Qwen3.6-27B-Mtp-GGUF",
|
||||
"unsloth/Qwen3.6-27B-MTP-GGUF:UD-Q4_K_XL",
|
||||
],
|
||||
)
|
||||
def test_is_mtp_model_name_detects_marker_in_identifier(identifier):
|
||||
assert _is_mtp_model_name(identifier) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"identifier",
|
||||
[
|
||||
"unsloth/Qwen3-27B-GGUF",
|
||||
"unsloth/Llama-3.1-8B-Instruct-GGUF",
|
||||
"google/gemma-3-4b-it",
|
||||
# mtp inside an org name should not match.
|
||||
"mtp-research/foo",
|
||||
"MTPower/bar",
|
||||
],
|
||||
)
|
||||
def test_is_mtp_model_name_does_not_overmatch(identifier):
|
||||
assert _is_mtp_model_name(identifier) is False
|
||||
|
||||
|
||||
def test_is_mtp_model_name_handles_none():
|
||||
assert _is_mtp_model_name(None) is False
|
||||
assert _is_mtp_model_name(None, None) is False
|
||||
assert _is_mtp_model_name("", "") is False
|
||||
|
||||
|
||||
def test_is_mtp_model_name_detects_marker_in_filename(tmp_path):
|
||||
gguf = tmp_path / "Qwen3.6-27B-MTP-Q4_K_M.gguf"
|
||||
gguf.write_bytes(b"")
|
||||
assert _is_mtp_model_name("local-model", str(gguf)) is True
|
||||
|
||||
|
||||
def test_is_mtp_model_name_filename_case_insensitive(tmp_path):
|
||||
gguf = tmp_path / "qwen3.6-35b-a3b-mtp-q4_k_m.gguf"
|
||||
gguf.write_bytes(b"")
|
||||
assert _is_mtp_model_name(None, str(gguf)) is True
|
||||
|
||||
|
||||
def test_is_mtp_model_name_ignores_non_mtp_filename(tmp_path):
|
||||
gguf = tmp_path / "Qwen3.6-27B-Q4_K_M.gguf"
|
||||
gguf.write_bytes(b"")
|
||||
assert _is_mtp_model_name("local-model", str(gguf)) is False
|
||||
|
||||
|
||||
# _already_in_target_state MTP promotion.
|
||||
|
||||
|
||||
class _FakeProcess:
|
||||
"""Minimal stand-in so is_loaded returns True."""
|
||||
|
||||
def terminate(self):
|
||||
pass
|
||||
|
||||
def wait(self, timeout = None):
|
||||
return 0
|
||||
|
||||
def kill(self):
|
||||
pass
|
||||
|
||||
def poll(self):
|
||||
return 0
|
||||
|
||||
|
||||
def _mtp_backend(**overrides):
|
||||
"""MTP-named GGUF backend that's already running with draft-mtp."""
|
||||
backend = LlamaCppBackend()
|
||||
backend._process = _FakeProcess()
|
||||
backend._healthy = True
|
||||
backend._model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF"
|
||||
backend._hf_variant = "Q4_K_M"
|
||||
backend._requested_n_ctx = 8192
|
||||
backend._cache_type_kv = None
|
||||
backend._speculative_type = "draft-mtp"
|
||||
backend._chat_template_override = None
|
||||
backend._is_vision = False
|
||||
backend._extra_args = None
|
||||
backend._extra_args_source = None
|
||||
backend._gguf_path = None
|
||||
for key, value in overrides.items():
|
||||
setattr(backend, key, value)
|
||||
return backend
|
||||
|
||||
|
||||
def test_already_in_target_state_matches_when_request_omits_spec_for_mtp_model():
|
||||
# Duplicate /load with no spec must match a running draft-mtp backend.
|
||||
backend = _mtp_backend()
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_matches_when_request_uses_default_for_mtp_model():
|
||||
backend = _mtp_backend()
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = "default",
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_non_mtp_model_unaffected():
|
||||
# Promotion is gated on the name; non-MTP must still mismatch req=None.
|
||||
backend = _mtp_backend(_model_identifier = "unsloth/Qwen3.6-27B-GGUF")
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "unsloth/Qwen3.6-27B-GGUF",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_explicit_off_still_mismatches_mtp_backend():
|
||||
backend = _mtp_backend()
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = "off",
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
# User override via extra_args (unsloth run / unsloth studio run).
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra_args",
|
||||
[
|
||||
["--spec-type", "none"],
|
||||
["--spec-type", "ngram-mod"],
|
||||
["--spec-type", "draft-mtp"],
|
||||
["--spec-type=none"],
|
||||
["--top-k", "20", "--spec-type", "ngram-simple", "--seed", "42"],
|
||||
["--spec-default"],
|
||||
],
|
||||
)
|
||||
def test_extra_args_set_spec_type_detects_user_override(extra_args):
|
||||
assert _extra_args_set_spec_type(extra_args) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra_args",
|
||||
[
|
||||
None,
|
||||
[],
|
||||
# Scalar tuning knobs compose safely with auto-emitted --spec-type.
|
||||
["--spec-draft-n-max", "4"],
|
||||
["--spec-ngram-mod-n-match", "32"],
|
||||
["--draft-max", "32"],
|
||||
["--top-k", "20", "--seed", "42"],
|
||||
],
|
||||
)
|
||||
def test_extra_args_set_spec_type_passes_on_non_spec_type_args(extra_args):
|
||||
assert _extra_args_set_spec_type(extra_args) is False
|
||||
|
||||
|
||||
def test_already_in_target_state_user_spec_type_override_matches_clean_backend():
|
||||
# User --spec-type none suppressed auto-MTP; repeat /load must not re-promote.
|
||||
backend = _mtp_backend(
|
||||
_speculative_type = None,
|
||||
_extra_args = ["--spec-type", "none"],
|
||||
)
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = None,
|
||||
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
|
||||
hf_variant = "Q4_K_M",
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = ["--spec-type", "none"],
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_already_in_target_state_local_file_mtp_match(tmp_path):
|
||||
# Local-file load: -MTP marker comes from the filename.
|
||||
gguf = tmp_path / "Qwen3.6-35B-A3B-MTP-Q4_K_M.gguf"
|
||||
gguf.write_bytes(b"")
|
||||
backend = _mtp_backend(
|
||||
_model_identifier = "local-qwen-mtp",
|
||||
_gguf_path = str(gguf),
|
||||
_hf_variant = None,
|
||||
)
|
||||
assert (
|
||||
backend._already_in_target_state(
|
||||
gguf_path = str(gguf),
|
||||
model_identifier = "local-qwen-mtp",
|
||||
hf_variant = None,
|
||||
n_ctx = 8192,
|
||||
cache_type_kv = None,
|
||||
speculative_type = None,
|
||||
chat_template_override = None,
|
||||
extra_args = None,
|
||||
is_vision = False,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
# GGUF-metadata-based detection (nextn_predict_layers).
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arch, nextn",
|
||||
[
|
||||
# Verified against real Unsloth MTP GGUFs (qwen35 / qwen35moe).
|
||||
("qwen35", 1),
|
||||
("qwen35moe", 1),
|
||||
# Future-proofing: any arch + n>0 should match.
|
||||
("qwen3moe", 2),
|
||||
("hypothetical_future_arch", 4),
|
||||
],
|
||||
)
|
||||
def test_read_gguf_metadata_captures_nextn_predict_layers(tmp_path, arch, nextn):
|
||||
gguf = _write_minimal_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
arch = arch,
|
||||
nextn = nextn,
|
||||
extra_uint32 = {f"{arch}.block_count": 4},
|
||||
)
|
||||
backend = LlamaCppBackend()
|
||||
backend._read_gguf_metadata(str(gguf))
|
||||
assert backend._nextn_predict_layers == nextn
|
||||
|
||||
|
||||
def test_read_gguf_metadata_leaves_nextn_unset_for_non_mtp_arch(tmp_path):
|
||||
gguf = _write_minimal_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
arch = "qwen3",
|
||||
nextn = None,
|
||||
extra_uint32 = {"qwen3.block_count": 4},
|
||||
)
|
||||
backend = LlamaCppBackend()
|
||||
backend._read_gguf_metadata(str(gguf))
|
||||
assert backend._nextn_predict_layers is None
|
||||
|
||||
|
||||
def test_read_gguf_metadata_zero_nextn_is_falsy(tmp_path):
|
||||
# bool(0) is False, so the spec block short-circuits.
|
||||
gguf = _write_minimal_gguf(
|
||||
tmp_path / "model.gguf",
|
||||
arch = "qwen35",
|
||||
nextn = 0,
|
||||
extra_uint32 = {"qwen35.block_count": 4},
|
||||
)
|
||||
backend = LlamaCppBackend()
|
||||
backend._read_gguf_metadata(str(gguf))
|
||||
assert backend._nextn_predict_layers == 0
|
||||
assert bool(backend._nextn_predict_layers) is False
|
||||
|
||||
|
||||
def test_unload_resets_nextn_predict_layers():
|
||||
# MTP state from a previous load must not bleed into the next load.
|
||||
backend = LlamaCppBackend()
|
||||
backend._nextn_predict_layers = 1
|
||||
backend.unload_model()
|
||||
assert backend._nextn_predict_layers is None
|
||||
|
||||
|
||||
# llama-server capability probe.
|
||||
|
||||
|
||||
def _make_fake_llama_server(path: Path, help_text: str) -> Path:
|
||||
"""Bash stub that prints `help_text` on --help."""
|
||||
path.write_text("#!/usr/bin/env bash\n" f"cat <<'EOF'\n{help_text}\nEOF\n")
|
||||
path.chmod(0o755)
|
||||
return path
|
||||
|
||||
|
||||
def _clear_caps_cache():
|
||||
LlamaCppBackend._capability_cache.clear()
|
||||
|
||||
|
||||
def test_probe_server_capabilities_detects_draft_mtp(tmp_path):
|
||||
# Original naming from llama.cpp #22673.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type none,draft-simple,draft-eagle3,draft-mtp,"
|
||||
"ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps["found"] is True
|
||||
assert caps["mtp_token"] == "draft-mtp"
|
||||
assert caps["supports_mtp"] is True
|
||||
|
||||
|
||||
def test_probe_server_capabilities_detects_renamed_mtp(tmp_path):
|
||||
# Renamed upstream: draft-mtp -> mtp.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type [none|mtp|ngram-cache|ngram-simple|ngram-map-k|"
|
||||
"ngram-map-k4v|ngram-mod]",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps["mtp_token"] == "mtp"
|
||||
assert caps["supports_mtp"] is True
|
||||
|
||||
|
||||
def test_probe_server_capabilities_reports_outdated_binary(tmp_path):
|
||||
# Pre-MTP llama.cpp: only ngram variants.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type none,ngram-simple,ngram-mod",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps["found"] is True
|
||||
assert caps["mtp_token"] is None
|
||||
assert caps["supports_mtp"] is False
|
||||
|
||||
|
||||
def test_probe_server_capabilities_handles_missing_binary():
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities("/no/such/llama-server")
|
||||
assert caps["found"] is False
|
||||
assert caps["supports_mtp"] is False
|
||||
|
||||
|
||||
def test_probe_server_capabilities_caches_by_mtime(tmp_path):
|
||||
# Same (path, mtime) -> cache hit. Bumped mtime -> re-probe.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type none,ngram-mod",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps1 = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps1["supports_mtp"] is False
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
_make_fake_llama_server(
|
||||
fake,
|
||||
"--spec-type none,draft-mtp,ngram-mod",
|
||||
)
|
||||
new_mtime = int(time.time()) + 2
|
||||
os.utime(fake, (new_mtime, new_mtime))
|
||||
caps2 = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps2["mtp_token"] == "draft-mtp"
|
||||
assert caps2["supports_mtp"] is True
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue