diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs deleted file mode 100644 index 17d96cd0f5..0000000000 --- a/.git-blame-ignore-revs +++ /dev/null @@ -1,8 +0,0 @@ -# Commits listed here are skipped by `git blame` so that bulk, whitespace-only -# changes don't obscure the real authorship of a line. -# -# GitHub honors this file automatically. To use it locally, run once: -# git config blame.ignoreRevsFile .git-blame-ignore-revs - -# chore(studio/frontend): normalize line endings to LF -c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh new file mode 100755 index 0000000000..3c7cea919c --- /dev/null +++ b/.github/scripts/agent-guides-drive.sh @@ -0,0 +1,502 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Drive one coding agent against the running `unsloth run` server for the +# Local Agent Guides CI. All failures from here are failure class (c) +# "guide drift": the server preflight already passed and the agent CLI +# already installed, so a failure here means the documented recipe in +# unsloth_cli/commands/connect.py no longer produces a working flow. +# +# Self-updating: for the 5 agents with a connect.py recipe we obtain the +# exact env + command from `unsloth connect --no-launch` and run +# THAT, so a recipe change is exercised automatically. Pi (no connect.py +# command at HEAD) is driven by a hand-written recipe. +# +# Every agent invocation is wrapped in `timeout` so a headless-TTY prompt +# can never hang the runner -- a timeout is reported as guide drift with a +# distinct message. +# +# Usage: +# agent-guides-drive.sh connection +# agent-guides-drive.sh file-edit +# agent-guides-drive.sh attribution-ab claude +# +# Required env (exported by serve-unsloth-run.sh): +# UNSLOTH_BASE_URL UNSLOTH_API_KEY UNSLOTH_MODEL_ID +# UNSLOTH_LLAMA_LOG_DIR AGENT_INVOKE_TIMEOUT UNSLOTH_SEED +set -uo pipefail + +MODE="${1:?usage: agent-guides-drive.sh }" +AGENT="${2:?usage: agent-guides-drive.sh }" + +: "${UNSLOTH_BASE_URL:?serve step did not export UNSLOTH_BASE_URL}" +: "${UNSLOTH_API_KEY:?serve step did not export UNSLOTH_API_KEY}" +: "${UNSLOTH_MODEL_ID:?serve step did not export UNSLOTH_MODEL_ID}" +# Determinism (seed/temp) is applied at the server level by +# serve-unsloth-run.sh --extra; agents inherit it through the API. +TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}" + +# Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner +# IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless +# to the other agents, which ignore it. +export IS_SANDBOX=1 + +# Absolute paths anchored at the repo root (this script lives in +# .github/scripts/). Everything writes here regardless of the current working +# directory, so the file-edit mode can `cd` into a scratch work dir without +# breaking log/redaction writes. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +LOGS_DIR="$REPO_ROOT/logs" +REDACTED_DIR="$REPO_ROOT/redacted-configs" +WORKDIR_BASE="$REPO_ROOT/agent-workdir" +CACHE_HELPER="$SCRIPT_DIR/assert-prompt-cache.sh" +mkdir -p "$LOGS_DIR" "$REDACTED_DIR" +CONNECT_REF="unsloth_cli/commands/connect.py" + +# Prefill-shrinking flags for Claude Code. The heavyweight agents send +# multi-thousand-token system prompts + full tool schemas, which on a CPU-only +# runner is minutes of prefill per model round-trip (~16 tok/s for a 4B model). +# Replacing the ~5.7k default system prompt with a tiny one (--system-prompt-file) +# and restricting tools cuts the prefill to a few hundred tokens so it completes +# quickly on CPU. These only shape the request size; the connect.py recipe +# (endpoint, auth, model) is still exercised end to end. +# +# The bulk of Claude Code's prompt is the built-in tool JSON schemas: measured +# via `claude -p /context`, the default prompt is ~28k tokens of which ~18k is +# "System tools" alone. --allowedTools/--disallowedTools only gate PERMISSION to +# call a tool; they do NOT remove its schema from what is sent to the model, so +# the earlier whitelist left the full ~18k in the prompt and CPU prefill +# (~16 tok/s) overran claude's own request timeout into a retry loop. --tools is +# the flag that restricts which schemas are sent. (The ~8k "Memory files" chunk +# is auto-loaded CLAUDE.md; the unsloth repo ships none, so it is 0 in CI.) +# +# Connection probe: --tools "" sends ZERO tool schemas, leaving ~20 tokens total +# (a one-line --system-prompt-file + the user turn), which prefills instantly. +CLAUDE_CONNECT_FLAGS=( + --system-prompt-file "$SCRIPT_DIR/ci-connect-prompt.txt" + --tools "" +) +# File-edit: the task needs the file/shell tools, so send only those schemas +# (~2.3k tokens vs ~18k for the full set). +CLAUDE_EDIT_FLAGS=( + --system-prompt-file "$SCRIPT_DIR/ci-min-system-prompt.txt" + --tools "Bash,Edit,Write,Read" +) + +guide_fail() { + echo "::error::[guide drift] agent=${AGENT}: $* (preflight passed + install OK, so the documented flow in ${CONNECT_REF} drifted)." >&2 + exit 1 +} + +# Redact the API key from any file we are about to keep as an artifact. +# Portable across GNU sed (Linux runners) and BSD sed (macOS), so the +# redaction is never silently skipped. +redact() { + local f + for f in "$@"; do + [ -f "$f" ] || continue + if sed --version >/dev/null 2>&1; then + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + else + sed -i '' "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + fi + done +} + +# A reply must be non-empty and free of connection/auth errors. +assert_reply() { + local out="$1" + if [ ! -s "$out" ]; then + guide_fail "agent produced an EMPTY reply" + fi + if grep -qiE 'connection refused|connection error|econnrefused|fetch failed|http 4[0-9][0-9]|unauthorized|invalid api key|authentication failed' "$out"; then + guide_fail "agent reply contained a connection/auth error: $(grep -iE 'connection|unauthorized|auth|http 4' "$out" | head -1)" + fi + echo "[$AGENT] reply (first 20 lines):" + head -20 "$out" +} + +# Run a command under a hard timeout; map 124 to a guide-drift hang message. +run_timed() { # $1=outfile, rest=command + local out="$1"; shift + timeout "$TIMEOUT" "$@" > "$out" 2>&1 + local rc=$? + if [ "$rc" -eq 124 ]; then + redact "$out" # guide_fail exits below, so scrub the transcript here too + echo "[$AGENT] last 40 lines before timeout:"; tail -40 "$out" 2>/dev/null || true + guide_fail "invoke timed out after ${TIMEOUT}s (headless-TTY hang -- the recipe likely needs a non-interactive/print flag)" + fi + return "$rc" +} + +# ── Pi: no connect.py command at HEAD -> hand-written recipe ────────────── +write_pi_config() { + if unsloth connect pi --help >/dev/null 2>&1; then + # Tripwire: once a real recipe exists, the hand-written config would mask any + # drift in it, defeating the point of this CI. Fail hard so the cell is + # migrated to the self-updating `unsloth connect pi --no-launch` path. + guide_fail "connect.py now ships a 'pi' command -- migrate this CI cell to the 'unsloth connect pi --no-launch' path so the documented recipe is exercised (the hand-written Pi config no longer reflects it)" + fi + mkdir -p "$HOME/.pi/agent" + python3 - "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" "$UNSLOTH_MODEL_ID" <<'PY' +import json, os, sys +base, key, model = sys.argv[1], sys.argv[2], sys.argv[3] +cfg = {"providers": {"unsloth": { + "api": "openai-completions", + "baseUrl": f"{base}/v1", + "apiKey": key, + "models": [{"id": model}], +}}} +path = os.path.expanduser("~/.pi/agent/models.json") +with open(path, "w") as fh: + json.dump(cfg, fh, indent=2) +PY + cp "$HOME/.pi/agent/models.json" "$REDACTED_DIR/pi-models.json" 2>/dev/null || true + redact "$REDACTED_DIR/pi-models.json" +} + +# ── 5-agent connect.py path: parse env + command from --no-launch ───────── +# Populates globals CONNECT_ENV (export/unset lines) and CONNECT_CMD (the +# launch command on the last printed line), and runs connect.py's config +# writers as a side effect (it writes ~/.codex, ~/.claude, etc.). +parse_connect() { + local raw="$LOGS_DIR/connect-${AGENT}.txt" + if ! unsloth connect "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + cat "$raw" + guide_fail "'unsloth connect ${AGENT} --no-launch' exited non-zero" + fi + echo "[$AGENT] connect --no-launch printed:"; cat "$raw" + CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)" + # The launch command is the last non-export, non-status line. connect.py + # prints "Studio · model " and "Updated ..." status lines first. + CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \ + | grep -E '[^[:space:]]' | tail -1)" + [ -n "$CONNECT_CMD" ] || guide_fail "could not parse a launch command from connect --no-launch output" + redact "$raw" +} + +# Cross-check the documented contract knobs so silent connect.py changes +# (env-var rename, wire_api flip, attribution setting drop) also fail/flag. +crosscheck_contract() { + local raw="$LOGS_DIR/connect-${AGENT}.txt" + case "$AGENT" in + codex) + grep -q 'UNSLOTH_STUDIO_AUTH_TOKEN' "$raw" \ + || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (connect.py _CODEX_ENV_KEY)" + if [ -f "$HOME/.codex/config.toml" ]; then + grep -q 'wire_api = "responses"' "$HOME/.codex/config.toml" \ + || guide_fail "Codex wire_api is no longer \"responses\" in ~/.codex/config.toml" + cp "$HOME/.codex/config.toml" "$REDACTED_DIR/codex-config.toml" + fi + grep -q 'codex --oss --profile unsloth_api' "$raw" \ + || echo "::warning::Codex launch command changed from 'codex --oss --profile unsloth_api'" + ;; + claude) + grep -q 'ANTHROPIC_AUTH_TOKEN' "$raw" \ + || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (connect.py claude())" + if [ -f "$HOME/.claude/settings.json" ]; then + grep -q '"CLAUDE_CODE_ATTRIBUTION_HEADER"' "$HOME/.claude/settings.json" \ + || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER not written to ~/.claude/settings.json (ensure_claude_attribution_header)" + cp "$HOME/.claude/settings.json" "$REDACTED_DIR/claude-settings.json" + fi + ;; + hermes) + grep -q 'UNSLOTH_API_KEY' "$raw" \ + || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (connect.py _HERMES_ENV_KEY)" + [ -f "$HOME/.hermes/config.yaml" ] && cp "$HOME/.hermes/config.yaml" "$REDACTED_DIR/hermes-config.yaml" + ;; + openclaw) + if [ -f "$HOME/.openclaw/openclaw.json" ]; then + grep -q '"openai-completions"' "$HOME/.openclaw/openclaw.json" \ + || echo "::warning::OpenClaw provider api is no longer 'openai-completions' (write_openclaw_config)" + cp "$HOME/.openclaw/openclaw.json" "$REDACTED_DIR/openclaw.json" + fi + ;; + opencode) + [ -f "$HOME/.config/opencode/opencode.json" ] && cp "$HOME/.config/opencode/opencode.json" "$REDACTED_DIR/opencode.json" + ;; + esac + redact "$REDACTED_DIR"/* 2>/dev/null || true +} + +# Heavyweight agents (hermes, openclaw) bake a large system prompt + tool JSON +# schemas into every request, which a CPU runner cannot prefill before the invoke +# timeout. As with claude's --tools, we shrink the request from the agent's own +# config: zero tools for the connection probe collapses the prompt to a few +# hundred tokens, since both CLIs gate the bulk of their prompt on having tools. + +# Hermes: an explicit empty cli toolset disables all tools (and drops the +# tool-gated guidance blocks), so -z sends ~300 tokens instead of thousands. +# hermes ships a DEFAULT config.yaml that already has a populated +# platform_toolsets, and `unsloth connect` merges into it, so we must override +# cli (not just append). That needs a YAML parser, and the runner's bare +# python3 has no PyYAML -- but the venv that ships `unsloth` does (connect.py +# imports yaml), so run the patch with that interpreter. +# (-z reads platform_toolsets.cli; --ignore-rules is a no-op under -z.) +patch_hermes_tools() { # $1 = none|default + # Find a python that can import yaml. The runner's bare python3 cannot, but the + # interpreter in the `unsloth` console-script shebang provably can (it runs + # connect.py's write_hermes_config, which imports yaml). Try that first, then + # any python on PATH, then the venv sibling, picking the first with PyYAML. + local cand py="" shebang + shebang="$(head -1 "$(command -v unsloth)" 2>/dev/null | sed -n 's/^#![[:space:]]*//p' | awk '{print $1}')" + for cand in "$shebang" python3 python "$(dirname "$(command -v unsloth)")/python"; do + [ -n "$cand" ] || continue + { [ -x "$cand" ] || command -v "$cand" >/dev/null 2>&1; } || continue + if "$cand" -c 'import yaml' 2>/dev/null; then py="$cand"; break; fi + done + [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch ~/.hermes/config.yaml" + echo "[hermes] patching config with $py" + "$py" - "$1" <<'PY' +import os, sys +import yaml +mode = sys.argv[1] +p = os.path.expanduser("~/.hermes/config.yaml") +cfg = (yaml.safe_load(open(p)) or {}) if os.path.exists(p) else {} +ts = cfg.get("platform_toolsets") +if not isinstance(ts, dict): + ts = cfg["platform_toolsets"] = {} +if mode == "none": + ts["cli"] = [] # explicit empty list -> zero tools (not "defaults") +else: + ts.pop("cli", None) # file-edit needs real tools -> restore defaults +with open(p, "w") as fh: + yaml.safe_dump(cfg, fh, sort_keys=False) +print(f"[hermes] platform_toolsets.cli = {ts.get('cli', 'default')}") +PY +} + +# OpenClaw: 'openclaw agent' has no tool/prompt flags, so we define a 'ci' agent +# in openclaw.json. tools.deny ["*"] sends zero tool schemas (deny always wins) +# for the connection probe; contextInjection "never" + defaults.skipBootstrap +# drop the auto-injected AGENTS.md/SOUL.md bootstrap (the bulk of the prompt) for +# both modes. --agent must reference a defined agent, so write it before invoking. +patch_openclaw_agent() { # $1 = notools|tools + python3 - "$1" <<'PY' +import os, sys, json +mode = sys.argv[1] +p = os.path.expanduser("~/.openclaw/openclaw.json") +cfg = json.load(open(p)) if os.path.exists(p) else {} +agents = cfg.setdefault("agents", {}) +agents.setdefault("defaults", {})["skipBootstrap"] = True +lst = [a for a in agents.get("list", []) if a.get("id") != "ci"] +agent = {"id": "ci", "contextInjection": "never"} +if mode == "notools": + agent["tools"] = {"deny": ["*"]} +lst.append(agent) +agents["list"] = lst +with open(p, "w") as fh: + json.dump(cfg, fh, indent=2) +print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}") +PY +} + +# Build an invoke script that applies connect.py's env then runs the launch +# command (with extra args appended) under bash. We do NOT eval connect's env +# into this shell; we write it into a one-shot script so the export/unset +# semantics are exactly what connect.py printed. The script path is absolute +# so it is valid even when the caller has cd'd into a scratch work dir. +invoke_via_connect() { # $1=outfile, rest=extra args appended to the command + local out="$1"; shift + local script="$LOGS_DIR/invoke-${AGENT}.sh" + local real; real="$(mktemp)" + { + echo "set -uo pipefail" + echo "$CONNECT_ENV" + # Append extra args (the prompt / flags) to the launch command verbatim. + printf '%s' "$CONNECT_CMD" + local a + for a in "$@"; do printf ' %q' "$a"; done + printf '\n' + } > "$real" + # Upload a REDACTED copy of the script, but EXECUTE the un-redacted one from a + # temp path outside the artifact dir. Redacting the script we run would turn + # the real `export TOKEN=sk-...` line into `export TOKEN=`, which is + # invalid bash (the `<`/`>` are redirections) and silently breaks every agent. + # Writing the redacted copy up front keeps the key out of the artifact even if + # the run times out (run_timed exits before returning here). + cp "$real" "$script"; redact "$script" + echo "[$AGENT] invoking (timeout ${TIMEOUT}s): $CONNECT_CMD $*" + run_timed "$out" bash "$real" + local rc=$? + rm -f "$real" + redact "$out" # the transcript can echo the token; scrub before upload + return "$rc" +} + +# ═════════════════════════════════════════════════════════════════════════ +case "$MODE" in + # ── connection: trivial prompt, assert a non-empty, error-free reply ──── + connection) + PROMPT='Reply with exactly the single word: pong' + OUT="$LOGS_DIR/${AGENT}-connection.txt" + if [ "$AGENT" = "pi" ]; then + write_pi_config + run_timed "$OUT" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$PROMPT" + else + parse_connect + crosscheck_contract + # claude/codex run in print mode via the flags connect.py emits + # (claude -p / codex exec). For agents whose default subcommand prints + # to stdout we pass the prompt through ctx.args. + case "$AGENT" in + claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; + codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; + opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; + hermes) patch_hermes_tools none + invoke_via_connect "$OUT" -z "$PROMPT" ;; + openclaw) patch_openclaw_agent notools + invoke_via_connect "$OUT" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; + *) invoke_via_connect "$OUT" "$PROMPT" ;; + esac + fi + # A non-zero exit from the documented launch command is drift even if it + # printed something: a benign-looking "command not found" / usage dump would + # otherwise slip past assert_reply (which only flags empty/error-keyword text). + rc=$? + [ "$rc" -eq 0 ] || guide_fail "the documented launch command exited non-zero (rc=$rc) -- see the transcript above" + assert_reply "$OUT" + echo "[$AGENT] connection OK" + ;; + + # ── file-edit: deterministic 2-turn hello.py test (Qwen3.5-2B) ────────── + file-edit) + WORK="$WORKDIR_BASE/${AGENT}" + rm -rf "$WORK"; mkdir -p "$WORK" + OUT1="$LOGS_DIR/${AGENT}-fileedit-turn1.txt" + OUT2="$LOGS_DIR/${AGENT}-fileedit-turn2.txt" + T1='Create a file named hello.py in the current directory whose entire contents are a single line: print("Hello"). Do not run it.' + T2='Run hello.py with python and show me the exact output.' + + # The connect.py recipe writers + crosscheck must see the repo; run them + # from the repo root BEFORE cd-ing into the scratch work dir. + if [ "$AGENT" != "pi" ]; then + parse_connect + crosscheck_contract + # File-edit needs real tools, so we cannot zero them as in connection. + # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md + # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work + # dir is empty, so no project context files are auto-loaded either. + case "$AGENT" in + hermes) patch_hermes_tools default ;; + openclaw) patch_openclaw_agent tools ;; + esac + else + write_pi_config + fi + + # Drive from inside the work dir so the agent edits files there. All log + # writes use absolute $LOGS_DIR, so cwd does not matter for them. + cd "$WORK" || guide_fail "could not enter work dir $WORK" + + invoke_turn() { # $1=outfile $2=continue? $3=prompt + local out="$1" cont="$2" prompt="$3" + case "$AGENT" in + pi) run_timed "$out" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$prompt" ;; + claude) + # --dangerously-skip-permissions lets headless claude actually use the + # Write/Bash tools (otherwise it blocks on an approval prompt and emits + # nothing). IS_SANDBOX=1 (exported above) authorizes it. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" "${CLAUDE_EDIT_FLAGS[@]}" --dangerously-skip-permissions -p --continue "$prompt" + else + invoke_via_connect "$out" "${CLAUDE_EDIT_FLAGS[@]}" --dangerously-skip-permissions -p "$prompt" + fi ;; + codex) + # --dangerously-bypass-approvals-and-sandbox gives codex exec + # workspace-write (default is read-only -> cannot create hello.py) and + # skips the bubblewrap sandbox that the runner lacks. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" exec --dangerously-bypass-approvals-and-sandbox resume --last "$prompt" + else + invoke_via_connect "$out" exec --dangerously-bypass-approvals-and-sandbox "$prompt" + fi ;; + opencode) invoke_via_connect "$out" run "$prompt" ;; + hermes) invoke_via_connect "$out" -z "$prompt" ;; + openclaw) invoke_via_connect "$out" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;; + *) invoke_via_connect "$out" "$prompt" ;; + esac + } + + # Turn 1: create hello.py. + invoke_turn "$OUT1" fresh "$T1" + # Fail on a non-zero agent exit before trusting side effects: an agent can + # error out (API/tool failure) yet leave a plausible file/transcript behind, + # which would otherwise slip past the assertions below (mirrors connection). + rc=$? + [ "$rc" -eq 0 ] || { echo "[$AGENT] turn-1 transcript:"; tail -40 "$OUT1" 2>/dev/null || true; \ + guide_fail "turn 1 (create hello.py) exited non-zero (rc=$rc)"; } + + # Hard assertions on the side effect (the real test): file + content + run. + if [ ! -f hello.py ]; then + echo "[$AGENT] turn-1 transcript:"; tail -40 "$OUT1" 2>/dev/null || true + guide_fail "turn 1 did not create hello.py" + fi + grep -q 'Hello' hello.py || guide_fail "hello.py does not contain 'Hello'" + RUN_OUT="$(python3 hello.py 2>&1 || true)" + [ "$RUN_OUT" = "Hello" ] || guide_fail "python3 hello.py printed '$RUN_OUT', expected exactly 'Hello'" + echo "[$AGENT] turn 1 OK (file created, prints 'Hello')" + + # Turn 2: same cwd + session continuation; assert the agent's run output + # contains Hello. Narration drift is WARN-only, missing output is a hard fail. + invoke_turn "$OUT2" continue "$T2" + rc=$? + [ "$rc" -eq 0 ] || { echo "[$AGENT] turn-2 transcript:"; tail -60 "$OUT2" 2>/dev/null || true; \ + guide_fail "turn 2 (run hello.py) exited non-zero (rc=$rc)"; } + if grep -q 'Hello' "$OUT2"; then + echo "[$AGENT] turn 2 OK (run output contains 'Hello')" + else + echo "[$AGENT] turn-2 transcript:"; tail -60 "$OUT2" 2>/dev/null || true + guide_fail "turn 2 run/bash output did not contain 'Hello'" + fi + cd "$REPO_ROOT" || true + echo "[$AGENT] file-edit OK" + ;; + + # ── attribution-ab: Claude Code KV-cache HIT vs MISS ──────────────────── + attribution-ab) + [ "$AGENT" = "claude" ] || guide_fail "attribution-ab only applies to claude" + # The llama-server log filename uses the INTERNAL random llama.cpp port, + # not STUDIO_PORT, so we never glob by port: assert-prompt-cache.sh picks + # the newest llama-*.log and we slice it by a byte offset (`mark`) captured + # right before the measured turn, so an earlier turn's reuse can't leak in. + LLAMA_LOG_DIR="${UNSLOTH_LLAMA_LOG_DIR:-$HOME/.unsloth/studio/logs/llama-server}" + export LLAMA_LOG_DIR + parse_connect # writes ~/.claude/settings.json (header=0) + env + crosscheck_contract + PROMPT='Reply with exactly the single word: pong' + + # Phase A: header DISABLED (=0, the documented setting) -> expect a HIT on + # the continued turn. connect.py's ensure_claude_attribution_header() set 0. + invoke_via_connect "$LOGS_DIR/claude-ab-hit-1.txt" -p "$PROMPT" # turn 1 primes + FROM_HIT="$(bash "$CACHE_HELPER" mark)" # offset before turn 2 + invoke_via_connect "$LOGS_DIR/claude-ab-hit-2.txt" -p --continue "$PROMPT again" + CACHE_LOG_FROM="$FROM_HIT" bash "$CACHE_HELPER" log HIT + + # Phase B: header ENABLED -> expect a MISS. The header prepends a + # per-request-changing attribution line to the system prompt, so the shared + # prefix changes every turn and the KV cache is invalidated (~90% slower); + # this is exactly what the guide flag prevents. + python3 - <<'PY' +import json, os +p = os.path.expanduser("~/.claude/settings.json") +s = json.load(open(p)) if os.path.exists(p) else {} +s.setdefault("env", {})["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "1" +json.dump(s, open(p, "w"), indent=2) +PY + invoke_via_connect "$LOGS_DIR/claude-ab-miss-1.txt" -p "$PROMPT" + FROM_MISS="$(bash "$CACHE_HELPER" mark)" + invoke_via_connect "$LOGS_DIR/claude-ab-miss-2.txt" -p --continue "$PROMPT again" + CACHE_LOG_FROM="$FROM_MISS" bash "$CACHE_HELPER" log MISS + echo "[claude] attribution A/B OK (header=0 HIT, header=1 MISS)" + ;; + + *) + echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2 + exit 2 + ;; +esac diff --git a/.github/scripts/agent-guides-install.sh b/.github/scripts/agent-guides-install.sh new file mode 100755 index 0000000000..dfab8aec80 --- /dev/null +++ b/.github/scripts/agent-guides-install.sh @@ -0,0 +1,105 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Install one coding-agent CLI for the Local Agent Guides CI. Isolated as +# failure class (b) "agent package install failed": npm/curl flakiness here +# is the single biggest source of false reds, so installs retry with +# backoff and the only ::error:: this script can emit is class (b). The +# install recipes mirror the install_hint strings in +# unsloth_cli/commands/connect.py at HEAD. +# +# Usage: agent-guides-install.sh +# agent in: claude codex hermes openclaw opencode pi +set -uo pipefail + +AGENT="${1:?usage: agent-guides-install.sh }" +mkdir -p logs +LOG="logs/install-${AGENT}.log" + +install_fail() { + echo "::error::[agent install failed] agent=${AGENT}: $* (class (b): the agent CLI did not install; not a server or guide problem)." >&2 + echo "---- tail $LOG ----" >&2 + tail -60 "$LOG" 2>/dev/null || true + exit 1 +} + +# npm registry flakiness is common in CI; retry 3x with linear backoff. +npm_retry() { + local pkg="$1" i + for i in 1 2 3; do + if npm install -g "$pkg" >> "$LOG" 2>&1; then + return 0 + fi + echo "[install] npm install -g $pkg attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + sleep "$((i * 10))" + done + return 1 +} + +# curl|bash installers, retried at the curl layer. We download to a temp file +# first and only execute on a fully successful fetch, so a truncated download +# (network hiccup mid-stream) can never run a half-written installer. +curl_bash() { + local url="$1"; shift + local i tmp + tmp="$(mktemp)" + for i in 1 2 3; do + if curl -fsSL --retry 3 --retry-delay 5 "$url" -o "$tmp" 2>>"$LOG" \ + && bash "$tmp" "$@" >> "$LOG" 2>&1; then + rm -f "$tmp" + return 0 + fi + echo "[install] curl|bash $url attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + sleep "$((i * 10))" + done + rm -f "$tmp" + return 1 +} + +echo "[install] agent=$AGENT (log=$LOG)" +case "$AGENT" in + claude) + # connect.py install_hint: curl -fsSL https://claude.ai/install.sh | bash + curl_bash "https://claude.ai/install.sh" || install_fail "claude installer failed" + # The installer drops the binary under ~/.local/bin. + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + ;; + codex) + # connect.py install_hint: npm install -g @openai/codex + npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed" + ;; + opencode) + # connect.py install_hint: npm install -g opencode-ai + npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed" + ;; + openclaw) + # connect.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash + # npm is the more deterministic path in CI and matches the agent's docs; + # fall back to the connect.py curl installer if the npm tag is missing. + if ! npm_retry "openclaw@latest"; then + curl_bash "https://openclaw.ai/install.sh" || install_fail "openclaw install failed (npm + curl)" + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + fi + ;; + hermes) + # connect.py install_hint: + # curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash + curl_bash "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh" \ + --non-interactive --skip-setup --skip-browser --no-skills \ + || install_fail "hermes installer failed" + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + ;; + pi) + # No connect.py recipe; the agent's documented package name. The CLI moved + # from the now-deprecated @mariozechner scope to @earendil-works (the old + # scope is frozen, so installing it would test a stale Pi against the API). + npm_retry "@earendil-works/pi-coding-agent" \ + || install_fail "npm install -g @earendil-works/pi-coding-agent failed" + ;; + *) + install_fail "unknown agent '$AGENT'" + ;; +esac + +echo "[install] OK for $AGENT" diff --git a/.github/scripts/assert-prompt-cache.sh b/.github/scripts/assert-prompt-cache.sh new file mode 100755 index 0000000000..f5b6b075eb --- /dev/null +++ b/.github/scripts/assert-prompt-cache.sh @@ -0,0 +1,238 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Prompt-cache (KV-cache prefix reuse) detection, two strategies in one helper: +# +# mode=api A 2-turn /v1/chat/completions probe. Turn 2 prepends turn 1 + +# its reply, so the shared prefix must be served from llama.cpp's +# KV cache. Asserts usage.prompt_tokens_details.cached_tokens > 0 +# on turn 2. This is the OpenAI-dialect server cache sanity. +# WHY this works on chat completions: the chat path forwards +# llama-server's real cached_tokens through +# studio/backend/routes/inference.py:482-489 (_prompt_tokens_details) +# into prompt_tokens_details (inference.py:519). +# +# mode=log Read the llama-server log and decide HIT vs MISS from the +# prompt-reprocessing trace. WHY the log (not the API field): +# the Anthropic /v1/messages path builds AnthropicUsage( +# input_tokens=..., output_tokens=...) at inference.py:8787-8790 +# / :8829-8832 and NEVER sets cache_read_input_tokens, which +# therefore stays at its model default of 0 +# (studio/backend/models/inference.py:1655). So an Anthropic-path +# client (Claude Code, OpenClaw is openai-completions but Claude +# Code is the canonical Anthropic agent) can get a real KV-cache +# hit that the API usage field reports as 0. The only ground +# truth for the Anthropic path is the llama-server log. +# +# Log location (verified): studio/backend/core/inference/llama_cpp.py:4363-4365 +# _swa_cache_path().parent/"logs"/"llama-server"/llama-[label]-port-

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

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

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

-try.log`). Filenames are + # tool-generated timestamps, so ls -t is safe here. + # shellcheck disable=SC2012 + ls -1t "$dir"/llama-*.log 2>/dev/null | head -1 +} + +case "$MODE" in + # ------------------------------------------------------------------------- + # mark: emit the current byte size of the newest llama log so a caller can + # scan only the slice a single turn produced (set CACHE_LOG_FROM to it). + # ------------------------------------------------------------------------- + mark) + log="$(_newest_log || true)" + if [ -n "$log" ] && [ -f "$log" ]; then + wc -c < "$log" | tr -d ' ' + else + echo 0 + fi + exit 0 + ;; + + # ------------------------------------------------------------------------- + # api: 2-turn /v1/chat/completions, assert turn-2 cached_tokens > 0. + # ------------------------------------------------------------------------- + api) + BASE_URL="${2:?usage: assert-prompt-cache.sh api BASE_URL API_KEY}" + API_KEY="${3:?usage: assert-prompt-cache.sh api BASE_URL API_KEY}" + + # A deliberately long, fixed system prompt makes the shared prefix big so a + # KV-cache hit is unambiguous (cached_tokens grows with the reused prefix). + SYS='You are a meticulous assistant. Always answer concisely and correctly. This is a fixed system preamble that exists only to create a large, identical prompt prefix across both turns so the KV cache has something substantial to reuse on the second request. Do not mention this preamble.' + + turn1_body() { + jq -n --arg sys "$SYS" '{ + model: "default", + messages: [ + {role:"system", content:$sys}, + {role:"user", content:"What is the capital of France?"} + ], + temperature: 0.0, seed: 3407, max_tokens: 40, stream: false, + enable_thinking: false + }' + } + + echo "[cache/api] turn 1 (prime the KV cache)" + R1="$(curl -fs -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Authorization: Bearer ${API_KEY}" -H 'content-type: application/json' \ + --max-time 240 -d "$(turn1_body)")" || { + echo "::error::[cache/api] turn-1 /v1/chat/completions request failed. Unsloth server/API regression." + exit 1 + } + A1="$(echo "$R1" | jq -r '.choices[0].message.content // ""')" + + turn2_body() { + jq -n --arg sys "$SYS" --arg a1 "$A1" '{ + model: "default", + messages: [ + {role:"system", content:$sys}, + {role:"user", content:"What is the capital of France?"}, + {role:"assistant", content:$a1}, + {role:"user", content:"And the capital of Germany?"} + ], + temperature: 0.0, seed: 3407, max_tokens: 40, stream: false, + enable_thinking: false + }' + } + + echo "[cache/api] turn 2 (expect cached_tokens > 0)" + R2="$(curl -fs -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Authorization: Bearer ${API_KEY}" -H 'content-type: application/json' \ + --max-time 240 -d "$(turn2_body)")" || { + echo "::error::[cache/api] turn-2 /v1/chat/completions request failed. Unsloth server/API regression." + exit 1 + } + + CACHED="$(echo "$R2" | jq -r '.usage.prompt_tokens_details.cached_tokens // 0')" + PROMPT_TOK="$(echo "$R2" | jq -r '.usage.prompt_tokens // 0')" + echo "[cache/api] turn-2 usage: prompt_tokens=${PROMPT_TOK} cached_tokens=${CACHED}" + + if [ -z "$CACHED" ] || ! [ "$CACHED" -gt 0 ] 2>/dev/null; then + echo "::error::[cache/api] turn-2 usage.prompt_tokens_details.cached_tokens=${CACHED}, expected > 0. The server is not surfacing llama.cpp KV-cache hits on /v1/chat/completions. Check studio/backend/routes/inference.py:482-489 (_prompt_tokens_details) and :519. Full turn-2 usage:" + echo "$R2" | jq -c '.usage' 2>/dev/null || echo "$R2" + exit 1 + fi + echo "[cache/api] PASS server cache sanity (cached_tokens=${CACHED} > 0)" + exit 0 + ;; + + # ------------------------------------------------------------------------- + # log: classify the newest llama-server log (from CACHE_LOG_FROM bytes on) + # as HIT or MISS and compare to EXPECT. + # ------------------------------------------------------------------------- + log) + EXPECT="${2:?usage: assert-prompt-cache.sh log HIT|MISS}" + FROM="${CACHE_LOG_FROM:-0}" + + log="$(_newest_log || true)" + if [ -z "$log" ] || [ ! -f "$log" ]; then + echo "::error::[cache/log] no llama-server log under ${LLAMA_LOG_DIR:-$(_default_log_dir)}. Cannot read KV-cache trace. (Path contract: studio/backend/core/inference/llama_cpp.py:4363-4365.)" + exit 1 + fi + echo "[cache/log] reading $log from byte $FROM" + + # Scan only the slice produced after FROM. + slice="$(tail -c "+$((FROM + 1))" "$log" 2>/dev/null || cat "$log")" + + # ---- HIT detectors (most-specific first) ----------------------------- + # 1. Modern + legacy "re-used N tokens" / "reused N" (N>0). Primary signal + # per the design brief. + reused_n="$(printf '%s\n' "$slice" \ + | grep -aoiE 're-?used[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 2. "kv cache rm [START, end)" with START>0 => prefix [0,START) reused. + cache_rm_start="$(printf '%s\n' "$slice" \ + | grep -aoiE 'kv cache rm \[[0-9]+' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 3. "n_past = N" with N>0 after a prompt-processing line (prefix kept). + n_past_n="$(printf '%s\n' "$slice" \ + | grep -aoiE 'n_past[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 4. tokens_cached / tokens from cache (some builds). + tok_cached="$(printf '%s\n' "$slice" \ + | grep -aoiE 'tokens_cached[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + + # ---- MISS detectors -------------------------------------------------- + # Explicit forced full re-processing (SWA / recurrent) or kv cache rm [0,. + forced_full=0 + if printf '%s\n' "$slice" | grep -aqiE 'forcing full prompt re-?processing|kv cache rm \[0,'; then + forced_full=1 + fi + + HIT=0 + why="" + if [ -n "$reused_n" ] && [ "$reused_n" -gt 0 ] 2>/dev/null; then + HIT=1; why="re-used=$reused_n" + elif [ -n "$cache_rm_start" ] && [ "$cache_rm_start" -gt 0 ] 2>/dev/null; then + HIT=1; why="kv-cache-rm-start=$cache_rm_start" + elif [ -n "$tok_cached" ] && [ "$tok_cached" -gt 0 ] 2>/dev/null; then + HIT=1; why="tokens_cached=$tok_cached" + elif [ "$forced_full" = "0" ] && [ -n "$n_past_n" ] && [ "$n_past_n" -gt 0 ] 2>/dev/null; then + # n_past>0 is the weakest signal; only trust it if nothing forced a full + # reprocess. (On a cold slot n_past tracks total processed, so it is a + # last-resort fallback per the brief.) + HIT=1; why="n_past=$n_past_n(fallback)" + fi + [ "$HIT" = "1" ] || why="${why:-no-reuse-markers (forced_full=$forced_full)}" + + OBSERVED="MISS"; [ "$HIT" = "1" ] && OBSERVED="HIT" + echo "[cache/log] observed=$OBSERVED expected=$EXPECT ($why)" + + if [ "$OBSERVED" != "$EXPECT" ]; then + echo "::error::[cache/log] KV-cache observed=$OBSERVED but expected=$EXPECT ($why). See the attribution A/B note in the workflow." + echo "---- llama-server log slice (last 60 lines) ----" + printf '%s\n' "$slice" | tail -60 + exit 1 + fi + echo "[cache/log] PASS ($OBSERVED == $EXPECT)" + exit 0 + ;; + + *) + echo "::error::unknown mode '$MODE' (want api|log|mark)" + exit 1 + ;; +esac diff --git a/.github/scripts/ci-connect-prompt.txt b/.github/scripts/ci-connect-prompt.txt new file mode 100644 index 0000000000..2d96f2b1a8 --- /dev/null +++ b/.github/scripts/ci-connect-prompt.txt @@ -0,0 +1 @@ +You are a helpful assistant in a CI connectivity check. Answer the user directly in plain text. Do not use any tools, do not take any actions, and do not explain. Just reply with the answer. diff --git a/.github/scripts/ci-min-system-prompt.txt b/.github/scripts/ci-min-system-prompt.txt new file mode 100644 index 0000000000..d55b828bbd --- /dev/null +++ b/.github/scripts/ci-min-system-prompt.txt @@ -0,0 +1 @@ +You are a coding assistant running non-interactively in a CI smoke test. Use the available file-editing and shell tools to complete the user's request directly and concisely. Do not ask questions or explain; just do the task. diff --git a/.github/scripts/hf-download-with-retry.sh b/.github/scripts/hf-download-with-retry.sh index c5ee013c80..013a459f46 100755 --- a/.github/scripts/hf-download-with-retry.sh +++ b/.github/scripts/hf-download-with-retry.sh @@ -1,4 +1,6 @@ #!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 # # Download a single file from a Hugging Face repo with a stall-retry # watchdog. Used by the Studio CI workflows so a hung hf-xet transfer diff --git a/.github/scripts/serve-unsloth-run.sh b/.github/scripts/serve-unsloth-run.sh new file mode 100755 index 0000000000..34b8b962c6 --- /dev/null +++ b/.github/scripts/serve-unsloth-run.sh @@ -0,0 +1,172 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Boot `unsloth run --disable-tools` in the background, wait for it to be +# healthy, parse the minted API key from the banner, and resolve the +# /v1/models id. Exports everything downstream steps need into $GITHUB_ENV +# (or prints it when run outside Actions). Factored out of the workflow so +# the failure-isolation logic lives in one shellcheck-clean place. +# +# Usage: +# serve-unsloth-run.sh --model REPO --gguf-variant VAR --port PORT \ +# [--gguf-file PATH] [--extra "--seed 3407 --temp 0"] \ +# [--log-dir logs] [--health-timeout 300] +# +# Why a helper and not inline YAML +# -------------------------------- +# * Every `unsloth run` invocation here is the *Unsloth server* under test. +# A failure to come up healthy is class (a) "server/API regression" and +# must be reported with a distinct `::error::` BEFORE any agent runs. +# * The banner is the documented contract a human copies from. We parse the +# exact `API Key:` line printed by unsloth_cli/commands/studio.py +# (` API Key: ` non-silent, `API Key: ` silent) so a +# silent change to that line is also caught. +# * `unsloth run` re-execs into the studio venv ($STUDIO_HOME/unsloth_studio), +# so in CI after `install.sh --local` it runs the PR's repo code. +# +# Outputs written to $GITHUB_ENV (and echoed): +# UNSLOTH_API_KEY the sk-unsloth-* key minted on the banner +# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth connect` +# finds THIS server, not the hardcoded :8888) +# UNSLOTH_BASE_URL same as UNSLOTH_STUDIO_URL (alias for clarity) +# UNSLOTH_MODEL_ID the canonical id reported by /v1/models +# UNSLOTH_SERVER_PID pid of the backgrounded `unsloth run` +# UNSLOTH_LLAMA_LOG_DIR ~/.unsloth/studio/logs/llama-server + +set -uo pipefail + +# ── arg parse ──────────────────────────────────────────────────────────── +MODEL="" +GGUF_VARIANT="" +GGUF_FILE="" +PORT="" +EXTRA="" +LOG_DIR="logs" +HEALTH_TIMEOUT="300" + +while [ "$#" -gt 0 ]; do + case "$1" in + --model) MODEL="$2"; shift 2 ;; + --gguf-variant) GGUF_VARIANT="$2"; shift 2 ;; + --gguf-file) GGUF_FILE="$2"; shift 2 ;; + --port) PORT="$2"; shift 2 ;; + --extra) EXTRA="$2"; shift 2 ;; + --log-dir) LOG_DIR="$2"; shift 2 ;; + --health-timeout) HEALTH_TIMEOUT="$2"; shift 2 ;; + *) echo "serve-unsloth-run.sh: unknown arg '$1'" >&2; exit 2 ;; + esac +done + +[ -n "$PORT" ] || { echo "serve-unsloth-run.sh: --port is required" >&2; exit 2; } +if [ -z "$MODEL" ] && [ -z "$GGUF_FILE" ]; then + echo "serve-unsloth-run.sh: one of --model or --gguf-file is required" >&2 + exit 2 +fi + +mkdir -p "$LOG_DIR" +SERVER_LOG="$LOG_DIR/unsloth-run-${PORT}.log" +BASE_URL="http://127.0.0.1:${PORT}" +STUDIO_HOME_DIR="${STUDIO_HOME:-$HOME/.unsloth/studio}" +LLAMA_LOG_DIR="${STUDIO_HOME_DIR}/logs/llama-server" + +# Emit a key=value pair to $GITHUB_ENV when set, always echo for local runs. +emit() { + echo "$1=$2" + if [ -n "${GITHUB_ENV:-}" ]; then + echo "$1=$2" >> "$GITHUB_ENV" + fi +} + +server_fail() { + echo "::error::Unsloth server/API regression: $*" >&2 + echo "---- last 200 lines of $SERVER_LOG ----" >&2 + tail -200 "$SERVER_LOG" 2>/dev/null || true + exit 1 +} + +# ── port collision guard ───────────────────────────────────────────────── +# A leftover listener (or a parallel matrix cell that wandered onto our port) +# would make us attach to the wrong server and mask a real regression. Fail +# fast instead. +if command -v ss >/dev/null 2>&1; then + if ss -tln 2>/dev/null | grep -q ":${PORT}\b"; then + server_fail "port ${PORT} already has a listener before we started (collision)" + fi +fi + +# ── build the command ──────────────────────────────────────────────────── +# `unsloth run` == alias of `unsloth studio run`. --disable-tools is REQUIRED +# (passthrough mode) so the agent's own tools relay instead of the server's. +# --no-cloudflare keeps us off the network (loopback bind, no tunnel attempt). +CMD=(unsloth run -H 127.0.0.1 -p "$PORT" --disable-tools --no-cloudflare) +if [ -n "$GGUF_FILE" ]; then + CMD+=(--model "$GGUF_FILE") +else + CMD+=(--model "$MODEL") + [ -n "$GGUF_VARIANT" ] && CMD+=(--gguf-variant "$GGUF_VARIANT") +fi +# Determinism knobs + any caller passthrough (e.g. --seed 3407 --temp 0). +# shellcheck disable=SC2206 # intentional word-split of caller-controlled flags +[ -n "$EXTRA" ] && CMD+=($EXTRA) + +echo "[serve] launching: ${CMD[*]}" +echo "[serve] server log: $SERVER_LOG" + +# Run detached, no controlling TTY (setsid avoids any TTY-prompt hang and +# detaches from this step's process group so the job's teardown is clean). +setsid "${CMD[@]}" > "$SERVER_LOG" 2>&1 < /dev/null & +SERVER_PID=$! +emit UNSLOTH_SERVER_PID "$SERVER_PID" + +# ── wait for /api/health == healthy ────────────────────────────────────── +HEALTHY=0 +for _ in $(seq 1 "$HEALTH_TIMEOUT"); do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + server_fail "process exited before becoming healthy (pid $SERVER_PID)" + fi + if curl -fs "${BASE_URL}/api/health" -o "$LOG_DIR/health-${PORT}.json" 2>/dev/null; then + if jq -e '.status == "healthy"' "$LOG_DIR/health-${PORT}.json" >/dev/null 2>&1; then + HEALTHY=1 + break + fi + fi + sleep 1 +done +[ "$HEALTHY" = "1" ] || server_fail "did not report /api/health healthy within ${HEALTH_TIMEOUT}s" +echo "[serve] /api/health healthy" + +# ── parse the API key from the banner ──────────────────────────────────── +# Match both the non-silent " API Key: " and silent "API Key: " +# forms. We do NOT trust a fixed column count; we take the sk-unsloth-* token. +API_KEY="" +for _ in $(seq 1 30); do + API_KEY="$(grep -aoE 'sk-unsloth-[A-Za-z0-9_-]+' "$SERVER_LOG" 2>/dev/null | head -1 || true)" + [ -n "$API_KEY" ] && break + sleep 1 +done +if [ -z "$API_KEY" ]; then + # Fallback: take whatever follows an "API Key:" label, in case the key + # prefix scheme changes. Still a parse-fragility guard, not silent. + API_KEY="$(grep -aE 'API Key:' "$SERVER_LOG" 2>/dev/null \ + | sed -E 's/.*API Key:[[:space:]]*//' | head -1 || true)" +fi +[ -n "$API_KEY" ] || server_fail "could not parse an API key from the banner (banner-parse fragility -- check the 'API Key:' line in unsloth_cli/commands/studio.py)" +echo "::add-mask::${API_KEY}" +emit UNSLOTH_API_KEY "$API_KEY" + +# ── resolve /v1/models id ──────────────────────────────────────────────── +if ! curl -fs "${BASE_URL}/v1/models" \ + -H "Authorization: Bearer ${API_KEY}" -o "$LOG_DIR/models-${PORT}.json" 2>/dev/null; then + server_fail "/v1/models did not respond (or rejected the banner key)" +fi +MODEL_ID="$(jq -r '.data[0].id // empty' "$LOG_DIR/models-${PORT}.json" 2>/dev/null || true)" +[ -n "$MODEL_ID" ] || server_fail "/v1/models returned no model id (model failed to load)" +echo "[serve] resolved model id: $MODEL_ID" + +emit UNSLOTH_MODEL_ID "$MODEL_ID" +emit UNSLOTH_STUDIO_URL "$BASE_URL" +emit UNSLOTH_BASE_URL "$BASE_URL" +emit UNSLOTH_LLAMA_LOG_DIR "$LLAMA_LOG_DIR" + +echo "[serve] server is up: ${BASE_URL} (model ${MODEL_ID})" diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index f7c338d76b..b56a6c2615 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -2204,12 +2204,13 @@ jobs: pip install -e "$RUNNER_TEMP/unsloth-zoo" --no-deps pip show unsloth_zoo - - name: llama.cpp install via unsloth_zoo.llama_cpp + `llama-cli --help` smoke + - name: llama.cpp install via unsloth_zoo.llama_cpp + CLI `--help` smoke # Exercise the canonical `unsloth_zoo.llama_cpp.install_llama_cpp` # flow that GGUF export uses at runtime: clone ggml-org/llama.cpp # into ~/.unsloth/llama.cpp, build the LLAMA_CPP_TARGETS list # (llama-quantize, llama-cli, llama-mtmd-cli, llama-gguf-split, - # llama-server) via cmake, then run `llama-cli --help`. + # llama-server) via cmake, then run `--help` on whichever CLI + # inference binary the build actually produced. # # This replaces the previous "download upstream prebuilt zip" # approach, which silently exited 0 with the message @@ -2218,6 +2219,18 @@ jobs: # matched their current asset names). The build path is the same # one Unsloth users hit in production via `model.save_pretrained_gguf`. # + # We do NOT hard-require `llama-cli` specifically: upstream + # ggml-org/llama.cpp moved the cli/server/ui targets behind the + # `LLAMA_BUILD_SERVER` cmake option (tools/CMakeLists.txt) and the + # set of binaries that survive a given checkout drifts over time + # (e.g. a recent build root shipped llama-server + llama-quantize + # + llama-diffusion-cli but no llama-cli). The durable contract is + # "install_llama_cpp produced a working CLI inference binary AND a + # working quantizer", so we --help-probe the first of + # llama-cli / llama-mtmd-cli / llama-server that exists. If a + # future llama.cpp restores llama-cli it is first in the list and + # is preferred, so this stays backwards compatible. + # # Wall-time budget: ~3-5 min cold, dominated by cmake build of # 5 targets on the runner's 4 cores. Apt-package install is # handled by `install_llama_cpp` itself via its @@ -2252,8 +2265,9 @@ jobs: print(f"Build targets: {LLAMA_CPP_TARGETS}") # install_llama_cpp returns (quantizer_path, converter_script_path). # The quantizer's directory is the `llama.cpp` install root, which - # also holds llama-cli after build/bin/llama-* gets copied up - # (llama_cpp.py:867-871). + # also holds the CLI inference binaries after build/bin/llama-* gets + # copied up (llama_cpp.py:1450-1454; on Windows they stay in + # build/bin/Release/). quantizer, converter = install_llama_cpp(print_output=True) assert quantizer and os.path.exists(quantizer), ( f"install_llama_cpp returned quantizer={quantizer!r} but file missing" @@ -2262,25 +2276,54 @@ jobs: f"install_llama_cpp returned converter={converter!r} but missing" ) install_root = os.path.dirname(quantizer) - cli = os.path.join(install_root, "llama-cli") - assert os.path.exists(cli), ( - f"llama-cli not found at {cli!r} after build. Build root contents: " - f"{sorted(p for p in os.listdir(install_root) if p.startswith('llama-'))[:20]}" - ) - assert os.access(cli, os.X_OK), f"{cli!r} not executable" - # `llama-cli --help` exits non-zero on some builds; the contract - # is that recognizable help text appears on stdout/stderr. + is_windows = sys.platform == "win32" + exe = ".exe" if is_windows else "" + # Search both the copied-up root and the Windows build/bin/Release/ + # location the quantizer might already live in. + search_dirs = [install_root] + win_release = os.path.join(install_root, "build", "bin", "Release") + if win_release not in search_dirs: + search_dirs.append(win_release) + # Any of these proves a working llama.cpp CLI inference binary was + # built. Order = preference: llama-cli is canonical (restored first + # if upstream brings it back), then the multimodal CLI, then the + # server (always built whenever cli would be, behind LLAMA_BUILD_SERVER). + cli_names = [f"llama-cli{exe}", f"llama-mtmd-cli{exe}", f"llama-server{exe}"] + cli = None + cli_name = None + for name in cli_names: + for d in search_dirs: + candidate = os.path.join(d, name) + if os.path.exists(candidate) and (is_windows or os.access(candidate, os.X_OK)): + cli, cli_name = candidate, name + break + if cli is not None: + break + if cli is None: + found = [] + for d in search_dirs: + if os.path.isdir(d): + found += [p for p in os.listdir(d) if p.startswith("llama-")] + raise AssertionError( + f"No CLI inference binary ({', '.join(cli_names)}) found after " + f"build in {search_dirs}. Build root contents: {sorted(set(found))[:20]}" + ) + print(f"Using CLI inference binary: {cli_name} -> {cli}") + # `--help` exits non-zero on some builds; the contract is that + # recognizable help text appears on stdout/stderr. llama-server + # exposes a different flag set than llama-cli, so accept its + # tokens too (e.g. --host / --port / "server"). proc = subprocess.run( [cli, "--help"], capture_output=True, text=True, timeout=30, ) combined = (proc.stdout or "") + (proc.stderr or "") - print("--- llama-cli --help (first 30 lines) ---") + print(f"--- {cli_name} --help (first 30 lines) ---") print("\n".join(combined.splitlines()[:30])) assert any( tok in combined.lower() - for tok in ("usage", "--help", "--model", "-m,") + for tok in ("usage", "--help", "--model", "-m,", "--host", "--port", "server") ), ( - f"llama-cli --help produced no recognizable help text. " + f"{cli_name} --help produced no recognizable help text. " f"exit={proc.returncode}\nstdout: {proc.stdout[:400]!r}\n" f"stderr: {proc.stderr[:400]!r}" ) @@ -2296,7 +2339,7 @@ jobs: f"stderr: {q.stderr[:400]!r}" ) print( - f"\nOK: install_llama_cpp produced a working llama-cli at {cli} " + f"\nOK: install_llama_cpp produced a working {cli_name} at {cli} " f"and llama-quantize at {quantizer}." ) PY diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml new file mode 100644 index 0000000000..299ee3f18b --- /dev/null +++ b/.github/workflows/local-agent-guides-ci.yml @@ -0,0 +1,611 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Local Agent Guides CI +# ===================== +# Detects when our local-agent setup recipes drift out of sync with +# `unsloth run`. Boots a real `unsloth run --disable-tools` server and +# drives the coding agents end to end through the *exact* recipes defined +# in unsloth_cli/commands/connect.py (the in-repo source of truth -- there +# is no docs/ tree). Wherever connect.py has a recipe we drive the agent +# via `unsloth connect --no-launch` and execute what it prints, so +# the test self-updates against connect.py and catches silent recipe drift. +# +# Source-of-truth files this workflow guards: +# unsloth_cli/commands/connect.py the `unsloth connect ` recipes +# unsloth_cli/commands/studio.py the `unsloth run` banner (API Key line) +# +# Failure taxonomy (each surfaced with a distinct ::error:: + the agent name +# + the connect.py location, so a red X is immediately triageable): +# (a) Unsloth server/API regression -- the dialect HTTP preflight fails +# BEFORE the agent runs (or the server never becomes healthy). +# (b) Agent package install failed -- npm/curl install of the CLI failed. +# (c) Guide drift -- preflight passed + install ok, but +# the documented `unsloth connect` flow produced no/garbled output. +# +# Agents covered (6): claude, codex, hermes, openclaw, opencode, pi. +# - claude/codex/hermes/openclaw/opencode have a connect.py recipe. +# - pi has NO `unsloth connect pi` command in connect.py at HEAD; it is +# driven by a hand-written recipe and the matrix cell asserts that the +# missing connect recipe is the (known) reason, so the day connect.py +# grows a `pi` command this cell flips to the self-updating path. + +name: Local Agent Guides CI + +on: + # Off-peak weekly, deliberately a NON-:00 minute to dodge the top-of-hour + # GitHub-hosted-runner stampede. + schedule: + - cron: '37 7 * * 1' + workflow_dispatch: + pull_request: + paths: + - 'unsloth_cli/**' + - 'studio/backend/routes/**' + # Contracts this workflow asserts that live outside routes/**: the + # /api/health endpoint, the llama-server KV-cache log behavior, and the + # request/response schemas the agent dialects depend on. + - 'studio/backend/main.py' + - 'studio/backend/core/inference/llama_cpp.py' + - 'studio/backend/models/**' + - 'install.sh' + - '.github/workflows/local-agent-guides-ci.yml' + - '.github/scripts/serve-unsloth-run.sh' + - '.github/scripts/assert-prompt-cache.sh' + - '.github/scripts/agent-guides-install.sh' + - '.github/scripts/agent-guides-drive.sh' + - '.github/scripts/ci-connect-prompt.txt' + - '.github/scripts/ci-min-system-prompt.txt' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +# Secret handling on pull_request: these jobs check out and run PR-controlled code +# (install.sh, .github/scripts/**), so HF_TOKEN (an external HF credential) is gated +# off pull_request at each step below -- public GGUF repos still download anonymously. +# GH_TOKEN (GITHUB_TOKEN) is kept: it is the job-scoped contents:read token and +# install_llama_prebuilt.py needs it for the GitHub releases API (else 403s). + +env: + # Determinism precedent (studio-inference-smoke.yml): temp 0 + fixed seed. + UNSLOTH_SEED: '3407' + # A single invoke must never hang the runner on a headless TTY prompt. With + # prefill-shrinking flags (minimal system prompt + restricted tools) a turn on + # a 4B model finishes in a couple of minutes on CPU; this also caps how long a + # still-large-prompt agent burns before failing. Well under the 6h job cap. + AGENT_INVOKE_TIMEOUT: '600' + +jobs: + # ═════════════════════════════════════════════════════════════════════ + # Job 1: connection + # Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect, + # install the agent, run `unsloth connect --no-launch`, execute + # the emitted recipe with a trivial prompt, assert a non-empty reply. + # Runs on PR + weekly + dispatch. Each matrix cell is its own runner so + # it serves exactly one model on its own port. + # ═════════════════════════════════════════════════════════════════════ + connection: + name: connection (${{ matrix.agent }}) + runs-on: ubuntu-latest + timeout-minutes: 40 + strategy: + fail-fast: false + matrix: + agent: [claude, codex, hermes, openclaw, opencode, pi] + include: + # OpenClaw needs Node 24; everything else is happy on 22. + - agent: openclaw + node: '24' + env: + # gemma-4-E4B (128K context, capable enough to drive every agent for a + # trivial reply; the 270m model produced empty/failed responses for + # codex/openclaw and is below hermes' 64K context floor). Served as a flat + # GGUF file (the -MTP- repo ships no separate draft, so this is plain 4B). + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18901' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: ${{ matrix.node || '22' }} + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + # ── boot the server under test (factored helper) ────────────────── + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + # ── (a) server/API preflight: prove the dialect works BEFORE the agent ─ + # Distinct error class. If this step fails it is a SERVER regression, + # not the agent's or the guide's fault, and the agent steps never run. + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect; this is class (a), not guide drift). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + # Anthropic Messages dialect. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + # Codex always streams /v1/responses. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). + # OpenClaw's connect.py recipe writes an "openai-completions" + # provider (write_openclaw_config), so it uses this path, not + # /v1/messages. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + # ── (b) install the agent CLI (hardened npm/curl, retried) ───────── + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + # ── (c) drive the agent via connect.py and assert a reply ────────── + # For the 5 agents with a connect.py recipe we run + # `unsloth connect --no-launch`, eval its env/unset exports, + # then run the printed command with a hard timeout (no headless-TTY + # hang). Pi has no connect recipe, so it is driven by hand and the + # cell asserts that absence is the (known) reason. + - name: Drive ${{ matrix.agent }} via unsloth connect (class-c isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: connection-${{ matrix.agent }}-log + path: | + logs/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job 2: file-edit + # The deterministic 2-turn hello.py test on Qwen3.5-4B (smaller models + # can't reliably drive the heavyweight agents' edit flows). Weekly + + # dispatch only -- it is the slow, model-heavy job and must not gate PRs. + # ═════════════════════════════════════════════════════════════════════ + file-edit: + name: file-edit (${{ matrix.agent }}) + if: github.event_name != 'pull_request' + runs-on: ubuntu-latest + timeout-minutes: 60 + # hermes and openclaw drive a multi-turn tool loop that a CPU-only runner + # cannot finish in time (e.g. openclaw holds its 300s session-write-lock past + # expiry; each turn re-prefills the tool prompt at ~16 tok/s). Their endpoint + # wiring + generation are already hard-gated by the connection job, so the + # file-edit cell is best-effort here -- it still runs and uploads logs, but a + # timeout does not fail the workflow. Drop best_effort (or move e2e to a GPU + # runner) to make it blocking again. + continue-on-error: ${{ matrix.best_effort || false }} + strategy: + fail-fast: false + matrix: + agent: [claude, codex, hermes, openclaw, opencode, pi] + include: + - agent: openclaw + node: '24' + best_effort: true + - agent: hermes + best_effort: true + env: + # gemma-4-E4B served as a flat GGUF file (cache size tracks the .gguf 1:1, + # no xet-chunk inflation; the -MTP- repo ships no separate draft file). + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18902' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: ${{ matrix.node || '22' }} + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect; this is class (a), not guide drift). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + # Probe the same dialect the agent will use, so a streaming/messages + # regression in the weekly run is reported as class (a) here instead of + # surfacing later as guide drift (mirrors the connection job). + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: 2-turn hello.py test (class-c isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh file-edit "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: file-edit-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job 3: prompt-cache + # (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0 + # (server prompt-cache sanity). + # (b) Claude Code attribution A/B: with CLAUDE_CODE_ATTRIBUTION_HEADER=0 + # expect a llama-server KV-cache HIT on turn 2; without it expect a + # MISS. If it inverts, the guide flag is stale. + # PR + weekly + dispatch (cheap, gemma-3-270m). + # ═════════════════════════════════════════════════════════════════════ + prompt-cache: + name: prompt-cache (gemma-3-270m) + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18903' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-3-270m) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --model "$GGUF_REPO" --gguf-variant "$GGUF_VARIANT" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" + + # (a) server prompt-cache sanity on the OpenAI chat path. The helper runs + # the 2-turn probe internally (turn 2 reuses turn 1's prefix) and asserts + # turn-2 usage.prompt_tokens_details.cached_tokens > 0. This is the hard + # gate -- it proves llama.cpp KV reuse is surfaced on /v1/chat/completions. + - name: Server prompt-cache sanity (cached_tokens > 0) + run: bash .github/scripts/assert-prompt-cache.sh api "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" + + - name: Install Claude Code (class-b isolation) + env: + AGENT: claude + run: bash .github/scripts/agent-guides-install.sh claude + + # (b) Claude attribution A/B against the llama-server log. This is the most + # environment-sensitive check (it depends on the bundled llama.cpp's + # slot-reuse log wording and on claude --continue reusing the prefix), so + # it is non-blocking until calibrated on the first scheduled run; the + # server cache sanity above is the hard gate. The step still prints the + # observed HIT/MISS so drift is visible in the log + artifacts. + - name: Claude attribution A/B (HIT with header=0, MISS without) + continue-on-error: true + run: bash .github/scripts/agent-guides-drive.sh attribution-ab claude + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: prompt-cache-log + path: | + logs/ + redacted-configs/ + retention-days: 7 diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 221e86f235..864630f9f0 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -241,7 +241,8 @@ jobs: # non-zero binary exit is an Unsloth/Studio bug. - name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1) env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} # 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 @@ -332,7 +333,8 @@ jobs: # 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 }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} UNSLOTH_COMPILE_DISABLE: '1' run: | mkdir -p mlx_workdir @@ -348,7 +350,8 @@ jobs: # the saved dir. - name: MLX export round-trip — RELOAD LoRA (fresh process) env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} UNSLOTH_COMPILE_DISABLE: '1' run: | python tests/studio/run_real_mlx_smoke.py reload \ @@ -357,7 +360,8 @@ jobs: - name: MLX export round-trip — RELOAD merged_16bit (fresh process) env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} UNSLOTH_COMPILE_DISABLE: '1' run: | python tests/studio/run_real_mlx_smoke.py reload \ @@ -372,7 +376,8 @@ jobs: # 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 }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then python tests/studio/run_real_mlx_smoke.py reload \ diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml index e747605322..188e078f90 100644 --- a/.github/workflows/release-desktop.yml +++ b/.github/workflows/release-desktop.yml @@ -353,7 +353,7 @@ jobs: if: matrix.platform == 'ubuntu-22.04' run: | sudo apt-get update - sudo apt-get install -y libwebkit2gtk-4.1-dev libayatana-appindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf + sudo apt-get install -y libwebkit2gtk-4.1-dev libappindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf # ── Node.js ── - name: Setup Node.js @@ -406,9 +406,65 @@ jobs: if (config.bundle?.linux?.rpm) { throw new Error('bundle.linux.rpm must not be configured'); } + if (config.bundle?.linux?.appimage?.bundleMediaFramework !== false) { + throw new Error('Linux AppImage bundleMediaFramework must stay false'); + } const workflow = readFileSync('.github/workflows/release-desktop.yml', 'utf8'); const lines = workflow.split(/\r?\n/); + const linuxInstallLines = lines.filter((line) => line.includes('sudo apt-get install')); + const ayatanaPackage = ['libayatana', 'appindicator3-dev'].join('-'); + if (linuxInstallLines.some((line) => line.includes(ayatanaPackage))) { + throw new Error('Desktop Linux release must not install the Ayatana appindicator dev package'); + } + if (!linuxInstallLines.some((line) => line.includes('libappindicator3-dev'))) { + throw new Error('Desktop Linux release must install libappindicator3-dev'); + } + const linuxdeployLines = lines.filter((line) => line.includes('github.com/linuxdeploy/linuxdeploy/releases/download')); + if (!linuxdeployLines.some((line) => line.includes('1-alpha-20250213-2/linuxdeploy-x86_64.AppImage'))) { + throw new Error('Desktop Linux release must pin linuxdeploy 1-alpha-20250213-2'); + } + // A pinned version/path is reproducibility, not integrity: the asset + // can be replaced after upload. Require the immutable SHA-256 digest + // to be pinned AND verified before chmod +x. Scope every check to the + // real "Pin linuxdeploy for AppImage" step so this guard cannot + // satisfy itself; a file-wide scan would match the guard's own code. + const expectedLinuxdeployDigest = '4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a'; + const isComment = (line) => { + const trimmed = line.trim(); + return trimmed.startsWith('#') || trimmed.startsWith('//'); + }; + const stepStart = lines.findIndex((line) => /^\s*- name: Pin linuxdeploy for AppImage\s*$/.test(line)); + if (stepStart === -1) { + throw new Error('Desktop Linux release must keep the "Pin linuxdeploy for AppImage" step'); + } + const stepIndent = lines[stepStart].search(/\S/); + let stepEnd = lines.length; + for (let i = stepStart + 1; i < lines.length; i += 1) { + const line = lines[i]; + if (line.trim() === '') continue; + const indent = line.search(/\S/); + // The next sibling step ('- ...') at the same indent, or any dedent + // below the step, ends this step's block. + if (indent < stepIndent || (indent === stepIndent && /^\s*-\s/.test(line))) { + stepEnd = i; + break; + } + } + const stepLines = lines.slice(stepStart, stepEnd); + const digestEnvRe = /^\s*LINUXDEPLOY_SHA256:\s*["']([0-9a-f]{64})["']\s*$/; + const digestEnvLine = stepLines.find((line) => digestEnvRe.test(line)); + if (!digestEnvLine || digestEnvLine.match(digestEnvRe)[1] !== expectedLinuxdeployDigest) { + throw new Error('Desktop Linux release must pin the linuxdeploy SHA-256 digest in the LINUXDEPLOY_SHA256 env'); + } + const sha256Idx = stepLines.findIndex((line) => !isComment(line) && line.includes('sha256sum -c')); + if (sha256Idx === -1) { + throw new Error('Desktop Linux release must verify the linuxdeploy digest with sha256sum -c before use'); + } + const chmodIdx = stepLines.findIndex((line) => !isComment(line) && /chmod\s+\+x/.test(line)); + if (chmodIdx !== -1 && sha256Idx > chmodIdx) { + throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x'); + } const releaseBodies = []; for (let i = 0; i < lines.length; i += 1) { const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/); @@ -438,6 +494,12 @@ jobs: if (/\brpm\b|\.rpm/i.test(body)) { throw new Error('Desktop release body must not advertise RPM packages'); } + if (/AppImage.*universal|universal.*AppImage/i.test(body)) { + throw new Error('Desktop release body must not advertise AppImage as universal'); + } + if (!/AppImage.*experimental/i.test(body)) { + throw new Error('Desktop release body must mark AppImage as experimental'); + } } JS @@ -562,6 +624,33 @@ jobs: Get-Command trusted-signing-cli -ErrorAction SilentlyContinue || Write-Output "trusted-signing-cli NOT in PATH" trusted-signing-cli --version || Write-Output "trusted-signing-cli failed to run" + # ── Linux: pin AppImage packaging toolchain ── + - name: Pin linuxdeploy for AppImage + if: matrix.platform == 'ubuntu-22.04' + shell: bash + env: + # Pinning the versioned release path is reproducibility, not + # integrity: a GitHub release asset can be replaced (or its delivery + # path compromised) after upload. The SHA-256 below is the immutable + # digest of this exact asset and is the integrity gate. If linuxdeploy + # publishes a new build under this tag, this run fails closed and the + # digest must be re-pinned deliberately. + LINUXDEPLOY_URL: "https://github.com/linuxdeploy/linuxdeploy/releases/download/1-alpha-20250213-2/linuxdeploy-x86_64.AppImage" + LINUXDEPLOY_SHA256: "4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a" + run: | + set -euo pipefail + tools_dir="$RUNNER_TEMP/tauri-tools-cache/tauri" + mkdir -p "$tools_dir" + dest="$tools_dir/linuxdeploy-x86_64.AppImage" + curl -fsSL "$LINUXDEPLOY_URL" -o "$dest" + # Verify the digest BEFORE the binary is ever marked executable. The + # next step builds the AppImage with the Tauri signing key and a + # contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy + # that ran here could exfiltrate signing material or tamper with + # published release artifacts. Fail closed on any mismatch. + echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c - + chmod +x "$dest" + # ── Linux: build + sign + upload ── - name: Build Linux app if: matrix.platform == 'ubuntu-22.04' @@ -570,6 +659,7 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} + XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache with: projectPath: studio tauriScript: npx --prefix . tauri @@ -580,9 +670,10 @@ jobs: **macOS**: Download the Apple Silicon `.dmg`. **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal). + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} @@ -611,9 +702,10 @@ jobs: **macOS**: Download the Apple Silicon `.dmg`. **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal). + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} @@ -643,9 +735,10 @@ jobs: **macOS**: Download the Apple Silicon `.dmg`. **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal). + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 33ac3b9bd8..0ef2ad1e9d 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -434,7 +434,7 @@ jobs: # ───────────────────────────────────────────────────────────── # Semgrep: design-flaw detection (catches what regex-pattern - # scanning of malicious authors cannot — first-party logic bugs + # scanning of malicious authors cannot, e.g. first-party logic bugs # like langchain-core CVE-2025-68664 dumps/dumpd injection, # n8n CVE-2025-68668 _pyodide.eval_code sandbox escape, marimo # CVE-2026-39987 unauth WebSocket). @@ -849,10 +849,13 @@ jobs: grep -q "Standalone pre-install package scanner" scripts/scan_packages.py - name: Scan declared + transitive Python deps - # scan_packages.py exits 1 on CRITICAL/HIGH findings, 0 on - # clean. We swallow the exit because the baseline isn't - # triaged yet; surface the findings in the workflow summary. - # Drop continue-on-error after the first clean run on main. + # scan_packages.py exits 1 on NON-baselined CRITICAL/HIGH + # findings, 0 otherwise. It scans code-only (docstrings and + # comments are blanked first) and suppresses reviewed + # known-good findings via scripts/scan_packages_baseline.json, + # so legitimate-library noise no longer red-fails the gate. + # The step stays advisory until SCAN_ENFORCE=1 (see env below); + # then PIPESTATUS propagates the scanner's exit code. # # `--with-deps` walks PyPI metadata to enumerate every # transitive dep the declared set would install, then scans @@ -869,6 +872,14 @@ jobs: # downloads in exchange for wall-clock parallelism. env: SHARD_FILES: ${{ matrix.shard.files }} + # Enforcement switch. "1" = blocking: a non-baselined CRITICAL/HIGH + # fails the build. scan_packages.py scans code-only (docstrings/comments + # stripped), fetches sdist-only packages directly from PyPI (no build) + # so every shard resolves, and honors the reviewed allowlist at + # scripts/scan_packages_baseline.json, so only NON-baselined + # CRITICAL/HIGH cause its exit 1. The committed baseline makes all three + # shards exit 0 today; set this back to "0" to return to advisory. + SCAN_ENFORCE: "1" run: | set +e mkdir -p logs @@ -884,12 +895,14 @@ jobs: fi done echo "::endgroup::" + rc=0 if [ ${#REQ_ARGS[@]} -eq 0 ]; then echo "[security-audit] shard ${{ matrix.shard.id }}: no PyPI specs, nothing to scan" \ | tee "$LOG" else python scripts/scan_packages.py --with-deps "${REQ_ARGS[@]}" \ 2>&1 | tee "$LOG" + rc=${PIPESTATUS[0]} fi { echo "## scan_packages :: shard ${{ matrix.shard.id }}" @@ -897,11 +910,19 @@ jobs: echo "### Files in this shard" for f in $SHARD_FILES; do echo "- audit-reqs/$f.txt"; done echo + echo "scan_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)" + echo echo '### Findings (tail)' echo '```' tail -200 "$LOG" echo '```' } >> "$GITHUB_STEP_SUMMARY" + # Advisory by default; blocking once SCAN_ENFORCE=1 and the baseline + # is committed. PIPESTATUS is captured above so `tee` does not mask the + # scanner's exit code. + if [ "$SCAN_ENFORCE" = "1" ]; then + exit "$rc" + fi - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 if: always() @@ -975,24 +996,37 @@ jobs: python3 -c "import ast; ast.parse(open('scripts/scan_npm_packages.py').read())" - name: Scan npm tarballs (declared + transitive, no install) - # The script exits 1 on HIGH/CRITICAL findings; we capture the - # full log and surface it in the step summary either way. It - # never runs `npm install`, never executes anything from a - # downloaded tarball, and only fetches from registry.npmjs.org. - # Initially non-blocking so the baseline can settle; drop - # continue-on-error once the baseline is clean for a week. + # scan_npm_packages.py exits 1 on NON-baselined HIGH/CRITICAL + # findings, 0 otherwise. It scans code-only (JS/TS comments are + # blanked first) and honors a reviewed allowlist at + # scripts/scan_npm_packages_baseline.json. It never runs + # `npm install`, never executes anything from a downloaded + # tarball, and only fetches from registry.npmjs.org. The npm + # corpus is clean (the baseline is empty), so the gate is + # enforcing (SCAN_ENFORCE=1) and any new finding fails the build. + env: + SCAN_ENFORCE: "1" run: | - set -o pipefail + set +e LOG=logs-scan-npm.txt python3 scripts/scan_npm_packages.py 2>&1 | tee "$LOG" + rc=${PIPESTATUS[0]} { echo "## scan_npm_packages" echo + echo "scan_npm_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)" + echo echo '### Findings (tail)' echo '```' tail -300 "$LOG" echo '```' } >> "$GITHUB_STEP_SUMMARY" + # Blocking: the npm corpus is clean, so any non-baselined + # HIGH/CRITICAL is new and should fail the build. PIPESTATUS is + # captured above so `tee` does not mask the scanner's exit code. + if [ "$SCAN_ENFORCE" = "1" ]; then + exit "$rc" + fi - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 if: always() diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index b196805cf7..15efee382e 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -83,7 +83,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -100,7 +101,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 88c7344683..ea60252cf6 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -222,9 +222,14 @@ jobs: for s in \ tests/sh/test_get_torch_index_url.sh \ tests/sh/test_mac_intel_compat.sh \ + tests/sh/test_node_decision.sh \ + tests/sh/test_studio_home_node_dir.sh \ + tests/sh/test_system_node_readonly.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ + tests/sh/test_resolve_cuda_archs.sh \ tests/sh/test_tauri_install_exit_order.sh \ - tests/sh/test_torch_constraint.sh; do + tests/sh/test_torch_constraint.sh \ + tests/sh/test_torch_flavor.sh; do echo "::group::$s" bash "$s" echo "::endgroup::" diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index cffb33f71d..aebf90380a 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -97,7 +97,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -114,7 +115,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -317,7 +319,7 @@ jobs: 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 + # (Qwen3.5-2B at Q4_K_XL = ~1.28 GiB). Caching HF_HOME would # store xet chunks + blobs + snapshots = ~4 GiB compressed -- # 4-5x file-size inflation, dominated by xet chunks. Use main's # `--local-dir gguf-cache` pattern to cache the flat .gguf only. @@ -326,8 +328,11 @@ jobs: # path keeps the test off HF_HOME entirely so the cache size # tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images # jobs still cover the gguf_variant resolution path. + # Q4_K_XL, not IQ3_XXS: at IQ3_XXS this model emits malformed + # tool calls that llama-server's peg-native parser rejects with a + # 500. Mac/Windows already use Q4_K_XL for the same reason. GGUF_REPO: unsloth/Qwen3.5-2B-GGUF - GGUF_FILE: Qwen3.5-2B-UD-IQ3_XXS.gguf + GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf STUDIO_PORT: '18889' steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -361,7 +366,8 @@ jobs: id: download-gguf if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -377,7 +383,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -772,6 +779,9 @@ jobs: kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 ss -tln | grep ":${STUDIO_PORT}" || true + # Capture backend + llama-server logs so a 500 has a server-side traceback. + mkdir -p logs/server-logs + cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true - name: Upload logs # Always upload so green runs are still reviewable. @@ -784,6 +794,7 @@ jobs: path: | logs/studio.log logs/install.log + logs/server-logs/ retention-days: 7 # ───────────────────────────────────────────────────────────────────── @@ -838,7 +849,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -856,7 +868,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index 412726538c..617ce189dc 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -68,7 +68,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -85,7 +86,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index c794a34acd..d562294d42 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -91,7 +91,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -110,7 +111,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -346,7 +348,8 @@ jobs: id: download-gguf if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -363,7 +366,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -725,7 +729,8 @@ jobs: # Authenticated + parallel: shared macos-14 NAT egress stalls # multi-GB anonymous downloads. env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -752,7 +757,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml index da944d4b5c..362305cdd4 100644 --- a/.github/workflows/studio-mac-install-matrix.yml +++ b/.github/workflows/studio-mac-install-matrix.yml @@ -63,7 +63,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 4f9f94b534..512af54d53 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -68,7 +68,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -85,7 +86,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml index f554a16415..d104306c7e 100644 --- a/.github/workflows/studio-mac-update-smoke.yml +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -62,7 +62,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -74,7 +75,8 @@ jobs: - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -93,7 +95,8 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml index 1156c264ae..018857de68 100644 --- a/.github/workflows/studio-tauri-smoke.yml +++ b/.github/workflows/studio-tauri-smoke.yml @@ -47,7 +47,7 @@ jobs: run: | sudo apt-get update sudo apt-get install -y \ - libwebkit2gtk-4.1-dev libayatana-appindicator3-dev \ + libwebkit2gtk-4.1-dev libappindicator3-dev \ librsvg2-dev libxdo-dev libssl-dev patchelf - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index de106e201f..297a585430 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -82,7 +82,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -99,7 +100,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -276,6 +278,10 @@ jobs: run: | kill "${STUDIO_IME_PID}" 2>/dev/null || true sleep 2 + # Capture backend + llama-server logs (all three Studios share this + # dir) so a stray 500 has a server-side traceback. + mkdir -p logs/server-logs + cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true - name: Upload Playwright artifacts # Always upload so a green run's screenshots stay reviewable -- @@ -289,6 +295,7 @@ jobs: logs/studio_extra.log logs/studio_ime.log logs/install.log + logs/server-logs/ logs/playwright logs/playwright_extra logs/playwright_ime diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml index 307bb51972..08a79afacd 100644 --- a/.github/workflows/studio-update-smoke.yml +++ b/.github/workflows/studio-update-smoke.yml @@ -71,7 +71,8 @@ jobs: # prebuilt path falls back to source build. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -86,7 +87,8 @@ jobs: # idempotency regressed. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -109,7 +111,8 @@ jobs: # the first one. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index 78efe918ac..e9abd2d669 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -75,7 +75,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -124,7 +125,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index a772a6d102..c44c68278d 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -26,6 +26,7 @@ on: - 'unsloth_cli/**' - 'install.ps1' - 'pyproject.toml' + - 'tests/studio_setup_ps1/**' - '.github/workflows/studio-windows-inference-smoke.yml' push: branches: [main, pip] @@ -65,17 +66,34 @@ jobs: with: persist-credentials: false - # Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit - # test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke. - - name: setup.ps1 unit test (Resolve-CudaToolkit) + # Fast GPU-free gate: parse install.ps1 + setup.ps1 and run the PowerShell + # unit tests (CUDA-toolkit + torch-flavor helpers) before the heavy GGUF smoke. + - name: PowerShell installer unit tests + shell: pwsh + run: | + foreach ($f in @('install.ps1', 'studio/setup.ps1')) { + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path $f).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "$f parsed with no errors" + } + pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 + pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1 + pwsh -NoProfile -File tests/studio/test_node_decision.ps1 + pwsh -NoProfile -File tests/studio/test_node_probe_guard.ps1 + + # uninstall.ps1: native uninstall must keep the shared unsloth.ico while a + # WSL shortcut still references it (dual install), else that shortcut blanks. + - name: uninstall.ps1 unit test (dual-install icon preserve) shell: pwsh run: | $errs = $null [void][System.Management.Automation.Language.Parser]::ParseFile( - (Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs) + (Resolve-Path scripts/uninstall.ps1).Path, [ref]$null, [ref]$errs) if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } - Write-Host "setup.ps1 parsed with no errors" - pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 + Write-Host "uninstall.ps1 parsed with no errors" + pwsh -NoProfile -File tests/studio/test_uninstall_dual_install_icon.ps1 - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: @@ -109,7 +127,8 @@ jobs: # described above (outcome != success). if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -161,7 +180,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -458,7 +478,8 @@ jobs: id: download-gguf if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -506,7 +527,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -888,7 +910,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -938,7 +961,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -1239,3 +1263,539 @@ jobs: logs/install.log logs/llama-server/*.log retention-days: 7 + + # ── folded from studio-windows-no-vs-smoke.yml: install + run with no Visual Studio ── + no-vs-cpu: + name: Studio install + inference without Visual Studio + runs-on: windows-latest + timeout-minutes: 35 + defaults: + run: + shell: bash + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18820' + HF_HOME: ${{ github.workspace }}/hf-cache + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + run: | + $ProgressPreference = 'SilentlyContinue' + npm install -g 'npm@^11' 2>&1 | Out-Host + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { } + } + + - name: Hide Visual Studio + CMake (simulate a host with no build tools) + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + # Rename the Visual Studio install roots (incl. the Installer that holds + # vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss. + foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { + if (Test-Path -LiteralPath $d) { + Rename-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff') + Write-Host "Hid VS: $d" + } + } + # Surgically rename each cmake executable on PATH (not its parent dir -- + # cmake can share a dir with other shims) so Get-Command cmake fails. + $hidden = @() + foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) { + if ($c.Source -and (Test-Path -LiteralPath $c.Source)) { + Rename-Item -LiteralPath $c.Source -NewName ((Split-Path $c.Source -Leaf) + '.off') + $hidden += $c.Source + Write-Host "Hid cmake: $($c.Source)" + } + } + ("HIDDEN_CMAKE=" + ($hidden -join '|')) | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + + - name: Assert Visual Studio + CMake are genuinely undetectable + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + . (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1') + $setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1' + foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Find-VsBuildTools')) { + . ([scriptblock]::Create((Get-FunctionSource -Path $setup -Name $fn))) + } + $vs = Find-VsBuildTools + if ($vs) { Write-Error "Find-VsBuildTools still detects VS: $($vs.Generator) @ $($vs.InstallPath)"; exit 1 } + if (Get-Command cmake -ErrorAction SilentlyContinue) { Write-Error "cmake is still on PATH"; exit 1 } + if (Get-Command cl.exe -ErrorAction SilentlyContinue) { Write-Error "cl.exe is still on PATH"; exit 1 } + Write-Host "Confirmed: no Visual Studio, no cmake, no cl.exe." + + - name: PyTorch CPU wheel installs and imports (no Visual Studio) + run: | + python -m pip install --upgrade pip + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())" + + - name: Install Studio (--local, --no-torch) with no build tools present + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + New-Item -ItemType Directory -Force -Path logs | Out-Null + $ProgressPreference = 'SilentlyContinue' + & ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log + + - name: Assert prebuilt used AND no build tools were installed + run: | + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + fail=0 + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp without VS."; fail=1 + fi + # The deferred build-tool installs must NOT run on the prebuilt path. + for pat in "Kitware.CMake" "Microsoft.VisualStudio.2022.BuildTools" "installing via winget"; do + if grep -qi "$pat" logs/install.log; then + echo "::error::unexpected build-tool install on the prebuilt path: '$pat'"; fail=1 + fi + done + [ -f "$INFO" ] || { echo "::error::no UNSLOTH_PREBUILT_INFO.json"; ls -la "$LLAMA_DIR" || true; fail=1; } + [ -f "$BIN" ] || { echo "::error::no llama-server.exe"; ls -la "$LLAMA_DIR/build/bin" || true; fail=1; } + if [ "$fail" != "0" ]; then grep -iE "cmake|visual studio|prebuilt|source build" logs/install.log | tail -60; exit 1; fi + echo "Prebuilt installed with no build tools:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + run: | + SHIM_DIR=~/.unsloth/studio/bin + [ -f "$SHIM_DIR/unsloth.exe" ] || { echo "::error::unsloth.exe shim not found"; ls -la ~/.unsloth/studio/ || true; exit 1; } + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, load the GGUF + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json || { tail -200 logs/studio.log; exit 1; } + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CINoVS-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP"; cat /tmp/load.json || true; sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_gguf}' /tmp/load.json + + - name: Inference works via the prebuilt llama.cpp (no VS) + run: | + RESP=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/v1/chat/completions" \ + -H "Authorization: Bearer $API_KEY" -H 'content-type: application/json' \ + --max-time 240 \ + -d '{"model":"default","messages":[{"role":"user","content":"What is 1+1? Answer briefly."}],"temperature":0,"max_tokens":32,"stream":false}') + echo "$RESP" | jq '.choices[0].message' || { echo "$RESP"; exit 1; } + CONTENT=$(echo "$RESP" | jq -r '.choices[0].message.content') + [ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; } + echo "Inference OK without Visual Studio: $CONTENT" + + - name: Restore Visual Studio + CMake + if: always() + shell: pwsh + run: | + foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { + $off = "$d.vsoff" + if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } + } + if ($env:HIDDEN_CMAKE) { + foreach ($src in ($env:HIDDEN_CMAKE -split '\|')) { + if ($src -and (Test-Path -LiteralPath "$src.off")) { Rename-Item -LiteralPath "$src.off" -NewName (Split-Path $src -Leaf) } + } + } + + - name: Stop Studio + if: always() + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + + - name: Collect llama-server logs + if: always() + continue-on-error: true + run: | + mkdir -p logs/llama-server + cp -v ~/.unsloth/studio/logs/llama-server/*.log logs/llama-server/ 2>/dev/null || echo "no llama-server logs" + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-no-vs-cpu-log + path: | + logs/install.log + logs/studio.log + logs/llama-server/*.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job B: the GPU (CUDA) prebuilt path is also VS-free (resolve/availability) + # ───────────────────────────────────────────────────────────────────── + no-vs-gpu-resolve: + name: GPU prebuilt resolves without Visual Studio + runs-on: windows-latest + timeout-minutes: 15 + defaults: + run: + shell: bash + env: + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Hide Visual Studio + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { + if (Test-Path -LiteralPath $d) { Rename-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" } + } + + - name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + curl -fsSL -H "Authorization: Bearer $GH_TOKEN" \ + "https://api.github.com/repos/unslothai/llama.cpp/releases/latest" > /tmp/rel.json + echo "release: $(jq -r .tag_name /tmp/rel.json)" + ASSETS=$(jq -r '.assets[].name' /tmp/rel.json) + echo "$ASSETS" | grep -iE 'windows-x64-cuda[0-9]' || { + echo "::error::no Windows x64 CUDA prebuilt asset found in unslothai/llama.cpp latest release" + echo "$ASSETS"; exit 1; } + # AMD parity: hosted runners have no AMD GPU, so the resolver step below + # can't exercise the ROCm path (it resolves to CPU). Pin the per-gfx + # Windows ROCm bundles here so a release that drops them fails loudly -- + # the AMD no-VS guarantee otherwise rides only on shared resolver code. + echo "$ASSETS" | grep -iE 'windows-x64-rocm-gfx' || { + echo "::error::no Windows x64 ROCm (per-gfx) prebuilt asset found in unslothai/llama.cpp latest release" + echo "$ASSETS"; exit 1; } + echo "Windows CUDA and ROCm prebuilts are available -- GPU users get them without compiling." + + - name: The prebuilt resolver runs without Visual Studio + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + # Resolver-only (no GPU on hosted runners, so the host resolves to the + # CPU bundle). The point is that resolution needs no compiler/VS. + python -m pip install --upgrade huggingface_hub + python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > /tmp/resolve.json || { + echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; } + cat /tmp/resolve.json + echo "Prebuilt resolver ran with no Visual Studio present." + + - name: Restore Visual Studio + if: always() + shell: pwsh + run: | + foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { + $off = "$d.vsoff" + if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } + } + + # ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ── + pester: + name: setup.ps1 unit tests (VS 2026 / CMake guard) + runs-on: windows-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Install Pester v5 + shell: pwsh + run: | + Set-PSRepository PSGallery -InstallationPolicy Trusted + Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser + Import-Module Pester -MinimumVersion 5.5.0 + Get-Module Pester | Select-Object Name, Version | Format-Table + + - name: Run Pester suite + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $testDir = Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1' + if (-not (Test-Path $testDir)) { + Write-Error "Test directory not found: $testDir" + exit 1 + } + $cfg = New-PesterConfiguration + $cfg.Run.Path = $testDir + $cfg.Run.Exit = $true # non-zero exit => job fails + $cfg.Run.Throw = $true # also throw on test failure / 0 tests + $cfg.TestResult.Enabled = $true + $cfg.TestResult.OutputFormat = 'NUnitXml' + $cfg.TestResult.OutputPath = Join-Path $env:GITHUB_WORKSPACE 'pester-results.xml' + $cfg.Output.Verbosity = 'Detailed' + Invoke-Pester -Configuration $cfg + + - name: Upload Pester results + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: pester-results-setup-ps1 + path: pester-results.xml + if-no-files-found: warn + + vs-integration: + # Real detection against the VS installed on the runner image (no mocks). + name: real-VS detection (${{ matrix.label }}) + strategy: + fail-fast: false + matrix: + include: + - { os: windows-2022, label: 'VS 2022', expectGen: 'Visual Studio 17 2022', expectToolset: 'v170' } + - { os: windows-2025-vs2026, label: 'VS 2026', expectGen: 'Visual Studio 18 2026', expectToolset: 'v180' } + runs-on: ${{ matrix.os }} + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Detect the real Visual Studio with setup.ps1 functions + shell: pwsh + env: + EXPECT_GEN: ${{ matrix.expectGen }} + EXPECT_TOOLSET: ${{ matrix.expectToolset }} + run: | + $ErrorActionPreference = 'Stop' + . (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1') + $setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1' + foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Get-VcBuildCustomizationsDir', 'Find-VsBuildTools')) { + . ([scriptblock]::Create((Get-FunctionSource -Path $setup -Name $fn))) + } + + # Ground truth from the real vswhere (independent of our code), for visibility. + $vsw = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" + if (Test-Path $vsw) { + $year = (& $vsw -latest -property catalog_productLineVersion 2>$null | Select-Object -First 1) + $path = (& $vsw -latest -property installationPath 2>$null | Select-Object -First 1) + Write-Host "Real vswhere: productLineVersion='$year' installPath='$path'" + } else { + Write-Host "vswhere not present at $vsw (relying on filesystem fallback)" + } + + # Our detection must find the real VS and report the expected generator. + $r = Find-VsBuildTools + if (-not $r) { throw "Find-VsBuildTools returned null on a host with real $env:EXPECT_GEN" } + Write-Host "Find-VsBuildTools -> Generator='$($r.Generator)' Source='$($r.Source)' InstallPath='$($r.InstallPath)'" + if ($r.Generator -ne $env:EXPECT_GEN) { + throw "Detection mismatch: got '$($r.Generator)', expected '$env:EXPECT_GEN'" + } + if (-not (Test-Path $r.InstallPath)) { throw "Detected InstallPath does not exist: $($r.InstallPath)" } + + # Toolset path derivation must match the expected v-number... + $bc = Get-VcBuildCustomizationsDir -VsInstallPath $r.InstallPath -Generator $r.Generator + $derived = Split-Path (Split-Path $bc -Parent) -Leaf # e.g. v170 / v180 + Write-Host "Get-VcBuildCustomizationsDir -> '$bc' (toolset='$derived')" + if ($derived -ne $env:EXPECT_TOOLSET) { + throw "Toolset mismatch: derived '$derived', expected '$env:EXPECT_TOOLSET'" + } + + # ...and that v-number is a real folder on the VS install (where CUDA's + # BuildCustomizations would land). + $vcRoot = Join-Path $r.InstallPath 'MSBuild\Microsoft\VC' + if (Test-Path $vcRoot) { + $realToolsets = @((Get-ChildItem -Path $vcRoot -Directory -ErrorAction SilentlyContinue).Name) + Write-Host "Real VC toolset dirs: $($realToolsets -join ', ')" + if ($realToolsets -notcontains $derived) { + throw "Derived toolset '$derived' is not present on the real $env:EXPECT_GEN install (have: $($realToolsets -join ', '))" + } + Write-Host "OK: toolset '$derived' exists on the real VS install." + } else { + Write-Warning "VC MSBuild root absent ($vcRoot) - C++ workload not installed; skipping on-disk toolset check." + } + + Write-Host "PASS: real $env:EXPECT_GEN detected correctly with toolset '$derived'." + + vcredist-clean-box: + # Validate Test-VCRedistInstalled + Ensure-VCRedist on a throwaway runner: + # present on the stock image, fires on a clean box (signals removed restorably), + # then a literal uninstall/reinstall round trip. Always restored before the end. + name: VC++ runtime detect + install round-trip (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [windows-latest, windows-2025-vs2026] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Detect present, fire on a clean box, and round-trip the install + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + . (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1') + $setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1' + # Dot-source the guard + the logging closure it reaches + # (step/substep -> Write-StudioStdoutMirror / Get-StudioAnsi). + $script:StudioVtOk = $false + $script:UnslothVerbose = $false + foreach ($fn in @('Get-StudioAnsi', 'Write-StudioStdoutMirror', 'step', 'substep', + 'Invoke-SetupCommand', 'Refresh-Environment', + 'Test-VCRedistInstalled', 'Ensure-VCRedist')) { + $src = Get-FunctionSource -Path $setup -Name $fn + if (-not $src) { throw "Function '$fn' not found in setup.ps1" } + . ([scriptblock]::Create($src)) + } + + $regKeys = @( + 'HKLM\SOFTWARE\Microsoft\VisualStudio\14.0\VC\Runtimes\x64', + 'HKLM\SOFTWARE\WOW6432Node\Microsoft\VisualStudio\14.0\VC\Runtimes\x64' + ) + function Show-GroundTruth { + $dll = Join-Path $env:SystemRoot 'System32\vcruntime140_1.dll' + Write-Host (" System32\vcruntime140_1.dll present: {0}" -f (Test-Path $dll)) + foreach ($k in $regKeys) { + $r = Get-ItemProperty -Path "HKLM:\$($k.Substring(5))" -ErrorAction SilentlyContinue + if ($r) { Write-Host (" {0}: Installed={1} {2}.{3}" -f $k, $r.Installed, $r.Major, $r.Minor) } + else { Write-Host (" {0}: (absent)" -f $k) } + } + } + + Write-Host '== A. Detection on the stock runner (expect present) ==' + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'Test-VCRedistInstalled reported ABSENT on a stock runner that ships the VC++ runtime (detection regression).' } + Write-Host ' Test-VCRedistInstalled -> present OK' + + Write-Host '== B. Genuinely clean box (restorable): detection must FIRE ==' + $scratch = Join-Path $env:RUNNER_TEMP 'cleanwin' + New-Item -ItemType Directory -Force -Path (Join-Path $scratch 'System32') | Out-Null + $backup = Join-Path $env:RUNNER_TEMP 'vcreg_backup' + New-Item -ItemType Directory -Force -Path $backup | Out-Null + $origSysRoot = $env:SystemRoot + try { + for ($i = 0; $i -lt $regKeys.Count; $i++) { + reg query $regKeys[$i] *> $null + if ($LASTEXITCODE -eq 0) { + reg export $regKeys[$i] (Join-Path $backup "$i.reg") /y *> $null + reg delete $regKeys[$i] /f *> $null + } + } + $env:SystemRoot = $scratch + if (Test-VCRedistInstalled) { throw 'Detection still PRESENT after both signals were removed (it would never trigger an install on a clean box).' } + Write-Host ' Test-VCRedistInstalled -> absent OK (detection fires on a clean box)' + } finally { + $env:SystemRoot = $origSysRoot + for ($i = 0; $i -lt $regKeys.Count; $i++) { + $f = Join-Path $backup "$i.reg" + if (Test-Path $f) { reg import $f *> $null } + } + } + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'Detection did not recover after restoring the registry (test restore bug).' } + + Write-Host '== C. Literal uninstall on this throwaway VM (official installer), observe detection ==' + $exe = Join-Path $env:RUNNER_TEMP 'vc_redist.x64.exe' + Invoke-WebRequest -Uri 'https://aka.ms/vs/17/release/vc_redist.x64.exe' -OutFile $exe + Start-Process -FilePath $exe -ArgumentList '/uninstall', '/quiet', '/norestart' -Wait + Show-GroundTruth + Write-Host (" Test-VCRedistInstalled after uninstall -> {0}" -f (Test-VCRedistInstalled)) + if (Test-VCRedistInstalled) { + Write-Host ' Note: the Visual Studio on this image ref-counts the runtime, so the package' + Write-Host ' uninstall is a no-op here; section B already proved detection on a clean box.' + } + + Write-Host '== D. Restore via Ensure-VCRedist (winget product path), installer fallback if needed ==' + Ensure-VCRedist + if (-not (Test-VCRedistInstalled)) { + Write-Host ' winget path did not restore it; using the official installer to close the round trip.' + Start-Process -FilePath $exe -ArgumentList '/install', '/quiet', '/norestart' -Wait + } + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'VC++ runtime could not be restored after the uninstall round-trip.' } + Write-Host ' Test-VCRedistInstalled -> present OK' + Write-Host 'PASS: detection is correct on a real install, fires on a clean box, and the install round-trip restores the runtime.' diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index 40d8e530cd..405309916a 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -91,7 +91,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -136,6 +137,17 @@ jobs: } } + - name: Seed a legacy launch-studio.vbs (upgrade-cleanup check) + # Simulate a pre-hardening install so the post-install assertion below + # proves the installer DELETES an existing launch-studio.vbs (the exact + # Kaspersky-flagged file), not merely stops generating it. + shell: pwsh + run: | + $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio' + New-Item -ItemType Directory -Force -Path $appDir | Out-Null + Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode + Write-Host "seeded legacy launch-studio.vbs at $appDir" + - name: Install Studio (--local, --no-torch) # install.ps1 is the supported Windows installer. install.sh # has no Windows branch (apt-get / brew calls). The PS1 @@ -144,7 +156,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 redirects ALL PowerShell streams (stdout, stderr, @@ -192,6 +205,69 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" + - name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut) + # The shortcut launch path is otherwise untested here (the steps below + # boot `unsloth studio` directly). Guard against re-introducing the VBS + # that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk + # pointing anywhere other than hidden PowerShell over launch-studio.ps1. + shell: pwsh + run: | + $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio' + if (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.vbs')) { + throw "regression: launch-studio.vbs exists (the Kaspersky VBS-FP shape)" + } + if (-not (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.ps1'))) { + throw "missing launch-studio.ps1 in $appDir" + } + $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk' + if (-not (Test-Path -LiteralPath $lnk)) { + $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk' + } + if (-not (Test-Path -LiteralPath $lnk)) { throw "no Unsloth Studio.lnk on Desktop or Start Menu" } + $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk) + Write-Host "shortcut target: $($sc.TargetPath)" + Write-Host "shortcut args: $($sc.Arguments)" + if ($sc.TargetPath -match 'wscript\.exe$') { throw "shortcut still targets wscript.exe (VBS host)" } + if ($sc.TargetPath -notmatch 'powershell\.exe$') { throw "unexpected shortcut target: $($sc.TargetPath)" } + if ($sc.Arguments -notmatch '-WindowStyle Hidden') { + throw "shortcut must launch windowless (-WindowStyle Hidden)" + } + Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)" + + - name: Launch Studio via the shortcut and assert health + # Run the exact command the .lnk stores (hidden PowerShell over + # launch-studio.ps1) and confirm it brings the backend up. This is the + # only step that proves the shortcut launch is not silently broken. + # Default port range is 8888-8908; the later UI tests use 18896/18897, so + # there is no conflict, and we tear this server down before they boot. + shell: pwsh + run: | + $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk' + if (-not (Test-Path -LiteralPath $lnk)) { + $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk' + } + $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk) + Write-Host "launching: $($sc.TargetPath) $($sc.Arguments)" + Start-Process -FilePath $sc.TargetPath -ArgumentList $sc.Arguments -WorkingDirectory $sc.WorkingDirectory + $foundPort = 0 + foreach ($i in 1..180) { + foreach ($port in 8888..8908) { + try { + $r = Invoke-RestMethod -Uri "http://127.0.0.1:$port/api/health" -TimeoutSec 1 + if ($r.status -eq 'healthy' -and $r.service -eq 'Unsloth UI Backend') { $foundPort = $port; break } + } catch {} + } + if ($foundPort) { break } + Start-Sleep -Seconds 1 + } + # Tear down the shortcut-launched server before the main UI tests boot. + try { + $owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess + if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null } + } catch {} + if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" } + Write-Host "Studio healthy on port $foundPort (launched via the shortcut)" + - name: Add Studio shim to GITHUB_PATH # install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe # and adds that dir to the User PATH via the Windows registry. diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 4a4806cfb1..888b3d70a3 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -79,15 +79,15 @@ jobs: # 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. + # (1) npm. setup.ps1's Get-NodeDecision requires Node 22.12+ + # (or 20.19+ / 23+) AND npm >=11 because Vite 8 needs both. # actions/setup-node@v4 with `node-version: '22'` lands - # Node 22.22.2 + the npm 10.9.7 it bundles, so the 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. + # Node 22.22.2 + the npm 10.9.7 it bundles, so the decision + # is "bundled" and setup.ps1 downloads an isolated Node (~30 + # MB) we don't need on a runner that already has a fine Node. + # `npm install -g npm@^11` updates the runner's npm in-place + # in ~5 s, flipping the decision to "system" so setup.ps1 + # reuses the existing Node with no download. # # (2) Defender. windows-latest's real-time scan opens / hashes # every file Studio writes during install (Vite output = @@ -133,7 +133,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -180,7 +181,8 @@ jobs: - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -199,7 +201,8 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.gitignore b/.gitignore index 2caf92e546..7926cb3989 100644 --- a/.gitignore +++ b/.gitignore @@ -237,3 +237,5 @@ package-lock.json llama.cpp/ async_task_outputs/ individual_reviews/ +# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo. +/~/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cffbf73cd5..8dcb9130b3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.16 + rev: v0.15.18 hooks: - id: ruff args: diff --git a/README.md b/README.md index b6a4b836a4..9162d29b1c 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,8 @@ unsloth studio -p 8888 ``` For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally. +To reach Studio over HTTPS, use `unsloth studio --secure`. Studio stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Studio reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below). + #### Docker Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run: ```bash @@ -162,13 +164,19 @@ Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Ad ## 📥 Advanced Installation The below advanced instructions are for Unsloth Studio. For Unsloth Core advanced installation, [view our docs](https://unsloth.ai/docs/get-started/install/pip-install#advanced-pip-installation). -#### Developer installs: macOS, Linux, WSL: +#### Developer / Nightly / Experimental installs: macOS, Linux, WSL: +The developer install builds from the `main` branch, which is the latest (nightly) source. ```bash git clone https://github.com/unslothai/unsloth cd unsloth ./install.sh --local unsloth studio -p 8888 ``` +To install into an isolated location (its own virtual env, `auth/`, `studio.db`, cache and llama.cpp build), set `UNSLOTH_STUDIO_HOME` and pass it again at launch: +```bash +UNSLOTH_STUDIO_HOME="$PWD/.studio" ./install.sh --local +UNSLOTH_STUDIO_HOME="$PWD/.studio" unsloth studio -p 8888 +``` Then to update : ```bash cd unsloth && git pull @@ -176,7 +184,8 @@ cd unsloth && git pull unsloth studio -p 8888 ``` -#### Developer installs: Windows PowerShell: +#### Developer / Nightly / Experimental installs: Windows PowerShell: +The developer install builds from the `main` branch, which is the latest (nightly) source. ```powershell git clone https://github.com/unslothai/unsloth.git cd unsloth @@ -184,40 +193,31 @@ Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass .\install.ps1 --local unsloth studio -p 8888 ``` +To install into an isolated location (its own virtual env, `auth/`, `studio.db`, cache and llama.cpp build), set `UNSLOTH_STUDIO_HOME` and pass it again at launch: +```powershell +$env:UNSLOTH_STUDIO_HOME="$PWD\.studio"; .\install.ps1 --local +$env:UNSLOTH_STUDIO_HOME="$PWD\.studio"; unsloth studio -p 8888 +``` Then to update : -```bash -cd unsloth && git pull -./install.sh --local -unsloth studio -p 8888 -``` - -#### Nightly: MacOS, Linux, WSL: -```bash -git clone https://github.com/unslothai/unsloth -cd unsloth -git checkout nightly -./install.sh --local -unsloth studio -p 8888 -``` -Then to launch every time: -```bash -unsloth studio -p 8888 -``` - -#### Nightly: Windows: -Run in Windows Powershell: ```powershell -git clone https://github.com/unslothai/unsloth.git -cd unsloth -git checkout nightly -Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass +cd unsloth; git pull .\install.ps1 --local unsloth studio -p 8888 ``` -Then to launch every time: + +#### Remote access: `--secure` (HTTPS tunnel) vs raw port +By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach it from another device, pick one of: + +- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Studio stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed. ```bash -unsloth studio -p 8888 +unsloth studio --secure -p 8888 ``` +- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network. +```bash +unsloth studio -H 0.0.0.0 -p 8888 +``` + +Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Studio. #### Advanced launch options Installer options can be passed as environment variables. On macOS, Linux and WSL place the variable after the pipe so the shell passes it to `sh`; on Windows set it with `$env:` before piping to `iex`. @@ -246,6 +246,15 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh $env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex ``` +Point the frontend build at a corporate npm mirror/proxy with `UNSLOTH_NPM_REGISTRY` (for the developer install behind a firewall that blocks `registry.npmjs.org`): +```bash +UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh --local +``` +```powershell +$env:UNSLOTH_NPM_REGISTRY='https://artifactory.example.com/api/npm/npm/'; .\install.ps1 --local +``` +It is threaded as `--registry` into the Studio frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force. + Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`. #### Uninstall diff --git a/build.sh b/build.sh index 1558dca240..dc272f0de1 100644 --- a/build.sh +++ b/build.sh @@ -1,4 +1,6 @@ #!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 set -euo pipefail @@ -33,10 +35,19 @@ _restore_gitignores() { } trap _restore_gitignores EXIT +# Corporate-mirror / proxy escape hatch (#6491). When UNSLOTH_NPM_REGISTRY is set we +# thread it as `--registry ` into the installs (overrides frontend/.npmrc's pinned +# registry for both bun and npm; min-release-age / save-exact stay in force). Empty +# array (the default) expands to nothing under `set -u`. +_NPM_REGISTRY_ARGS=() +if [ -n "${UNSLOTH_NPM_REGISTRY:-}" ]; then + _NPM_REGISTRY_ARGS=(--registry "$UNSLOTH_NPM_REGISTRY") +fi + # Use bun for install if available (faster), fall back to npm. _install_ok=false if command -v bun &>/dev/null; then - if bun install; then + if bun install "${_NPM_REGISTRY_ARGS[@]+"${_NPM_REGISTRY_ARGS[@]}"}"; then _install_ok=true else echo "⚠ bun install failed, falling back to npm" @@ -44,8 +55,10 @@ if command -v bun &>/dev/null; then fi fi if [ "$_install_ok" != "true" ]; then - if ! npm install; then + if ! npm install "${_NPM_REGISTRY_ARGS[@]+"${_NPM_REGISTRY_ARGS[@]}"}"; then echo "❌ ERROR: package install failed" >&2 + echo " If you are behind a corporate firewall/proxy, set UNSLOTH_NPM_REGISTRY to your mirror and retry, e.g.:" >&2 + echo " UNSLOTH_NPM_REGISTRY=https://your-mirror.example/api/npm/ ./build.sh" >&2 exit 1 fi fi diff --git a/install.ps1 b/install.ps1 index f036f7f8fc..33cf103317 100644 --- a/install.ps1 +++ b/install.ps1 @@ -482,6 +482,37 @@ function Install-UnslothStudio { } } + # Retry Invoke-InstallCommand on transient uv download failures with backoff. + # Returns the last exit code on permanent failure so rollback still fires. + function Invoke-InstallCommandRetry { + param( + [Parameter(Mandatory = $true, Position = 0)][ScriptBlock]$Command, + [string]$Label = "install step" + ) + # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables). + # TryParse with bounds avoids an Int32 overflow throw. Bounds: 1..100 retries, 0..3600s. + $maxAttempts = 3 + $parsedAttempts = 0 + if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRIES, [ref]$parsedAttempts) -and $parsedAttempts -ge 1 -and $parsedAttempts -le 100) { + $maxAttempts = $parsedAttempts + } + $delay = 3 + $parsedDelay = 0 + if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRY_DELAY, [ref]$parsedDelay) -and $parsedDelay -ge 0 -and $parsedDelay -le 3600) { + $delay = $parsedDelay + } + $attempt = 1 + while ($true) { + $code = Invoke-InstallCommand $Command + if ($code -eq 0) { return 0 } + if ($attempt -ge $maxAttempts) { return $code } + substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow" + Start-Sleep -Seconds $delay + $attempt++ + $delay = $delay * 2 + } + } + function New-StudioShortcuts { param( [Parameter(Mandatory = $true)][string]$UnslothExePath @@ -506,7 +537,6 @@ function Install-UnslothStudio { } $appDir = $StudioDataDir $launcherPs1 = Join-Path $appDir "launch-studio.ps1" - $launcherVbs = Join-Path $appDir "launch-studio.vbs" $desktopDir = [Environment]::GetFolderPath("Desktop") $desktopLink = if ($desktopDir -and $desktopDir.Trim()) { Join-Path $desktopDir "Unsloth Studio.lnk" @@ -799,19 +829,30 @@ exit 0 # even when install.ps1 is executed from PowerShell 7. $utf8Bom = New-Object System.Text.UTF8Encoding($true) [System.IO.File]::WriteAllText($launcherPs1, $launcherContent, $utf8Bom) - # shell.Run(cmd, 0, ...) already hides the window, so -WindowStyle Hidden - # is redundant; omitting it trims an AV-heuristic token (Kaspersky FP). - $vbsContent = @" -Set shell = CreateObject("WScript.Shell") -cmd = "powershell -NoProfile -ExecutionPolicy Bypass -File ""$launcherPs1""" -shell.Run cmd, 0, False -"@ - # WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths. - Set-Content -LiteralPath $launcherVbs -Value $vbsContent -Encoding Unicode -Force + # No .vbs launcher is written. A WScript.Shell .vbs that spawns a hidden + # ExecutionPolicy-Bypass PowerShell is exactly the shape VBS-dropper + # heuristics score (e.g. Kaspersky HEUR:Trojan.VBS.Agent.gen). The .lnk + # shortcuts instead point straight at powershell.exe running + # launch-studio.ps1 with a hidden window (selected below). + + # Delete any launch-studio.vbs left by a pre-hardening install. New + # installs no longer generate it, but an upgrade that merely stopped + # generating it would leave the exact file AV flags on disk, so remove + # it explicitly. Covers default and env-mode installs (same $appDir). + $legacyLauncherVbs = Join-Path $appDir "launch-studio.vbs" + if (Test-Path -LiteralPath $legacyLauncherVbs) { + Remove-Item -LiteralPath $legacyLauncherVbs -Force -ErrorAction SilentlyContinue + } # 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. + # Snapshot the existing icon first so we can tell whether it actually + # changed and gate the heavier icon-cache refresh on a real change. + $preIconHash = $null + if (Test-Path -LiteralPath $iconPath) { + try { $preIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash } catch {} + } $hasValidIcon = $false if ($bundledIcon -and (Test-Path -LiteralPath $bundledIcon)) { try { @@ -847,6 +888,24 @@ shell.Run cmd, 0, False } } + # Did the icon content actually change vs the previous install? + # Only a real change (or a first/removed icon) should trigger the heavy + # refresh; a no-op reinstall with no icon at all must not. + $iconChanged = $false + if ($hasValidIcon) { + if (-not $preIconHash) { + $iconChanged = $true + } else { + try { + $postIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash + $iconChanged = ($postIconHash -ne $preIconHash) + } catch { $iconChanged = $true } + } + } elseif ($preIconHash) { + # A previously present icon was removed or invalidated. + $iconChanged = $true + } + # Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts # that may point at a deleted workspace; launcher + icon stay. if ($StudioRedirectMode -eq 'env') { @@ -854,8 +913,22 @@ shell.Run cmd, 0, False return } - $wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe" - $shortcutArgs = "//B //Nologo `"$launcherVbs`"" + # Whether this is effectively a first install (no pre-existing .lnk). + # Used to gate the heavier icon-cache refresh below so a no-op reinstall + # does not repeatedly clear caches / restart StartMenuExperienceHost -- + # a behavioral cluster AV heuristics can score as dropper-like. + $firstInstall = -not ( + ($desktopLink -and (Test-Path -LiteralPath $desktopLink)) -or + ($startMenuLink -and (Test-Path -LiteralPath $startMenuLink)) + ) + + # Launch transport for the shortcuts: powershell.exe runs + # launch-studio.ps1 with a hidden window. We deliberately avoid a + # .vbs/WScript.Shell wrapper -- that script-engine shape is what AV + # VBS-dropper heuristics score (Kaspersky HEUR:Trojan.VBS.Agent.gen). + $powershellForLnk = Join-Path $env:SystemRoot "System32\WindowsPowerShell\v1.0\powershell.exe" + $shortcutTarget = $powershellForLnk + $shortcutArgs = "-NoProfile -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$launcherPs1`"" try { $wshell = New-Object -ComObject WScript.Shell @@ -865,9 +938,11 @@ shell.Run cmd, 0, False if (-not $linkPath -or [string]::IsNullOrWhiteSpace($linkPath)) { continue } try { $shortcut = $wshell.CreateShortcut($linkPath) - $shortcut.TargetPath = $wscriptExe + $shortcut.TargetPath = $shortcutTarget $shortcut.Arguments = $shortcutArgs $shortcut.WorkingDirectory = $appDir + # Start minimized so the brief PowerShell console flash is muted. + $shortcut.WindowStyle = 7 $shortcut.Description = "Launch Unsloth Studio" if ($hasValidIcon) { $shortcut.IconLocation = "$iconPath,0" @@ -881,15 +956,13 @@ shell.Run cmd, 0, False } if ($createdShortcutCount -gt 0) { substep "Created Unsloth Studio shortcut" - # Force Explorer to re-read each new shortcut's icon so it renders - # immediately instead of a stale/generic entry (a same-name .lnk - # recreated across reinstalls keeps Explorer's cached per-item icon). - # The reliable, non-disruptive fix (no explorer restart) is a per-item - # SHChangeNotify SHCNE_UPDATEITEM + SHCNF_PATHW per .lnk; the global - # SHCNE_ASSOCCHANGED broadcast alone does NOT recover a stale item. - # Also clear the on-disk icon cache (covers heavier staleness). - try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {} - try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {} + # Always do the cheap, non-disruptive per-item refresh so a + # rewritten same-name .lnk renders with its new target/icon + # immediately (a same-name .lnk recreated across reinstalls keeps + # Explorer's cached per-item icon). The reliable fix (no explorer + # restart) is a per-item SHChangeNotify SHCNE_UPDATEITEM + + # SHCNF_PATHW per .lnk; the global SHCNE_ASSOCCHANGED broadcast + # alone does NOT recover a stale item. try { Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int eventId, uint flags, string item1, System.IntPtr item2);' -ErrorAction SilentlyContinue # SHCNE_UPDATEITEM (0x00002000) + SHCNF_PATHW (0x0005) per shortcut @@ -899,21 +972,31 @@ shell.Run cmd, 0, False # SHCNE_ASSOCCHANGED (0x08000000) global refresh (belt-and-suspenders) [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero) } catch {} - # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN - # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT - # invalidate, so a rewritten same-name shortcut shows the old tile - # until the host restarts. Drop only the render caches (NEVER - # start2.bin -- the pinned layout) and let the host rebuild. - # Best-effort; Win10 has no such host (Test-Path skips it). - try { - $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState" - if (Test-Path -LiteralPath $smehTemp) { - Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue | - Remove-Item -Force -ErrorAction SilentlyContinue - Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue - Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue - } - } catch {} + # Heavier on-disk icon-cache clear + StartMenuExperienceHost tile + # rebuild only when the icon actually changed or this is a first + # install. Running "clear icon cache + kill StartMenuExperienceHost" + # on every no-op reinstall is a dropper-like behavioral cluster and + # is unnecessary when the icon is unchanged (the per-item notify + # above already refreshes the rewritten shortcut). + if ($firstInstall -or $iconChanged) { + try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {} + try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {} + # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN + # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT + # invalidate, so a rewritten same-name shortcut shows the old tile + # until the host restarts. Drop only the render caches (NEVER + # start2.bin -- the pinned layout) and let the host rebuild. + # Best-effort; Win10 has no such host (Test-Path skips it). + try { + $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState" + if (Test-Path -LiteralPath $smehTemp) { + Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue | + Remove-Item -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue + Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue + } + } catch {} + } } else { substep "no Unsloth Studio shortcuts were created" "Yellow" } @@ -1181,7 +1264,7 @@ shell.Run cmd, 0, False # ── Install uv ── Write-TauriLog "STEP" "Installing uv package manager" - $UvMinVersion = "0.7.22" + $UvMinVersion = "0.8.16" function Test-UvVersionOk { $cmd = Get-Command uv -ErrorAction SilentlyContinue if (-not $cmd) { return $false } @@ -1252,6 +1335,15 @@ shell.Run cmd, 0, False $env:UV_COMPILE_BYTECODE_TIMEOUT = "180" } + # uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read + # timeout for large wheel downloads. User-provided values are preserved. + if (-not $env:UV_HTTP_RETRIES) { + $env:UV_HTTP_RETRIES = "5" + } + if (-not $env:UV_HTTP_TIMEOUT) { + $env:UV_HTTP_TIMEOUT = "180" + } + # ── Create venv (migrate old layout if possible, otherwise fresh) ── # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. @@ -1531,29 +1623,87 @@ shell.Run cmd, 0, False if (-not $HasNvidiaSmi) { # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution). # AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type. - $hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue - if (-not $hipinfoExe) { - $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } - $hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } - if ($hipRoot) { - $hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe" - if (Test-Path $hipinfoCandidate) { - Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow - Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow - Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow - $hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate } - } else { - Write-Host " [WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" -ForegroundColor Yellow - Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow - Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow + # Ignore the venv hipInfo.exe (AMD wheel, on PATH): not a HIP SDK, so + # amd-smi would still auto-elevate. Cf. _path_inside_venv(). + function Test-HipinfoIsVenvInternal { + param([AllowNull()][string]$HipinfoPath) + if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false } + # Also derive the venv from the setup python + default Studio home, so + # the venv hipInfo is caught when VenvDir/VIRTUAL_ENV are unset. + $venvRoots = @() + if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV } + $vd = Get-Variable -Name VenvDir -ValueOnly -ErrorAction SilentlyContinue + if ($vd) { $venvRoots += $vd } + if ($env:UNSLOTH_SETUP_PYTHON) { + try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {} + } + if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") } + # A custom Studio home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the + # venv off the default path; seed it too or its hipInfo escapes the filter. + $studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null } + if ($studioHomeEnv) { + # Expand a leading ~ like the canonical resolver; else GetFullPath + # keeps the literal ~ (cwd-relative) and the hipInfo escapes the filter. + if (($studioHomeEnv -eq "~" -or $studioHomeEnv -like "~/*" -or $studioHomeEnv -like "~\*") -and -not [string]::IsNullOrWhiteSpace($env:USERPROFILE)) { + # A bare "~" leaves an empty child path; Join-Path rejects that on + # PS 5.1, so use USERPROFILE directly and only join a real remainder. + $studioHomeRest = $studioHomeEnv.Substring(1).TrimStart('/', '\') + $studioHomeEnv = if ($studioHomeRest) { Join-Path $env:USERPROFILE $studioHomeRest } else { $env:USERPROFILE } } + $venvRoots += (Join-Path $studioHomeEnv "unsloth_studio") + } + try { $hip = [System.IO.Path]::GetFullPath($HipinfoPath).TrimEnd('\', '/') } catch { return $false } + foreach ($root in $venvRoots) { + if ([string]::IsNullOrWhiteSpace($root)) { continue } + try { $r = [System.IO.Path]::GetFullPath($root).TrimEnd('\', '/') } catch { continue } + # Skip a bare drive root (e.g. a non-venv UNSLOTH_SETUP_PYTHON like + # C:\Python311\python.exe yields C:) -- it would match every path on that drive. + if ($r -match '^[a-zA-Z]:$') { continue } + if ($hip.Equals($r, [System.StringComparison]::OrdinalIgnoreCase) -or + $hip.StartsWith($r + [System.IO.Path]::DirectorySeparatorChar, [System.StringComparison]::OrdinalIgnoreCase)) { + return $true + } + } + return $false + } + # Scan all hipinfo and keep the first non-venv one (the venv copy from the + # bnb fix could shadow a real HIP SDK's). -CommandType Application matches + # only real executables, not a user alias/function named hipinfo. + $hipinfoExe = Get-Command hipinfo -CommandType Application -All -ErrorAction SilentlyContinue | + Where-Object { -not (Test-HipinfoIsVenvInternal $_.Source) } | + Select-Object -First 1 + if (-not $hipinfoExe) { + # Iterate the env roots (mirrors the Python list) and take the first non-venv + # bin\hipinfo.exe, so a venv-internal HIP_PATH can't mask a real SDK in ROCM_PATH. + $hipMissingLabel = $null; $hipMissingRoot = $null; $hipMissingCandidate = $null + foreach ($hipEnvLabel in @("HIP_PATH", "HIP_PATH_57", "ROCM_PATH")) { + $hipRoot = [Environment]::GetEnvironmentVariable($hipEnvLabel) + if ([string]::IsNullOrWhiteSpace($hipRoot)) { continue } + $hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe" + if (-not (Test-Path $hipinfoCandidate)) { + if (-not $hipMissingLabel) { $hipMissingLabel = $hipEnvLabel; $hipMissingRoot = $hipRoot; $hipMissingCandidate = $hipinfoCandidate } + continue + } + if (Test-HipinfoIsVenvInternal $hipinfoCandidate) { continue } # venv copy (AMD wheel): not a HIP SDK + Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow + Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow + Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow + $hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate } + break + } + if ((-not $hipinfoExe) -and $hipMissingLabel) { + Write-Host " [WARN] ${hipMissingLabel}=$hipMissingRoot is set but hipinfo.exe not found at $hipMissingCandidate" -ForegroundColor Yellow + Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow + Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow } } if ($hipinfoExe) { $HipSdkInstalled = $true # binary found → SDK is installed regardless of device state try { $hipOut = & $hipinfoExe.Source 2>&1 | Out-String - if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { + if ($hipOut -match "(?i)gcnArchName") { + # hipinfo can crash after printing gcnArchName (#6043). + # Once the arch is printed, keep the ROCm wheel path. $HasROCm = $true $_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() }) $_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } @@ -1563,8 +1713,13 @@ shell.Run cmd, 0, False } else { $ROCmGpuLabel = "AMD ROCm" } + if ($LASTEXITCODE -ne 0) { + Write-Host " [INFO] hipinfo exited with code $LASTEXITCODE but reported gcnArchName -- treating as ROCm-capable (see #6043)" -ForegroundColor Cyan + } } elseif ($LASTEXITCODE -ne 0) { - # hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected") + # hipinfo ran but returned a HIP runtime error without any gcnArchName + # output (e.g. "no ROCm-capable device detected"), or crashed before + # printing device info. $firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1) Write-Host " [WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" -ForegroundColor Yellow Write-Host " $firstLine" -ForegroundColor Yellow @@ -1625,11 +1780,10 @@ shell.Run cmd, 0, False } catch {} } # ── Arch resolution: env-var override → name inference ────────────── - # Runs even when the hipinfo/amd-smi probe could NOT confirm a runtime - # ($HasROCm false): the gfx arch inferred from the WMI GPU name lets the - # studio setup forward --rocm-gfx and pull a GPU-accelerated ROCm - # llama.cpp, which bundles its own ROCm runtime. PyTorch's ROCm wheels - # still require a confirmed HIP SDK -- they stay gated on $HasROCm below. + # Runs even when the probe can't confirm a runtime ($HasROCm false): the + # WMI-name gfx arch drives both ROCm llama.cpp and torch. repo.amd.com + # wheels bundle their own runtime (no HIP SDK), so a mapped arch installs + # ROCm torch directly below -- no wasted CPU base. if (-not $ROCmGfxArch) { # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. if ($env:UNSLOTH_ROCM_GFX_ARCH) { @@ -1809,6 +1963,64 @@ shell.Run cmd, 0, False substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow" return "$baseUrl/cu126" } + + # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── + # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, + # matching setup.ps1's stale-venv parse. + function ConvertTo-TorchFlavorTag { + param([string]$TorchVersion) + if (-not $TorchVersion) { return $null } + if ($TorchVersion -match '\+(cu\d+)') { return $Matches[1] } + if ($TorchVersion -match '\+rocm') { return 'rocm' } + if ($TorchVersion -match '\+cpu') { return 'cpu' } + return 'cpu' + } + + # Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a + # gfx* leaf -> rocm). $null on an unknown leaf (odd mirror) so repair no-ops. + function Get-ExpectedTorchFlavorTag { + param([string]$TorchIndexUrl, [string]$ROCmIndexUrl) + if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' } + if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null } + $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if ($leaf -match '^cu\d+$') { return $leaf } + if ($leaf -eq 'cpu') { return 'cpu' } + if ($leaf -match '^rocm') { return 'rocm' } + if ($leaf -match '^gfx') { return 'rocm' } + return $null + } + + # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses + # ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference. + function Get-InstalledTorchTag { + param([string]$PythonExe) + if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $PythonExe + $psi.Arguments = '-c "import torch; print(torch.__version__)"' + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.UseShellExecute = $false + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + # Drain BOTH streams async, then WaitForExit. A synchronous ReadToEnd() + # before the wait would block forever if a wedged "import torch" never + # closes stdout; leaving the redirected stderr undrained would deadlock a + # child that floods it past the pipe buffer. Async reads let a noisy-but- + # exiting probe finish, while a truly hung one still hits the 30s timeout + # and is killed -- bounded either way. + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + $finished = $proc.WaitForExit(30000) + if (-not $finished) { try { $proc.Kill() } catch {}; return $null } + $torchVer = $outTask.GetAwaiter().GetResult().Trim() + [void]$errTask.GetAwaiter().GetResult() + if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null } + return ConvertTo-TorchFlavorTag $torchVer + } catch { return $null } + } + $TorchIndexUrl = Get-TorchIndexUrl # ── GPU arch → newest compatible Windows ROCm wheel release ── @@ -1820,7 +2032,7 @@ shell.Run cmd, 0, False # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null - if ($HasROCm -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { + if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 @@ -1843,6 +2055,17 @@ shell.Run cmd, 0, False "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" } + # Companion ranges track the torch ceiling so pip resolves a consistent + # trio on AMD's per-arch index (each published independently). Mirrors + # setup.ps1 / install_python_stack.py; bump all three together for 2.12.x. + $torchvisionFloorMap = @{ + "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" + "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" + } + $torchaudioFloorMap = @{ + "gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0" + "gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0" + } $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } if ($archFamily) { $ROCmIndexUrl = "$amdIndexBase/$archFamily/" @@ -1871,10 +2094,10 @@ shell.Run cmd, 0, False if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") { Write-Host "" if ($ROCmGfxArch) { - # Known AMD arch: install.ps1 lays down CPU PyTorch as a base, then - # setup.ps1 swaps in AMD's bundled-runtime GPU ROCm wheels (no HIP SDK). - substep "Installing CPU PyTorch as a base -- Studio setup installs GPU ROCm" "Cyan" - substep "wheels for $ROCmGfxArch next (bundled runtime; HIP SDK not required)." "Cyan" + # Only an unmapped arch reaches here (a mapped one set $ROCmIndexUrl + # above). No ROCm torch wheels for this arch (e.g. RDNA2 gfx103X) -> CPU. + substep "Installing CPU PyTorch -- no ROCm PyTorch wheels are available for $ROCmGfxArch." "Yellow" + substep "PyTorch (training and Transformers inference) runs on CPU on this GPU." "Yellow" } else { if ($HipSdkInstalled -and -not $HasROCm) { substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow" @@ -1928,21 +2151,21 @@ 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.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below # is --no-deps). All transitive deps are torch-free. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1956,7 +2179,7 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } substep "overlaying unsloth-zoo from git main..." - $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } + $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } if ($zooOverlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) @@ -1969,15 +2192,34 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" substep "installing PyTorch from $ROCmIndexUrl..." $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } - $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio } + # Pin the companions to match $torchSpec; bare names can resolve an + # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. + $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { - Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red - return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit) + # Transient AMD-index failure: fall back to a CPU base so the install + # still completes; Studio setup retries ROCm afterwards. + substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Studio setup retries ROCm." "Yellow" + # --force-reinstall: a failed ROCm install can leave an unpinned ROCm + # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU + # torch>= range, so without it uv would keep the ROCm build and only swap + # the companions -- a mismatched venv the flavor-repair block won't fix. + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + if ($torchInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) + } + # CPU base is in; drop the ROCm expectation so the flavor-repair + # block below won't retry the just-failed index and abort. setup.ps1 + # reinstalls ROCm afterwards (recomputes its own index URL). + $ROCmIndexUrl = $null + $ROCmTorchFloor = $null } } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -1989,21 +2231,21 @@ 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.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2018,7 +2260,7 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } substep "overlaying unsloth-zoo from git main..." - $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } + $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } if ($zooOverlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) @@ -2029,7 +2271,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.6.7" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --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) @@ -2041,13 +2283,13 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } substep "overlaying unsloth-zoo from git main..." - $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } + $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } if ($zooOverlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) @@ -2055,6 +2297,54 @@ shell.Run cmd, 0, False } } + # ── Enforce the installed torch flavor matches the detected GPU build ── + # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv + # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on + # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is + # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* + # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install + # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. + if (-not $SkipTorch) { + $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl + if ($expectedTorchTag -and $expectedTorchTag -ne 'cpu') { + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { + if ($expectedTorchTag -eq 'rocm' -and $ROCmIndexUrl) { + # AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path + # would have force-reinstalled. Repair from the same repo.amd.com index. + $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } + # Pin companions like the fresh ROCm path (bare names can pull an + # ABI-incompatible torchvision/torchaudio from the per-arch index). + $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + if ($torchFixExit -ne 0) { + Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) + } + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + } elseif ($expectedTorchTag -ne 'rocm') { + # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. + substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + if ($torchFixExit -ne 0) { + Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) + } + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + } + } + # Safety net (incl. AMD): GPU build expected but still CPU -> warn loudly. + if ($installedTorchTag -eq 'cpu') { + Write-Host "" + Write-Host " [WARN] PyTorch is CPU-only but a $expectedTorchTag GPU build was expected for this machine." -ForegroundColor Yellow + Write-Host " [WARN] Training and GPU inference will run on CPU until this is fixed." -ForegroundColor Yellow + Write-Host " [WARN] Re-run this installer, or reinstall the GPU build manually for your GPU." -ForegroundColor Yellow + } + } + } + # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped # for --local: the editable install above already makes _PACKAGE_ROOT in # unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__). @@ -2102,7 +2392,9 @@ shell.Run cmd, 0, False # ── Run studio setup ── # setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools, - # CUDA Toolkit, Node.js, and other dependencies automatically via winget. + # CUDA Toolkit, and other dependencies automatically via winget. Node.js is + # NOT installed via winget -- setup.ps1 uses an isolated Node it manages and + # never touches the system Node/npm. Write-TauriLog "STEP" "Running studio setup" step "setup" "running unsloth studio setup..." $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" @@ -2308,6 +2600,7 @@ shell.Run cmd, 0, False step "launch" "to start later, run:" substep "unsloth studio -p 8888" substep "(add -H 0.0.0.0 to allow network / cloud access)" + substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)" Write-Host "" } } else { @@ -2328,6 +2621,7 @@ shell.Run cmd, 0, False substep "unsloth studio -p 8888" } substep "(add -H 0.0.0.0 to allow network / cloud access)" + substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)" Write-Host "" } } diff --git a/install.sh b/install.sh index 52324d1ef5..318b03e054 100755 --- a/install.sh +++ b/install.sh @@ -163,6 +163,40 @@ run_install_cmd() { return $_rc } +# Retry run_install_cmd on transient uv download failures with backoff. Returns +# the last exit code on permanent failure so the set -e rollback trap still fires. +: "${UNSLOTH_INSTALL_RETRIES:=3}" +: "${UNSLOTH_INSTALL_RETRY_DELAY:=3}" +run_install_cmd_retry() { + _ricr_label="$1" + # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables). + # Length guard precedes the numeric test so a huge value can't overflow `[ -ge ]`. + # 0?* rejects leading-zero delays ("08"/"09" break the later $((delay*2)) as octal); + # bare "0" stays valid. Bounds: 1..100 retries, 0..3600s base delay. + case "$UNSLOTH_INSTALL_RETRIES" in + ''|*[!0-9]*|0) _ricr_max=3 ;; + *) if [ "${#UNSLOTH_INSTALL_RETRIES}" -le 3 ] && [ "$UNSLOTH_INSTALL_RETRIES" -ge 1 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRIES" -le 100 ] 2>/dev/null; then _ricr_max=$UNSLOTH_INSTALL_RETRIES; else _ricr_max=3; fi ;; + esac + case "$UNSLOTH_INSTALL_RETRY_DELAY" in + ''|*[!0-9]*|0?*) _ricr_delay=3 ;; + *) if [ "${#UNSLOTH_INSTALL_RETRY_DELAY}" -le 4 ] && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -ge 0 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -le 3600 ] 2>/dev/null; then _ricr_delay=$UNSLOTH_INSTALL_RETRY_DELAY; else _ricr_delay=3; fi ;; + esac + _ricr_attempt=1 + while :; do + # AND-OR (not `if`) preserves the real failure code: $? after a non-taken + # `if` is 0 in sh/dash/bash, which would break the rollback path. + run_install_cmd "$@" && return 0 + _ricr_rc=$? + if [ "$_ricr_attempt" -ge "$_ricr_max" ]; then + return "$_ricr_rc" + fi + substep "retrying \"$_ricr_label\" after transient failure (attempt $((_ricr_attempt + 1))/$_ricr_max, waiting ${_ricr_delay}s)..." "$C_WARN" + sleep "$_ricr_delay" || true + _ricr_attempt=$((_ricr_attempt + 1)) + _ricr_delay=$((_ricr_delay * 2)) + done +} + # Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main # wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2 # NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the @@ -413,8 +447,12 @@ _on_install_exit() { if [ "$_status" -ne 0 ]; then _restore_studio_venv_replacement fi + [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true exit "$_status" } +# Empty so an inherited value can never reach the trap's rm; only a temp dir +# this script creates below (Apple Silicon, spaced path) is ever removed. +_UV_OVERRIDE_TMPDIR="" trap _on_install_exit EXIT # ── Helper: download a URL to a file (supports curl and wget) ── @@ -1227,6 +1265,10 @@ if (-not \$targetExe) { exit 1 } # native install if one exists) so the WSL shortcut shows the proper icon. \$iconDir = Join-Path \$env:LOCALAPPDATA 'Unsloth Studio' \$iconPath = Join-Path \$iconDir 'unsloth.ico' +\$preIconHash = \$null +if (Test-Path -LiteralPath \$iconPath) { + try { \$preIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash } catch {} +} if (-not (Test-Path -LiteralPath \$iconPath)) { try { New-Item -ItemType Directory -Force -Path \$iconDir | Out-Null @@ -1242,9 +1284,11 @@ if (Test-Path -LiteralPath \$iconPath) { (Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs') ) \$created = @() +\$firstShortcut = \$false foreach (\$dir in \$locations) { if (-not \$dir -or -not (Test-Path \$dir)) { continue } \$linkPath = Join-Path \$dir '$_css_lnk_name_ps' + if (-not (Test-Path -LiteralPath \$linkPath)) { \$firstShortcut = \$true } \$shortcut = \$WshShell.CreateShortcut(\$linkPath) \$shortcut.TargetPath = \$targetExe \$shortcut.Arguments = '$_css_sc_args_ps' @@ -1253,27 +1297,43 @@ foreach (\$dir in \$locations) { \$shortcut.Save() \$created += \$linkPath } -# Force Explorer to re-read EACH new shortcut's icon so it renders immediately -# instead of a stale/blank (generic) icon. The reliable, NON-disruptive fix -# (no explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM, -# SHCNF_PATHW, ) -- the global SHCNE_ASSOCCHANGED alone does not recover a -# stale item. Also clear the on-disk icon cache for heavier staleness. -try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {} -try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {} +\$iconChanged = \$false +if (\$hasIcon) { + if (-not \$preIconHash) { + \$iconChanged = \$true + } else { + try { + \$postIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash + \$iconChanged = (\$postIconHash -ne \$preIconHash) + } catch { \$iconChanged = \$true } + } +} elseif (\$preIconHash) { + \$iconChanged = \$true +} +# Per-item refresh always (cheap, non-disruptive) so the rewritten .lnk renders +# immediately instead of a stale/blank (generic) icon. The reliable fix (no +# explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM, SHCNF_PATHW, +# ) -- the global SHCNE_ASSOCCHANGED alone does not recover a stale item. try { Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int e, uint f, string a, System.IntPtr b);' -ErrorAction SilentlyContinue foreach (\$p in \$created) { try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, \$p, [System.IntPtr]::Zero) } catch {} } [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, \$null, [System.IntPtr]::Zero) } catch {} -# Win11 Start Menu keeps its own tile-icon cache (preserve start2.bin). -try { - \$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState' - if (Test-Path -LiteralPath \$smeh) { - Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue - Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue - Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue - } -} catch {} +# Heavier on-disk icon-cache clear + StartMenuExperienceHost tile rebuild +# (preserve start2.bin) only on first install or a real icon change, so a no-op +# WSL reinstall does not run a dropper-like clear-cache + kill cluster each time. +if (\$created.Count -gt 0 -and (\$firstShortcut -or \$iconChanged)) { + try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {} + try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {} + try { + \$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState' + if (Test-Path -LiteralPath \$smeh) { + Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue + Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue + } + } catch {} +} WSLPS1_EOF # Convert WSL path to Windows path for powershell.exe @@ -1371,6 +1431,25 @@ fi if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then _OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt" if [ -f "$_OVERRIDES_FILE" ]; then + # uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace + # truncates it and aborts every later uv call (issue #6503). Hand uv a copy. + case "$_OVERRIDES_FILE" in + *[[:space:]]*) + _UV_OVERRIDE_TMPDIR=$(mktemp -d 2>/dev/null) || _UV_OVERRIDE_TMPDIR="" + case "$_UV_OVERRIDE_TMPDIR" in + "") ;; + *[[:space:]]*) rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true; _UV_OVERRIDE_TMPDIR="" ;; + *) + if cp "$_OVERRIDES_FILE" "$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt" 2>/dev/null; then + _OVERRIDES_FILE="$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt" + else + rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + _UV_OVERRIDE_TMPDIR="" + fi + ;; + esac + ;; + esac export UV_OVERRIDE="$_OVERRIDES_FILE" fi fi @@ -1383,18 +1462,110 @@ elif [ "$OS" = "macos" ]; then fi tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none" -# ── Check system dependencies ── -# cmake and git are needed by unsloth studio setup to build the GGUF inference -# engine (llama.cpp). build-essential and libcurl-dev are also needed on Linux. -tauri_log "STEP" "Checking system dependencies" -MISSING="" +# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04) +# with a 24.04 distro present, re-run the install there and stop; else fall through +# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before +# the STUDIO_HOME mkdir/venv so the origin distro is untouched. +_maybe_reroute_strixhalo_to_2404() { + [ "${OS:-}" = "wsl" ] || return 0 + [ "${SKIP_TORCH:-false}" = "false" ] || return 0 + [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 + [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 + [ -e /dev/dxg ] || return 0 + grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # Already ROCm-on-WSL? leave a working GPU alone, whatever the version. + if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then + return 0 + fi + _rr_ver="" + [ -r /etc/os-release ] && _rr_ver=$(. /etc/os-release 2>/dev/null; printf '%s' "${VERSION_ID:-}") + # The bootstrap (scripts/install_rocm_wsl_strixhalo.sh) dies on any VERSION_ID but + # 24.04 and pins the noble repo, so 24.04 is the sole GPU-supported target; leave a + # 24.04 user alone. (Working ROCm on other versions was caught by librocdxg above.) + case "$_rr_ver" in 24.04) return 0 ;; esac + # Distro is now unsupported. If we can't reroute to a 24.04 target, stay CPU-only + # AND skip the later origin-distro ROCm bootstrap (it ignores distro version, so it + # would otherwise install ROCm into 26.04 etc.). + command -v wsl.exe >/dev/null 2>&1 || { UNSLOTH_SKIP_ROCM_WSL_SETUP=1; return 0; } + # Route only to an installed Ubuntu-24.04 (bootstrap's only target). Match the whole + # line (one distro per line from wsl.exe -l -q), not a substring, so "Ubuntu-24.04-test" + # can't masquerade as it and then fail `wsl -d`. + # || true: no match is expected, not an error (script runs under set -e). + _rr_distros=$(wsl.exe -l -q 2>/dev/null | tr -d '\000\r') + _rr_target=$(printf '%s\n' "$_rr_distros" | grep -ixF "Ubuntu-24.04" | head -n1) || true + [ -n "$_rr_target" ] || { + substep "ROCm-on-WSL (GPU) needs Ubuntu 24.04; this distro is Ubuntu ${_rr_ver:-unknown}." "$C_WARN" + substep "No Ubuntu-24.04 WSL distro found; staying CPU-only. Install Ubuntu-24.04 and re-run there for GPU." "$C_WARN" + UNSLOTH_SKIP_ROCM_WSL_SETUP=1 + return 0 + } -command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake" -command -v git >/dev/null 2>&1 || MISSING="$MISSING git" + echo "" + substep "ROCm-on-WSL (GPU) needs Ubuntu 24.04; this distro is Ubuntu ${_rr_ver:-unknown}." "$C_WARN" + substep "Found an existing $_rr_target distro -- continuing the GPU install there." "$C_OK" + # A --local checkout can't be replayed via curl|sh (the repo isn't in the target + # distro), so tell the user to re-run there rather than silently run a different install. + if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + substep "This is a --local install; re-run it from $_rr_target instead:" "$C_WARN" + substep " wsl -d $_rr_target -- bash -lc 'cd && ./install.sh --local'" "$C_WARN" + substep "Continuing CPU-only in Ubuntu ${_rr_ver:-this distro} for now." "$C_WARN" + # Unsupported distro, can't reroute a --local checkout: skip the origin ROCm bootstrap. + UNSLOTH_SKIP_ROCM_WSL_SETUP=1 + return 0 + fi + # Forward the caller's options/env (custom package/python/home) so the rerouted + # install matches what was asked for, not a default install. + _rr_q() { printf "'%s'" "$(printf '%s' "$1" | sed "s/'/'\\\\''/g")"; } + _rr_exports="set -o pipefail; export UNSLOTH_WSL_REROUTED=1" + [ "$_STUDIO_HOME_REDIRECT" = "env" ] && _rr_exports="$_rr_exports; export UNSLOTH_STUDIO_HOME=$(_rr_q "$STUDIO_HOME")" + # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the + # GPU instead of falling back to the desktop-app prompt path. + [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1" + _rr_args="" + [ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")" + [ -n "$_USER_PYTHON" ] && _rr_args="$_rr_args --python $(_rr_q "$_USER_PYTHON")" + [ "$_VERBOSE" = true ] && _rr_args="$_rr_args --verbose" + [ "$TAURI_MODE" = true ] && _rr_args="$_rr_args --tauri" + if [ -n "${UNSLOTH_WSL_REROUTE_CMD:-}" ]; then + _rr_cmd="$UNSLOTH_WSL_REROUTE_CMD" # user took full control + elif [ -n "$_rr_args" ]; then + _rr_cmd="curl -fsSL https://unsloth.ai/install.sh | sh -s --$_rr_args" + else + _rr_cmd="curl -fsSL https://unsloth.ai/install.sh | sh" + fi + # pipefail so a failed curl in `curl | sh` isn't masked by sh exiting 0 on empty + # input (which would wrongly report success and exit 0 the parent installer). + _rr_rc=0 + wsl.exe -d "$_rr_target" -- bash -lc "$_rr_exports; $_rr_cmd" || _rr_rc=$? + if [ "$_rr_rc" -eq 0 ]; then + exit 0 + fi + # In Tauri mode the child uses exit 2 ([TAURI:NEED_SUDO]) to ask the desktop app to + # elevate for the target distro; the child already printed the NEED_SUDO line, so + # propagate the code instead of masking it as a reroute failure and dropping to CPU. + if [ "$TAURI_MODE" = true ] && [ "$_rr_rc" -eq 2 ]; then + exit 2 + fi + substep "Could not auto-continue in $_rr_target; run it yourself:" "$C_WARN" + substep " wsl -d $_rr_target -- bash -lc 'curl -fsSL https://unsloth.ai/install.sh | sh'" + substep "Continuing CPU-only in Ubuntu ${_rr_ver:-this distro} for now." "$C_WARN" + # Reroute failed; don't let the later bootstrap install ROCm into this unsupported + # distro -- stay CPU-only. + UNSLOTH_SKIP_ROCM_WSL_SETUP=1 + return 0 +} +_maybe_reroute_strixhalo_to_2404 || true + +# ── Check system dependencies ── +# cmake/git are only needed to *build* llama.cpp from source. Studio downloads a +# prebuilt by default, and setup.sh self-skips the source build when they're +# absent -- so macOS doesn't block on cmake (requiring it would force a manual +# Homebrew install). Linux keeps requiring them; its package manager has them. +tauri_log "STEP" "Checking system dependencies" case "$OS" in macos) - # Xcode Command Line Tools provide the C/C++ compiler + # Xcode Command Line Tools provide the C/C++ compiler and git. if ! xcode-select -p >/dev/null 2>&1; then echo "" echo "==> Xcode Command Line Tools are required." @@ -1403,8 +1574,19 @@ case "$OS" in echo " After the installation completes, please re-run this script." exit 1 fi + # cmake is only needed for a source build; the default prebuilt path + # doesn't use it, so its absence is not fatal -- no Homebrew prerequisite. + if command -v cmake >/dev/null 2>&1; then + step "deps" "all system dependencies found" + else + step "deps" "using prebuilt llama.cpp (cmake not found)" "$C_WARN" + substep "Install cmake only if you want a source build: brew install cmake" + fi ;; linux|wsl) + MISSING="" + command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake" + command -v git >/dev/null 2>&1 || MISSING="$MISSING git" # curl or wget is needed for downloads; check both if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then MISSING="$MISSING curl" @@ -1412,27 +1594,12 @@ case "$OS" in command -v gcc >/dev/null 2>&1 || MISSING="$MISSING build-essential" # libcurl dev headers for llama.cpp HTTPS support command -v curl-config >/dev/null 2>&1 || MISSING="$MISSING libcurl4-openssl-dev" - ;; -esac -MISSING=$(echo "$MISSING" | sed 's/^ *//') - -if [ -n "$MISSING" ]; then - echo "" - step "deps" "missing: $MISSING" "$C_WARN" - substep "These are needed to build the GGUF inference engine." - - case "$OS" in - macos) - if ! command -v brew >/dev/null 2>&1; then - echo "" - echo " Homebrew is required to install them." - echo " Install Homebrew from https://brew.sh then re-run this script." - exit 1 - fi - brew install $MISSING /dev/null 2>&1; then _smart_apt_install $MISSING else @@ -1447,21 +1614,28 @@ if [ -n "$MISSING" ]; then echo " openSUSE: sudo zypper install cmake git gcc gcc-c++ make libcurl-devel" exit 1 fi - ;; - esac - echo "" -else - step "deps" "all system dependencies found" -fi + echo "" + else + step "deps" "all system dependencies found" + fi + ;; +esac # ── Install uv ── tauri_log "STEP" "Installing uv package manager" -UV_MIN_VERSION="0.7.22" +UV_MIN_VERSION="0.8.16" # When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables). : "${UV_COMPILE_BYTECODE_TIMEOUT:=180}" export UV_COMPILE_BYTECODE_TIMEOUT +# uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read +# timeout for large wheel downloads. ":=" preserves any user override. +: "${UV_HTTP_RETRIES:=5}" +export UV_HTTP_RETRIES +: "${UV_HTTP_TIMEOUT:=180}" +export UV_HTTP_TIMEOUT + version_ge() { # returns 0 if $1 >= $2 _a=$1 @@ -1937,6 +2111,45 @@ get_torch_index_url() { else echo "$_base/cpu"; fi } +# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── +# torch.__version__ ($1) -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu. +_torch_flavor_tag() { + case "$1" in + *+cu[0-9]*) printf '%s\n' "$1" | sed -n 's/.*+\(cu[0-9][0-9]*\).*/\1/p' ;; + *+rocm*) echo "rocm" ;; + *+cpu*) echo "cpu" ;; + "") echo "" ;; + *) echo "cpu" ;; + esac +} + +# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> +# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. +_expected_torch_flavor_tag() { + _u="${1%/}" + _leaf="${_u##*/}" + case "$_leaf" in + cu[0-9]*) echo "$_leaf" ;; + cpu) echo "cpu" ;; + rocm*|gfx*) echo "rocm" ;; + *) echo "" ;; + esac +} + +# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv +# resolves (torch + every transitive dep) via --index-url -- the same URLs the +# fresh-install paths above already use -- so a stale wheel is auto-repairable. +# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. +_torch_index_repairable() { + _u="${1%/}" + _leaf="${_u##*/}" + case "$_leaf" in + cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; + *) echo "no" ;; + esac +} + get_radeon_wheel_url() { # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/, @@ -2368,6 +2581,9 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then substep "ROCm: $_rocm_root" [ -n "$_gpu_rocm_ver" ] && substep "hipconfig: $_gpu_rocm_ver" [ -n "$_gpu_disp_mkt" ] && [ -n "$_gpu_disp_gfx" ] && substep "GPU: $_gpu_disp_mkt" +elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + # Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only. + step "gpu" "Apple Silicon (Metal, unified memory)" else step "gpu" "none (CPU-only)" "$C_WARN" fi @@ -2430,28 +2646,28 @@ if [ "$_MIGRATED" = true ]; then # PyPI metadata still declares torch as a hard dep), then install # runtime deps (typer, safetensors, transformers, etc.) with --no-deps # to prevent transitive torch resolution. - run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ + run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.7" unsloth-zoo + "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. - run_install_cmd "install pydantic (with deps for compatible core)" \ + run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _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" + run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi else - run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.7" unsloth-zoo + "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then 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..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" fi @@ -2460,13 +2676,13 @@ if [ "$_MIGRATED" = true ]; then # fresh reinstall. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" # Repair ROCm torch if overwritten during migrated install _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." - run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" \ --force-reinstall @@ -2592,7 +2808,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \ [ "$_radeon_versions_match" != true ]; then substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" else @@ -2604,30 +2820,30 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # filelock / sympy / networkx which are not in the # Radeon listing. if [ -n "$_tri_whl" ]; then - run_install_cmd "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ --find-links "$_RADEON_BASE_URL" \ "$_tri_whl" "$_torch_whl" "$_tv_whl" "$_ta_whl" else - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ --find-links "$_RADEON_BASE_URL" \ "$_torch_whl" "$_tv_whl" "$_ta_whl" fi fi else substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). @@ -2636,7 +2852,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # which is only useful once torch is present for training. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" ;; esac @@ -2647,46 +2863,46 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, 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 \ + run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.7" unsloth-zoo + "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" # Same pydantic-with-deps trick as the migrated branch. - run_install_cmd "install pydantic (with deps for compatible core)" \ + run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _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" + run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then 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..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo + run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ + --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" 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..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else - run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ --upgrade-package unsloth -- "$PACKAGE_NAME" fi # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." - run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" \ --force-reinstall @@ -2699,15 +2915,50 @@ 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.6.7" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --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..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME" + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME" + fi +fi + +# ── Enforce the installed torch flavor matches the detected GPU build ── +# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv +# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on +# CPU. Reinstall the right wheel triplet when a GPU build is expected; if it +# can't be reinstalled, warn loudly. --no-torch / CPU-only / macOS: no-op. +if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then + _expected_torch_tag=$(_expected_torch_flavor_tag "$TORCH_INDEX_URL") + # Only act when a GPU build is expected (cuXXX / rocm); cpu and unknown skip. + if [ -n "$_expected_torch_tag" ] && [ "$_expected_torch_tag" != "cpu" ]; then + _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) + _installed_torch_tag="" + [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") + # Repair when flavor is wrong AND the index is plain --index-url reinstallable + # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. + if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ + && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then + substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." + run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" \ + --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio + _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) + _installed_torch_tag="" + [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") + fi + # Safety net (incl. AMD/WSL): GPU build expected but still CPU -> warn loudly. + if [ "$_installed_torch_tag" = "cpu" ]; then + substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" + substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" + substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + fi fi fi @@ -2902,10 +3153,11 @@ echo "" if [ -t 1 ]; then echo "" printf " Start Unsloth Studio now? [Y/n] " + # No readable answer (closed/EOF tty) defaults to no; Enter is still yes. if [ -r /dev/tty ]; then - read -r _reply =2026.6.5", + "unsloth_zoo>=2026.6.7", "wheel>=0.42.0", "packaging", "numpy", @@ -92,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.7", "torchvision", "unsloth[triton]", ] @@ -582,7 +584,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.7", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index fe90afa7e6..c1d156d40a 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -412,7 +412,7 @@ BLOCKED_NPM_VERSIONS: dict[str, set[str]] = { "@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 + # 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"}, @@ -916,6 +916,204 @@ def _evidence( LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare") +# ───────────────────────────────────────────────────────────────────── +# Code-only scanning for JS/TS sources. Blank `//` and `/* */` comments +# before matching (the top FP source: scary strings in JSDoc/changelog +# comments), tracking string/template/regex context so a `//` inside +# "http://..." is not mistaken for a comment. Strings are NOT blanked +# (droppers hide payloads there). Fail open on lexer confusion: the raw +# text is still scanned. JS sibling of scan_packages.py::_strip_noncode. +# ───────────────────────────────────────────────────────────────────── +_JS_FAMILY_SUFFIXES = (".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx") + +# Keywords after which a `/` begins a regex literal (not division). +_REGEX_PRECEDING_KEYWORDS = frozenset( + { + "return", + "typeof", + "instanceof", + "in", + "of", + "new", + "delete", + "void", + "throw", + "yield", + "await", + "do", + "else", + "case", + } +) +_IDENT_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$") + + +def _slash_is_regex(prev_tok: str) -> bool: + """Disambiguate a lone ``/``: regex literal vs division operator. + + Biased toward regex when ambiguous -- regex state never blanks, so a + wrong guess only costs FP reduction (or a fail-open), never a missed + detection. + """ + if prev_tok == "": + return True # start of file -> expression position + if prev_tok in _REGEX_PRECEDING_KEYWORDS: + return True + last = prev_tok[-1] + if last.isalnum() or last in "_$)]": + return False # previous token ends a value -> division + return True # operators, punctuation, `{`, `}` -> regex (safe bias) + + +def _strip_js_noncode(text: str) -> str: + """Blank JS/TS comments, preserving byte geometry. Fail-open on confusion.""" + if "//" not in text and "/*" not in text: + return text # nothing to strip + n = len(text) + out = list(text) + nl = ("\n", "\r") + + def _blank(a: int, b: int) -> None: + for k in range(a, b): + if out[k] not in nl: + out[k] = " " + + state = "code" + prev_tok = "" + tmpl_stack: list[str] = [] + i = 0 + try: + while i < n: + c = text[i] + nxt = text[i + 1] if i + 1 < n else "" + if state == "code": + if c == "/" and nxt == "/": + start = i + i += 2 + while i < n and text[i] not in nl: + i += 1 + _blank(start, i) + continue + if c == "/" and nxt == "*": + start = i + i += 2 + closed = False + while i < n: + if text[i] == "*" and i + 1 < n and text[i + 1] == "/": + i += 2 + closed = True + break + i += 1 + if not closed: + return text # unterminated block comment + _blank(start, i) + continue + if c == "'": + state = "sq" + i += 1 + continue + if c == '"': + state = "dq" + i += 1 + continue + if c == "`": + state = "tmpl" + i += 1 + continue + if c == "/": + if _slash_is_regex(prev_tok): + state = "regex" + i += 1 + continue + prev_tok = "/" + i += 1 + continue + if c.isspace(): + i += 1 + continue + if c in _IDENT_CHARS: + j = i + while j < n and text[j] in _IDENT_CHARS: + j += 1 + prev_tok = text[i:j] + i = j + continue + if c == "}" and tmpl_stack: + state = tmpl_stack.pop() + i += 1 + continue + prev_tok = c + i += 1 + continue + elif state in ("sq", "dq"): + q = "'" if state == "sq" else '"' + if c == "\\": + i += 2 + continue + if c == q: + state = "code" + prev_tok = "_v" + i += 1 + continue + if c in nl: + return text # unterminated string literal + i += 1 + continue + elif state == "tmpl": + if c == "\\": + i += 2 + continue + if c == "`": + state = "code" + prev_tok = "_v" + i += 1 + continue + if c == "$" and nxt == "{": + tmpl_stack.append("tmpl") + state = "code" + prev_tok = "{" + i += 2 + continue + i += 1 + continue + elif state == "regex": + if c == "\\": + i += 2 + continue + if c == "[": + state = "regex_cc" + i += 1 + continue + if c == "/": + state = "code" + prev_tok = "_v" + i += 1 + continue + if c in nl: + return text # unterminated regex literal + i += 1 + continue + elif state == "regex_cc": + if c == "\\": + i += 2 + continue + if c == "]": + state = "regex" + i += 1 + continue + if c in nl: + return text + i += 1 + continue + else: + return text + if state != "code" or tmpl_stack: + return text # unterminated construct -> fail open + except Exception: + return text + return "".join(out) + + def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: findings: list[Finding] = [] try: @@ -1042,6 +1240,14 @@ def _host_in_outbound_context(text: str, host: str) -> bool: def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: findings: list[Finding] = [] + # Code-only scanning for JS/TS sources: blank comments before matching so + # an IOC host / `eval(atob)` example / campaign marker quoted in a comment + # cannot manufacture a false positive. Assigned string literals (where real + # droppers hide base64 payloads) are preserved. Non-JS text (json/yaml/sh/ + # py/html) is scanned as-is -- this lexer only understands JS comments. + if rel.lower().endswith(_JS_FAMILY_SUFFIXES): + text = _strip_js_noncode(text) + # IOC substrings (literal, case-sensitive). for needle, (sev, why) in KNOWN_IOC_STRINGS.items(): if needle in text: @@ -1236,6 +1442,131 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N pass +# ───────────────────────────────────────────────────────────────────── +# Baseline allowlist: triaged known-good HIGH/CRITICAL findings so the gate +# can enforce without red-failing on rare legitimate-library behavior. +# Matched on ``(normalized package, package-relative path, pattern)`` -- not +# evidence text -- so a version bump does not reopen a finding, but a *new* +# kind of finding in a listed file is a different pattern and still fails. +# Mirrors scan_packages.py. Regenerate with ``--write-baseline``. +# ───────────────────────────────────────────────────────────────────── + +_DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json") + +# Bumped when the entry-key semantics change. v2 keys on the package-relative +# path; v1 stored only a basename, so a v1 entry could suppress a same-named file +# in a different directory. A pre-v2 baseline with entries is ignored (fail +# closed) rather than mis-applied. +_BASELINE_SCHEMA_VERSION = 2 + + +def _norm_pkg_name(display: str) -> str: + """``@scope/pkg@1.2.3`` / ``pkg@1.2.3`` -> name without the version. + + The version is the LAST ``@``-separated field; a leading ``@`` (scope) + is preserved. Lower-cased (npm names are case-insensitive). Sentinels + like ```` / ```` pass through unchanged. + """ + s = (display or "").strip() + at = s.rfind("@") + if at > 0: # >0 so a leading @scope is not treated as the version sep + s = s[:at] + return s.lower() + + +_NPM_TARBALL_ROOT = "package/" + + +def _relpath_in_package(filename: str) -> str: + """Path within the published package, stable across version bumps. npm + tarballs root every file at ``package/``; strip it so the key is the real + source path (``dist/index.js``) and a new file with the same basename in a + different directory is not silently suppressed.""" + f = (filename or "").replace("\\", "/") + return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f + + +def _finding_key(f: Finding) -> tuple[str, str, str]: + """Stable allowlist key: normalized package, package-relative path, pattern.""" + return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern) + + +def _load_baseline(path: str) -> set[tuple[str, str, str]]: + """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" + try: + with open(path, "r", encoding = "utf-8") as fh: + data = json.load(fh) + except FileNotFoundError: + return set() + except (OSError, json.JSONDecodeError) as exc: + print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) + return set() + entries = data.get("entries", []) + if entries and data.get("version") != _BASELINE_SCHEMA_VERSION: + print( + f" [WARN] baseline schema v{data.get('version')} predates package-relative " + f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.", + file = sys.stderr, + ) + return set() + keys: set[tuple[str, str, str]] = set() + for e in entries: + try: + keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"])) + except (KeyError, TypeError): + continue + return keys + + +def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int: + """Persist at-or-above-threshold findings as an allowlist for triage.""" + entries = [] + seen: set[tuple[str, str, str]] = set() + for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)): + if _SEVERITY_RANK[f.severity] > threshold_rank: + continue + key = _finding_key(f) + if key in seen: + continue + seen.add(key) + entries.append( + { + "package": _norm_pkg_name(f.package), + "file": _relpath_in_package(f.filename), + "pattern": f.pattern, + "severity": f.severity, + "evidence": (f.evidence or f.detail)[:240], + } + ) + doc = { + "_comment": ( + "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL " + "finding manually judged benign. Matched on (package, " + "package-relative path, pattern); evidence/severity are for review " + "only. Regenerate with --write-baseline AFTER reviewing every line." + ), + "version": _BASELINE_SCHEMA_VERSION, + "entries": entries, + } + with open(path, "w", encoding = "utf-8") as fh: + json.dump(doc, fh, indent = 2, sort_keys = False) + fh.write("\n") + print(f" Wrote {len(entries)} baseline entr(y/ies) to {path}") + return len(entries) + + +def _partition_baseline( + findings: list[Finding], baseline: set[tuple[str, str, str]] +) -> tuple[list[Finding], list[Finding]]: + """Split findings into (active, suppressed) by allowlist membership.""" + if not baseline: + return list(findings), [] + active, suppressed = [], [] + for f in findings: + (suppressed if _finding_key(f) in baseline else active).append(f) + return active, suppressed + + def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description = "Pre-install npm tarball content scanner.", @@ -1263,6 +1594,30 @@ def main(argv: list[str] | None = None) -> int: "Medium and below print but exit 0." ), ) + parser.add_argument( + "--baseline", + metavar = "FILE", + default = None, + help = ( + "Allowlist JSON of triaged known-good findings to suppress. " + "Defaults to scan_npm_packages_baseline.json next to this script " + "if present." + ), + ) + parser.add_argument( + "--no-baseline", + action = "store_true", + help = "Ignore the auto-discovered baseline allowlist.", + ) + parser.add_argument( + "--write-baseline", + metavar = "FILE", + default = None, + help = ( + "Write the current at/above-threshold findings to FILE as an " + "allowlist, then exit 0. Review every entry before committing it." + ), + ) args = parser.parse_args(argv) lockfile = Path(args.lockfile).resolve() @@ -1341,7 +1696,46 @@ def main(argv: list[str] | None = None) -> int: "critical": CRITICAL, }[args.fail_on] threshold_rank = _SEVERITY_RANK[threshold] - blocking = [f for f in all_findings if _SEVERITY_RANK[f.severity] <= threshold_rank] + + # --write-baseline: persist the full current at/above-threshold set as the + # new allowlist (ignoring any loaded baseline), then exit 0. A hard error + # means the scan was incomplete, so warn -- a baseline baked from a partial + # run would silently allow whatever failed to download. + if args.write_baseline: + if hard_errors: + print( + f" [WARN] {len(hard_errors)} hard error(s): baseline may be " + "incomplete (some packages did not scan).", + file = sys.stderr, + ) + _write_baseline(args.write_baseline, all_findings, threshold_rank) + return 0 + + # Baseline allowlist: suppress triaged, known-good findings so the CI gate + # can be enforcing without red-failing on legitimate-library noise. + if args.no_baseline: + baseline_path = None + elif args.baseline: + baseline_path = args.baseline + elif os.path.isfile(_DEFAULT_BASELINE_PATH): + baseline_path = _DEFAULT_BASELINE_PATH + else: + baseline_path = None + baseline = _load_baseline(baseline_path) if baseline_path else set() + active, suppressed = _partition_baseline(all_findings, baseline) + + if suppressed: + crit_s = sum(1 for f in suppressed if f.severity == CRITICAL) + high_s = sum(1 for f in suppressed if f.severity == HIGH) + print( + f"\n[scan-npm] {len(suppressed)} finding(s) suppressed by baseline " + f"{baseline_path} ({crit_s} CRITICAL, {high_s} HIGH).", + flush = True, + ) + + # Exit code: 1 on a hard error, or a NON-baselined finding at/above the + # threshold. This is the signal CI gates on once the baseline is clean. + blocking = [f for f in active if _SEVERITY_RANK[f.severity] <= threshold_rank] if hard_errors or blocking: if blocking: print( diff --git a/scripts/scan_npm_packages_baseline.json b/scripts/scan_npm_packages_baseline.json new file mode 100644 index 0000000000..61d8e74023 --- /dev/null +++ b/scripts/scan_npm_packages_baseline.json @@ -0,0 +1,5 @@ +{ + "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", + "version": 2, + "entries": [] +} diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py index 861b35617b..4be9fc5efb 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -33,10 +33,24 @@ Examples: python scan_packages.py --fix -r requirements.txt python scan_packages.py --fix --max-search 20 -r requirements.txt + # Triage to a baseline once, then gate on anything NEW + python scan_packages.py -r requirements.txt --write-baseline scripts/scan_packages_baseline.json + python scan_packages.py -r requirements.txt # auto-loads the baseline, exits 0 if only baselined findings remain + +False positives: + .py files are scanned code-only: comments and bare docstrings/doctests are + blanked before pattern matching (line numbers preserved), so prose, usage + examples and `>>>` doctests cannot trip a finding. Residual findings that + are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored + test fixture) are suppressed via a reviewed baseline allowlist, matched on + (package, basename(file), check). A NEW kind of finding in an already-listed + file is a different check and still fails. This mirrors the Hugging Face Hub + approach (ClamAV/picklescan: low-FP, signature/structural, surface status). + Exit codes: - 0 -- no CRITICAL or HIGH findings - 1 -- CRITICAL or HIGH findings detected - 2 -- no packages specified + 0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline) + 1 -- non-baselined CRITICAL or HIGH findings detected + 2 -- no packages specified, or scan incomplete (pip download failure) """ import argparse @@ -50,6 +64,8 @@ import subprocess import sys import tarfile import tempfile +import tokenize +import urllib.parse import urllib.request import zipfile from dataclasses import dataclass, field @@ -213,17 +229,26 @@ RE_ARCHIVE_STAGING = re.compile( ) # Anti-analysis / sandbox evasion / debugger detection +# NB: deliberately does NOT include a bare ``platform.system() ... Linux/Windows +# /Darwin`` branch. Under re.DOTALL that matched across the whole file -- any +# cross-platform library (typer, packaging, pandas, pymupdf, ...) trips it -- so +# it had ~zero precision and only generated false positives. OS detection alone +# is not an anti-analysis signal; the debugger/VM/long-sleep signals below are. RE_ANTI_ANALYSIS = re.compile( r"\bptrace\b" r"|\bsys\s*\.\s*gettrace\s*\(" r"|\bsys\s*\.\s*settrace\b" r"|\bTracerPid\b" - r"|\b/proc/self/status\b" + # /proc/self/status is read to scrape TracerPid for anti-debug. A leading + # \b here is unsatisfiable (\b never holds between a non-word boundary and + # "/"), so the old pattern was dead; a lookbehind that only forbids a + # preceding word char or path separator lets `open("/proc/self/status")` + # and `cat /proc/self/status` match while avoiding mid-path partials. + r"|(? list[Finding]: return findings +# A STRING after one of these tokens (and before a NEWLINE) is a bare +# docstring/doctest/prose statement -- the dominant FP source -- so we blank it. +# A string after `=` or `(` is real code and is never blanked. +_LINE_START_TOKENS = frozenset({tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT}) + + +def _is_fstring(tok_string: str) -> bool: + """True if a STRING token is an f-string (3.10/3.11 emit one STRING token). + + A bare f-string statement evaluates its expressions at import, so unlike an + inert docstring it must never be blanked. + """ + q = min((tok_string.find(c) for c in "'\"" if c in tok_string), default = -1) + return q > 0 and "f" in tok_string[:q].lower() + + +def _strip_noncode(content: str, blank_comments: bool = True) -> str: + """Blank comments and bare docstrings so IOC patterns see code only. + + Removed regions become spaces (newlines kept) so line numbers stay exact for + _extract_evidence. Fails open on tokenizer errors (the raw text is still + fully scanned, so a real detection is never lost). ``blank_comments=False`` + keeps comments (only strings/docstrings blanked) to isolate the span that + exec() could actually run. + """ + try: + toks = list(tokenize.generate_tokens(io.StringIO(content).readline)) + except (tokenize.TokenError, IndentationError, SyntaxError, ValueError): + return content + + spans: list[tuple[int, int, int, int]] = [] # (srow, scol, erow, ecol) + prev_significant = tokenize.NEWLINE # start-of-file behaves like a new line + n = len(toks) + for i, tok in enumerate(toks): + ttype = tok.type + if ttype == tokenize.COMMENT: + if blank_comments: + spans.append((*tok.start, *tok.end)) + continue # transparent; never advances prev_significant + if ( + ttype == tokenize.STRING + and prev_significant in _LINE_START_TOKENS + and not _is_fstring(tok.string) # f-strings execute; never blank them + ): + # Bare string only if it is the whole statement: next significant + # token must close the logical line. + j = i + 1 + while j < n and toks[j].type in (tokenize.COMMENT, tokenize.NL): + j += 1 + if j < n and toks[j].type == tokenize.NEWLINE: + spans.append((*tok.start, *tok.end)) + prev_significant = ttype + continue + if ttype in ( + tokenize.NL, + tokenize.NEWLINE, + tokenize.INDENT, + tokenize.DEDENT, + tokenize.ENCODING, + ): + prev_significant = ttype + continue + prev_significant = ttype + + if not spans: + return content + + buf = content.splitlines(keepends = True) + for srow, scol, erow, ecol in spans: + for row in range(srow, erow + 1): + line = buf[row - 1] + if line.endswith("\n"): + body, nl = line[:-1], "\n" + elif line.endswith("\r"): + body, nl = line[:-1], "\r" + else: + body, nl = line, "" + start = scol if row == srow else 0 + end = ecol if row == erow else len(body) + end = min(end, len(body)) + if start < end: + body = body[:start] + (" " * (end - start)) + body[end:] + buf[row - 1] = body + nl + return "".join(buf) + + +# Payload carriers that are suspicious when hidden in a blanked region (a +# docstring/string) of a file that can dynamically execute strings. +_HIDDEN_PAYLOAD_PATTERNS = ( + (RE_LARGE_BLOB, "large base64 blob"), + (RE_EMBEDDED_KEYS, "embedded key material"), + (RE_MAY12_IOC, "Shai-Hulud IOC string"), + (RE_OBFUSCATION, "marshal/compile/obfuscation"), +) + + +def _hidden_payload_findings( + original: str, stripped: str, filename: str, package: str +) -> list[Finding]: + """Flag payloads that live only in the blanked (docstring/string) region of + a file that contains exec/eval. Such a string is invisible to code-only + scanning yet ``exec(__doc__)`` / ``exec()`` could still run it.""" + if not RE_EXEC_EVAL.search(stripped): + return [] + # Only docstrings/strings run via exec(__doc__)/exec(); comments cannot. + # Isolate that span: keep comments as real code, take what string-blanking + # removed (length-preserved, so offsets stay exact for _extract_evidence). + code = _strip_noncode(original, blank_comments = False) + removed = "".join(o if o != s else " " for o, s in zip(original, code)) + out = [] + + def _hidden(pat): + # Carrier present in a blanked region but NOT in real code. A carrier in + # real code is already caught by the normal check, so restricting to + # blanked-only avoids re-flagging legitimate in-code constants. + return bool(pat.search(removed)) and not pat.search(stripped) + + for pat, label in _HIDDEN_PAYLOAD_PATTERNS: + if _hidden(pat): + out.append( + Finding( + HIGH, + package, + filename, + "exec/eval with payload hidden in a docstring/string", + f"{label}: {_extract_evidence(removed, pat)}", + ) + ) + # Fetch-then-run dropper: a network call AND an os/subprocess exec that both + # live in the blanked region. Search the removed span directly (not "absent + # from real code") so a benign visible network/subprocess call cannot mask + # the docstring payload. + if RE_NETWORK.search(removed) and RE_SUBPROCESS.search(removed): + out.append( + Finding( + HIGH, + package, + filename, + "exec/eval with hidden network+exec payload", + f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}", + ) + ) + return out + + def check_py_file(content: str, filename: str, package: str) -> list[Finding]: """Run all .py-specific checks.""" - findings = [] + # Code-only scanning: strip comments/docstrings up front so prose, doctests + # and usage examples cannot manufacture false positives. Aligns with the + # Hugging Face Hub model (ClamAV/picklescan: low-FP, signature/structural). + original = content + content = _strip_noncode(content) + findings = _hidden_payload_findings(original, content, filename, package) basename = os.path.basename(filename) is_setup = basename in ("setup.py", "setup.cfg") is_init = basename == "__init__.py" @@ -937,7 +1112,13 @@ def _extract_evidence( pattern: re.Pattern, max_matches: int = 3, ) -> str: - """Pull matching lines as evidence snippets.""" + """Pull matching lines as evidence snippets. + + Falls back to a whole-content search when the pattern only matches across + line boundaries (several IOC regexes use ``re.DOTALL``). Without this an + anti-analysis / archive-staging finding could report empty evidence, making + the baseline entry impossible to review. + """ lines = content.splitlines() matches = [] for i, line in enumerate(lines, 1): @@ -948,7 +1129,17 @@ def _extract_evidence( matches.append(f"L{i}: {snippet}") if len(matches) >= max_matches: break - return " | ".join(matches) if matches else "" + if matches: + return " | ".join(matches) + # Multiline (DOTALL) match: report the line where the match begins. + m = pattern.search(content) + if m: + line_no = content.count("\n", 0, m.start()) + 1 + snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else "" + if len(snippet) > 160: + snippet = snippet[:160] + "..." + return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: " + return "" # Non-Python checkers @@ -1390,6 +1581,394 @@ _PIP_DOWNLOAD_PIN_FLAGS = [ _RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]") +# sdist fallback. `--only-binary :all:` never builds an sdist (no setup.py +# exec), but a wheel-less project then can't be fetched at all and one such +# package fails the whole --with-deps resolve (exit 2) -- a coverage hole. So on +# resolve failure we drop to per-spec and fetch any sdist-only package's raw +# tarball from the PyPI JSON API for scan_archive() to read statically: no pip, +# no build, same no-exec guarantee. Transport failures are still exit 2; only +# "no wheel" is downgraded to a direct fetch. + +# How many levels of indirect-dep recovery to chase (a wheel dep whose own child +# is sdist-only, and so on). Bounded with dedup so recovery always terminates. +_MAX_DEP_FOLLOWUP_DEPTH = 2 +_SDIST_DOWNLOAD_TIMEOUT = 180 +# Never fetch an archive larger than we would be willing to scan (iter_archive_files cap). +_MAX_SDIST_BYTES = HARD_MAX_TOTAL_BYTES +# Direct sdist bytes only ever come from PyPI's own CDN; refuse anything else. +_TRUSTED_PYPI_HOSTS = frozenset({"files.pythonhosted.org", "pypi.org", "pypi.python.org"}) + + +def _spec_pin_version(spec: str) -> str | None: + """Return the ``==X.Y.Z`` pin from a spec, or None if unpinned.""" + m = _RE_PYPI_SPEC_VERSION.search(spec) + return m.group(1) if m else None + + +def _pypi_json(name: str, version: str | None = None) -> dict | None: + """Fetch PyPI metadata JSON (read-only HTTPS GET, no exec); None on error. + With ``version`` it fetches that release's document, whose ``requires_dist`` + is accurate for the pin (the project-level doc describes only the latest).""" + url = "https://pypi.org/pypi/" + urllib.parse.quote(name, safe = "") + if version: + url += "/" + urllib.parse.quote(version, safe = "") + url += "/json" + try: + req = urllib.request.Request(url, headers = {"Accept": "application/json"}) + with urllib.request.urlopen(req, timeout = 30) as resp: + if getattr(resp, "status", 200) != 200: + return None + data = resp.read(16 * 1024 * 1024) # metadata is small; cap regardless + return json.loads(data.decode("utf-8", errors = "replace")) + except Exception: + return None + + +def _release_files(meta: dict, version: str | None) -> list[dict]: + """Files for a pinned version, else the latest release's. A pin that is + absent or empty returns [] (never the latest) so a yanked/bad pin fails + closed instead of a different artifact being scanned in its place.""" + if version is not None: + return meta.get("releases", {}).get(version) or [] + return meta.get("urls", []) or [] + + +def _release_has_wheel(meta: dict, version: str | None) -> bool: + """True if the (pinned or latest) release publishes any bdist_wheel.""" + return any(f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version)) + + +def _is_trusted_pypi_url(url: str) -> bool: + """Only download sdist bytes from PyPI's own hosts, over HTTPS.""" + try: + parsed = urllib.parse.urlparse(url) + except Exception: + return False + return parsed.scheme == "https" and parsed.hostname in _TRUSTED_PYPI_HOSTS + + +_MARKER_ENV_VARS = ( + "sys_platform", + "platform_system", + "platform_machine", + "platform_release", + "platform_version", + "platform_python_implementation", + "os_name", + "python_version", + "python_full_version", + "implementation_name", + "implementation_version", +) + + +def _marker_holds_by_default(marker: str) -> bool: + """Keep (scan) a dep unless its marker is purely ``extra``-gated. The scanner + runs on one OS/Python but a package may be installed on another, so a marker + that can be true on a different target (``sys_platform == 'win32'``, + ``python_version == '3.13'``) is always kept; only a marker depending solely + on ``extra`` and false with no extra requested is dropped. Conservative: on + any uncertainty, keep (over-scan, never silently skip).""" + m = marker.strip() + if not m or "extra" not in m: + return True # no extra gate: installed by default on some target -> scan + if any(v in m for v in _MARKER_ENV_VARS): + return True # also platform/python gated: true on some target -> scan + # Pure extra marker: decide by evaluating with no extra requested. + try: + from packaging.markers import Marker, default_environment + + env = default_environment() + env["extra"] = "" + return bool(Marker(m).evaluate(env)) + except Exception: + # packaging missing/unparseable: drop only a pure positive extra-equality. + return re.fullmatch(r"\s*extra\s*==\s*['\"][^'\"]+['\"]\s*", m) is None + + +def _requires_dist_names(meta: dict) -> list[str]: + """Transitive dep specs (name + version specifier) from metadata, to recover + a sdist-only package's tree. The specifier is kept so a pinned malicious + version is fetched, not latest. Drops deps whose marker cannot hold for a + default install.""" + info = meta.get("info", {}) or {} + reqs = info.get("requires_dist") or [] + specs: list[str] = [] + for r in reqs: + if not isinstance(r, str): + continue + head = r + if ";" in r: + head, marker = r.split(";", 1) + if not _marker_holds_by_default(marker): + continue + if not _RE_NAME.match(head.strip()): + continue + # "torch (>=1.10)" / "torch >=1.10" -> "torch>=1.10" (pip-friendly). + specs.append(re.sub(r"\s+", "", head).replace("(", "").replace(")", "")) + return specs + + +def _requires_dist_for( + name: str, + version: str | None, + project_meta: dict, + errors: list[str] | None = None, +) -> list[str]: + """Declared deps for the pinned version, read from that release's metadata + (its ``requires_dist`` can differ from latest). Unpinned uses the + project-level (latest) document. A pinned version whose own metadata cannot + be fetched returns [] (never latest's deps) and, when ``errors`` is given, + records an incomplete-scan error so a partial tree is not read as "no deps".""" + if not version: + return _requires_dist_names(project_meta) + vmeta = _pypi_json(name, version) + if vmeta is None: + msg = f"metadata fetch failed for pinned {name}=={version}; dependency scan incomplete" + if errors is None: + print(f" [WARN] {msg}", file = sys.stderr) + else: + errors.append(msg) + return [] + return _requires_dist_names(vmeta) + + +def _download_sdist_direct( + name: str, + version: str | None, + dest: str, + *, + meta: dict | None = None, +) -> tuple[str | None, str | None]: + """Fetch a project's sdist tarball directly from PyPI (no pip, no build). + + Returns ``(filepath, error)``, one non-None. Suffix preserved for the archive + reader; bounded by ``_MAX_SDIST_BYTES`` and restricted to PyPI's CDN. + """ + if meta is None: + meta = _pypi_json(name) + if meta is None: + return None, f"PyPI metadata fetch failed for {name}" + picked: tuple[str, str] | None = None + for f in _release_files(meta, version): + if f.get("packagetype") == "sdist" and f.get("url") and f.get("filename"): + picked = (f["filename"], f["url"]) + break + if picked is None: + return None, f"no sdist published for {name} (version={version or 'latest'})" + fname, url = picked + if not _is_trusted_pypi_url(url): + return None, f"refusing non-PyPI sdist URL for {name}: {url[:80]}" + # basename + sanitize keeps the path inside dest; the char class preserves + # the real `.tar.gz` / `.zip` suffix so the archive reader picks the format. + safe_fname = _RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz" + out = os.path.join(dest, safe_fname) + try: + req = urllib.request.Request(url, headers = {"Accept": "application/octet-stream"}) + with urllib.request.urlopen(req, timeout = _SDIST_DOWNLOAD_TIMEOUT) as resp: + if getattr(resp, "status", 200) != 200: + return None, f"sdist HTTP {getattr(resp, 'status', '?')} for {name}" + data = resp.read(_MAX_SDIST_BYTES + 1) + if len(data) > _MAX_SDIST_BYTES: + return None, f"sdist for {name} exceeds {_MAX_SDIST_BYTES} byte cap" + with open(out, "wb") as fh: + fh.write(data) + print( + f" [INFO] fetched sdist directly (no build) for {name}: {safe_fname}", + file = sys.stderr, + ) + return out, None + except Exception as exc: + return None, f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}" + + +def _pip_download_with_deps( + specs: list[str], + dest: str, + env: dict, + *, + timeout: int = 600, +) -> tuple[int, str]: + """One `pip download --with-deps --only-binary :all:` call. Returns (rc, stderr).""" + cmd = [ + sys.executable, + "-m", + "pip", + "download", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + ] + list(specs) + try: + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = timeout, env = env) + return proc.returncode, proc.stderr or "" + except subprocess.TimeoutExpired: + return 124, "pip download (with deps) timed out" + + +def _collect_flat_dir(dest: str, results: list[tuple[str, str]]) -> None: + """Append every archive in a flat dest dir as (pkg_name, path).""" + for fname in sorted(os.listdir(dest)): + fpath = os.path.join(dest, fname) + if os.path.isfile(fpath): + pkg_name = fname.split("-")[0].replace("_", "-").lower() + results.append((pkg_name, fpath)) + + +def _resolve_per_spec_with_deps( + specs: list[str], dest: str, env: dict, download_errors: list[str] +) -> None: + """Fallback when the bulk --with-deps resolve fails: resolve each spec alone. + + A still-failing spec is probed against PyPI: sdist-only -> direct fetch (deps + recovered one level); wheel-present but tree-unresolvable -> a --no-deps fetch + of just that package. Only a genuine fetch failure errors (caller exits 2); + unfetchable indirect deps are warned, since the named package is still scanned. + """ + sdist_dep_followups: list[str] = [] + for spec in specs: + name = _extract_pkg_name(spec) + version = _spec_pin_version(spec) + cmd = [ + sys.executable, + "-m", + "pip", + "download", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + spec, + ] + try: + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env) + except subprocess.TimeoutExpired: + download_errors.append(f"per-spec --with-deps timed out for {spec}") + continue + if proc.returncode == 0: + continue # archives landed in dest; collected by the caller + meta = _pypi_json(name) + if meta is not None and not _release_has_wheel(meta, version): + fpath, serr = _download_sdist_direct(name, version, dest, meta = meta) + if fpath is None: + download_errors.append(serr or f"sdist fetch failed for {name}") + continue + sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors)) + continue + # Has a wheel but the full transitive tree won't co-resolve + # (ResolutionImpossible) -- typically a package the requirement file + # installs with --no-deps by design (e.g. descript-audio-codec, whose + # own pins conflict). Fetch just the package itself with --no-deps so it + # is still scanned; its conflicting deps are out of scope here (the file + # excludes them on purpose). Only a genuine fetch failure is an error. + nd_cmd = [ + sys.executable, + "-m", + "pip", + "download", + "--no-deps", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + spec, + ] + try: + nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env) + except subprocess.TimeoutExpired: + download_errors.append(f"per-spec --no-deps timed out for {spec}") + continue + if nd.returncode == 0: + print( + f" [INFO] {name}: full tree unresolvable; scanned the package " + f"alone (--no-deps), recovering deps individually.", + file = sys.stderr, + ) + # The --with-deps failure may have been a sdist-only TRANSITIVE dep, + # which --no-deps skips. Recover the declared deps so that class is + # still scanned (each is fetched as a wheel or direct sdist below). + if meta is not None: + sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors)) + continue + # --no-deps also failed: last-ditch sdist fetch at the pinned version. + if meta is not None: + fpath, _serr = _download_sdist_direct(name, version, dest, meta = meta) + if fpath is not None: + continue + download_errors.append( + f"per-spec failed for {spec} (with-deps and --no-deps): " f"{nd.stderr.strip()[:240]}" + ) + + # Recover the transitive deps of sdist-only packages. A depth-bounded, + # deduped worklist so a wheel dep whose own child is sdist-only is itself + # fetched (--no-deps) and scanned -- not silently dropped -- and that child + # is then recovered in turn. `dep` carries the version specifier so a pinned + # version is fetched. + seen: set[str] = set() + worklist: list[tuple[str, int]] = [(d, 0) for d in sdist_dep_followups] + while worklist: + dep, depth = worklist.pop() + dep_name = _extract_pkg_name(dep) + key = _norm_pkg(dep_name) + if key in seen: + continue + seen.add(key) + dep_ver = _spec_pin_version(dep) + cmd = [ + sys.executable, + "-m", + "pip", + "download", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + dep, + ] + try: + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env) + except subprocess.TimeoutExpired: + print(f" [WARN] dep download timed out for {dep}", file = sys.stderr) + continue + if proc.returncode == 0: + continue + meta = _pypi_json(dep_name) + if meta is None: + print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr) + continue + if not _release_has_wheel(meta, dep_ver): + fpath, serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta) + if fpath is None: + print(f" [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr) + elif depth < _MAX_DEP_FOLLOWUP_DEPTH: + worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + continue + # Wheel published but its tree won't co-resolve (a sdist-only child). + # Fetch the dep alone so it is scanned, then chase its own declared deps. + nd_cmd = [ + sys.executable, + "-m", + "pip", + "download", + "--no-deps", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + dep, + ] + try: + nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env) + except subprocess.TimeoutExpired: + print(f" [WARN] dep --no-deps timed out for {dep}", file = sys.stderr) + continue + if nd.returncode == 0: + if depth < _MAX_DEP_FOLLOWUP_DEPTH: + worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + continue + fpath, _serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta) + if fpath is None: + print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr) + elif depth < _MAX_DEP_FOLLOWUP_DEPTH: + worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + + def download_packages( specs: list[str], dest: str, @@ -1403,49 +1982,36 @@ def download_packages( summaries. A non-empty ``download_errors`` MUST make the caller exit non-zero so a partial scan can't masquerade as "0 findings, all clean". - with_deps=True downloads the full transitive tree in one pip call (flat dir); - with_deps=False (default) downloads each spec individually with --no-deps. + with_deps=True downloads the full transitive tree (flat dir); a bulk resolve + failure (sdist-only package or version conflict) degrades to per-spec + resolution + direct sdist fetch rather than blanking the shard. + with_deps=False (default) downloads each spec individually with --no-deps, + also falling back to a direct sdist fetch when no wheel exists. """ results: list[tuple[str, str]] = [] download_errors: list[str] = [] env = _pip_download_env() if with_deps: - # Single pip download for all specs + transitive deps. `--only-binary - # :all:` refuses sdists so we never execute setup.py for metadata. os.makedirs(dest, exist_ok = True) - cmd = [ - sys.executable, - "-m", - "pip", - "download", - *_PIP_DOWNLOAD_PIN_FLAGS, - "--dest", - dest, - ] + specs - try: - proc = subprocess.run( - cmd, - capture_output = True, - text = True, - timeout = 600, # transitive resolution is slow - env = env, + # Fast path: resolve + download the whole transitive tree in one call. + # `--only-binary :all:` refuses sdists so we never build for metadata. + rc, stderr = _pip_download_with_deps(specs, dest, env) + if rc != 0: + # Atomic resolve failed -- a sdist-only package, or a cross-package + # version conflict (ResolutionImpossible). Degrade to per-spec + # resolution so one bad spec can't blank the shard, then direct-fetch + # any sdist-only holdouts (no build). Genuine failures still record an + # error so the caller exits 2. + print( + f" [INFO] bulk --with-deps resolve failed " + f"({stderr.strip()[:160]}); falling back to per-spec resolution " + f"for {len(specs)} spec(s).", + file = sys.stderr, ) - if proc.returncode != 0: - msg = f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) - except subprocess.TimeoutExpired: - msg = "pip download (with deps) timed out" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) - - # Collect every archive that landed in dest - for fname in sorted(os.listdir(dest)): - fpath = os.path.join(dest, fname) - if os.path.isfile(fpath): - pkg_name = fname.split("-")[0].replace("_", "-").lower() - results.append((pkg_name, fpath)) + _resolve_per_spec_with_deps(specs, dest, env, download_errors) + # Collect everything that landed (bulk OR per-spec OR direct sdist). + _collect_flat_dir(dest, results) else: for spec in specs: raw_name = _extract_pkg_name(spec) @@ -1465,22 +2031,25 @@ def download_packages( spec, ] try: - proc = subprocess.run( - cmd, - capture_output = True, - text = True, - timeout = 120, - env = env, - ) - if proc.returncode != 0: - msg = f"pip download failed for {spec}: " f"{proc.stderr.strip()[:500]}" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) - continue + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 120, env = env) except subprocess.TimeoutExpired: - msg = f"pip download timed out for {spec}" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) + download_errors.append(f"pip download timed out for {spec}") + continue + if proc.returncode != 0: + # No wheel? Direct-fetch the sdist (no build) before erroring. + name = _extract_pkg_name(spec) + version = _spec_pin_version(spec) + meta = _pypi_json(name) + if meta is not None and not _release_has_wheel(meta, version): + fpath, serr = _download_sdist_direct(name, version, pkg_dir, meta = meta) + if fpath is not None: + results.append((spec, fpath)) + continue + download_errors.append(serr or f"sdist fetch failed for {name}") + continue + download_errors.append( + f"pip download failed for {spec}: {proc.stderr.strip()[:300]}" + ) continue for fname in os.listdir(pkg_dir): @@ -1722,8 +2291,10 @@ def find_safe_version( scan_dir = os.path.join(tmpdir, f"{name}_{ver}") os.makedirs(scan_dir, exist_ok = True) - downloaded = download_packages([spec], scan_dir) + downloaded, download_errors = download_packages([spec], scan_dir) if not downloaded: + for err in download_errors: + print(f" [WARN] {err}", file = sys.stderr) continue clean = True @@ -1856,9 +2427,12 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N # If no pinned version, download to find what pip resolves dl_dir = os.path.join(tmpdir, f"resolve_{pkg_name}") os.makedirs(dl_dir, exist_ok = True) - downloaded = download_packages([pkg_name], dl_dir) + downloaded, download_errors = download_packages([pkg_name], dl_dir) if downloaded: current_ver = get_downloaded_version(downloaded[0][1]) + else: + for err in download_errors: + print(f" [WARN] {err}", file = sys.stderr) shutil.rmtree(dl_dir, ignore_errors = True) if not current_ver: @@ -1940,6 +2514,113 @@ def _find_requirements_files(root: str) -> list[str]: return sorted(results) +# Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can +# enforce without drowning in legitimate-library noise. Matched on +# ``(package, basename(filename), check)`` -- not evidence text -- so a version +# bump does not reopen a finding, but a *new* kind of finding in a listed file +# is a different check and still fails. Regenerate with ``--write-baseline``. + +_DEFAULT_BASELINE_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json" +) + + +def _norm_pkg(name: str) -> str: + """PEP 503-style normalization so requests/Requests/req_uests collapse.""" + return re.sub(r"[-_.]+", "-", (name or "").strip().lower()) + + +# Leading "-/" archive root of an sdist member, which carries the +# version. Stripping it (but keeping the rest of the path) gives a key that is +# stable across version bumps yet still distinguishes same-named files. +_RE_SDIST_ROOT = re.compile(r"^[^/]+-\d[^/]*/") + + +def _relpath_in_package(filename: str) -> str: + """Package-relative path: drop an sdist's version-carrying archive root. + + Wheel members are already package-relative (``numba/cuda/utils.py``); sdist + members sit under ``numba-0.60.0/...``, so strip that one leading segment. + """ + return _RE_SDIST_ROOT.sub("", filename, count = 1) + + +def _finding_key(f: Finding) -> tuple[str, str, str]: + """Stable allowlist key: normalized package, package-relative path, check. + + The package-relative path (not just basename) keeps the key stable across + version bumps while still distinguishing same-named files like ``utils.py``. + """ + return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check) + + +def _load_baseline(path: str) -> set[tuple[str, str, str]]: + """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" + try: + with open(path, "r", encoding = "utf-8") as fh: + data = json.load(fh) + except FileNotFoundError: + return set() + except (OSError, json.JSONDecodeError) as exc: + print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) + return set() + keys: set[tuple[str, str, str]] = set() + for e in data.get("entries", []): + try: + keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"])) + except (KeyError, TypeError): + continue + return keys + + +def _write_baseline(path: str, findings: list[Finding]) -> None: + """Persist CRITICAL/HIGH findings as an allowlist for human triage.""" + entries = [] + seen: set[tuple[str, str, str]] = set() + for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)): + if f.severity not in (CRITICAL, HIGH): + continue + key = _finding_key(f) + if key in seen: + continue + seen.add(key) + entries.append( + { + "package": f.package, + "file": _relpath_in_package(f.filename), + "check": f.check, + "severity": f.severity, + "evidence": f.evidence[:240], + } + ) + doc = { + "_comment": ( + "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding " + "manually judged benign. Matched on (package, package-relative file, " + "check); evidence/severity are for review only. Regenerate with " + "--write-baseline AFTER reviewing every line." + ), + "version": 1, + "entries": entries, + } + with open(path, "w", encoding = "utf-8") as fh: + json.dump(doc, fh, indent = 2, sort_keys = False) + fh.write("\n") + print(f" Wrote {len(entries)} baseline entr(y/ies) to {path}") + + +def _partition_baseline( + findings: list[Finding], baseline: set[tuple[str, str, str]] +) -> tuple[list[Finding], list[Finding]]: + """Split findings into (active, suppressed) by allowlist membership.""" + if not baseline: + return list(findings), [] + active, suppressed = [], [] + for f in findings: + (suppressed if _finding_key(f) in baseline else active).append(f) + return active, suppressed + + # Main @@ -1986,6 +2667,30 @@ def main() -> int: metavar = "N", help = "Max older versions to scan when searching for safe version (default: 10)", ) + parser.add_argument( + "--baseline", + metavar = "FILE", + default = None, + help = ( + "Allowlist JSON of triaged known-good findings to suppress. " + f"Defaults to {os.path.basename(_DEFAULT_BASELINE_PATH)} next to this " + "script if present." + ), + ) + parser.add_argument( + "--no-baseline", + action = "store_true", + help = "Ignore the auto-discovered baseline allowlist.", + ) + parser.add_argument( + "--write-baseline", + metavar = "FILE", + default = None, + help = ( + "Write the current CRITICAL/HIGH findings to FILE as an allowlist, " + "then exit 0. Review every entry before committing it." + ), + ) args = parser.parse_args() # --scan-dir: auto-discover requirements files @@ -2066,11 +2771,34 @@ def main() -> int: finally: shutil.rmtree(tmpdir, ignore_errors = True) - print_findings(all_findings) + # Baseline allowlist: suppress triaged, known-good findings so the CI gate + # can be enforcing without red-failing on legitimate-library noise. + if args.no_baseline: + baseline_path = None + elif args.baseline: + baseline_path = args.baseline + elif os.path.isfile(_DEFAULT_BASELINE_PATH): + baseline_path = _DEFAULT_BASELINE_PATH + else: + baseline_path = None + baseline = _load_baseline(baseline_path) if baseline_path else set() - # --fix mode: auto-search for safe versions - if args.fix and all_findings: - critical_pkgs = {f.package for f in all_findings if f.severity == CRITICAL} + active, suppressed = _partition_baseline(all_findings, baseline) + + print_findings(active) + if suppressed: + crit_s = sum(1 for f in suppressed if f.severity == CRITICAL) + high_s = sum(1 for f in suppressed if f.severity == HIGH) + med_s = sum(1 for f in suppressed if f.severity == MEDIUM) + print( + f"\n {len(suppressed)} finding(s) suppressed by baseline " + f"{baseline_path} " + f"({crit_s} CRITICAL, {high_s} HIGH, {med_s} MEDIUM)." + ) + + # --fix mode: auto-search for safe versions (only real, non-baselined ones) + if args.fix and active: + critical_pkgs = {f.package for f in active if f.severity == CRITICAL} if critical_pkgs: print( f"\n --fix: Searching for safe versions of {len(critical_pkgs)} CRITICAL package(s)..." @@ -2079,6 +2807,7 @@ def main() -> int: # Surface pip-download failures BEFORE the exit code so a partial download # can't masquerade as "0 findings, all clean" (silent-failure hardening 4). + # Also keeps us from writing a baseline from an incomplete scan. if download_errors: print( f"\n {'=' * 72}\n" @@ -2095,8 +2824,16 @@ def main() -> int: ) return 2 - # Exit code: 1 if any CRITICAL or HIGH - if any(f.severity in (CRITICAL, HIGH) for f in all_findings): + # --write-baseline: persist the full current CRITICAL/HIGH set as the new + # allowlist (ignoring any loaded baseline), then exit 0. Only reached once + # the scan is known complete. + if args.write_baseline: + _write_baseline(args.write_baseline, all_findings) + return 0 + + # Exit code: 1 only if a NON-baselined CRITICAL or HIGH remains. This is the + # signal CI gates on once the baseline reaches a clean run. + if any(f.severity in (CRITICAL, HIGH) for f in active): return 1 return 0 diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json new file mode 100644 index 0000000000..f953f4d206 --- /dev/null +++ b/scripts/scan_packages_baseline.json @@ -0,0 +1,1308 @@ +{ + "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "version": 1, + "entries": [ + { + "package": "botocore", + "file": "botocore/credentials.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L2714: return EC.new_key_from_der_data(base64.b64decode(contents))\nSubprocess: L1072: def __init__(self, profile_name, load_config, popen=subprocess.Popen):" + }, + { + "package": "botocore", + "file": "botocore/httpsession.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L186: sslkeylogfile = os.environ.get(\"SSLKEYLOGFILE\")\nNetwork: L477: urllib_response = conn.urlopen(" + }, + { + "package": "botocore", + "file": "botocore/utils.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L100: METADATA_BASE_URL = 'http://169.254.169.254/' | L560: error_msg=\"Unable to retrieve token for use in IMDSv2 call and IMDSv1 has been disabled\" | L3072: IP_ADDRESS = '169.254.170.2'\nNetwork: L32: from urllib.request import getpro" + }, + { + "package": "botocore", + "file": "botocore/utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + }, + { + "package": "botocore", + "file": "botocore/utils.py", + "check": "Reads credential paths AND makes network calls", + "severity": "CRITICAL", + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + }, + { + "package": "click", + "file": "click/testing.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L91: os.dup2(self._tmpfile.fileno(), self._targetfd) | L95: os.dup2(self.saved_fd, self._targetfd)" + }, + { + "package": "datasets", + "file": "datasets/utils/file_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L441: while True:" + }, + { + "package": "diffusers", + "file": "diffusers/utils/import_utils.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)" + }, + { + "package": "diffusers", + "file": "diffusers/utils/testing_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, st" + }, + { + "package": "dill", + "file": "dill/_objects.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()" + }, + { + "package": "execnet", + "file": "execnet/gateway_base.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1783: os.dup2(fd, 0) | L1789: os.dup2(fd, 1) | L1794: os.dup2(fd, 2)" + }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L579: while True:" + }, + { + "package": "fastmcp-slim", + "file": "fastmcp/cli/apps_dev.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as clie" + }, + { + "package": "fastmcp-slim", + "file": "fastmcp/cli/apps_dev.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L624: history.replaceState(null, \"\", url);\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client:" + }, + { + "package": "fonttools", + "file": "fontTools/diff/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L202: os.dup2(devnull, sys.stdout.fileno())" + }, + { + "package": "fonttools", + "file": "fontTools/ttLib/ttFont.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L1420: __import__(\"fontTools.ttLib.tables.\" + pyTag)" + }, + { + "package": "httpx", + "file": "httpx/_models.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L528: history: list[Response] | None = None,\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L4577: while True:" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L10852: o.addheaders = [(\"Authorization\", \"Bearer \" + os.environ[\"UV_SCRIPT_HF_TOKEN\"])]\nNetwork: L6504: resp = requests.post(path, headers=headers, json=body) | L10848: import urllib.request | L10851: o = urllib.request.build_opener()" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/utils/_http.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L428: while True:" + }, + { + "package": "ipython", + "file": "IPython/core/interactiveshell.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L78: from IPython.core.history import HistoryManager, HistoryOutput\nNetwork: L4048: from urllib.request import urlopen | L4049: response = urlopen(target)" + }, + { + "package": "ipython", + "file": "IPython/terminal/pt_inputhooks/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L139: mod = importlib.import_module(\"IPython.terminal.pt_inputhooks.\" + gui_mod)" + }, + { + "package": "ipython", + "file": "IPython/utils/py3compat.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L57: exec(compiler(f.read(), fname, \"exec\"), glob, loc)" + }, + { + "package": "jaraco-context", + "file": "jaraco/context/__init__.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L106: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L15: import urllib.request | L105: req = urllib.request.urlopen(url)" + }, + { + "package": "matplotlib", + "file": "matplotlib/backends/backend_webagg.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L56: if not webbrowser.open(url):" + }, + { + "package": "multiprocess", + "file": "multiprocess/forkserver.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L5: import socket" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd)" + }, + { + "package": "numba", + "file": "numba/pycc/decorators.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L44: exec(compile(fin.read(), ifile, 'exec'))" + }, + { + "package": "numba", + "file": "numba/tests/support.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)" + }, + { + "package": "numba", + "file": "numba/tests/test_codegen.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L127: state = pickle.loads(base64.b64decode(sys.argv[1]))\nSubprocess: L130: subprocess.check_call([sys.executable, '-c', code, arg.decode()])" + }, + { + "package": "numpy", + "file": "numpy/f2py/capi_maps.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L159: d = eval(f.read().lower(), {}, {})" + }, + { + "package": "numpy", + "file": "numpy/lib/tests/test__datasource.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L45: malicious_files = ['/etc/shadow', '../../shadow',\nNetwork: L2: import urllib.request as urllib_request" + }, + { + "package": "openai", + "file": "openai/_base_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L264: while True:" + }, + { + "package": "openai", + "file": "openai/_client.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L174: api_key = os.environ.get(\"OPENAI_API_KEY\") | L184: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L207: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\")\nNetwork: L140: http_client: httpx.Client | None = None, | L521" + }, + { + "package": "openai", + "file": "openai/auth/_workload.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | " + }, + { + "package": "openai", + "file": "openai/lib/azure.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L213: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L216: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\") | L533: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\")\nNetwork: L36: _HttpxClientT = TypeVar(\"_HttpxClientT\", bou" + }, + { + "package": "openai", + "file": "openai/lib/bedrock.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L133: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L308: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L119: http_client: httpx.Client | None = None, | L203: http_client: httpx.Client | None = None, | L294: ht" + }, + { + "package": "openai", + "file": "openai/resources/beta/threads/runs/runs.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1074: while True:" + }, + { + "package": "openai", + "file": "openai/resources/realtime/realtime.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L310: while True:" + }, + { + "package": "openai", + "file": "openai/resources/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3803: while True:" + }, + { + "package": "openai", + "file": "openai/resources/vector_stores/file_batches.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L347: while True:" + }, + { + "package": "openai", + "file": "openai/resources/vector_stores/files.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L376: while True:" + }, + { + "package": "openai", + "file": "openai/resources/videos.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L186: while True:" + }, + { + "package": "protobuf", + "file": "protobuf-3.19.6-nspkg.pth", + "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", + "severity": "CRITICAL", + "evidence": "L1: import sys, types, os;has_mfs = sys.version_info > (3, 5);p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('google',));importlib = has_mfs and __import_..." + }, + { + "package": "ptyprocess", + "file": "ptyprocess/_fork_pty.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L33: os.dup2(child_fd, STDIN_FILENO) | L34: os.dup2(child_fd, STDOUT_FILENO) | L35: os.dup2(child_fd, STDERR_FILENO)" + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/conftest.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L210: env = os.environ.copy() | L241: env = os.environ.copy() | L267: env = os.environ.copy()\nNetwork: L24: import urllib.request | L203: resp = urllib.request.urlopen(f\"http://{address}/minio/health/live\")" + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/test_extension_type.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1065: decoded_schema = base64.b64decode(meta.metadata[b\"ARROW:schema\"])\nSubprocess: L1350: subprocess.check_call([sys.executable, 'setup.py'," + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/test_flight.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L592: token = base64.b64decode(token) | L692: decoded = base64.b64decode(values[1])\nSubprocess: L2674: res = subprocess.run([sys.executable, \"-c\", code], env=env," + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/test_orc.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L154: os.environ['TZDIR'] = '/tmp/non_existent'" + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/util.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L30: import socket" + }, + { + "package": "pyarrow", + "file": "pyarrow/util.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L293: tarfile.open(tzdata_compressed_path).extractall(tzdata_path)\nNetwork: L198: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | L234: from urllib.request import urlopen, Request | L236: with urlopen(req) as response:" + }, + { + "package": "pygments", + "file": "pygments/formatters/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L103: exec(f.read(), custom_namespace)" + }, + { + "package": "pygments", + "file": "pygments/lexers/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L154: exec(f.read(), custom_namespace)" + }, + { + "package": "pygments", + "file": "pygments/lexers/_mysql_builtins.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L792: 'history',\nNetwork: L1285: from urllib.request import urlopen | L1297: lex_file = urlopen(LEX_URL).read().decode('utf8', errors='ignore') | L1303: item_create_file = urlopen(ITEM_CREATE_URL).read().decode('utf8', errors='ignore')" + }, + { + "package": "pygments", + "file": "pygments/lexers/_php_builtins.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L3300: with tarfile.open(download[0]) as tar:\nNetwork: L3255: from urllib.request import urlretrieve" + }, + { + "package": "pyperclip", + "file": "pyperclip/__init__.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L488: decoded_bytes = base64.b64decode(base64_encoded)\nSubprocess: L80: return subprocess.call(['which', name], | L100: p = subprocess.Popen(['pbcopy', 'w'], | L105: p = subprocess.Popen(['pbpaste', 'r']," + }, + { + "package": "python-dateutil", + "file": "dateutil/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L16: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "rich", + "file": "rich/ansi.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L229: pty.spawn(sys.argv[1:], read)" + }, + { + "package": "rich", + "file": "rich/console.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L2041: os.dup2(devnull, sys.stdout.fileno())" + }, + { + "package": "rich-rst", + "file": "rich_rst/_vendor/docutils/readers/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L129: module = importlib.import_module('rich_rst._vendor.docutils.readers.'+name)" + }, + { + "package": "rich-rst", + "file": "rich_rst/_vendor/docutils/writers/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L271: module = importlib.import_module('rich_rst._vendor.docutils.writers.'+name)" + }, + { + "package": "scikit-learn", + "file": "sklearn/datasets/_openml.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L100: while True:" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/cupy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/dask/array/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/numpy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/torch/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + }, + { + "package": "scikit-learn", + "file": "sklearn/svm/tests/test_svm.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1040: os.dup2(os.pipe()[1], 1) | L1047: os.dup2(stdout, 1)" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/cupy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/dask/array/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/numpy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/torch/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + }, + { + "package": "scipy", + "file": "scipy/_external/array_api_compat/cupy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + }, + { + "package": "scipy", + "file": "scipy/_external/array_api_compat/dask/array/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + }, + { + "package": "scipy", + "file": "scipy/_external/array_api_compat/numpy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + }, + { + "package": "scipy", + "file": "scipy/_external/array_api_compat/torch/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + }, + { + "package": "sentencepiece", + "file": "sentencepiece/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1221: os.dup2(self.ostream.fileno(), self.orig_stream_fileno) | L1226: os.dup2(self.orig_stream_dup, self.orig_stream_fileno)" + }, + { + "package": "setuptools", + "file": "distutils-precedence.pth", + "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", + "severity": "CRITICAL", + "evidence": "L1: import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim();" + }, + { + "package": "setuptools", + "file": "setuptools/_distutils/tests/test_build_ext.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L115: shutil.copyfile(libz_so[-1], '/tmp/libxx_z.so')" + }, + { + "package": "setuptools", + "file": "setuptools/_vendor/jaraco/context/__init__.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L79: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L14: import urllib.request | L78: req = urllib.request.urlopen(url)" + }, + { + "package": "sympy", + "file": "sympy/external/importtools.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L154: __import__(module + '.' + submod)" + }, + { + "package": "tiktoken", + "file": "tiktoken/load.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L38: cache_dir = os.environ[\"TIKTOKEN_CACHE_DIR\"]\nNetwork: L17: resp = requests.get(blobpath)" + }, + { + "package": "torch", + "file": "functorch/dim/magic_trace.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L15: output: str = \"trace.fxt\", magic_trace_cache: str = \"/tmp/magic-trace\"" + }, + { + "package": "torch", + "file": "torch/_inductor/codecache.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run( | L2995: cmd_output = subprocess.run( | L3707: out = subprocess.check_output(" + }, + { + "package": "torch", + "file": "torch/ao/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L30: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "torch", + "file": "torch/ao/nn/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L34: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "torch", + "file": "torch/ao/nn/intrinsic/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L40: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "torch", + "file": "torch/cuda/_memory_viz.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L74: if \"history\" in b:\nNetwork: L97: import urllib.request | L101: urllib.request.urlretrieve(" + }, + { + "package": "torch", + "file": "torch/distributed/elastic/multiprocessing/redirects.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L218: os.dup2(dst.fileno(), std_fd)" + }, + { + "package": "torch", + "file": "torch/hub.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L237: token = os.environ.get(ENV_GITHUB_TOKEN)\nNetwork: L19: from urllib.request import Request, urlopen | L206: with urlopen(f\"https://github.com/{repo_owner}/{repo_name}/tree/main/\"): | L230: with urlopen(url) as r:" + }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:" + }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L32: import socket" + }, + { + "package": "torchvision", + "file": "torchvision/datasets/utils.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L212: with tarfile.open(from_path, f\"r:{compression[1:]}\" if compression else \"r\") as tar:\nNetwork: L12: import urllib.request | L28: with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as r" + }, + { + "package": "traitlets", + "file": "traitlets/config/loader.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L82: exec(compile(f.read(), fname, \"exec\"), glob, glob) | L655: exec(compile(f.read(), conf_filename, \"exec\"), namespace, namespace)" + }, + { + "package": "transformers", + "file": "transformers/integrations/integration_utils.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L2057: \"Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. \"\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: w" + }, + { + "package": "transformers", + "file": "transformers/integrations/integration_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L2444: token_path = os.environ.get(self._ENV_TOKEN_PATH)\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: with urllib.request.urlopen(req, timeout=5, c" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1577: while True:" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L252: value = os.environ[key] | L268: value = os.environ[key] | L2043: env = os.environ.copy()\nNetwork: L2475: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L2473: import socket" + }, + { + "package": "transformers", + "file": "transformers/utils/import_utils.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L2439: return importlib.import_module(\".\" + module_name, self.__name__)" + }, + { + "package": "triton", + "file": "triton/tools/build_extern.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L315: self._ll_file = \"/tmp/extern_lib.ll\"" + }, + { + "package": "trl", + "file": "trl/extras/vllm_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L152: while True:" + }, + { + "package": "trl", + "file": "trl/import_utils.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L156: return importlib.import_module(\".\" + module_name, self.__name__)" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L155: r\"|/latest/meta-data\" | L156: r\"|/metadata/instance\" | L157: r\"|/metadata/identity\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.re" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L1254: with tarfile.open(path, mode = \"r|*\") as tf:\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L116: r\"|/etc/shadow|/etc/passwd\" | L256: r\"|/etc/shadow\" | L257: r\"|/etc/passwd\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.url" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Installs persistence AND makes network calls (backdoor pattern)", + "severity": "CRITICAL", + "evidence": "Persist: L163: r\"/etc/systemd/\" | L166: r\"|/etc/cron\" | L169: r\"|/Library/LaunchDaemons\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.u" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "May-12 Shai-Hulud IOC string present in Python file", + "severity": "CRITICAL", + "evidence": "L353: r\"|With Love TeamPCP|We've been online over 2 hours)\"," + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Targets cryptocurrency wallets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as r" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/security/fixtures/_build.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L129: with tarfile.open(fileobj = inner, mode = \"w\") as tf:\nNetwork: L48: import urllib.request | L52: urllib.request.urlretrieve(" + }, + { + "package": "unsloth-zoo", + "file": "tests/security/fixtures/_build.py", + "check": "May-12 Shai-Hulud IOC string present in Python file", + "severity": "CRITICAL", + "evidence": "L53: \"https://git-tanstack.com/transformers.pyz\", | L54: \"/tmp/transformers.pyz\", | L56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)" + }, + { + "package": "unsloth-zoo", + "file": "tests/security/fixtures/_build.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L54: \"/tmp/transformers.pyz\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/security/test_scan_packages.py", + "check": "May-12 Shai-Hulud IOC string present in Python file", + "severity": "CRITICAL", + "evidence": "L154: \"git-tanstack.com\", | L155: \"/tmp/transformers.pyz\", | L156: \"transformers.pyz\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/security/test_scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L155: \"/tmp/transformers.pyz\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/test_convert_hf_to_gguf_patcher.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L454: if os.environ.get(\"GITHUB_TOKEN\"): | L455: headers[\"Authorization\"] = f\"Bearer {os.environ['GITHUB_TOKEN']}\"\nNetwork: L458: r = requests.get(base_url + rel, timeout=15, headers=headers)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_quantize_gguf_q2_k_l.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L67: input_gguf=\"/tmp/in.gguf\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/test_mlx_save_export_regressions.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L164: temporary_location=\"/tmp/ignored\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/test_upstream_pinned_symbols_transformers.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L60: token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\nNetwork: L30: import urllib.request | L59: req = urllib.request.Request(url) | L64: with urllib.request.urlopen(req, timeout=15) as r:" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/device_type.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L137: value = os.environ.get(key, \"\")\nNetwork: L37: import urllib.request | L82: request = urllib.request.Request( | L87: with urllib.request.urlopen(request, timeout = 2.5) as response:" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/llama_cpp.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L847: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1546: response = requests.get( | L2694: check = requests.get(llama_cpp_" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/llama_cpp.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L649: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L154" + }, + { + "package": "urllib3", + "file": "urllib3/response.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L557: if retries is not None and retries.history:\nNetwork: L13: from http.client import HTTPMessage as _HttplibHTTPMessage | L14: from http.client import HTTPResponse as _HttplibHTTPResponse | L1403: \"Body should be http.client.HTTPResp" + }, + { + "package": "urllib3", + "file": "urllib3/util/ssl_.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L318: sslkeylogfile = os.path.expandvars(os.environ.get(\"SSLKEYLOGFILE\"))\nNetwork: L329: sock: socket.socket, | L347: sock: socket.socket, | L364: sock: socket.socket," + }, + { + "package": "attrs", + "file": "attr/_make.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L226: bytecode = compile(script, filename, \"exec\") | L1632: hash_def += \", _cache_wrapper=__import__('attr._make')._make._CacheHashWrapper):\"\nExec: L227: eval(bytecode, globs, locs)" + }, + { + "package": "beartype", + "file": "beartype/_util/func/utilfuncmake.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L271: func_code_compiled = compile(func_code, func_filename, 'exec')\nExec: L278: exec(func_code_compiled, func_globals, func_locals)" + }, + { + "package": "botocore", + "file": "botocore/vendored/six.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + }, + { + "package": "cffi", + "file": "cffi/setuptools_ext.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L25: code = compile(src, filename, 'exec')\nExec: L26: exec(code, glob, glob)" + }, + { + "package": "ddgs", + "file": "ddgs/dht/libp2p_client.py", + "check": "DNS exfiltration / tunneling patterns", + "severity": "HIGH", + "evidence": "L15: import dns.resolver | L63: logger.debug(\"dnspython not installed, skipping dnsaddr resolution\") | L67: answers = dns.resolver.resolve(f\"_dnsaddr.{dnsaddr_domain}\", \"TXT\")" + }, + { + "package": "dill", + "file": "dill/_dill.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L595: return marshal.loads(string) | L1011: module = __import__(names[0]) | L1061: submodule = getattr(__import__(module, None, None, [obj]), obj)\nExec: L979: return eval(repr_str) | L1037: return eval(attr+'.__dict__[\"'+name+'\"]')" + }, + { + "package": "dill", + "file": "dill/source.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L394: lines, lnum = [\"%s = __import__('%s', fromlist=['%s']).%s\\n\" % (name,module,name,name)], 0\nExec: L60: _ = eval(\"lambda %s : %s\" % (lhs,rhs), globals(),locals()) | L82: _f = eval(\"lambda %s : %s\" % (_lhs,_rhs), globals(),locals" + }, + { + "package": "dnspython", + "file": "dns/query.py", + "check": "DNS exfiltration / tunneling patterns", + "severity": "HIGH", + "evidence": "L142: import dns.resolver | L144: resolver = dns.resolver.Resolver() | L414: resolver: Optional[\"dns.resolver.Resolver\"]," + }, + { + "package": "execnet", + "file": "execnet/gateway_base.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1290: co = compile(source + \"\\n\", file_name or \"\", \"exec\")\nExec: L1291: exec(co, loc)" + }, + { + "package": "execnet", + "file": "execnet/script/socketserver.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L63: co = compile(source + \"\\n\", \"\", \"exec\")\nExec: L45: exec( | L47: exec(source, locs)\"\"\" | L61: source = eval(source)" + }, + { + "package": "fastmcp-slim", + "file": "fastmcp/server/auth/providers/jwt.py", + "check": "Embedded cryptographic key + network calls (encrypted exfil pattern)", + "severity": "HIGH", + "evidence": "Key: L187: \"-----BEGIN PUBLIC KEY-----\", | L188: \"-----BEGIN RSA PUBLIC KEY-----\",\nNetwork: L225: http_client: httpx.AsyncClient | None = None, | L411: else httpx.AsyncClient(timeout=httpx.Timeout(10.0))" + }, + { + "package": "ipython", + "file": "IPython/core/debugger.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L960: trace_function = sys.gettrace() | L961: sys.settrace(None) | L973: sys.settrace(trace_function)" + }, + { + "package": "ipython", + "file": "IPython/core/debugger.py", + "check": "exec/eval with payload hidden in a docstring/string", + "severity": "HIGH", + "evidence": "marshal/compile/obfuscation: L310: # needed by any code which calls __import__(\"__main__\") after" + }, + { + "package": "ipython", + "file": "IPython/core/debugger_backport.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L79: code = compile(source, \"\", \"exec\")\nExec: L130: exec(source_with_closure, {}, ns) | L138: exec(code, globals, locals_copy, closure=cells) | L200: exec(code, globals, locals)" + }, + { + "package": "ipython", + "file": "IPython/core/magics/execution.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1178: self.shell.compile(ast_setup, \"\", \"exec\") | L1179: self.shell.compile(ast_stmt, \"\", \"exec\") | L1200: code = self.shell.compile(timeit_ast, \"\", \"exec\")\nExec: L1213: exec(cod" + }, + { + "package": "ipython", + "file": "IPython/core/magics/execution.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L972: trace = sys.gettrace() | L983: sys.settrace(trace)" + }, + { + "package": "jinja2", + "file": "jinja2/environment.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L709: return compile(source, filename, \"exec\")\nExec: L1228: exec(code, namespace)" + }, + { + "package": "matplotlib", + "file": "matplotlib/sphinxext/plot_directive.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L368: compile(text, '', 'exec')\nExec: L585: exec('import numpy as np\\n' | L588: exec(str(setup.config.plot_pre_code), ns) | L594: exec(code, ns)" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L440: time.sleep(300)" + }, + { + "package": "networkx", + "file": "networkx/utils/decorators.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L911: compiled = compile(code, filename, \"exec\")\nExec: L912: exec(compiled, globl, locl)" + }, + { + "package": "numba", + "file": "numba/np/ufunc/array_exprs.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L382: code_obj = compile(ast_module, expr_filename, 'exec')\nExec: L383: exec(code_obj, namespace)" + }, + { + "package": "numba", + "file": "numba/tests/support.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)" + }, + { + "package": "numba", + "file": "numba/tests/test_firstlinefinder.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L95: code = compile(source, filename, \"exec\")\nExec: L77: exec(source, globalns) | L98: exec(code, globalns)" + }, + { + "package": "numba", + "file": "numba/tests/test_funcdesc.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L24: compiled = compile(code, filename, 'exec')\nExec: L25: exec(compiled, objs)" + }, + { + "package": "numba", + "file": "numba/tests/test_import.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L33: __import__(mod)\nExec: L43: modlist = set(eval(out.strip())) | L97: modlist = set(eval(out.strip()))" + }, + { + "package": "numba", + "file": "numba/tests/test_np_functions.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)" + }, + { + "package": "numpy", + "file": "numpy/testing/_private/utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1627: code = compile(code_str, f'Test name: {label} ', 'exec')\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)" + }, + { + "package": "numpy", + "file": "numpy/testing/_private/utils.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)" + }, + { + "package": "numpy", + "file": "numpy/tests/test_public_api.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L543: core_submodule = __import__(\nExec: L405: eval(module_name)" + }, + { + "package": "pillow", + "file": "PIL/Image.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:" + }, + { + "package": "protobuf", + "file": "protobuf-3.19.6-nspkg.pth", + "check": "Unusually large executable .pth (539 bytes)", + "severity": "HIGH", + "evidence": "1 import line(s) in 539-byte .pth file" + }, + { + "package": "pygments", + "file": "pygments/formatters/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L38: mod = __import__(module_name, None, None, ['__all__'])\nExec: L103: exec(f.read(), custom_namespace)" + }, + { + "package": "pygments", + "file": "pygments/lexers/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L45: mod = __import__(module_name, None, None, ['__all__'])\nExec: L154: exec(f.read(), custom_namespace)" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/torch/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')\nExec: L12: exec(f\"{n} = torch.{n}\")" + }, + { + "package": "scipy", + "file": "scipy/optimize/_optimize.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L4155: __import__(mod_name)\nExec: L323: def eval(x):" + }, + { + "package": "setuptools", + "file": "pkg_resources/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L423: __import__(moduleOrReq) | L1739: code = compile(source, script_filename, 'exec') | L1750: script_code = compile(script_text, script_filename, 'exec')\nExec: L1740: exec(code, namespace, namespace) | L1751: exec(script_code, nam" + }, + { + "package": "setuptools", + "file": "setuptools/_distutils/compilers/C/base.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):" + }, + { + "package": "setuptools", + "file": "setuptools/launch.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L31: code = compile(norm_script, script_name, 'exec')\nExec: L32: exec(code, namespace)" + }, + { + "package": "setuptools", + "file": "setuptools/tests/config/test_pyprojecttoml.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\"," + }, + { + "package": "setuptools", + "file": "setuptools/tests/test_editable_install.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)" + }, + { + "package": "setuptools", + "file": "setuptools/wheel.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L35: NAMESPACE_PACKAGE_INIT = \"__import__('pkg_resources').declare_namespace(__name__)\\n\"\nExec: L191: def eval(req, **env): | L212: (req for req in reqs if for_extra(req) and eval(req, extra=extra))," + }, + { + "package": "six", + "file": "six.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L87: __import__(name)\nExec: L740: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + }, + { + "package": "sympy", + "file": "sympy/external/importtools.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L145: mod = __import__(module, **import_kwargs) | L154: __import__(module + '.' + submod)\nExec: L21: return eval(debug_str)" + }, + { + "package": "sympy", + "file": "sympy/plotting/experimental_lambdify.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L249: namespace.update({'math': __import__('math')}) | L251: namespace.update({'cmath': __import__('cmath')}) | L254: namespace.update({'np': __import__('numpy')})\nExec: L268: exec(\"MYNEWLAMBDA = %s\" % eval_str, namespace)" + }, + { + "package": "sympy", + "file": "sympy/utilities/lambdify.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace)" + }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large (1918 KB) JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "" + }, + { + "package": "torch", + "file": "torch/_dynamo/bytecode_debugger.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)" + }, + { + "package": "torch", + "file": "torch/_functorch/_aot_autograd/subclass_codegen.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L342: code = compile(source, f\"<{artifact_name}>\", \"exec\")\nExec: L344: exec(code, globals_dict, local_dict)" + }, + { + "package": "torch", + "file": "torch/fx/experimental/rewriter.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)" + }, + { + "package": "torch", + "file": "torch/fx/graph_module.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L106: exec(compile(src, key, \"exec\"), globals)\nExec: L106: exec(compile(src, key, \"exec\"), globals)" + }, + { + "package": "torch", + "file": "torch/package/package_importer.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)" + }, + { + "package": "triton", + "file": "triton/runtime/interpreter.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1435: compiled_code = compile(transformed_ast, filename=self.filename, mode='exec')\nExec: L1441: exec(compiled_code, fn_globals, local_namespace)" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "exec/eval with payload hidden in a docstring/string", + "severity": "HIGH", + "evidence": "marshal/compile/obfuscation: L132: r\"|\\bbytearray\\s*\\(\\s*\\[.*?\\]\\s*\\)\" # bytearray([104,101,...]) | L135: r\"|\\bgetattr\\s*\\(\\s*__builtins__\" # getattr(__builtins__, ...)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_compiler_dynamic_exec.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L126: code = compile(source, f\"<{entry_point}>\", \"exec\")\nExec: L134: exec(code, sandbox)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_fused_forward_install.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L268: code = compile(src, fake_path, \"exec\")\nExec: L269: exec(code, namespace)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_upstream_pinned_symbols_trl_vllm.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L379: mod = __import__(modpath, fromlist=[\"Logprob\"])\nExec: L238: \"unsloth_zoo dispatch via `eval(f'trl.trainer.{trainer_file}.{name}')` breaks\"" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/compiler.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\ | L4292: f\"O^O/ {chr(92)}_/ {c" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/fused_losses/forward_install.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L274: code = compile(new_src, synthetic_path, \"exec\")\nExec: L275: exec(code, ns)" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/mlx/loader.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1739: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L141: mx.eval(model.parameters()) | L1543: model.eval() | L2126: mx.eval(model.parameters())" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/patching_utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L706: compile(new_source, '', 'exec')\nExec: L221: try: exec(_try_compile_argument) | L226: try: exec(_try_dynamo_argument) | L570: exec(\"from torch._dynamo.compiled_autograd import (\" + \", \".join(x for x in good_" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/saving_utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L3078: module = __import__('transformers', fromlist=[model_class_name])\nExec: L2960: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3006: exec(save_pretrained, globals(), functions)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_mlx_trainer_internals.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L430: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L408: def eval(self):" + }, + { + "package": "werkzeug", + "file": "werkzeug/routing/rules.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)" + } + ] +} diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index f0f6e4fddc..88defb9ea0 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -22,15 +22,64 @@ function Uninstall-UnslothStudio { param([string]$Path) if ([string]::IsNullOrWhiteSpace($Path)) { return } if (-not (Test-Path -LiteralPath $Path)) { return } - for ($attempt = 1; $attempt -le 3; $attempt++) { + for ($attempt = 1; $attempt -le 4; $attempt++) { try { Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop + } catch { + if ($attempt -lt 4) { Start-Sleep -Milliseconds 700; continue } + _Substep "could not remove: $Path ($($_.Exception.Message))" "Yellow" + return + } + # Remove-Item -Recurse can report success yet leave a transiently-locked + # child (e.g. unsloth.ico in Explorer's icon cache); verify + retry so we + # never falsely claim "removed" or orphan the dir. + if (-not (Test-Path -LiteralPath $Path)) { _Substep "removed: $Path" "Green" return - } catch { - if ($attempt -lt 3) { Start-Sleep -Milliseconds 700; continue } - _Substep "could not remove: $Path ($($_.Exception.Message))" "Yellow" } + if ($attempt -lt 4) { Start-Sleep -Milliseconds 700; continue } + _Substep "still present (files held open): $Path" "Yellow" + } + } + + # Remove the shared data dir, but keep unsloth.ico if a WSL shortcut still points + # at it (else that shortcut blanks); uninstall.sh drops it when WSL is removed. + function _RemoveDataDirKeepingWslIcon { + param( + [string]$DataDir, + # WSL-shortcut search dirs; default Start Menu + Desktop, overridable for tests. + [string[]]$ShortcutDirs = $null + ) + if ([string]::IsNullOrWhiteSpace($DataDir)) { return } + if (-not (Test-Path -LiteralPath $DataDir)) { return } + # $null = not passed (use defaults); test $null not truthiness so an explicit + # @() is honored (-not @() is $true). + if ($null -eq $ShortcutDirs) { + # Guard $env:APPDATA: it can be unset in service/CI Windows contexts, where + # an unguarded Join-Path emits a noisy parameter-binding error. + $ShortcutDirs = @() + if (-not [string]::IsNullOrWhiteSpace($env:APPDATA)) { + $ShortcutDirs += Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs" + } + try { + $desktop = [Environment]::GetFolderPath("Desktop") + if (-not [string]::IsNullOrWhiteSpace($desktop)) { $ShortcutDirs += $desktop } + } catch {} + } + $wslShortcuts = @() + foreach ($d in $ShortcutDirs) { + if ($d -and (Test-Path -LiteralPath $d)) { + $wslShortcuts += Get-ChildItem -LiteralPath $d -Filter "Unsloth Studio (WSL*.lnk" -ErrorAction SilentlyContinue + } + } + if (@($wslShortcuts).Count -eq 0) { + _RemovePath $DataDir + return + } + # A WSL shortcut survives: drop everything except its shared icon. + _Substep "keeping $(Join-Path $DataDir 'unsloth.ico') for the WSL shortcut" "Gray" + Get-ChildItem -LiteralPath $DataDir -Force -ErrorAction SilentlyContinue | ForEach-Object { + if ($_.Name -ne "unsloth.ico") { _RemovePath $_.FullName } } } @@ -287,6 +336,9 @@ function Uninstall-UnslothStudio { $defaultUnslothHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth" } else { $null } $defaultLlamaCpp = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "llama.cpp" } else { $null } $defaultCache = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".cache" } else { $null } + # Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in + # default mode. No-op in env/custom mode (nested under the custom root) and absent. + $defaultNode = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "node" } else { $null } # llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging, # sibling of the install dir). Usually pruned after activate, but an interrupted # build can leave a ".staging-XXXX" tree; removing it lets the empty-dir @@ -310,7 +362,7 @@ function Uninstall-UnslothStudio { _StopStudioProcesses -KnownRoots $knownRoots # Also stop anything holding a handle on the exact paths we delete (llama-server, # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused. - _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache)) + _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode)) # ── Remove custom-root install trees ── _Step "Removing data and install directories..." @@ -328,12 +380,18 @@ function Uninstall-UnslothStudio { # Default install dir (always at %USERPROFILE%\.unsloth\studio when present). if ($defaultStudioHome) { _RemovePath $defaultStudioHome } # Default data dir. - if ($defaultDataDir) { _RemovePath $defaultDataDir } + if ($defaultDataDir) { _RemoveDataDirKeepingWslIcon $defaultDataDir } # Default-mode shared llama.cpp build + cache (siblings of studio under # ~/.unsloth). No-op in env/custom mode and when absent. if ($defaultLlamaCpp) { _RemovePath $defaultLlamaCpp } if ($defaultCache) { _RemovePath $defaultCache } + # Isolated Node.js runtime (sibling of studio under ~/.unsloth). No-op in env/ + # custom mode (nested under the custom root, removed with it) and when absent. + if ($defaultNode) { _RemovePath $defaultNode } if ($defaultStaging) { _RemovePath $defaultStaging } + # llama.cpp install lock (serializes the shared build); a stray lock keeps + # ~/.unsloth from being pruned below. No-op in env/custom mode and when absent. + if ($defaultUnslothHome) { _RemovePath (Join-Path $defaultUnslothHome ".llama.cpp.install.lock") } # Drop ~/.unsloth itself, but ONLY if now empty -- never nuke unrelated content. if ($defaultUnslothHome -and (Test-Path -LiteralPath $defaultUnslothHome) -and -not (Get-ChildItem -LiteralPath $defaultUnslothHome -Force -ErrorAction SilentlyContinue)) { @@ -362,6 +420,11 @@ function Uninstall-UnslothStudio { } } catch { } + # Re-sweep: the first pass may have left unsloth.ico locked by Explorer/SMEH for + # the native shortcut; that handle is now freed. (A surviving WSL shortcut still + # keeps the icon -- see the helper.) + if ($defaultDataDir -and (Test-Path -LiteralPath $defaultDataDir)) { _RemoveDataDirKeepingWslIcon $defaultDataDir } + # ── Clean user PATH and registry backup ── _Step "Cleaning user PATH and registry..." try { diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index e97b28799b..31e851fcbb 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -217,10 +217,16 @@ _remove_path "$HOME/.unsloth/studio" # when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept. _remove_path "$HOME/.unsloth/llama.cpp" _remove_path "$HOME/.unsloth/.cache" +# Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in +# default mode. No-op in env/custom mode (nested under the custom root) and absent. +_remove_path "$HOME/.unsloth/node" # llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging). # Normally pruned after activate, but an interrupted build can leave it behind; # removing it lets the rmdir below succeed. No-op in env/custom mode and absent. _remove_path "$HOME/.unsloth/.staging" +# llama.cpp install lock (serializes the shared build); a stray one keeps ~/.unsloth +# from being pruned below. No-op in env/custom mode and when absent. +_remove_path "$HOME/.unsloth/.llama.cpp.install.lock" # ROCm-on-WSL helper artifacts (librocdxg build clone + smoke-test venv). No-op # where they don't exist; removing them lets the rmdir below succeed. _remove_path "$HOME/.unsloth/librocdxg" @@ -298,11 +304,50 @@ case "$_os" in Remove-Item -LiteralPath $_.FullName -Force -ErrorAction SilentlyContinue } catch { } } + } + # Keep the shared icon while any Unsloth shortcut still uses it (native + # install or another WSL distro); drop it only with the last one. + $iconInUse = $false; + foreach ($d in $dirs) { + if (-not $d -or -not (Test-Path -LiteralPath $d)) { continue } + if (Get-ChildItem -LiteralPath $d -Filter "Unsloth Studio*.lnk" -ErrorAction SilentlyContinue) { $iconInUse = $true; break } + } + # Guard LOCALAPPDATA: empty on a service/SYSTEM account makes + # Join-Path throw, aborting the icon cleanup (mirror uninstall.ps1). + if (-not [string]::IsNullOrWhiteSpace($env:LOCALAPPDATA)) { + $iconDir = Join-Path $env:LOCALAPPDATA "Unsloth Studio"; + $ico = Join-Path $iconDir "unsloth.ico"; + if ((-not $iconInUse) -and (Test-Path -LiteralPath $ico)) { Remove-Item -LiteralPath $ico -Force -ErrorAction SilentlyContinue } + if ((Test-Path -LiteralPath $iconDir) -and -not (Get-ChildItem -LiteralPath $iconDir -Force -ErrorAction SilentlyContinue)) { Remove-Item -LiteralPath $iconDir -Recurse -Force -ErrorAction SilentlyContinue } }' >/dev/null 2>&1 || true fi - # Fallback when powershell.exe can't run (interop disabled): remove the - # WSL .lnk files via drvfs. The "Unsloth Studio (WSL..." name is - # WSL-specific, so a native install's "Unsloth Studio.lnk" never matches. + # Remove $1's shared unsloth.ico only if no Unsloth shortcut (native install + # or another WSL distro) still uses it, then drop the dir if empty. Reciprocal + # of uninstall.ps1's _RemoveDataDirKeepingWslIcon (keeps the icon for a + # surviving WSL shortcut when the native side is removed). + _drop_shared_icon_if_unused() { + _du="$1" + _icodir="$_du/AppData/Local/Unsloth Studio" + _icon_in_use=0 + for _sd in \ + "$_du/Desktop" \ + "$_du/OneDrive/Desktop" \ + "$_du"/OneDrive*/Desktop \ + "$_du/AppData/Roaming/Microsoft/Windows/Start Menu/Programs"; do + [ -d "$_sd" ] || continue + for _any in "$_sd"/"Unsloth Studio"*.lnk; do + [ -e "$_any" ] && { _icon_in_use=1; break; } + done + [ "$_icon_in_use" = "1" ] && break + done + if [ "$_icon_in_use" = "0" ]; then + [ -f "$_icodir/unsloth.ico" ] && rm -f "$_icodir/unsloth.ico" 2>/dev/null || true + fi + [ -d "$_icodir" ] && rmdir "$_icodir" 2>/dev/null || true + } + # Fallback when powershell.exe can't run (interop disabled): remove WSL .lnk + # files via drvfs. The "Unsloth Studio (WSL..." name is WSL-specific, so a + # native install's "Unsloth Studio.lnk" never matches. if [ "$_ps_ran" = "0" ]; then for _drive in /mnt/c /mnt/d /mnt/e; do [ -d "$_drive/Users" ] || continue @@ -325,6 +370,8 @@ case "$_os" in done fi done + # Drop the shared icon only when no shortcut still needs it. + _drop_shared_icon_if_unused "$_udir" done done fi diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b8cb0573fe..b4c908b0cb 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -581,9 +581,16 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] ) # 3. TARGET-CHANGED (same scope+name resolves to a different import target) + # Only a *swap* is dangerous: a BEFORE target that is no longer reachable in + # AFTER means a reference was silently re-pointed. A pure superset growth + # (tbefore <= tafter) is the benign `import pkg.subA` + `import pkg.subB` + # case: both statements bind the same top-level name `pkg` to the same + # package object and only *add* submodule attributes (e.g. adding + # `import urllib.error` next to `import urllib.request`). Nothing the name + # resolved to before is lost, so no reference is re-pointed -- skip it. for key, tafter in b["target_by_use"].items(): tbefore = a["target_by_use"].get(key) - if tbefore and tbefore != tafter: + if tbefore and tbefore != tafter and (tbefore - tafter): findings.append( ( "BLOCKER", diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1b10b557e4..1c7a409bc1 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -277,6 +277,13 @@ "min_p": 0.01, "repetition_penalty": 1.0 }, + "minimax-m2.7": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 40, + "min_p": 0.01, + "repetition_penalty": 1.0 + }, "minimax-m2.5": { "temperature": 1.0, "top_p": 0.95, @@ -390,7 +397,7 @@ "deepseek-r1", "deepseek-v3", "deepseek-ocr", "glm-5", "glm-4", "nemotron", - "minimax-m2.5", "minimax", + "minimax-m2.7", "minimax-m2.5", "minimax", "gpt-oss", "granite-4", "kimi-k2", "kimi", "lfm2", "smollm", "olmo", "falcon", "ernie", "seed", "grok", "mimo" diff --git a/studio/backend/assets/configs/model_defaults/default.yaml b/studio/backend/assets/configs/model_defaults/default.yaml index 12566019b8..841e8ba166 100644 --- a/studio/backend/assets/configs/model_defaults/default.yaml +++ b/studio/backend/assets/configs/model_defaults/default.yaml @@ -2,7 +2,6 @@ # Used for models without specific configurations training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -48,7 +47,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 top_p: 0.95 top_k: -1 diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml index 52511c6eaf..734115ec41 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/ERNIE-4.5-21B-A3B-PT training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml index 524a723dc2..1032449e8c 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml index fa7bd8c1ea..c8e5f35841 100644 --- a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: tiiuae/Falcon-H1-0.5B-Instruct, unsloth/Falcon-H1-0.5B-Instruct training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml index 62836dc0cd..251409c29d 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -45,6 +44,5 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0 top_p: 0.9 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml index f97a842d2a..89b1d7f938 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml index 56f10cdc4f..e3292b5972 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml @@ -2,7 +2,6 @@ # Based on Gemma2_(9B)-Alpaca.ipynb (same defaults for larger models) training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml index a4acbe9262..98fe497912 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/gemma-2-2b-bnb-4bit, google/gemma-2-2b training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml index 455407abf8..bda5471643 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml index 2bcdf67c15..18392568bd 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml index 7c123da0b8..434ac41b46 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml index 492c42812e..5f0a7b26ce 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 2 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml index 23d00df752..dd5ae51ab0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 1024 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: audio_input: true inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml index bf5e111b7d..e53e163a04 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 2 num_epochs: 0 @@ -45,7 +44,6 @@ logging: audio_input: true inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml index c80506d9f5..ebe344e382 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-26B-A4B-it, unsloth/gemma-4-26B-A4B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml index 9e579be503..fb89a07133 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-26B-A4B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml index cec4ea95e1..4a089992ac 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-31B-it, unsloth/gemma-4-31B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml index 717cdd5e63..ae7524b7c6 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-31B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml index 43e3d78a23..10c1abd8a5 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E2B-it, unsloth/gemma-4-E2B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml index bd86cef751..fb5c1d9dea 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E2B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml index a8ef51836b..189e5dc6b2 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E4B-it, unsloth/gemma-4-E4B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml index 740cc99df5..aa51440b6a 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E4B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml index bd39e70a96..e2d67bcb0b 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml index 839e9a5b75..aa436117a1 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 1024 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml index 9557fc296f..3f2cb84a94 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -47,7 +46,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml index ce73c6a8ee..ab756fe764 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -47,7 +46,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml index d9a75c391d..1a7a91e56f 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml index 2bc3f6f871..7c7bb8dc3e 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Llama-3.2-1B-Instruct-unsloth-bnb-4bit, meta-llama/Llama-3.2-1B-Instruct, unsloth/Llama-3.2-1B-Instruct-bnb-4bit, RedHatAI/Llama-3.2-1B-Instruct-FP8, unsloth/Llama-3.2-1B-Instruct-FP8-Block, unsloth/Llama-3.2-1B-Instruct-FP8-Dynamic training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 5 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml index 82091c7d35..f73b0c09b6 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml index 5a014a63bf..ffefb29e24 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml index 885f7b47fd..cd986a6da1 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Meta-Llama-3.1-8B-bnb-4bit, unsloth/Meta-Llama-3.1-8B-unsloth-bnb-4bit, meta-llama/Meta-Llama-3.1-8B, unsloth/Meta-Llama-3.1-8B, unsloth/Meta-Llama-3.1-70B, meta-llama/Meta-Llama-3.1-70B, unsloth/Meta-Llama-3.1-405B-bnb-4bit, meta-llama/Meta-Llama-3.1-405B training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml index 1ff06cca6f..55dd3144c6 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit", "meta-llama/Meta-Llama-3.1-8B-Instruct", "unsloth/Meta-Llama-3.1-8B-Instruct","RedHatAI/Llama-3.1-8B-Instruct-FP8","unsloth/Llama-3.1-8B-Instruct-FP8-Block","unsloth/Llama-3.1-8B-Instruct-FP8-Dynamic" training: - trust_remote_code: false max_seq_length: 8192 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml index 95ee5ead5c..8c9cb07fb9 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/llama-3-8b-Instruct, meta-llama/Meta-Llama-3-8B-Instruct training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml index a05ac86f43..32441c5674 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/llama-3-8b, meta-llama/Meta-Llama-3-8B training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml index 1f473c3af1..6bba9c9633 100644 --- a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml +++ b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.2 top_p: 1.2 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml index 5a53bb52eb..f9833ce705 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 min_p: 0.01 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml index b84f7e1abb..0ba857cd40 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.15 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml index abdac62c0c..3476f2dd6d 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Mistral-Nemo-Base-2407", "mistralai/Mistral-Nemo-Base-2407", "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", "unsloth/Mistral-Nemo-Instruct-2407", "mistralai/Mistral-Nemo-Instruct-2407", training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml index 149f2a24f1..eda04d21f9 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Mistral-Small-Instruct-2409-bnb-4bit, mistralai/Mistral-Small-Instruct-2409 training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml index 3976cd0aa0..bcd0d20c8c 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml index 55d5dd289b..34a033e32f 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/mistral-7b-instruct-v0.3, mistralai/Mistral-7B-Instruct-v0.3 training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml index 5b24f5b581..98105eaf38 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml @@ -2,7 +2,6 @@ # Based on Mistral_v0.3_(7B)-Alpaca.ipynb # Also applies to: "unsloth/mistral-7b-v0.3", "mistralai/Mistral-7B-v0.3", training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml index 87b94ce67c..72b5b018e1 100644 --- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml @@ -6,7 +6,6 @@ audio_type: dac training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.4 top_k: 40 top_p: 0.9 diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml index 03748cd5fd..d20751b0c7 100644 --- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml +++ b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml @@ -6,7 +6,6 @@ audio_type: bicodec training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -48,7 +47,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.8 top_k: 50 top_p: 1.0 diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml index 5c1e180f8c..8a80282a2a 100644 --- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml +++ b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml @@ -5,7 +5,6 @@ audio_type: csm training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -45,6 +44,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml index 6d8be3656f..a973c2d4e4 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/GLM-4.7-Flash-unsloth-bnb-4bit, unsloth/GLM-4.7-Flash-bnb-4bit, THUDM/GLM-4.7-Flash training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 0.7 top_p: 0.8 top_k: 20 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml index 39a2fe0a5b..b0feafbd6e 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -39,7 +38,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.3 min_p: 0.15 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml index 663ce87d5f..2c44c91eab 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -47,7 +46,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 1.0 top_p: 1.0 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml index b7587bbd91..e1fbc08e4d 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml index cc5d130bfa..2abdfd8ac3 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml @@ -2,7 +2,6 @@ # Based on bert_classification.ipynb training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 1 num_epochs: 0 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml index 883761675f..5a3c4abb48 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml @@ -6,7 +6,6 @@ audio_type: snac training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -48,7 +47,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml index 35c850c71f..a6ce27620f 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 1 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index 9140878e0e..050774a8cd 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -6,7 +6,6 @@ audio_type: whisper audio_input: true training: - trust_remote_code: false eval_steps: 5 max_seq_length: 448 # num_epochs: 4 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml index 1088df7796..c574714d78 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Phi-3-medium-4k-instruct-bnb-4bit", "microsoft/Phi-3-medium-4k-instruct", training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml index 79812a74c4..e803c842b3 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Phi-3.5-mini-instruct-bnb-4bit", "microsoft/Phi-3.5-mini-instruct" training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml index aaa4feac45..4de3d9437d 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.8 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml index fa7b9c4e8b..bb75b3ce52 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml @@ -4,7 +4,6 @@ # MoE model - includes gate_up_proj for MoE layers training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -46,7 +45,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml index 3e64a6ca48..c305d328c2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2-7B-bnb-4bit, Qwen/Qwen2-7B training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml index 894751bed1..6cee3d0949 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml index 1d37cc9829..20ba81df2c 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-1.5B-Instruct-unsloth-bnb-4bit, Qwen/Qwen2.5-1.5B-Instruct, unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml index 99f3a66e23..9930786c24 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-7B-unsloth-bnb-4bit, Qwen/Qwen2.5-7B, unsloth/Qwen2.5-7B-bnb-4bit training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml index c48b943cba..775c7ce08f 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit, Qwen/Qwen2.5-Coder-1.5B-Instruct training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml index 830bfcf1cb..856db0c1b3 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml index db88c3b033..5900392547 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-Coder-7B-Instruct, Qwen/Qwen2.5-Coder-7B-Instruct training: - trust_remote_code: false max_seq_length: 32768 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml index cb9bcb104b..bd54b1d015 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml index 13f066a27d..9feb6dcaae 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 1024 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml index 87c042705b..a40eace253 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml index a8ecbb4365..c130771c32 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml index 485dd7a111..2fb3a95c30 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml @@ -4,7 +4,6 @@ # MoE model - includes gate_up_proj for MoE layers training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -46,7 +45,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml index 0de64d50ae..152f4ae06a 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml index dc5940d58c..94fe000708 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 top_p: 0.80 top_k: 20 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml index 6392ee0ae9..3c325485d2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_p: 0.95 top_k: 20 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml index ef52fad763..5b47c3bdd2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 top_p: 0.8 top_k: 20 diff --git a/studio/backend/assets/preview_page.html b/studio/backend/assets/preview_page.html new file mode 100644 index 0000000000..36483a824c --- /dev/null +++ b/studio/backend/assets/preview_page.html @@ -0,0 +1,403 @@ + + + + + + __TITLE__ - Unsloth + + + +

+ Unsloth__TITLE__ +
+
+
+

Chat with your model

+

Fine-tuned with Unsloth

+
+
+
+
+
+
+ + +
+
Served by Unsloth Studio
+
+
+ + + diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index fa5b985513..1f153699d7 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -4,9 +4,12 @@ """SQLite storage for auth data (user credentials + JWT secret).""" import hashlib +import hmac +import ipaddress import os import secrets import sqlite3 +import threading from datetime import datetime, timezone from typing import Optional, Tuple @@ -98,6 +101,14 @@ def get_connection() -> sqlite3.Connection: """Get a connection to the auth database, creating tables if needed.""" ensure_dir(DB_PATH.parent) conn = sqlite3.connect(DB_PATH) + # Keep the auth dir + DB private (they hold the JWT/identity secrets and + # password hashes); sqlite3.connect would otherwise create the DB 0644 under + # a 022 umask, letting another OS user read the identity secret and forge proofs. + for _path, _mode in ((DB_PATH.parent, 0o700), (DB_PATH, 0o600)): + try: + os.chmod(_path, _mode) + except OSError: + pass conn.row_factory = sqlite3.Row conn.execute( """ @@ -208,6 +219,114 @@ def _get_or_create_api_key_pbkdf2_salt() -> bytes: return salt +# Secret answering the /api/auth/identity challenge (HMAC(secret, nonce)). Lives +# in this same-user DB so a port squatter or remote/fake server can't forge a +# proof. Separate from the per-user JWT secret. +_IDENTITY_SECRET_DB_KEY = "studio_identity_secret" +_identity_secret_cache: Optional[bytes] = None + + +def get_or_create_identity_secret() -> bytes: + """Return the identity secret (hex 32-byte row in app_secrets), creating it once.""" + global _identity_secret_cache + if _identity_secret_cache is not None: + return _identity_secret_cache + + conn = get_connection() + try: + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + if row is None: + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + (_IDENTITY_SECRET_DB_KEY, secrets.token_hex(32)), + ) + conn.commit() + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + secret = bytes.fromhex(row["value"]) + finally: + conn.close() + + _identity_secret_cache = secret + return secret + + +def compute_identity_proof(nonce: bytes, host: str, port: int) -> str: + """HMAC-SHA256 proof that the caller holds this install's identity secret, + bound to the loopback address and port the connection landed on. A proof + relayed from a Studio on a different address/port (a squatter proxying to the + real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was + computed for that other endpoint and won't match the one the client dialed.""" + try: + host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms + except ValueError: + host = (host or "").lower() + msg = b"|".join([nonce, host.encode(), str(int(port)).encode()]) + return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest() + + +# Capability secret for public ``/p`` preview share links. HMAC(secret, ref) +# turns the deterministic preview ref into an unguessable bearer capability, so a +# guessed run/checkpoint name can't reach inference. Dedicated (not the per-user +# JWT secret) so rotating it revokes every shared link without touching logins. +_PREVIEW_LINK_SECRET_DB_KEY = "preview_link_secret" +_preview_link_secret_cache: Optional[bytes] = None + + +def get_or_create_preview_link_secret() -> bytes: + """Return the preview-link signing secret (hex 32-byte row in app_secrets), creating it once.""" + global _preview_link_secret_cache + if _preview_link_secret_cache is not None: + return _preview_link_secret_cache + + conn = get_connection() + try: + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_PREVIEW_LINK_SECRET_DB_KEY,), + ).fetchone() + if row is None: + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + (_PREVIEW_LINK_SECRET_DB_KEY, secrets.token_hex(32)), + ) + conn.commit() + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_PREVIEW_LINK_SECRET_DB_KEY,), + ).fetchone() + secret = bytes.fromhex(row["value"]) + finally: + conn.close() + + _preview_link_secret_cache = secret + return secret + + +def rotate_preview_link_secret() -> bytes: + """Rotate the preview-link secret, immediately revoking every outstanding ``/p`` share link.""" + global _preview_link_secret_cache + new_secret_hex = secrets.token_hex(32) + conn = get_connection() + try: + conn.execute( + "INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)", + (_PREVIEW_LINK_SECRET_DB_KEY, new_secret_hex), + ) + conn.commit() + finally: + conn.close() + + secret = bytes.fromhex(new_secret_hex) + _preview_link_secret_cache = secret + return secret + + _API_KEY_PBKDF2_ITERATIONS = 100_000 DESKTOP_SECRET_PREFIX = "desktop-" _DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash" @@ -236,6 +355,29 @@ def _pbkdf2_desktop_secret(raw_secret: str) -> str: return _pbkdf2_api_key(raw_secret) +# Memoize the deterministic raw-key -> PBKDF2-hash derivation so the 100k-round +# KDF runs once per key instead of on every authenticated request. Keyed by a +# salted HMAC of the key (not the key itself); revocation/expiry are still +# enforced by the SQLite read on every call, so a cache hit only skips the KDF. +# Only keys present in the DB are cached, so unknown-key spam can't grow it. +_api_key_hash_cache: dict[str, str] = {} +_API_KEY_HASH_CACHE_MAX = 4096 +_api_key_hash_cache_lock = threading.Lock() + + +def _api_key_cache_id(raw_key: str) -> str: + """Cache id for a raw key: salted HMAC-SHA256 (not the key itself).""" + return hmac.new( + _get_or_create_api_key_pbkdf2_salt(), raw_key.encode("utf-8"), hashlib.sha256 + ).hexdigest() + + +def _reset_api_key_hash_cache() -> None: + """Drop memoized derivations (tests / salt change).""" + with _api_key_hash_cache_lock: + _api_key_hash_cache.clear() + + def is_initialized() -> bool: """Check if auth is ready for login (at least one user exists in DB).""" conn = get_connection() @@ -704,7 +846,9 @@ def validate_api_key(raw_key: str) -> Optional[str]: Also updates ``last_used_at`` on success. """ - key_hash = _pbkdf2_api_key(raw_key) + cache_id = _api_key_cache_id(raw_key) + cached_hash = _api_key_hash_cache.get(cache_id) + key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key) conn = get_connection() try: cur = conn.execute( @@ -714,6 +858,12 @@ def validate_api_key(raw_key: str) -> Optional[str]: row = cur.fetchone() if row is None: return None + # Real key: memoize so later requests skip the KDF. Bounded; clear on overflow. + if cached_hash is None: + with _api_key_hash_cache_lock: + if len(_api_key_hash_cache) >= _API_KEY_HASH_CACHE_MAX: + _api_key_hash_cache.clear() + _api_key_hash_cache[cache_id] = key_hash if not row["is_active"]: return None if row["expires_at"] is not None: diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index e5dba69452..ef7bacba67 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -49,6 +49,16 @@ def _windows_hidden_kwargs() -> dict: return {"creationflags": flags} if flags else {} +def _lifetime_kwargs() -> dict: + """Bind cloudflared to the parent's lifetime (Linux PDEATHSIG). Lazy + + best-effort so this module still loads standalone (storage_roots-style).""" + try: + from utils.process_lifetime import child_popen_kwargs + return child_popen_kwargs() + except Exception: + return {} + + def _asset_name() -> Optional[Tuple[str, bool]]: """(release asset filename, is_tgz) for this OS/arch, or None if unsupported.""" system = platform.system().lower() @@ -233,6 +243,7 @@ class CloudflareTunnel: errors = "replace", bufsize = 1, **_windows_hidden_kwargs(), + **_lifetime_kwargs(), ) self._proc = proc threading.Thread( diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 32523e469e..c238c250bd 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -176,6 +176,9 @@ class JobManager: daemon = True, ) proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) self._mp_q = mp_q self._proc = proc diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index fbe847f9ce..ebb1d39dfb 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -14,6 +14,7 @@ from pathlib import Path from typing import Any from loggers import get_logger +from utils.node_runtime import resolve_node_executable from utils.paths import ensure_dir, oxc_validator_tmp_root logger = get_logger(__name__) @@ -231,6 +232,14 @@ def _run_oxc_batch( "code_shape": code_shape, "codes": code_values, } + # Resolve a usable Node (system or the isolated install, which is not on the + # user's PATH); a bare "node" would fail for isolated-Node users. + node_executable = resolve_node_executable() + if not node_executable: + return _fallback_results( + len(code_values), + "Node.js not found (install Node >= 20.19, or re-run Studio setup to provision it).", + ) try: tmp_dir = ensure_dir(oxc_validator_tmp_root()) env = child_env_without_native_path_secret() @@ -238,8 +247,13 @@ def _run_oxc_batch( env["TMPDIR"] = tmp_dir_str env["TMP"] = tmp_dir_str env["TEMP"] = tmp_dir_str + # Resolved node's dir first on the child PATH so it finds its own npm/npx. + node_bin_dir = os.path.dirname(node_executable) + if node_bin_dir: + env["PATH"] = node_bin_dir + os.pathsep + env.get("PATH", "") + env.pop("NODE_PATH", None) proc = subprocess.run( - ["node", str(_OXC_RUNNER_PATH)], + [node_executable, str(_OXC_RUNNER_PATH)], cwd = str(_OXC_TOOL_DIR), input = json.dumps(payload), text = True, diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index b28b61f088..a0959741a4 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -145,13 +145,19 @@ class ExportBackend: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """ Load a checkpoint for export. + ``hf_token`` authenticates the actual weight load for gated/private + checkpoints, matching the token the worker used for the security preflight + (otherwise a gated repo passes scanning then 401s at from_pretrained). + Returns: Tuple of (success: bool, message: str) """ + token = hf_token if hf_token and hf_token.strip() else None try: logger.info(f"Loading checkpoint: {checkpoint_path}") @@ -169,8 +175,10 @@ class ExportBackend: model_id = base_model or checkpoint_path - self._audio_type = detect_audio_type(model_id) - self.is_vision = not self._audio_type and is_vision_model(model_id) + # Token the type-detection probes too, else a gated multimodal base + # 404s here and falls through to the text loader. + self._audio_type = detect_audio_type(model_id, hf_token = token) + self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token) if self._audio_type == "csm": from unsloth import FastModel @@ -184,6 +192,7 @@ class ExportBackend: auto_model = CsmForConditionalGeneration, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "whisper": @@ -197,6 +206,7 @@ class ExportBackend: load_in_4bit = False, auto_model = WhisperForConditionalGeneration, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "snac": @@ -207,6 +217,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "bicodec": @@ -218,6 +229,7 @@ class ExportBackend: dtype = None if _IS_MLX else torch.float32, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "dac": @@ -228,6 +240,7 @@ class ExportBackend: max_seq_length = max_seq_length, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self.is_vision: @@ -238,6 +251,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) tokenizer = processor # vision: processor acts as tokenizer @@ -249,6 +263,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) if _IS_MLX: diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 20158d1891..478624b48e 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -63,6 +63,19 @@ class ExportOrchestrator: self._run_start_seq: int = 0 # True while an export op runs; SSE ends the stream 1s after this flips False. self._export_active: bool = False + # Set by cancel_export(); reset when a new load/export run starts. Lets the + # caller distinguish a user cancel from a genuine subprocess crash. + self._cancel_requested: bool = False + + # Last finished operation, so a client whose blocking POST was cut off by a + # Cloudflare tunnel timeout (524 at ~100s, while the op runs for minutes) can + # poll /api/export/status and still learn the real outcome. Guarded by + # _op_lock. `_op_seq` is a monotonic counter the client uses as a baseline to + # tell "my op finished" (seq grew) from a stale previous result. + self._op_lock = threading.Lock() + self._op_seq: int = 0 + self._active_op_kind: Optional[str] = None + self._last_op: Optional[Dict[str, Any]] = None atexit.register(self._cleanup) logger.info("ExportOrchestrator initialized (subprocess mode)") @@ -119,6 +132,72 @@ class ExportOrchestrator: """True while an export / load / cleanup command is running.""" return self._export_active + def was_cancelled(self) -> bool: + """True if the in-flight (or most recent) run was cancelled by the user.""" + return self._cancel_requested + + def _record_op_finished(self, success: bool, message: str, output_path: Optional[str]) -> None: + """Snapshot the just-finished op so status pollers can recover its outcome. + + Called from each op's ``finally`` (with ``_active_op_kind`` still set) BEFORE + ``_export_active`` is cleared, so a status read that observes the op as + inactive is guaranteed to also see this matching result. + """ + with self._op_lock: + self._op_seq += 1 + status = "cancelled" if self._cancel_requested else ("success" if success else "error") + self._last_op = { + "seq": self._op_seq, + "kind": self._active_op_kind, + "status": status, + "output_path": output_path if success else None, + "error": None if success else (message or None), + } + + def get_last_op(self) -> Optional[Dict[str, Any]]: + """Return the last finished op record (or None), for status recovery.""" + with self._op_lock: + return dict(self._last_op) if self._last_op is not None else None + + def get_active_op_kind(self) -> Optional[str]: + """Return the kind of the currently running op (or None when idle).""" + return self._active_op_kind + + def cancel_export(self) -> bool: + """Terminate the in-flight export subprocess immediately. + + An export op holds ``self._lock`` for its whole duration (blocked in + ``_wait_response``), so we deliberately do NOT take the lock here -- we + kill the worker process directly, which unblocks that wait and makes the + in-flight op return a failure the caller surfaces as "cancelled". + + Only the export subprocess is touched; training and inference run in + their own subprocesses and are left untouched. + + Returns True if a live subprocess was terminated, False if none ran. + """ + self._cancel_requested = True + proc = self._proc + if proc is None or not proc.is_alive(): + return False + logger.info( + "Export cancel requested: terminating export subprocess (pid=%s)", + proc.pid, + ) + try: + proc.terminate() + proc.join(timeout = 5) + except Exception: + pass + if proc.is_alive(): + logger.warning("Export subprocess survived terminate, killing") + try: + proc.kill() + proc.join(timeout = 3) + except Exception: + pass + return True + # ------------------------------------------------------------------ # Subprocess lifecycle # ------------------------------------------------------------------ @@ -147,6 +226,9 @@ class ExportOrchestrator: daemon = True, ) self._proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Export subprocess started (pid=%s)", self._proc.pid) def _shutdown_subprocess(self, timeout: float = 10.0) -> None: @@ -301,7 +383,9 @@ class ExportOrchestrator: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, hf_token: Optional[str] = None, + subject: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -312,13 +396,18 @@ class ExportOrchestrator: "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, + "subject": subject, "hf_token": hf_token, } with self._lock: # Fresh log buffer so the UI sees only this run's output. self.clear_logs() + self._cancel_requested = False + self._active_op_kind = "load_checkpoint" self._export_active = True + op_success, op_message = False, "" try: # Always kill any existing subprocess and spawn fresh. if self._ensure_subprocess_alive(): @@ -336,6 +425,7 @@ class ExportOrchestrator: self.current_checkpoint = None self.is_vision = False self.is_peft = False + op_success, op_message = False, str(exc) return False, str(exc) if resp.get("success"): @@ -343,15 +433,19 @@ class ExportOrchestrator: self.is_vision = resp.get("is_vision", False) self.is_peft = resp.get("is_peft", False) logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path) - return True, resp.get("message", "Loaded successfully") + op_success, op_message = True, resp.get("message", "Loaded successfully") + return True, op_message else: error = resp.get("message", "Failed to load checkpoint") logger.error("Failed to load checkpoint: %s", error) self.current_checkpoint = None self.is_vision = False self.is_peft = False + op_success, op_message = False, error return False, error finally: + self._record_op_finished(op_success, op_message, None) + self._active_op_kind = None self._export_active = False def export_merged_model( @@ -453,7 +547,10 @@ class ExportOrchestrator: ) self.clear_logs() + self._cancel_requested = False + self._active_op_kind = f"export_{export_type}" self._export_active = True + op_success, op_message, op_output_path = False, "", None try: cmd = {"type": "export", "export_type": export_type, **params} try: @@ -462,14 +559,16 @@ class ExportOrchestrator: f"export_{export_type}_done", timeout = 3600, # GGUF for 30B+ models can take 30+ min ) - return ( - resp.get("success", False), - resp.get("message", ""), - resp.get("output_path"), - ) + op_success = resp.get("success", False) + op_message = resp.get("message", "") + op_output_path = resp.get("output_path") + return op_success, op_message, op_output_path except RuntimeError as exc: + op_success, op_message = False, str(exc) return False, str(exc), None finally: + self._record_op_finished(op_success, op_message, op_output_path) + self._active_op_kind = None self._export_active = False def cleanup_memory(self) -> bool: @@ -481,7 +580,9 @@ class ExportOrchestrator: self.is_peft = False return True + self._active_op_kind = "cleanup" self._export_active = True + success = False try: try: self._send_cmd({"type": "cleanup"}) @@ -498,6 +599,8 @@ class ExportOrchestrator: self.is_peft = False return success finally: + self._record_op_finished(success, "", None) + self._active_op_kind = None self._export_active = False def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]: diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index fb2a893014..fdaa306e10 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -159,7 +159,7 @@ def _setup_log_capture(resp_queue: Any) -> None: t_err.start() -def _activate_transformers_version(model_name: str) -> None: +def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None: """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on sys.path for utils imports. backend_path = str(Path(__file__).resolve().parent.parent.parent) @@ -168,7 +168,7 @@ def _activate_transformers_version(model_name: str) -> None: from utils.transformers_version import activate_transformers_for_subprocess - activate_transformers_for_subprocess(model_name) + activate_transformers_for_subprocess(model_name, hf_token) def _send_response(resp_queue: Any, response: dict) -> None: @@ -188,10 +188,16 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: # Auto-enable trust_remote_code for NemotronH/Nano models. if not trust_remote_code: + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _cp_lower = checkpoint_path.lower() - if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( - _cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/") + if ( + any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) + and (_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")) + # Genuine first-party Hub repo only (not a local/spoof name starting + # with "unsloth/"); authenticated so private repos resolve. + and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token")) ): trust_remote_code = True logger.info( @@ -199,6 +205,82 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: checkpoint_path, ) + # Malware gate: a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) every + # load. Local checkpoints have no Hub scan and are skipped in the helper; a + # LoRA merges its base weights, so gate that repo too. + from utils.security import evaluate_file_security, security_load_subdirs + + malware_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + _hf_token = cmd.get("hf_token") + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token) + ) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + }, + ) + return + + # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless + # pinned-approved. A LoRA merges its base model, whose code runs, so gate it too. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a local or remote adapter's base so its base repo is gated too. + base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = cmd.get("hf_token"), + trust_remote_code = True, + approved_fingerprint = cmd.get("approved_remote_code_fingerprint"), + subject = cmd.get("subject"), + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Checkpoint '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review and " + f"approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + }, + ) + return + try: _send_response( resp_queue, @@ -214,6 +296,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: max_seq_length = max_seq_length, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + hf_token = cmd.get("hf_token"), ) _send_response( @@ -377,7 +460,7 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None # ── 1. Activate correct transformers version BEFORE any ML imports ── try: - _activate_transformers_version(checkpoint_path) + _activate_transformers_version(checkpoint_path, config.get("hf_token") or None) except Exception as exc: _send_response( resp_queue, @@ -424,6 +507,11 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None if backend_path not in sys.path: sys.path.insert(0, backend_path) + # Recover from any namespace-package shadow before importing Unsloth. + from core.import_guards import ensure_real_packages + + ensure_real_packages("unsloth_zoo", "unsloth") + from core.export.export import ExportBackend import transformers diff --git a/studio/backend/core/import_guards.py b/studio/backend/core/import_guards.py new file mode 100644 index 0000000000..5b85a96cd2 --- /dev/null +++ b/studio/backend/core/import_guards.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Recover `unsloth`/`unsloth_zoo` from a namespace-package shadow. Stdlib-only.""" + +from __future__ import annotations + +import os +import sys + + +def ensure_real_packages(*names: str) -> None: + """Drop sys.path entries where a bare `/` dir (no __init__.py) shadows + the installed package as a namespace, import the real packages, restore + sys.path. No-op without a shadow. Pass dependency-first (e.g. "unsloth_zoo", + "unsloth"); imports run dependency-last.""" + import importlib + import importlib.util + + bad: set = set() + shadowed: list = [] + for name in names: + try: + spec = importlib.util.find_spec(name) + except (ImportError, ValueError, AttributeError): + spec = None + # real package -> spec.origin is its __init__; namespace shadow -> None/"namespace" + if spec is None or spec.origin not in (None, "namespace"): + continue + dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])} + if not dirs: + continue + shadowed.append(name) + for entry in sys.path: + pkg = os.path.join(entry or os.getcwd(), name) + if os.path.realpath(pkg) in dirs and not os.path.isfile( + os.path.join(pkg, "__init__.py") + ): + bad.add(entry) + if not bad: + return + saved = list(sys.path) + sys.path[:] = [e for e in sys.path if e not in bad] + for name in shadowed: + for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]: + del sys.modules[cached] + try: + importlib.invalidate_caches() + # import unsloth before unsloth_zoo: unsloth.__init__ runs GPU/bnb fixes zoo relies on + for name in reversed(names): + importlib.import_module(name) + finally: + sys.path[:] = saved diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py new file mode 100644 index 0000000000..f76a38576f --- /dev/null +++ b/studio/backend/core/inference/api_monitor.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Small in-memory monitor for OpenAI-compatible API traffic.""" + +from __future__ import annotations + +import threading +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Optional + + +_MAX_ENTRIES = 50 +_MAX_PROMPT_CHARS = 12000 +_MAX_REPLY_CHARS = 12000 +_PREVIEW_CHARS = 360 + + +def _trim(text: Optional[str], limit: int) -> str: + if not text: + return "" + if len(text) <= limit: + return text + # Guard against limit < 3 (slice would underflow). + if limit <= 3: + return "..."[:limit] + return text[: limit - 3] + "..." + + +@dataclass +class ApiMonitorEntry: + id: str + endpoint: str + method: str + model: str + prompt: str + status: str + started_at: float + updated_at: float + subject: Optional[str] = None + # Monotonic anchors so duration math survives wall-clock steps (NTP). + started_monotonic: float = 0.0 + finished_monotonic: Optional[float] = None + reply: str = "" + finished_at: Optional[float] = None + context_length: Optional[int] = None + prompt_tokens: Optional[int] = None + completion_tokens: Optional[int] = None + total_tokens: Optional[int] = None + total_tokens_authoritative: bool = False + error: Optional[str] = None + + def snapshot(self, *, include_details: bool = True) -> dict[str, Any]: + duration_ms = None + if self.finished_monotonic is not None: + duration_ms = max( + 0, + int((self.finished_monotonic - self.started_monotonic) * 1000), + ) + elif self.finished_at is not None: + duration_ms = max(0, int((self.finished_at - self.started_at) * 1000)) + context_usage = None + if self.total_tokens is not None and self.context_length: + context_usage = min(1.0, max(0.0, self.total_tokens / self.context_length)) + payload = { + "id": self.id, + "endpoint": self.endpoint, + "method": self.method, + "model": self.model, + "prompt_preview": _trim(self.prompt, _PREVIEW_CHARS), + "reply_preview": _trim(self.reply, _PREVIEW_CHARS), + "prompt_truncated": len(self.prompt) > _PREVIEW_CHARS, + "reply_truncated": len(self.reply) > _PREVIEW_CHARS, + "status": self.status, + "started_at": self.started_at, + "updated_at": self.updated_at, + "finished_at": self.finished_at, + "duration_ms": duration_ms, + "context_length": self.context_length, + "context_usage": context_usage, + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, + "total_tokens": self.total_tokens, + "error": self.error, + } + if include_details: + payload["prompt"] = self.prompt + payload["reply"] = self.reply + return payload + + +class ApiMonitor: + def __init__(self, max_entries: int = _MAX_ENTRIES): + self._entries: deque[ApiMonitorEntry] = deque() + self._max_entries = max(0, max_entries) + self._lock = threading.Lock() + + def start( + self, + *, + endpoint: str, + method: str, + model: str, + prompt: str, + context_length: Optional[int] = None, + subject: Optional[str] = None, + ) -> str: + now = time.time() + entry = ApiMonitorEntry( + id = f"apireq_{uuid.uuid4().hex[:12]}", + endpoint = endpoint, + method = method, + model = model or "default", + prompt = _trim(prompt, _MAX_PROMPT_CHARS), + status = "running", + started_at = now, + updated_at = now, + subject = subject, + started_monotonic = time.monotonic(), + context_length = context_length, + ) + with self._lock: + self._entries.appendleft(entry) + self._trim_terminal_locked() + return entry.id + + def append_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id or not text: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + # Preview is capped: once the "..." marker is present the head is + # frozen, so skip the per-chunk re-concat (avoids O(n^2) on long + # generations). A reply that landed exactly on the cap has no marker + # yet, so let one more append record the truncation before freezing. + if len(entry.reply) >= _MAX_REPLY_CHARS: + if not entry.reply.endswith("..."): + entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + return + entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + entry.reply = _trim(text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_usage( + self, + entry_id: Optional[str], + *, + prompt_tokens: Optional[int] = None, + completion_tokens: Optional[int] = None, + total_tokens: Optional[int] = None, + context_length: Optional[int] = None, + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if prompt_tokens is not None: + entry.prompt_tokens = prompt_tokens + if completion_tokens is not None: + entry.completion_tokens = completion_tokens + if total_tokens is not None: + entry.total_tokens = total_tokens + entry.total_tokens_authoritative = True + elif not entry.total_tokens_authoritative and ( + prompt_tokens is not None or completion_tokens is not None + ): + # Derive only when no authoritative total has been set; + # a later partial chunk must not clobber a provider total. + entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0) + if context_length is not None: + entry.context_length = context_length + entry.updated_at = time.time() + + def finish( + self, + entry_id: Optional[str], + status: str = "completed", + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + # Idempotent: second call (e.g. [DONE] after the finally block + # already ran) must not move finished_*. + if entry.finished_at is not None: + return + now = time.time() + entry.status = status + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() + + def fail(self, entry_id: Optional[str], error: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if entry.finished_at is not None: + # Already terminal; refresh error text only. + if error: + entry.error = _trim(error, 1000) + return + now = time.time() + entry.status = "error" + entry.error = _trim(error, 1000) + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() + + def snapshot( + self, + *, + include_details: bool = True, + subject: Optional[str] = None, + ) -> list[dict[str, Any]]: + with self._lock: + return [ + entry.snapshot(include_details = include_details) + for entry in self._entries + if subject is None or entry.subject == subject + ] + + def get( + self, + entry_id: str, + *, + subject: Optional[str] = None, + ) -> Optional[dict[str, Any]]: + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return None + if subject is not None and entry.subject != subject: + return None + return entry.snapshot(include_details = True) + + def active_count(self, *, subject: Optional[str] = None) -> int: + with self._lock: + return sum( + 1 + for entry in self._entries + if entry.status == "running" and (subject is None or entry.subject == subject) + ) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + + def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]: + for entry in self._entries: + if entry.id == entry_id: + return entry + return None + + def _trim_terminal_locked(self) -> None: + terminal_seen = 0 + kept: deque[ApiMonitorEntry] = deque() + for entry in self._entries: + if entry.status == "running": + kept.append(entry) + continue + if terminal_seen < self._max_entries: + kept.append(entry) + terminal_seen += 1 + self._entries = kept + + +api_monitor = ApiMonitor() diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 2b9517692f..4dca4db768 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -26,6 +26,7 @@ from utils.hardware import ( ) from core.inference.audio_codecs import AudioCodecManager from core.inference.runtime_context import runtime_context_length +from core.inference.message_content import content_to_text from io import StringIO import structlog from loggers import get_logger @@ -1018,7 +1019,7 @@ class InferenceBackend: user_message = "" if messages and messages[-1]["role"] == "user": import re - user_message = messages[-1]["content"] + user_message = content_to_text(messages[-1]["content"]) user_message = re.sub(r"]*>", "", user_message).strip() if not user_message: @@ -1181,7 +1182,7 @@ class InferenceBackend: if messages: for msg in reversed(messages): if msg["role"] == "user" and msg.get("content"): - user_text = msg["content"] + user_text = content_to_text(msg["content"]) break # ASR-specific default system prompt if none set @@ -1713,7 +1714,7 @@ class InferenceBackend: for msg in messages: role = msg.get("role", "") - content = msg.get("content", "") + content = content_to_text(msg.get("content", "")) if role in ["system", "user", "assistant"] and content.strip(): if role == last_role: @@ -1801,7 +1802,7 @@ class InferenceBackend: for msg in messages: role = msg["role"] - content = msg["content"] + content = content_to_text(msg["content"]) formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>" formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n" @@ -1817,14 +1818,14 @@ class InferenceBackend: for msg in messages: if msg["role"] == "system": - system_msg = msg["content"] + system_msg = content_to_text(msg["content"]) else: conversation.append(msg) i = 0 while i < len(conversation): if conversation[i]["role"] == "user": - user_content = conversation[i]["content"] + user_content = content_to_text(conversation[i]["content"]) if system_msg and i == 0: user_content = f"{system_msg}\n\n{user_content}" @@ -1832,7 +1833,7 @@ class InferenceBackend: formatted += f"[INST] {user_content} [/INST]" if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant": - formatted += f" {conversation[i + 1]['content']}" + formatted += f" {content_to_text(conversation[i + 1]['content'])}" i += 2 else: formatted += " " @@ -1848,7 +1849,7 @@ class InferenceBackend: for msg in messages: role = msg["role"] - content = msg["content"] + content = content_to_text(msg["content"]) formatted += f"<|im_start|>{role}\n{content}<|im_end|>\n" formatted += "<|im_start|>assistant\n" @@ -1860,16 +1861,17 @@ class InferenceBackend: system_msg = None for msg in messages: + content = content_to_text(msg["content"]) if msg["role"] == "system": - system_msg = msg["content"] + system_msg = content elif msg["role"] == "user": if system_msg: - formatted += f"### Instruction:\n{system_msg}\n\n### Input:\n{msg['content']}\n\n### Response:\n" + formatted += f"### Instruction:\n{system_msg}\n\n### Input:\n{content}\n\n### Response:\n" system_msg = None else: - formatted += f"### Human:\n{msg['content']}\n\n### Assistant:\n" + formatted += f"### Human:\n{content}\n\n### Assistant:\n" elif msg["role"] == "assistant": - formatted += f"{msg['content']}\n\n" + formatted += f"{content}\n\n" return formatted @@ -1879,7 +1881,7 @@ class InferenceBackend: for msg in messages: role = msg["role"].title() - content = msg["content"] + content = content_to_text(msg["content"]) formatted += f"{role}: {content}\n" formatted += "Assistant: " diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 57cc97f3b5..152a3f19b2 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -13,7 +13,6 @@ import json import os import re import struct -import structlog from loggers import get_logger import shutil import signal @@ -23,37 +22,32 @@ import sys import threading import time from pathlib import Path -from typing import Callable, Generator, Iterable, List, Optional +from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union import httpx from core.inference.llama_server_args import ( + _effective_tensor_parallel, + _tensor_parallel_matches_loaded, extra_args_disable_mmproj, parse_cache_override, + parse_cache_override_per_axis, parse_ctx_override, parse_split_mode_override, - resolve_cache_type_kv, resolve_requested_ctx, - resolve_tensor_parallel, strip_shadowing_flags, strip_split_mode_only, ) from core.tool_healing import ( - _TC_END_TAG_RE, - _TC_FUNC_CLOSE_RE, - _TC_FUNC_START_RE, - _TC_JSON_START_RE, - _TC_PARAM_CLOSE_RE, - _TC_PARAM_START_RE, _TOOL_ALL_PATS, - _TOOL_CLOSED_PATS, - parse_tool_calls_from_text, strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret +from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs from core.inference.tool_call_parser import ( RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, @@ -75,6 +69,112 @@ from state.tool_approvals import ( logger = get_logger(__name__) +class LlamaServerNotFoundError(RuntimeError): + """GGUF model needs the llama.cpp runtime but no llama-server is installed. + Subclasses RuntimeError so existing handlers still catch it.""" + + +# Shared so the from_identifier preflight and the load-time raise stay in sync. +LLAMA_SERVER_NOT_FOUND_DETAIL = ( + "This is a GGUF model, but the llama.cpp runtime (llama-server) is not " + "installed. Run `unsloth studio setup` to download the prebuilt runtime, " + "then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)" +) + + +# llama-server can serve HTTP 200 while running a model entirely on CPU when a +# GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so +# Studio can warn. Priority: explicit "offloaded N/M layers to GPU" counts +# (authoritative), then GPU "model buffer size" lines (host-pinned _Host +# excluded), then the "device_info:" device table (disconfirm only). +_GPU_OFFLOAD_MARKERS = ( + "CUDA", + "ROCm", + "ROCM", + "HIP", + "Metal", + "Vulkan", + "OpenCL", + "SYCL", + "MUSA", + "CANN", +) +_OFFLOADED_LAYERS_RE = re.compile( + r"offloaded\s+(\d+)\s*/\s*(\d+)\s+layers?\s+to\s+gpu", re.IGNORECASE +) +_DEVICE_ROW_RE = re.compile( + r"-\s*(CUDA|ROCm|ROCM|HIP|Metal|Vulkan|SYCL|OpenCL|MUSA|CANN|CPU)\w*\s*:", + re.IGNORECASE, +) +_GPU_DEVICE_PREFIXES = ( + "cuda", + "rocm", + "hip", + "metal", + "vulkan", + "sycl", + "opencl", + "musa", + "cann", +) + + +def classify_gpu_offload_lines(lines: "list[str]") -> Optional[bool]: + """True if the model landed on a GPU, False if it stayed on CPU despite GPU + intent, None when the log has no usable signal.""" + # Counted offload is authoritative, keyed on the model with the most layers. + # A separate MTP/draft model logs its own (much smaller) "offloaded N/M" + # line, so decide on the largest-M line: a drafter that fits on GPU must not + # mask a main model running on CPU. N>0 on that model is True, 0 is False. + max_total = -1 + offloaded_at_max = 0 + for line in lines: + match = _OFFLOADED_LAYERS_RE.search(line) + if not match: + continue + offloaded, total = int(match.group(1)), int(match.group(2)) + if total > max_total or (total == max_total and offloaded > offloaded_at_max): + max_total, offloaded_at_max = total, offloaded + if max_total >= 0: + return offloaded_at_max > 0 + + # GPU marker on a *model* buffer; _Host buffers are CPU-pinned, not offload. + # Buffer lines are authoritative: present but none on a GPU means CPU-only, + # so do not let the device table below override that. + saw_model_buffer = False + for line in lines: + if "model buffer size" not in line: + continue + saw_model_buffer = True + if "_Host" not in line and any(m in line for m in _GPU_OFFLOAD_MARKERS): + return True + if saw_model_buffer: + return False + + # device_info: lists *available* devices (printed whenever a GPU backend is + # visible), not where the model loaded, so it can only disconfirm: an + # all-CPU table means no usable GPU. A visible GPU device is not proof the + # model used it, so it does not return True. Rows after the header only. + after_header = False + saw_device_row = False + saw_gpu_device = False + for line in lines: + if "device_info:" in line: + after_header = True + continue + if not after_header: + continue + match = _DEVICE_ROW_RE.search(line) + if not match: + continue + saw_device_row = True + if match.group(1).lower().startswith(_GPU_DEVICE_PREFIXES): + saw_gpu_device = True + if saw_device_row and not saw_gpu_device: + return False + return None + + def _wsl_system_rocm_lib_dirs() -> "list[str]": """System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL. @@ -126,6 +226,10 @@ _MAX_REPROMPTS = 1 # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. _DEFAULT_MAX_TOKENS_FLOOR = 32768 _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min + +# Only large streamed tool payloads get an early provisional card; render_html +# is exempt because it needs immediate artifact feedback. +_PROVISIONAL_ARGS_MIN_CHARS = 256 _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min _REPROMPT_MAX_CHARS = 2000 _FORCED_REPEAT_PLAN_SIGNAL = re.compile( @@ -154,8 +258,8 @@ def _should_suppress_forced_no_tool_output(text: str) -> bool: # ── Pre-compiled patterns for GGUF shard detection ─────────── -_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$") -_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") +_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", re.IGNORECASE) +_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$", re.IGNORECASE) # ── Sliding-window-pattern resolver ─────────────────────────── @@ -270,6 +374,13 @@ def _period_from_layer_types(layer_types: list) -> Optional[int]: return None +def _swa_entry_from_layer_types(lt) -> Optional[object]: + """Period int, or per-layer bool mask, from a transformers ``layer_types`` list.""" + if isinstance(lt, list) and lt: + return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] + return None + + def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: try: from huggingface_hub import hf_hub_download @@ -283,10 +394,7 @@ def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: period = src.get("sliding_window_pattern") if isinstance(period, int) and period > 0: return period - lt = src.get("layer_types") - if isinstance(lt, list) and lt: - return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] - return None + return _swa_entry_from_layer_types(src.get("layer_types")) def _arch_aliases(arch: str) -> tuple: @@ -303,10 +411,7 @@ def _swa_entry_from_config_obj(cfg) -> Optional[object]: period = getattr(src, "sliding_window_pattern", None) if isinstance(period, int) and period > 0: return period - lt = getattr(src, "layer_types", None) - if isinstance(lt, list) and lt: - return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] - return None + return _swa_entry_from_layer_types(getattr(src, "layer_types", None)) _SWA_PATTERN_SOURCE_RE = re.compile(r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)") @@ -443,6 +548,27 @@ _TOOL_TEMPLATE_MARKERS = ( ) +# Canonical reasoning_effort levels, weakest -> strongest. Used to read the +# discrete set a template branches on (e.g. GLM-5.2 uses 'high' | 'max') so we +# only ever offer levels the template actually understands. +_REASONING_EFFORT_SCALE = ("minimal", "low", "medium", "high", "max") + + +def _extract_reasoning_effort_levels(chat_template: str) -> list: + """Return the reasoning_effort levels a template references, in canonical + (weakest -> strongest) order. + + Looks for the quoted literals (e.g. ``'high'`` / ``"max"``) the template + compares ``reasoning_effort`` against, so we surface exactly the levels it + branches on and nothing else. + """ + return [ + level + for level in _REASONING_EFFORT_SCALE + if f"'{level}'" in chat_template or f'"{level}"' in chat_template + ] + + def detect_reasoning_flags( chat_template: Optional[str], model_identifier: Optional[str] = None, @@ -451,17 +577,20 @@ def detect_reasoning_flags( ) -> dict: """Classify a chat template's reasoning and tool-calling capabilities. - Returns the same five keys as the GGUF sniffer: ``supports_reasoning``, - ``reasoning_style`` (``"enable_thinking"`` | ``"reasoning_effort"``), - ``reasoning_always_on``, ``supports_preserve_thinking``, - ``supports_tools``. Used by both the llama-server backend at load time - and the safetensors/transformers paths in ``routes/inference`` so they - agree on what the frontend sees. + Returns the same six keys as the GGUF sniffer: ``supports_reasoning``, + ``reasoning_style`` (``"enable_thinking"`` | ``"reasoning_effort"`` | + ``"enable_thinking_effort"``), ``reasoning_always_on``, + ``reasoning_effort_levels``, ``supports_preserve_thinking``, + ``supports_tools``. A falsy ``chat_template`` yields the all-default dict. + Used by both the llama-server backend at load time and the + safetensors/transformers paths in ``routes/inference`` so they agree on + what the frontend sees. """ flags = { "supports_reasoning": False, "reasoning_style": "enable_thinking", "reasoning_always_on": False, + "reasoning_effort_levels": [], "supports_preserve_thinking": False, "supports_tools": False, } @@ -470,7 +599,25 @@ def detect_reasoning_flags( tpl = chat_template prefix = f"{log_source}: " if log_source else "" - if "enable_thinking" in tpl: + effort_levels = ( + _extract_reasoning_effort_levels(tpl) + if ("reasoning_effort" in tpl and "enable_thinking" in tpl) + else [] + ) + if effort_levels: + # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort + # level among a discrete set (e.g. 'high' | 'max'). Distinct from + # gpt-oss (reasoning_effort only, no on/off gate) and Qwen + # (enable_thinking only). Disabling is enable_thinking=false; the levels + # are the quoted effort literals the template actually branches on. + flags["supports_reasoning"] = True + flags["reasoning_style"] = "enable_thinking_effort" + flags["reasoning_effort_levels"] = effort_levels + logger.info( + f"{prefix}model supports reasoning " + f"(enable_thinking + reasoning_effort: {effort_levels})" + ) + elif "enable_thinking" in tpl: flags["supports_reasoning"] = True flags["reasoning_style"] = "enable_thinking" logger.info(f"{prefix}model supports reasoning (enable_thinking)") @@ -508,11 +655,32 @@ def detect_reasoning_flags( return flags +# Gemma 4 ships MTP as a separate drafter (no "-mtp" in the name). Gemma 3n +# ships no drafter, so it is excluded -- it takes the normal non-MTP path. +_GEMMA_MTP_FAMILY_RE = re.compile(r"gemma[-_]?4[-_]", re.IGNORECASE) + + +def _is_gemma_mtp_family(name: Optional[str]) -> bool: + """Match Gemma 4 by name.""" + return bool(name) and bool(_GEMMA_MTP_FAMILY_RE.search(name)) + + +def _is_gemma_mtp_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: + """Match Gemma 4 by id or GGUF filename.""" + return _is_gemma_mtp_family(model_identifier) or _is_gemma_mtp_family( + Path(gguf_path).name if gguf_path else None + ) + + def _is_mtp_model_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: """Name-based MTP detector. Fallback for the metadata signal.""" for cand in (model_identifier, Path(gguf_path).name if gguf_path else None): if cand and "-mtp" in cand.lower(): return True + # Recognise Gemma 4 too, so a failed drafter download surfaces a + # fallback reason instead of silently defaulting. + if cand and _is_gemma_mtp_family(cand): + return True return False @@ -533,22 +701,164 @@ def _is_companion_gguf_path(path: str) -> bool: return name.startswith("mtp-") or "/mtp/" in f"/{p}" +_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE) +_GGUF_KNOWN_QUANT_RE = re.compile( + r"(UD-)?" + r"(MXFP[0-9]+(?:_[A-Z0-9]+)*" + r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" + r"|TQ[0-9]+_[0-9]+" + r"|Q[0-9]+_K_[A-Z]+" + r"|Q[0-9]+_[0-9]+" + r"|Q[0-9]+_K" + r"|BF16|F16|F32)", + re.IGNORECASE, +) + + +def _is_big_endian_gguf_path(path: str, variant_key: str = "") -> bool: + normalized = path.replace("\\", "/") + name = normalized.rsplit("/", 1)[-1] + stem = name.rsplit(".", 1)[0].lower() + variant_key = variant_key.strip().lower() + variant_index = stem.find(variant_key) if variant_key else -1 + parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else "" + variant_in_parent_only = ( + bool(parent) + and variant_index < 0 + and ( + (variant_key and variant_key in parent) + or (not variant_key and _GGUF_KNOWN_QUANT_RE.search(parent) is not None) + ) + ) + for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem): + if variant_index >= 0 and variant_index < match.start(): + return True + tail = stem[match.end() :].lstrip("._-") + if not tail or _GGUF_KNOWN_QUANT_RE.search(tail) is None: + return not variant_in_parent_only + return False + + +def _gguf_snapshot_files(snapshot: Path) -> list[str]: + return [ + p.relative_to(snapshot).as_posix() + for p in snapshot.rglob("*") + if p.is_file() and p.name.lower().endswith(".gguf") + ] + + +def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: + m = _SHARD_FULL_RE.match(first_shard) + if not m: + return [] + prefix = m.group(1) + total = m.group(3) + sibling_pat = re.compile( + r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$", + re.IGNORECASE, + ) + return sorted(f for f in files if f != first_shard and sibling_pat.match(f)) + + +def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]: + """Return main GGUF files matching a requested variant. + + Prefer exact quant-label matches over loose substring matches so a request + for ``stories260K`` does not resolve to ``stories260K-be.gguf``. + """ + variant_key = variant.strip().lower() + main_files = [ + f + for f in files + if f.lower().endswith(".gguf") + and not _is_companion_gguf_path(f) + and not _is_big_endian_gguf_path(f, variant_key) + ] + if not variant_key: + return sorted(main_files) + + try: + from utils.models.model_config import _extract_quant_label + except Exception: + _extract_quant_label = None + + if _extract_quant_label is not None: + try: + exact = sorted(f for f in main_files if _extract_quant_label(f).lower() == variant_key) + if exact: + return exact + except Exception as e: + logger.warning("Failed to extract GGUF quant labels: %s", e) + + boundary = re.compile(r"(? float: + """Bytes per KV-cache element for a llama.cpp cache type (f16 default).""" + return { + "f32": 4.0, + "f16": 2.0, + "bf16": 2.0, + "q8_0": 34 / 32, + "q5_1": 0.75, + "q5_0": 0.6875, + "q4_1": 0.625, + "q4_0": 0.5625, + "iq4_nl": 0.5625, + }.get((cache_type or "f16").strip().lower(), 2.0) + + +def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]: + """Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it + exceeds the f16 default, else None. Studio emits --cache-type only for the + param/extras path, so a heavier env (f32) would otherwise reach the child + unbudgeted; quantized env types stay over-reserved by f16 (-> None).""" + e = os.environ if env is None else env + f16_bpe = _kv_bytes_per_elem("f16") + heaviest: Optional[str] = None + heaviest_bpe = f16_bpe + for var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + raw = (e.get(var) or "").strip().lower() + if not raw: + continue + bpe = _kv_bytes_per_elem(raw) + if bpe > heaviest_bpe: + heaviest, heaviest_bpe = raw, bpe + return heaviest + + +def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) -> Optional[str]: + """Heavier (max bytes/elem) of the explicit --cache-type-k/-v extras, or None. + + Extras are appended last and win per axis, so an asymmetric K=f32,V=f16 must be + budgeted by its heavier axis. resolve_cache_type_kv returns only the last-wins + single type, which under-reserves the heavier axis when the lighter one is last.""" + k, v = parse_cache_override_per_axis(extra_args) + candidates = [c for c in (k, v) if c] + if not candidates: + return None + return max(candidates, key = _kv_bytes_per_elem) + + def _auto_mode_drops_mtp( req_mode: Optional[str], size_b: Optional[float], @@ -564,21 +874,241 @@ def _auto_mode_drops_mtp( return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B +def _mla_mtp_auto_enabled() -> bool: + """Whether Auto may pick embedded MTP for an MLA model (GLM-5.2/DeepSeek/Kimi). + + Off by default: llama.cpp's MLA/DSA MTP path keeps a duplicated full target-KV + context and recomputes the sparse-attention indexer every draft step, so it runs + ~2x slower than no speculation (GLM-5.2 bench: 27 vs 45 tok/s, flat across draft + depth and 96-100% acceptance) -- the opposite of the vLLM/SGLang speedup on the + same model. Set UNSLOTH_MLA_MTP_ENABLED=1 to let Auto promote MLA MTP again once + that path is optimized upstream. Forced mtp / mtp+ngram ignore this gate.""" + return os.environ.get("UNSLOTH_MLA_MTP_ENABLED", "0").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: """User passed --spec-type / --spec-default? llama-server takes one --spec-type (comma-separated to chain), so suppress auto-emit.""" + return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"}) + + +_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}) +_THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"}) + + +def _extra_arg_flag_name(token: str) -> Optional[str]: + if not token.startswith("-") or token in {"-", "--"}: + return None + if len(token) >= 2 and (token[1].isdigit() or token[1] == "."): + return None + return token.split("=", 1)[0] + + +def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool: if not extra_args: return False for raw in extra_args: - tok = str(raw) - if not tok.startswith("--"): - continue - flag = tok.split("=", 1)[0] - if flag in ("--spec-type", "--spec-default"): + flag = _extra_arg_flag_name(str(raw)) + if flag in flags: return True return False +def _effective_spec_type( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """The --spec-type llama-server will use: the last CLI --spec-type (or + --spec-default, which resolves non-MTP), else LLAMA_ARG_SPEC_TYPE. A CLI flag + overrides the env (matching llama.cpp), so a stale MTP env can't make the + budget reserve a drafter the launch won't load. None if neither sets it.""" + args = [str(a) for a in extra_args] if extra_args else [] + cli_present = False + cli_value: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag == "--spec-default": + cli_present = True + cli_value = "default" + continue + if flag != "--spec-type": + continue + cli_present = True + cli_value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if cli_present: + return cli_value + return (os.environ if env is None else env).get("LLAMA_ARG_SPEC_TYPE") + + +def _extra_args_requests_mtp( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects MTP (mtp/draft-mtp), so the + budget must reserve for it.""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("mtp", "draft-mtp") for p in value.split(",")) + + +def _extra_args_requests_separate_draft( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects a non-MTP model draft mode + (draft-simple/draft-eagle3), which loads a separate draft model the budget + must reserve (draft-mtp -> _extra_args_requests_mtp; ngram-* load no model).""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(",")) + + +def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optional[int]: + """Draft depth from extras (``--spec-draft-n-max`` or legacy ``--draft-max``), else None.""" + if not extra_args: + return None + args = [str(a) for a in extra_args] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--spec-draft-n-max", "--draft-max"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + return found + + +def _extra_args_mtp_draft_path( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """Separate drafter path from extras (local --model-draft/-md or HF + --spec-draft-hf/-hfd/...), else the LLAMA_ARG_SPEC_DRAFT_MODEL/_HF_REPO env, + else None. An HF repo isn't a local file, so the budget can't size it (falls + back to the flat reserve), but recognizing it avoids sizing the wrong one.""" + flags = { + "--model-draft", + "--spec-draft-model", + "-md", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", + } + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if value and not value.startswith("-"): + found = value + if found is not None: + return found + e = os.environ if env is None else env + return e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") or None + + +def _extra_args_draft_cache_types( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> tuple[Optional[str], Optional[str]]: + """Draft KV cache types (k_type, v_type), each from extras else the + LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V env, else None (f16). K and V are + independent: a one-sided override must not apply to both.""" + args = [str(a) for a in extra_args] if extra_args else [] + k_flags = {"--cache-type-k-draft", "--spec-draft-type-k", "-ctkd"} + v_flags = {"--cache-type-v-draft", "--spec-draft-type-v", "-ctvd"} + k_type: Optional[str] = None + v_type: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in k_flags and flag not in v_flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if not value or value.startswith("-"): + continue + if flag in k_flags: + k_type = value + else: + v_type = value + e = os.environ if env is None else env + if k_type is None: + k_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K") or None + if v_type is None: + v_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V") or None + return k_type, v_type + + +def _extra_args_draft_offloaded_to_cpu( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the SEPARATE draft model is on CPU (so the budget must not charge + its weights+KV): --spec-draft-ngl 0, or --spec-draft-device naming only + cpu/none, else the LLAMA_ARG_N_GPU_LAYERS_DRAFT env the child honors (the + device flag has no env). An embedded MTP head follows the main -ngl, so these + draft-only flags don't move it. Last-wins, so only each flag's final value counts.""" + ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"} + dev_flags = {"--spec-draft-device", "-devd", "--device-draft"} + args = [str(a) for a in extra_args] if extra_args else [] + last_ngl: Optional[str] = None + last_dev: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if flag in ngl_flags: + last_ngl = value + elif flag in dev_flags: + last_dev = value + if last_ngl is None: + last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT") + if last_ngl is not None: + try: + if int(last_ngl) == 0: + return True + except (TypeError, ValueError): + pass + if last_dev is not None: + devs = [d.strip().lower() for d in last_dev.split(",") if d.strip()] + if devs and all(d in ("cpu", "none") for d in devs): + return True + return False + + +def _extra_args_n_ubatch( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[int]: + """Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH + env, else None. It sizes the compute-graph buffer, so an override must reach + the VRAM reserve.""" + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--ubatch-size", "-ub"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + if found is not None: + return found + raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH") + if raw: + try: + return int(raw) + except (TypeError, ValueError): + pass + return None + + def _build_ngram_mod_flags( caps: Optional[dict], n_match: int = 24, @@ -697,9 +1227,9 @@ class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. Lifecycle: - 1. load_model() — start llama-server with the GGUF file - 2. generate_chat_completion() — proxy to /v1/chat/completions, stream back - 3. unload_model() — terminate the subprocess + 1. load_model(): start llama-server with the GGUF file + 2. generate_chat_completion(): proxy to /v1/chat/completions, stream back + 3. unload_model(): terminate the subprocess """ def __init__(self): @@ -724,6 +1254,7 @@ class LlamaCppBackend: self._is_diffusion: bool = False self._diffusion_visual_bin: Optional[str] = None self._healthy = False + self._stats_logger = None # vLLM-style engine-stats poller, set on load # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None self._context_length: Optional[int] = None @@ -734,6 +1265,7 @@ class LlamaCppBackend: self._supports_reasoning: bool = False self._reasoning_always_on: bool = False self._reasoning_style: str = "enable_thinking" + self._reasoning_effort_levels: list = [] self._supports_preserve_thinking: bool = False self._supports_tools: bool = False self._cache_type_kv: Optional[str] = None @@ -754,6 +1286,9 @@ class LlamaCppBackend: self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None self._embedding_length: Optional[int] = None + # For the compute-graph buffer estimate; vocab from the tokens array len. + self._feed_forward_length: Optional[int] = None + self._vocab_size: Optional[int] = None # Architecture-aware KV fields for 5-path estimation self._kv_key_length: Optional[int] = None self._kv_value_length: Optional[int] = None @@ -773,8 +1308,12 @@ class LlamaCppBackend: self._nextn_predict_layers: Optional[int] = None self._lock = threading.Lock() # Wraps load_model() end-to-end so concurrent loads serialise and never - # coexist as two llama-server processes (#5401). - self._serial_load_lock = threading.Lock() + # coexist as two llama-server processes (#5401). RLock so MTP-crash + # recovery can re-acquire it for its nested load_model. + self._serial_load_lock = threading.RLock() + # Serialises mid-session respawns so many generations hitting a killed + # server trigger at most one reload (see _respawn_if_dead). + self._respawn_lock = threading.Lock() # Set by the in-app updater while it swaps prebuilt binaries; load_model() # rejects fast so no server starts from a half-swapped binary. self._llama_update_in_progress = False @@ -785,6 +1324,18 @@ class LlamaCppBackend: self._extra_args: Optional[List[str]] = None self._extra_args_source: Optional[tuple[str, Optional[str]]] = None self._requested_n_ctx: int = 0 + # Raw kwargs of the last healthy load, for the MTP-crash reload. Memory-only + # (carries hf_token, never logged); single-flight via the lock below. + self._last_load_kwargs: Optional[dict] = None + self._mtp_runtime_fallback_lock = threading.Lock() + self._mtp_runtime_fallback_in_progress = False + # Background watchdog so an MTP+tensor crash recovers even when no request + # observes it (direct proxy endpoints, or nothing in flight). + self._mtp_watchdog_thread: Optional[threading.Thread] = None + self._mtp_watchdog_stop = threading.Event() + # True when the launch actually runs MTP+tensor (Studio- or user/env-driven); + # gates the probe, watchdog, and recovery so pass-through MTP is covered. + self._mtp_runtime_fallback_active = False self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None # llama-server tee log (see _drain_stdout / _kill_process). @@ -803,7 +1354,11 @@ class LlamaCppBackend: # to decide whether to wait for the VRAM reclaim to finish. self._last_kill_monotonic: float = 0.0 - self._kill_orphaned_servers() + _reaped = self._kill_orphaned_servers() + if _reaped: + # Reaped VRAM frees lazily; arm the settle wait so the first load + # waits before ranking GPUs by free memory. + self._last_kill_monotonic = time.monotonic() atexit.register(self._cleanup) # ── Properties ──────────────────────────────────────────────── @@ -821,6 +1376,12 @@ class LlamaCppBackend: def base_url(self) -> str: return f"http://127.0.0.1:{self._port}" + @property + def _auth_headers(self) -> "Optional[dict[str, str]]": + """Bearer header matching the --api-key direct-stream mode uses, else + None (so unauthenticated llama-server calls don't get a spurious 401).""" + return {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + @property def model_identifier(self) -> Optional[str]: return self._model_identifier @@ -842,6 +1403,11 @@ class LlamaCppBackend: def gguf_path(self) -> Optional[str]: return self._gguf_path + @property + def hf_repo(self) -> Optional[str]: + """HF repo of the loaded model, or None for local/native file loads.""" + return self._hf_repo + @property def mtp_draft_path(self) -> Optional[str]: return self._mtp_draft_path @@ -936,12 +1502,13 @@ class LlamaCppBackend: m = _SHARD_RE.match(stem) prefix = m.group(1) if m else None if prefix and parent.is_dir(): + prefix_lower = prefix.lower() for sibling in parent.iterdir(): if ( sibling.is_file() - and sibling.name.startswith(prefix) + and sibling.name.lower().startswith(prefix_lower) and sibling.name != stem - and sibling.suffix == ".gguf" + and sibling.suffix.lower() == ".gguf" ): try: bytes_total += sibling.stat().st_size @@ -993,6 +1560,12 @@ class LlamaCppBackend: def reasoning_style(self) -> str: return self._reasoning_style + @property + def reasoning_effort_levels(self) -> list: + """Discrete reasoning_effort levels the template offers (e.g. GLM-5.2's + ['high', 'max']). Empty unless reasoning_style == 'enable_thinking_effort'.""" + return self._reasoning_effort_levels + @property def supports_preserve_thinking(self) -> bool: return self._supports_preserve_thinking @@ -1002,6 +1575,10 @@ class LlamaCppBackend: return self._reasoning_default def _reasoning_kwargs(self, enable_thinking: bool) -> dict: + if self._reasoning_style == "enable_thinking_effort": + # GLM-5.2-style: enable_thinking is the on/off gate; when on, leave + # the template's default effort (max) in place. + return {"enable_thinking": enable_thinking} if self._reasoning_style == "reasoning_effort": return {"reasoning_effort": "high" if enable_thinking else "low"} return {"enable_thinking": enable_thinking} @@ -1022,7 +1599,20 @@ class LlamaCppBackend: # Always-on reasoning models hardcode tags and don't consume # enable_thinking / reasoning_effort -- skip. if self._supports_reasoning and not self._reasoning_always_on: - if self._reasoning_style == "reasoning_effort": + if self._reasoning_style == "enable_thinking_effort": + # GLM-5.2-style: enable_thinking gates thinking on/off, and the + # reasoning_effort level (e.g. 'high' | 'max') is only meaningful + # while thinking is on. Disabling is enable_thinking=false; a raw + # API caller can also disable via the OpenAI-style + # reasoning_effort="none" sentinel. We never coerce off into a + # 'low' effort the way gpt-oss does (those models genuinely + # cannot disable). + thinking_off = enable_thinking is False or reasoning_effort == "none" + if enable_thinking is not None or reasoning_effort == "none": + kwargs["enable_thinking"] = not thinking_off + if not thinking_off and reasoning_effort in self._reasoning_effort_levels: + kwargs["reasoning_effort"] = reasoning_effort + elif self._reasoning_style == "reasoning_effort": if reasoning_effort in ("none", "low", "medium", "high"): kwargs["reasoning_effort"] = reasoning_effort elif reasoning_effort == "minimal": @@ -1070,6 +1660,30 @@ class LlamaCppBackend: # ── Binary discovery ────────────────────────────────────────── + @staticmethod + def _resolved_studio_root_and_is_legacy() -> "tuple[Optional[Path], bool]": + """Resolve the Studio install root and classify it as the legacy + ~/.unsloth/studio root vs. a custom (env/venv-inferred) root. + + Returns (resolved_root, is_legacy). On any import/resolution failure the + root is treated as legacy and resolved_root is None -- callers must read + resolved_root only when is_legacy is False. Shared by + _find_llama_server_binary (discovery) and _kill_orphaned_servers + (cleanup) so the two never disagree on which root is legacy. + """ + try: + from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 + + resolved = _sr() + legacy_studio = Path.home() / ".unsloth" / "studio" + try: + is_legacy = resolved.resolve() == legacy_studio.resolve() + except (OSError, ValueError): + is_legacy = resolved == legacy_studio + return (None if is_legacy else resolved), is_legacy + except (ImportError, OSError, ValueError): + return None, True + @staticmethod def _find_llama_server_binary(*, include_denied: bool = False) -> Optional[str]: """ @@ -1156,33 +1770,16 @@ class LlamaCppBackend: # 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp; # default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio). legacy_llama = Path.home() / ".unsloth" / "llama.cpp" - try: - from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 - - _resolved_sr = _sr() - _legacy_studio = Path.home() / ".unsloth" / "studio" - try: - _is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve() - except (OSError, ValueError): - _is_legacy = _resolved_sr == _legacy_studio - if _is_legacy: - search_roots = [legacy_llama] - else: - # _kill_orphaned_servers excludes the legacy root in custom - # mode; discovery must match so we never spawn a server we - # then refuse to clean up. UNSLOTH_LLAMA_CPP_PATH (handled - # earlier) is the explicit way to share a build across roots. - search_roots = [_resolved_sr / "llama.cpp"] - except (ImportError, OSError, ValueError): + _resolved_sr, _is_legacy = LlamaCppBackend._resolved_studio_root_and_is_legacy() + if _is_legacy: search_roots = [legacy_llama] - _seen_roots: set[str] = set() - _unique_roots: list[Path] = [] - for r in search_roots: - k = str(r) - if k not in _seen_roots: - _seen_roots.add(k) - _unique_roots.append(r) - for unsloth_home in _unique_roots: + else: + # _kill_orphaned_servers excludes the legacy root in custom mode; + # discovery must match so we never spawn a server we then refuse to + # clean up. UNSLOTH_LLAMA_CPP_PATH (handled earlier) is the explicit + # way to share a build across roots. + search_roots = [_resolved_sr / "llama.cpp"] + for unsloth_home in search_roots: hit, locked = _scan_pinned(_layout_candidates(unsloth_home)) if locked is not None: return _unavailable(locked) @@ -1217,7 +1814,7 @@ class LlamaCppBackend: def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]: """Parse `llama-server --help` for feature flags. Returns {found, mtp_token, supports_mtp, ngram_mod_flavor, - supports_ngram_mod, spec_draft_n_max_flag}. + supports_ngram_mod, spec_draft_n_max_flag, cache flag support}. ``ngram_mod_flavor``: ``"new"`` when the post-rename ``--spec-ngram-mod-n-match / -n-min / -n-max`` are real args; @@ -1242,6 +1839,10 @@ class LlamaCppBackend: "spec_draft_n_max_flag": None, "supports_kv_unified": False, "supports_fit_ctx": False, + "supports_cache_ram": False, + "supports_ctx_checkpoints": False, + "supports_no_cache_prompt": False, + "supports_metrics": False, } try: mtime = int(Path(bin_path).stat().st_mtime) @@ -1257,13 +1858,20 @@ class LlamaCppBackend: spec_draft_n_max_flag: Optional[str] = None supports_kv_unified = False supports_fit_ctx = False + supports_cache_ram = False + supports_ctx_checkpoints = False + supports_no_cache_prompt = False + supports_metrics = False try: + probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( [bin_path, "--help"], capture_output = True, text = True, + errors = "replace", timeout = 10, check = False, + env = probe_env, ) help_text = (result.stdout or "") + "\n" + (result.stderr or "") # Split into per-flag blocks (each --flag line + its indented @@ -1347,6 +1955,10 @@ class LlamaCppBackend: supports_kv_unified = _is_real("--kv-unified") supports_fit_ctx = _is_real("--fit-ctx") + supports_cache_ram = _is_real("--cache-ram") + supports_ctx_checkpoints = _is_real("--ctx-checkpoints") + supports_no_cache_prompt = _is_real("--no-cache-prompt") + supports_metrics = _is_real("--metrics") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") @@ -1359,6 +1971,10 @@ class LlamaCppBackend: "spec_draft_n_max_flag": spec_draft_n_max_flag, "supports_kv_unified": supports_kv_unified, "supports_fit_ctx": supports_fit_ctx, + "supports_cache_ram": supports_cache_ram, + "supports_ctx_checkpoints": supports_ctx_checkpoints, + "supports_no_cache_prompt": supports_no_cache_prompt, + "supports_metrics": supports_metrics, } cls._capability_cache[cache_key] = info return info @@ -1376,7 +1992,8 @@ class LlamaCppBackend: if m: prefix, _, num_total = m.group(1), m.group(2), m.group(3) sibling_pat = re.compile( - r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$" + r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$", + re.IGNORECASE, ) for sibling in main.parent.iterdir(): if sibling != main and sibling_pat.match(sibling.name): @@ -1385,11 +2002,42 @@ class LlamaCppBackend: return total @staticmethod - def _amd_apu_wants_unified_memory() -> bool: + def _resolve_visible_physical_ids() -> Optional[list[int]]: + """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on + ROCm, CUDA otherwise). None when no mask is set; empty list for an empty + mask. Shared by the APU / datacenter / free-memory probes so they agree + on the ordinal->physical mapping.""" + try: + import torch + is_rocm = getattr(torch.version, "hip", None) is not None + except Exception: + is_rocm = False + if is_rocm: + hip_v = os.environ.get("HIP_VISIBLE_DEVICES") + rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") + cvd = ( + hip_v + if hip_v is not None + else rocr_v + if rocr_v is not None + else os.environ.get("CUDA_VISIBLE_DEVICES") + ) + else: + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is None: + return None + try: + return [int(x.strip()) for x in cvd.split(",") if x.strip()] + except ValueError: + return None + + @staticmethod + def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: """True only for AMD unified-memory APUs (gfx1150/gfx1151), where - GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM. - False elsewhere (the env hurts discrete GPUs). ROCm reuses torch.cuda.*; - gcnArchName suffix is stripped.""" + GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it + hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the + selected GPUs, so a dGPU on a mixed host is not treated as unified-memory; + None means every visible GPU.""" try: import torch @@ -1397,12 +2045,25 @@ class LlamaCppBackend: return False if not (hasattr(torch, "cuda") and torch.cuda.is_available()): return False - for _i in range(torch.cuda.device_count()): + # Map visible ordinal -> physical id via the active ROCm mask (HIP, + # then ROCR, then CUDA), mirroring _get_gpu_memory's ROCm branch. + physical_ids = LlamaCppBackend._resolve_visible_physical_ids() + arch_by_id: dict[int, str] = {} + for ordinal in range(torch.cuda.device_count()): try: - _arch = getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "") or "" + _arch = ( + getattr(torch.cuda.get_device_properties(ordinal), "gcnArchName", "") or "" + ) except Exception: continue - if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}: + pid = ( + physical_ids[ordinal] + if physical_ids is not None and ordinal < len(physical_ids) + else ordinal + ) + arch_by_id[pid] = _arch.split(":")[0].strip().lower() + for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id): + if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}: return True except Exception: return False @@ -1438,13 +2099,7 @@ class LlamaCppBackend: # Mirror _get_gpu_free_memory: map visible ordinal -> physical id via # CUDA_VISIBLE_DEVICES; unset/unparsable leaves physical id == ordinal. - physical_ids: Optional[list[int]] = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] - except ValueError: - physical_ids = None + physical_ids = LlamaCppBackend._resolve_visible_physical_ids() pattern = LlamaCppBackend._DATACENTER_GPU_RE names_by_id: dict[int, str] = {} @@ -1509,7 +2164,42 @@ class LlamaCppBackend: @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: - """Query free memory per GPU. + """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by + index; empty if no supported GPU is reachable. Thin wrapper over + ``_get_gpu_memory`` for callers that only need free VRAM.""" + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + + @staticmethod + def _apple_metal_memory_budget_bytes() -> int: + """Unified-memory budget for GGUF context fitting on Apple Silicon. + + No GPU is enumerated on Metal, so the context would default to native and + over-commit unified memory ("Compute error." at decode, #5118/#6529). Use a + fraction of MLX's Metal working-set, else total RAM; 0 off Apple Silicon or + when unresolvable, so callers skip the cap. + """ + from utils.hardware import is_apple_silicon + + if not is_apple_silicon(): + return 0 + rec_bytes = 0 + try: + import mlx.core as mx + if mx.metal.is_available(): + rec_bytes = int(mx.device_info().get("max_recommended_working_set_size") or 0) + except Exception: + rec_bytes = 0 + if rec_bytes <= 0: + try: + import psutil + rec_bytes = int(psutil.virtual_memory().total) + except Exception: + return 0 + return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION) + + @staticmethod + def _get_gpu_memory() -> list[tuple[int, int, int]]: + """Query free AND total memory per GPU. Order: 1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects @@ -1520,15 +2210,15 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. - Returns list of (gpu_index, free_mib) sorted by index; empty if no - supported GPU is reachable. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no + supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. """ # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( [ "nvidia-smi", - "--query-gpu=index,memory.free", + "--query-gpu=index,memory.free,memory.total", "--format=csv,noheader,nounits", ], capture_output = True, @@ -1548,15 +2238,30 @@ class LlamaCppBackend: allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) except ValueError: pass - gpus: list[tuple[int, int]] = [] + gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): - parts = line.split(",") - if len(parts) == 2: - idx = int(parts[0].strip()) - free_mib = int(parts[1].strip()) - if allowed is not None and idx not in allowed: - continue - gpus.append((idx, free_mib)) + parts = [p.strip() for p in line.split(",")] + if len(parts) < 2: + continue + # Index and free required; skip a bad line rather than abandon + # the probe to the torch fallback. + try: + idx = int(parts[0]) + free_mib = int(parts[1]) + except ValueError: + continue + # Total parsed separately: a two-column line or a non-integer + # total ("N/A" on MIG/vGPU) keeps the GPU at total 0 (fit uses + # the free*frac fallback) instead of dropping it. + total_mib = 0 + if len(parts) >= 3 and parts[2]: + try: + total_mib = int(parts[2]) + except ValueError: + total_mib = 0 + if allowed is not None and idx not in allowed: + continue + gpus.append((idx, free_mib, total_mib)) # Match the docstring's sort-by-id guarantee (driver order isn't). gpus.sort(key = lambda g: g[0]) if gpus: @@ -1576,44 +2281,69 @@ class LlamaCppBackend: # feed these IDs back into the subprocess as CVD, so visible ordinals # must be translated to physical indices first; otherwise CVD=2,3 # gets rewritten to 0,1 and targets the wrong GPUs. - physical_ids: Optional[list[int]] = None # Match utils/hardware/hardware.py::_get_parent_visible_gpu_spec: # treat an empty mask (HIP_VISIBLE_DEVICES="") as "no GPUs" rather # than falling through. ``or`` would coerce "" to the wrong source. - if getattr(torch.version, "hip", None) is not None: - hip_v = os.environ.get("HIP_VISIBLE_DEVICES") - rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") - cvd = ( - hip_v - if hip_v is not None - else rocr_v - if rocr_v is not None - else os.environ.get("CUDA_VISIBLE_DEVICES") - ) - else: - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - # Empty mask (CVD="") yields an empty list -> no GPUs, - # consistent with the nvidia-smi path. - physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] - except ValueError: - physical_ids = None + # Empty mask (CVD="") yields an empty list -> no GPUs, consistent + # with the nvidia-smi path. + physical_ids = LlamaCppBackend._resolve_visible_physical_ids() gpus = [] for ordinal in range(torch.cuda.device_count()): - free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal) + free_bytes, total_bytes = torch.cuda.mem_get_info(ordinal) idx = ( physical_ids[ordinal] if physical_ids is not None and ordinal < len(physical_ids) else ordinal ) - gpus.append((idx, free_bytes // (1024 * 1024))) + gpus.append((idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024))) # Match the nvidia-smi path's docstring guarantee of sorted-by-id. return sorted(gpus, key = lambda g: g[0]) except Exception as e: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _available_system_memory_mib() -> Optional[int]: + """Available system RAM in MiB (psutil, then /proc/meminfo), or None if + neither is readable. On a unified-memory APU this, not the ROCm-reported + VRAM, is the real ceiling: the weights load into shared system RAM.""" + try: + import psutil + return int(psutil.virtual_memory().available // (1024 * 1024)) + except Exception: + pass + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) // 1024 # kB -> MiB + except Exception: + pass + return None + + @staticmethod + def _apu_ram_shortfall_message( + model_size_bytes: int, + avail_mib: Optional[int], + headroom_mib: int = 2048, + ) -> Optional[str]: + """On a unified-memory APU, return a user-facing refusal when the weights + cannot fit in available system RAM (else None). Weights only: KV/context + auto-reduce, so counting them too would refuse loads that would succeed. + None avail (unknown RAM) never refuses.""" + if avail_mib is None: + return None + need_mib = model_size_bytes / (1024 * 1024) + if need_mib <= avail_mib - headroom_mib: + return None + return ( + f"This model needs about {need_mib / 1024:.0f} GB but only about " + f"{avail_mib / 1024:.0f} GB of memory is available. On a unified-memory " + "APU the weights load into system RAM, so a larger model is stopped by " + "the OS mid-load. Use a smaller or more quantized GGUF, or free memory " + "(on WSL, raise the memory limit in .wslconfig)." + ) + # Skip the wait when the last kill is older than this; the driver has # already reclaimed the prior process's allocations. _VRAM_SETTLE_WINDOW_S: float = 15.0 @@ -1685,19 +2415,17 @@ class LlamaCppBackend: # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). _GPU_PIN_VRAM_FRACTION = 0.95 - # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the - # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes - # via graph_reserve -- it is roughly EQUAL on every device (not proportional - # to the tensor split) and independent of context. Measured ~2.3 GB - # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a - # conservative headroom above that. It is (a) subtracted from each GPU's free - # VRAM before computing --tensor-split, so the roomier GPU absorbs more - # weight and the smallest GPU keeps room for KV, and (b) reserved per device - # when capping context. The auto-fallback to layer split covers any - # underestimate. NOTE: scales with the model's vocab / batch size; tune if a - # large-vocab model OOMs at load. + # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF + # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived + # path) returns 0. _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120 + # Fixed per-device overhead on every GPU of a LAYER split (CUDA context + + # scratch), beyond the conserved slot-scaling buffer. ~0.9 GB/device measured + # (Qwen3.6-27B, b9625), independent of --parallel; reserved per extra GPU so a + # tight layer split can't advertise a context that OOMs at load. + _PIPELINE_PER_DEVICE_OVERHEAD_MIB = 1024 + # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) @@ -1766,11 +2494,81 @@ class LlamaCppBackend: path_dirs.append(cuda_bin_x64) return path_dirs + @staticmethod + def _llama_server_env_for_binary(binary: str) -> dict[str, str]: + """Build a subprocess env that lets llama-server resolve native libs.""" + env = child_env_without_native_path_secret() + binary_dir = str(Path(binary).parent) + + if sys.platform == "win32": + # Ordering: see _build_windows_path_dirs. #5106. + path_dirs = LlamaCppBackend._build_windows_path_dirs( + binary_dir, + sys.prefix, + os.environ.get("CUDA_PATH", ""), + ) + existing_path = env.get("PATH", "") + env["PATH"] = ";".join(path_dirs) + ";" + existing_path + + # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile + # kernel files (rocblas/library/*.dat + *.hsaco); the DLL searches + # /rocblas/library/ which doesn't exist. + _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", "")) + if _hip_path: + _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library") + if os.path.isdir(_rocblas_lib): + env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib) + else: + # Linux: LD_LIBRARY_PATH for shared libs next to the binary plus + # CUDA runtime libs (libcudart, libcublas, etc.) + import platform + + lib_dirs = [] + # WSL: system HIP before the bundle's (which segfaults on /dev/dxg). + lib_dirs.extend(_wsl_system_rocm_lib_dirs()) + if lib_dirs: + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + lib_dirs.append(binary_dir) + _arch = platform.machine() # x86_64, aarch64, etc. + + # Pip-installed nvidia CUDA runtime libs. The prebuilt binary links + # libcudart.so.13 / libcublas.so.13 which live here, not in + # /usr/local/cuda. + import glob as _glob + + for _nv_pattern in [ + os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", _sub, "lib") + for _sub in ("cu*", "cudnn", "nvjitlink") + ]: + for _nv_dir in _glob.glob(_nv_pattern): + if os.path.isdir(_nv_dir): + lib_dirs.append(_nv_dir) + + for cuda_lib in [ + "/usr/local/cuda/lib64", + f"/usr/local/cuda/targets/{_arch}-linux/lib", + # Fallback CUDA compat paths (e.g. binary built with CUDA 12 + # where default /usr/local/cuda is CUDA 13+). + "/usr/local/cuda-12/lib64", + "/usr/local/cuda-12.8/lib64", + f"/usr/local/cuda-12/targets/{_arch}-linux/lib", + f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", + ]: + if os.path.isdir(cuda_lib): + lib_dirs.append(cuda_lib) + existing_ld = env.get("LD_LIBRARY_PATH", "") + new_ld = ":".join(lib_dirs) + env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld + + return env + @staticmethod def _select_gpus( model_size_bytes: int, gpus: list[tuple[int, int]], usable_fraction: Optional[float] = None, + total_by_idx: Optional[dict[int, int]] = None, + per_device_overhead_bytes: int = 0, ) -> tuple[Optional[list[int]], bool]: """Pick GPU(s) for a model from estimated VRAM and free memory. @@ -1778,6 +2576,11 @@ class LlamaCppBackend: ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides headroom for compute buffers, CUDA context, and other runtime overhead; callers lower it when MTP reserves VRAM for a draft model. + ``total_by_idx`` (index -> total MiB) makes the headroom an ABSOLUTE + ``(1 - fraction) * total`` per GPU instead of a fraction of free. + ``per_device_overhead_bytes`` is the fixed layer-split cost per GPU beyond + the first; a k-GPU pin must hold ``model + (k-1) * overhead`` or it can OOM + a device after -ngl -1 (no --fit fallback). Single-GPU adds none. Returns (gpu_indices, use_fit): - ([1], False) fits on 1 GPU at the headroom threshold @@ -1791,20 +2594,31 @@ class LlamaCppBackend: if usable_fraction is None: usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION - # Sort GPUs by free memory descending - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # Per-GPU usable budget: free - (1-frac)*total when total is known, else + # the legacy free*frac (also covers a total-0 two-column probe). + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - usable_fraction) * t) + return free_mib * usable_fraction + + # Rank by usable budget (free - reserve), not raw free: a more-used large + # card can have less usable room than a less-used small one. + ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True) # Try 1 GPU at the usable-VRAM threshold. - if ranked[0][1] * usable_fraction >= model_size_mib: + if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: return [ranked[0][0]], False - # Try N GPUs (accumulate free memory from most-free) - cumulative = 0 + # Try N GPUs (accumulate usable memory from most-free). Each GPU past the + # first adds a fixed per-device overhead the pool must hold. + overhead_mib = per_device_overhead_bytes / (1024 * 1024) + cumulative = 0.0 selected = [] for idx, free_mib in ranked: selected.append(idx) - cumulative += free_mib * usable_fraction - if cumulative >= model_size_mib: + cumulative += _usable(idx, free_mib) + if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib: return sorted(selected), False # Too large even for all GPUs; let --fit handle it @@ -1839,6 +2653,12 @@ class LlamaCppBackend: return self._n_kv_heads_by_layer[layer_idx] return fallback + def _legacy_head_dim(self) -> int: + """Head-dim fallback for GGUFs without explicit key/value dims. Reached + only via the legacy branch of _can_estimate_kv(), so _embedding_length + is non-None here.""" + return self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + def _estimate_kv_cache_bytes( self, n_ctx: int, @@ -1859,14 +2679,10 @@ class LlamaCppBackend: 5. Legacy -- fallback using embed // n_heads Server-flag knobs (mirror llama-server's CLI): - swa_full -- ``--swa-full``: force SWA layers to cache full - ``n_ctx`` (collapses path 3 to path 4 for them). - n_parallel -- ``--parallel`` slots: non-SWA layers stay constant - (cells split across slots), SWA layers scale linearly. - kv_unified -- ``--kv-unified`` (default on): no-op for memory math; - kept for API forward-compat. - ctx_checkpoints -- ``--ctx-checkpoints`` (PR #15293): N SWA snapshots - per slot, one sliding-window of state per SWA layer. + swa_full -- --swa-full: SWA layers cache full n_ctx (path 3->4). + n_parallel -- --parallel slots: non-SWA constant, SWA scale linearly. + kv_unified -- --kv-unified: memory no-op (API forward-compat). + ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot. Returns 0 if metadata is insufficient. """ @@ -1881,17 +2697,7 @@ class LlamaCppBackend: n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization - bpe = { - "f32": 4.0, - "f16": 2.0, - "bf16": 2.0, - "q8_0": 34 / 32, - "q5_1": 0.75, - "q5_0": 0.6875, - "q4_1": 0.625, - "q4_0": 0.5625, - "iq4_nl": 0.5625, - }.get(cache_type_kv or "f16", 2.0) + bpe = _kv_bytes_per_elem(cache_type_kv) slots = max(1, n_parallel) @@ -1916,18 +2722,15 @@ class LlamaCppBackend: n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division if key_len is not None and val_len is not None: return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe) - head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + head_dim = self._legacy_head_dim() return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) - # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). - # Pattern filled by the resolver at parse time; if absent, falls through - # to the legacy 1/4-global heuristic below. Per-layer-type --parallel N - # accounting (verified against llama-server): - # * non-SWA layers: total cells = n_ctx split across slots -> CONSTANT. - # * SWA layers: per-slot cells = 2*sliding_window (capped at n_ctx - # and per_slot_ctx) -> grows LINEARLY in slots. - # --swa-full forces full n_ctx for SWA layers; --ctx-checkpoints N adds - # N snapshots per SWA layer per slot. + # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern + # from the resolver; if absent, falls through to the legacy 1/4-global + # heuristic. --parallel N accounting (verified against llama-server): + # non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells + # = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots. + # --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -1936,8 +2739,7 @@ class LlamaCppBackend: ): swa = self._sliding_window per_slot_ctx = max(1, n_ctx // slots) - # --swa-full caches full context like non-SWA (per-slot cells = - # per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA + # --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA # caches 2*sliding_window per slot, clamped at per-slot ctx. swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) key_len_swa = self._kv_key_length_swa or key_len @@ -1987,9 +2789,172 @@ class LlamaCppBackend: return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe) # Path 5: Legacy fallback (old GGUFs without explicit dimensions) - head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + head_dim = self._legacy_head_dim() return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) + def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]: + """Lightweight backend with a drafter GGUF's metadata, to size its own KV + via _estimate_kv_cache_bytes. Cached per path; None if unreadable.""" + cache = getattr(self, "_draft_backend_cache", None) + if cache is not None and cache[0] == drafter_path: + return cache[1] + db: Optional[LlamaCppBackend] = None + try: + db = LlamaCppBackend.__new__(LlamaCppBackend) + for attr in ( + "_context_length", + "_n_layers", + "_n_kv_heads", + "_n_heads", + "_embedding_length", + "_kv_key_length", + "_kv_value_length", + "_kv_lora_rank", + "_sliding_window", + "_sliding_window_pattern", + "_ssm_inner_size", + "_full_attention_interval", + "_key_length_mla", + "_n_kv_heads_by_layer", + "_kv_key_length_swa", + "_kv_value_length_swa", + "_shared_kv_layers", + "_nextn_predict_layers", + ): + setattr(db, attr, None) + db._model_identifier = "mtp-draft" + db._read_gguf_metadata(drafter_path) + except Exception as e: # unreadable drafter -> caller falls back + logger.debug(f"Could not read drafter GGUF for MTP budget: {e}") + db = None + self._draft_backend_cache = (drafter_path, db) + return db + + def _mtp_draft_kv_bytes( + self, + n_ctx: int, + *, + drafter_path: Optional[str] = None, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + n_parallel: int = 1, + ) -> Optional[int]: + """Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are + independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes + at the heavier type. Embedded head (Qwen): nextn_predict_layers attention + layers from the main dims. None when dims are missing (flat fallback).""" + if n_ctx <= 0: + return None + bpe_k = _kv_bytes_per_elem(draft_cache_type_k) + bpe_v = _kv_bytes_per_elem(draft_cache_type_v) + if drafter_path: + db = self._draft_backend_for(drafter_path) + if db is None or not db._can_estimate_kv(): + return None + heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v + # The drafter is served under the same --parallel slot count as the + # main model, so price its KV per slot too: a sliding-window drafter + # (Gemma) grows KV with slots and would otherwise be under-reserved. + kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel) + return kv or None + nextn = self._nextn_predict_layers or 0 + n_kv = self._n_kv_heads or self._n_heads + k_len = self._kv_key_length + v_len = self._kv_value_length + if not (nextn and n_kv and k_len and v_len): + return None + # The embedded MTP head is one draft layer, so a quantized draft KV can't + # amortize its overhead and fits *less* context than f16 (llama.cpp#24102). + # Floor it at f16: a quantized override is priced as f16, f32 keeps its 4 + # bytes. The separate-drafter branch is multi-layer, so it keeps its type. + f16_bpe = _kv_bytes_per_elem("f16") + bpe_k = max(bpe_k, f16_bpe) + bpe_v = max(bpe_v, f16_bpe) + return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx) + + def _estimate_mtp_overhead_bytes( + self, + n_ctx: int, + *, + spec_draft_n_max: int = 0, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + drafter_path: Optional[str] = None, + draft_weights_bytes: int = 0, + n_parallel: int = 1, + mtp_keeps_target_ctx: bool = True, + ) -> Optional[int]: + """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate- + drafter weights + (MTP + MLA only) a duplicated target KV context. The + verify buffer rides in the ctx-fit headroom (no tuned constant). None when + the draft KV can't be sized (caller keeps the flat fallback). + ``draft_weights_bytes`` is the drafter file size (0 for embedded). + ``mtp_keeps_target_ctx`` is True for MTP draft modes (which keep the + duplicated target context) and False for separate-drafter spec modes + (draft-simple/draft-eagle3), which do not.""" + draft_kv = self._mtp_draft_kv_bytes( + n_ctx, + drafter_path = drafter_path, + draft_cache_type_k = draft_cache_type_k, + draft_cache_type_v = draft_cache_type_v, + n_parallel = n_parallel, + ) + weights = max(0, draft_weights_bytes) + # MLA models (GLM-5.x, DeepSeek, Kimi-K2) under MTP keep a *second* full copy + # of the target model's KV context for draft verification -- llama.cpp's + # `ctx_tgt=yes` -- allocated at f16 regardless of the main cache type. It is + # ~the main KV again and dwarfs the embedded draft head (GLM-5.2 @ 1M ctx: + # a ~2 GiB head next to a ~89 GiB target copy), so omitting it lets auto-fit + # pick a context that fits on paper but OOMs cublasCreate at the first + # decode. Gated on both MLA (kv_lora_rank present) and the engaged mode + # actually being MTP: non-MLA MTP (Qwen/Gemma) keeps no such copy, and the + # separate-drafter spec modes (draft-simple/draft-eagle3) load a small + # distinct drafter with its own KV -- already counted in draft_kv/weights -- + # rather than duplicating the target, so they must not be charged for it. + target_ctx_copy = 0 + if mtp_keeps_target_ctx and self._kv_lora_rank is not None: + target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel) + if draft_kv is None: + # KV unsized (exotic/remote drafter): still reserve known weights + any + # MLA target copy so a large config can't launch over budget (the small + # unsized draft KV rides in the cushion). Nothing known -> None, so the + # caller keeps the flat fallback. + total = weights + target_ctx_copy + return total if total > 0 else None + return draft_kv + weights + target_ctx_copy + + _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it + _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + + def _estimate_compute_buffer_bytes( + self, + *, + n_ubatch: Optional[int] = None, + n_parallel: int = 1, + per_device_tensor: bool = False, + ) -> int: + """Per-device compute-graph buffer (bytes) from GGUF dims: a vocab-width + output buffer + activation scratch. Context-independent; scales with + ``--parallel`` (serving slots). Tensor mode materializes it on every device. + A slight upper bound over measured allocations; 0 when dims are missing.""" + n_vocab = self._vocab_size or 0 + n_embd = self._embedding_length or 0 + if n_vocab <= 0 or n_embd <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + par = max(1, int(n_parallel)) + out_buffer = n_vocab * ub * 4 # f32 output/logits buffer + act_scratch = 4 * n_embd * ub * 4 # a few resident hidden-width buffers + if per_device_tensor: + # Output + comm/staging materialized on every device, every slot. + compute = 2 * act_scratch + out_buffer * par + else: + # Each extra concurrent slot adds one output buffer (chat decode sizes + # ~one logit row per slot; would under-count embeddings/--logits-all, + # not run here). Matches measured {1:36,2:492,4:1388,8:3220} MiB. + compute = act_scratch + out_buffer * max(0, par - 1) + return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _fit_context_to_vram( self, requested_ctx: int, @@ -2004,13 +2969,15 @@ class LlamaCppBackend: ctx_checkpoints: int = 0, kv_on_gpu: bool = True, mtp_engaged: bool = False, + mtp_overhead_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, + total_mib: Optional[int] = None, ) -> int: """Return the largest context length that fits in GPU VRAM. - Uses 90% of available VRAM as the ctx-fit budget -- tighter than - ``_GPU_PIN_VRAM_FRACTION`` on purpose (over-promising context OOMs at - runtime). If the weights alone don't fit, returns ``requested_ctx``. + Budget caps occupancy at ``_CTX_FIT_VRAM_FRACTION`` of the card: an + absolute ``free - (1 - frac) * total`` when ``total_mib`` is given, else + ``free * frac``. Weights alone over budget returns ``requested_ctx``. ``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV cache lives in CPU RAM and the requested context is honored verbatim. @@ -2038,20 +3005,28 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - # MTP engaged: carve the drafter's reserve out of the fit budget. Callers - # can override outright (tensor-parallel mode passes a fatter margin), so - # only compute a default when none was supplied. + # byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback + # when dims can't size the draft KV); callers may override budget_frac. if budget_frac is None: - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) - budget_bytes = available_mib * 1024 * 1024 * budget_frac + flat_mtp = mtp_engaged and mtp_overhead_fn is None + budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0) + # Absolute reserve off total when known, else fraction-of-free; clamp >=0. + if total_mib is not None and total_mib > 0: + budget_mib = max(0.0, available_mib - (1.0 - budget_frac) * total_mib) + else: + budget_mib = available_mib * budget_frac + budget_bytes = budget_mib * 1024 * 1024 model_footprint = model_size_bytes + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: return requested_ctx - # Weights alone exceed budget -- reducing ctx can't help; --fit handles it. + # Weights + compute buffer alone exceed budget -- reducing ctx can't help. if model_footprint >= budget_bytes: logger.debug( "Model footprint exceeds GPU budget before KV cache", @@ -2061,7 +3036,7 @@ class LlamaCppBackend: ) return requested_ctx - # Binary search for max context that fits + # Binary search for max context that fits (KV + MTP draft reserve at that ctx) remaining = budget_bytes - model_footprint effective_min = min(min_ctx, requested_ctx) lo, hi = effective_min, requested_ctx @@ -2069,7 +3044,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv <= remaining: + if kv + _mtp_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -2101,7 +3076,11 @@ class LlamaCppBackend: files = list_repo_files(hf_repo, token = hf_token) gguf_files = [ - f for f in files if f.endswith(".gguf") and not _is_companion_gguf_path(f) + f + for f in files + if f.lower().endswith(".gguf") + and not _is_companion_gguf_path(f) + and not _is_big_endian_gguf_path(f) ] if not gguf_files: return None @@ -2236,6 +3215,7 @@ class LlamaCppBackend: self._supports_reasoning = False self._reasoning_always_on = False self._reasoning_style = "enable_thinking" + self._reasoning_effort_levels = [] self._reasoning_default = True self._supports_preserve_thinking = False self._supports_tools = False @@ -2244,6 +3224,8 @@ class LlamaCppBackend: self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None + self._feed_forward_length = None + self._vocab_size = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None @@ -2265,6 +3247,8 @@ class LlamaCppBackend: WANTED = { "general.architecture", "tokenizer.chat_template", + # Vocab size = tokens array length (no vocab_size key in many GGUFs). + "tokenizer.ggml.tokens", # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner. "diffusion.canvas_length", # Source-repo hints for the SWA resolver's HF fallback. @@ -2328,6 +3312,7 @@ class LlamaCppBackend: f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", + f"{arch}.feed_forward_length": "feed_forward_length", f"{arch}.attention.key_length": "kv_key_length", f"{arch}.attention.value_length": "kv_value_length", f"{arch}.attention.sliding_window": "sliding_window", @@ -2361,12 +3346,15 @@ class LlamaCppBackend: elif vtype == 9: # ARRAY atype = struct.unpack(" %s from local HF cache", hf_variant, @@ -2832,10 +3767,11 @@ class LlamaCppBackend: _m = _SHARD_RE.match(gguf_filename) _prefix = _m.group(1) if _m else None if _prefix: + prefix_lower = _prefix.lower() gguf_extra_shards = sorted( f for f in all_gguf_files - if f.startswith(_prefix) + if f.lower().startswith(prefix_lower) and f != gguf_filename and not _is_companion_gguf_path(f) ) @@ -2859,27 +3795,27 @@ class LlamaCppBackend: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") dl_start = time.monotonic() - local_path = hf_hub_download( - repo_id = hf_repo, - filename = gguf_filename, - token = hf_token, + # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. + local_path = hf_hub_download_with_xet_fallback( + hf_repo, + gguf_filename, + hf_token, + cancel_event = self._cancel_event, + on_status = lambda m: logger.info(m), ) for shard in gguf_extra_shards: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download( - repo_id = hf_repo, - filename = shard, - token = hf_token, + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = self._cancel_event, ) - except RuntimeError as e: - if "Cancelled" in str(e): - raise - raise RuntimeError( - f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" - ) except Exception as e: + if isinstance(e, RuntimeError) and "Cancelled" in str(e): + raise raise RuntimeError( f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" ) @@ -2910,17 +3846,37 @@ class LlamaCppBackend: return None target: Optional[str] = None - try: - from huggingface_hub import list_repo_files - target = pick(list_repo_files(hf_repo, token = hf_token)) - except Exception as e: - logger.debug(f"Could not list repo files for {label}: {e}") + from huggingface_hub import list_repo_files + + # Retry a transient listing blip; permanent repo/auth errors and offline + # mode are not retried (offline raises at once -> fall through to cache). + for attempt in range(3): + if self._cancel_event.is_set(): + return None + try: + target = pick(list_repo_files(hf_repo, token = hf_token)) + break + except Exception as e: + if type(e).__name__ in ( + "RepositoryNotFoundError", + "GatedRepoError", + "RevisionNotFoundError", + "EntryNotFoundError", + "OfflineModeIsEnabled", + ): + logger.debug(f"Could not list repo files for {label}: {e}") + break + logger.debug( + f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}" + ) + if attempt < 2: + self._cancel_event.wait(2**attempt) if target is None: try: from utils.models.model_config import _iter_hf_cache_snapshots for snap in _iter_hf_cache_snapshots(hf_repo): - rel_files = [p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")] + rel_files = _gguf_snapshot_files(snap) target = pick(rel_files) if target is not None: logger.info("Resolved %s %s from local HF cache", label, target) @@ -2932,12 +3888,13 @@ class LlamaCppBackend: return None try: - from huggingface_hub import hf_hub_download logger.info(f"Downloading {label}: {hf_repo}/{target}") - return hf_hub_download( - repo_id = hf_repo, - filename = target, - token = hf_token, + # Same policy; companions are best-effort (caller below swallows failures to None). + return hf_hub_download_with_xet_fallback( + hf_repo, + target, + hf_token, + cancel_event = self._cancel_event, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") @@ -3084,7 +4041,10 @@ class LlamaCppBackend: @staticmethod def _classify_llama_start_failure( - output: str, gguf_path: Optional[str], model_identifier: Optional[str] + output: str, + gguf_path: Optional[str], + model_identifier: Optional[str], + returncode: Optional[int] = None, ) -> str: """Explain *why* llama-server failed to start, from its output. @@ -3156,6 +4116,24 @@ class LlamaCppBackend: "Ollama instead." ) + # SIGKILL with no diagnostic output is the OOM killer (e.g. a model too + # large for the WSL VM's RAM cap); name it actionably. + if returncode == -9: + return ( + "llama-server was stopped by the operating system (signal 9), " + "most likely out of memory. Try a smaller or more quantized " + "GGUF, lower the context length, or free memory (on WSL, raise " + "the memory limit in .wslconfig)." + ) + # SIGTERM is also how an unload/cancel or a supervisor stops the server, + # so report it neutrally rather than blaming memory. + if returncode == -15: + return ( + "llama-server was terminated (signal 15) before it became " + "healthy. If you cancelled or unloaded the model this is " + "expected; otherwise check the llama-server log for the cause." + ) + # Fallback: genuinely unknown failure (OOM, missing binary ...). return ( "llama-server failed to start. " @@ -3170,7 +4148,11 @@ class LlamaCppBackend: cache_type_kv: Optional[str] = None, n_parallel: int = 1, mtp_engaged: bool = False, + mtp_overhead_fn: Optional[Callable[[int], int]] = None, + mtp_flat_reserve_bytes: int = 0, max_target_ctx: Optional[int] = None, + total_by_idx: Optional[dict[int, int]] = None, + n_ubatch: Optional[int] = None, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -3183,22 +4165,41 @@ class LlamaCppBackend: Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``). + one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. - ``tensor_split`` is None (llama.cpp's even default, safe for every arch incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even - share fits the smallest GPU; otherwise it is weighted by - ``(free - buffer)`` so the roomier GPU absorbs more weight and the - smallest GPU keeps room for KV. + share fits the smallest GPU; otherwise it is weighted by usable budget + so the roomier GPU absorbs more weight and the smallest keeps room for KV. + ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes + the compute buffer. """ - # Drop GPUs that can't hold the per-device compute-graph buffer; they'd - # OOM in tensor mode. load_model already filters before calling, so this - # is defense-in-depth that also keeps the pure function self-contained - # (and unit-testable without a GPU). - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - usable_gpus = [g for g in gpus if g[1] >= reserve_mib] + + # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a + # two-column probe) the legacy free*frac. Mirrors _select_gpus and + # _gpu_usable so the 5% cushion is kept on every path, not dropped here. + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - _CTX_FIT_VRAM_FRACTION) * t) + return max(0.0, free_mib * _CTX_FIT_VRAM_FRACTION) + + # Drop GPUs whose usable budget can't hold the per-device compute-graph + # buffer; they'd OOM in tensor mode. Admitting on raw free would let a + # partly-used big card in with no budget left. Defense-in-depth (load_model + # gates too). Derived per-device reserve; flat fallback. + _reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = n_parallel, per_device_tensor = True + ) + reserve_mib = ( + _reserve_bytes // (1024 * 1024) + if _reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + usable_gpus = [g for g in gpus if _usable(g[0], g[1]) >= reserve_mib] gpu_indices = sorted(idx for idx, _ in usable_gpus) if len(gpu_indices) < 2: # Tensor parallelism is meaningless on <2 GPUs (the caller drops the @@ -3210,21 +4211,50 @@ class LlamaCppBackend: None, ) free_by_idx = {idx: free for idx, free in usable_gpus} - pool_mib = sum(free_by_idx.values()) - kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - if mtp_engaged: - # MTP keeps a draft model + its own KV cache on GPU. - kv_budget_b -= 2 * 1024**3 + usable_by_idx = {idx: _usable(idx, free_by_idx[idx]) for idx in gpu_indices} + pool_mib = sum(usable_by_idx.values()) + # MTP reserve: byte-accurate per-ctx inside _fit_ctx (mtp_overhead_fn) plus + # a flat cushion that the byte fn can't size -- 2 GiB when dims are wholly + # unavailable (no fn), or mtp_flat_reserve_bytes when the fn is weights-only + # because the draft KV couldn't be sized (_mtp_kv_unsized). Without this the + # binary search spends the unsized-KV cushion on main context and OOMs. + flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) + if mtp_engaged and mtp_overhead_fn is None: + flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + kv_budget_b = ( + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + ) + + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 def _fit_ctx(ctx: int) -> int: - # Largest context whose KV fits the pooled budget. Floors small, but - # never raises an explicit ctx above what was asked. + # Largest context whose KV (+ MTP draft reserve) fits the pooled + # budget. Floors small, but never raises an explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: # Weights + buffers exceed the pool -> floor; the load then # falls back to layer split. return ctx_floor + if mtp_overhead_fn is not None: + # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + def _consumer(c: int) -> int: + return self._estimate_kv_cache_bytes( + c, cache_type_kv, n_parallel = n_parallel + ) + _mtp_at(c) + + if _consumer(ctx) <= kv_budget_b: + return ctx + lo, hi, best = ctx_floor, ctx, ctx_floor + while lo <= hi: + mid = (lo + hi) // 2 + if _consumer(mid) <= kv_budget_b: + best = mid + lo = mid + 1 + else: + hi = mid - 1 + return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) if kv_at <= kv_budget_b: return ctx @@ -3239,16 +4269,19 @@ class LlamaCppBackend: max_available_ctx = _fit_ctx(max_ctx_target) effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) - min_free_mib = min(free_by_idx.values()) + min_usable_mib = min(usable_by_idx.values()) kv_bytes = ( self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) if (self._can_estimate_kv() and effective_ctx > 0) else 0 ) - even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024) + # The MTP reserve also has to fit the even split (mirror the pooled budget): + # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. + mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes + even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) tensor_split: Optional[list[int]] = None - if even_share_mib > (min_free_mib - reserve_mib): - adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices] + if even_share_mib > (min_usable_mib - reserve_mib): + adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -3279,6 +4312,71 @@ class LlamaCppBackend: and ("unknown" in text or "unsupported" in text or "not supported" in text) ) + @staticmethod + def _output_has_nonprojector_diagnostic(output: str) -> bool: + """True when the output already names a concrete non-projector cause (out + of memory, an unsupported architecture, a tensor-parallel limit). A hard + crash carrying such a marker must surface that error, not be silently + retried text-only as if the vision projector were at fault; a bare crash + with no marker still gets the text-only retry. + """ + text = (output or "").lower() + return any( + m in text + for m in ( + "out of memory", + "failed to allocate", + "unknown model architecture", + "split_mode_tensor not implemented", + ) + ) + + @staticmethod + def _is_signal_crash(returncode: Optional[int]) -> bool: + """True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a + Windows 0xC0000000+ status), not SIGKILL/SIGTERM/SIGINT (OOM killer / + unload) nor a clean exit or still-running (None) process. + """ + if returncode is None: + return False + if returncode >= 0xC0000000: # Windows access violation / illegal instruction + return True + return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV + + @staticmethod + def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: + """Return cmd with flash attention forced off, or None when its effective + (last-wins) value is already off/absent so there is nothing to retry. FA + kernels hard-crash at startup on some ROCm builds; disabling FA keeps + vision and MTP, the least destructive rung. A bare --flash-attn/-fa reads + as on, so it counts toward the effective value and is neutralised too; + every form is flipped in place (length preserved for downstream slices).""" + out = list(cmd) + + def explicit(i): + nxt = out[i + 1] if i + 1 < len(out) else None + return nxt if nxt in ("on", "auto", "off") else None + + effective = None + for i, tok in enumerate(out): + if tok.startswith(("--flash-attn=", "-fa=")): + effective = tok.partition("=")[2] + elif tok in ("--flash-attn", "-fa"): + effective = explicit(i) or "on" + if effective not in ("on", "auto"): + return None + for i, tok in enumerate(out): + if tok.startswith(("--flash-attn=", "-fa=")): + flag, _, value = tok.partition("=") + if value in ("on", "auto"): + out[i] = f"{flag}=off" + elif tok in ("--flash-attn", "-fa"): + if explicit(i) in ("on", "auto"): + out[i + 1] = "off" + elif explicit(i) is None: # bare flag (reads as on) -> explicit off + out[i] = f"{tok}=off" + return out + @staticmethod def _strip_mmproj_args(cmd: list[str]) -> list[str]: """Return cmd without the '--mmproj ' pair (text-only retry). @@ -3296,6 +4394,16 @@ class LlamaCppBackend: out.append(tok) return out + @staticmethod + def _redacted_cmd_for_log(cmd: "list[str]") -> "list[str]": + """Copy of cmd with the value after --api-key replaced by .""" + out = list(cmd) + if "--api-key" in out: + ki = out.index("--api-key") + 1 + if ki < len(out): + out[ki] = "" + return out + def _start_llama_process(self, cmd: list[str], env: dict) -> None: """Spawn llama-server from cmd and start draining its output. @@ -3332,12 +4440,7 @@ class LlamaCppBackend: # Log the argv per attempt (the text-only mmproj retry re-enters here # with --mmproj stripped), redacting the API key. - _log_cmd = list(cmd) - if "--api-key" in _log_cmd: - _ki = _log_cmd.index("--api-key") + 1 - if _ki < len(_log_cmd): - _log_cmd[_ki] = "" - logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") + logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") self._process = subprocess.Popen( cmd, @@ -3346,7 +4449,11 @@ class LlamaCppBackend: text = True, env = env, **_windows_hidden_subprocess_kwargs(), + **_child_popen_kwargs(), ) + # Cross-session backstop: record the PID so a later startup can reap this + # server if parent-death cleanup did not run (macOS / best-effort failure). + self._record_server_pid(self._process.pid) # Start background thread to drain stdout and prevent pipe deadlock self._stdout_thread = threading.Thread( @@ -3389,6 +4496,28 @@ class LlamaCppBackend: Returns True if the server started and the health check passed. """ + # Raw load inputs so the runtime MTP-crash reload can replay this model + # without MTP. Committed to _last_load_kwargs only on a healthy load. + _pending_load_kwargs = { + "gguf_path": gguf_path, + "mmproj_path": mmproj_path, + "mtp_draft_path": mtp_draft_path, + "hf_repo": hf_repo, + "hf_variant": hf_variant, + "hf_token": hf_token, + "model_identifier": model_identifier, + "is_vision": is_vision, + "n_ctx": n_ctx, + "chat_template_override": chat_template_override, + "cache_type_kv": cache_type_kv, + "speculative_type": speculative_type, + "spec_draft_n_max": spec_draft_n_max, + "tensor_parallel": tensor_parallel, + "n_threads": n_threads, + "n_gpu_layers": n_gpu_layers, + "n_parallel": n_parallel, + "extra_args": list(extra_args) if extra_args is not None else None, + } # Serialise the whole load so concurrent /load calls never leave two # llama-server processes alive (#5401 / #5161). Doesn't block /unload. with self._serial_load_lock: @@ -3424,36 +4553,8 @@ class LlamaCppBackend: except Exception as exc: logger.debug("Fast-path audio probe failed: %s", exc) detected = None - if detected in ("snac", "bicodec", "dac"): - with self._lock: - if not self._healthy: - return False - try: - self.init_audio_codec(detected) - self._is_audio = True - self._audio_type = detected - except Exception as exc: - logger.warning( - "Failed to init audio codec '%s': %s", - detected, - exc, - ) - self._audio_probed = False - return False - elif detected: - # csm / whisper / audio_vlm: track type but keep - # _is_audio False -- GGUF TTS routing only fires for - # snac/bicodec/dac. - with self._lock: - if not self._healthy: - return False - self._audio_type = detected - # Re-derive after a retried probe (_mmproj_has_audio persists). - from utils.models.model_config import is_audio_input_type - - self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( - self._mmproj_has_audio - ) + if not self._apply_detected_audio(detected): + return False if not self._healthy: return False return True @@ -3548,11 +4649,11 @@ class LlamaCppBackend: "(access-denied; antivirus or an in-flight install). " "Retry the load once it is released." ) - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) + # Reached only after the diffusion early-return above, so this is a + # genuine llama-server-backed GGUF with no runtime. Raise the typed + # error so /load returns the actionable 400 (not a generic 500), the + # same message remote validation already shows. + raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL) # Outside ``self._lock`` so /unload, /cancel, /status aren't # blocked. ``unload_model`` also records the kill, so the @@ -3575,27 +4676,58 @@ class LlamaCppBackend: ctx_override = parse_ctx_override(extra_args) requested_ctx = resolve_requested_ctx(extra_args, n_ctx) cache_override = parse_cache_override(extra_args) - cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) - # A user --split-mode in extras last-wins-overrides the - # toggle, so reconcile it back into tensor_parallel state. + # Budget the heavier of asymmetric --cache-type-k/-v extras (they + # win per axis at launch, appended last); resolve_cache_type_kv only + # returns the last-wins type, which under-reserves the heavier axis. + # The user's extras still set the real (possibly asymmetric) child + # cache, so this only affects the reserve, not the emitted command. + _extras_cache = _extra_args_main_cache_type_for_budget(extra_args) + cache_type_kv = _extras_cache if _extras_cache is not None else cache_type_kv + _cache_type_from_env = False + if cache_type_kv is None: + # Param/extras set nothing, so the child inherits + # LLAMA_ARG_CACHE_TYPE_K/_V. Adopt a heavier env type (f32) for + # the reserve only; the launch does NOT re-emit it (that would + # rewrite an asymmetric K=f32,V=f16 env into symmetric flags), + # so _cache_type_from_env keeps it out of the emitted flags. + cache_type_kv = _env_main_cache_type_for_budget() + _cache_type_from_env = cache_type_kv is not None + # A user --split-mode in extras last-wins-overrides the toggle, and + # an inherited tensor LLAMA_ARG_SPLIT_MODE flips it on (the child + # would run tensor unbudgeted otherwise). The duplicate-load matchers + # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) - tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel) + tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type - # that would re-impose it when appended last). The layer-split - # fallback re-runs with tensor_parallel False and keeps the type. - if ( - tensor_parallel - and cache_type_kv - and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES - ): + # that would re-impose it when appended last). Layer split does + # support it, so remember the dropped type and the original extras + # to restore (verbatim, incl. an asymmetric K/V) if we later fall + # back to layer split below. + _tensor_dropped_cache_type_kv: Optional[str] = None + _tensor_dropped_extra_args: Optional[list] = None + # Tensor mode rejects any quantized axis. cache_type_kv is the + # heavier-by-bytes budget type, which can mask a quantized axis (an + # f16 budget hides a paired q4_0), so also test each explicit + # --cache-type-k/-v extra, not just the budget type. + _ck_extra, _cv_extra = parse_cache_override_per_axis(extra_args) + _cache_non_tensor_safe = any( + c and c.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + for c in (cache_type_kv, _ck_extra, _cv_extra) + ) + if tensor_parallel and _cache_non_tensor_safe: logger.info( "Tensor parallelism requires a non-quantized KV cache; " "ignoring cache type %s for the tensor attempt.", cache_type_kv, ) + _tensor_dropped_cache_type_kv = cache_type_kv cache_type_kv = None if extra_args: + # Keep the originals so a layer downgrade restores the real + # (possibly asymmetric) --cache-type-k/-v the layer path + # supports, not just the scalar heavier type. + _tensor_dropped_extra_args = list(extra_args) extra_args = strip_shadowing_flags( extra_args, strip_context = False, @@ -3604,10 +4736,24 @@ class LlamaCppBackend: strip_template = False, strip_split_mode = False, ) + # The launch keeps an inherited tensor-safe env cache type (the + # env cleanup only pops quantized ones), so re-adopt a heavier + # env type (f32) for the budget here too -- mirrors the initial + # adoption, which was skipped because the param/extras set the + # (now-dropped) quantized type. Else the child allocates f32 KV + # against an f16 budget. + _env_tensor_cache = _env_main_cache_type_for_budget() + if _env_tensor_cache is not None: + cache_type_kv = _env_tensor_cache + _cache_type_from_env = True if ctx_override is not None and ctx_override > 0: logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") if cache_override is not None: - logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + _ck, _cv = parse_cache_override_per_axis(extra_args) + logger.info( + f"User --cache-type-k/-v (k={_ck}, v={_cv}) honored; " + "KV estimate budgets the heavier axis" + ) if split_mode_override is not None: logger.info( f"User --split-mode {split_mode_override} honored; " @@ -3632,6 +4778,7 @@ class LlamaCppBackend: "Vision-capable GGUF loaded without a usable mmproj; " "image input will be disabled for this session" ) + model_size = None # set in the fit try; used by the APU RAM guard try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -3639,7 +4786,27 @@ class LlamaCppBackend: self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0 ) model_size = gguf_size + mmproj_size - gpus = self._get_gpu_free_memory() + # 2-tuple gpus for existing logic + a total map for the absolute + # per-GPU headroom (correct when the GPU is already partly used). + _gpu_mem = self._get_gpu_memory() + gpus = [(idx, free) for idx, free, _t in _gpu_mem] + total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + + def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): + # Per-GPU usable budget for ranking: free - (1-frac)*total. + # Callers pass the ACTIVE fraction so the ranking matches the + # budget the fit then tests (else mixed totals mis-order). + idx, free = g + t = total_by_idx.get(idx, 0) + if t > 0: + return free - (1.0 - frac) * t + return free * frac + + def _pool_budget_mib(subset, frac): + # Sum each GPU's own usable budget. Pooling free and total + # separately would let an unknown-total GPU (MIG/vGPU/N/A) + # add full free with no cushion among known-total GPUs. + return sum(max(0.0, _gpu_usable(g, frac)) for g in subset) # Resolve effective context: 0 means let llama-server use # the model's native length. Only expand to a known native @@ -3655,12 +4822,10 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx - # Will MTP engage on this load? If so, auto-fit reserves - # extra VRAM for the draft model. Mirrors - # _build_speculative_flags' resolver: forced mtp / mtp+ngram - # always engage; auto only on an MTP model >= 3B; ngram / - # ngram-simple / off never engage MTP. A separate drafter - # (Gemma) counts as an MTP model just like a baked-in head. + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. + # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always + # engage; auto only on an MTP model >= 3B; ngram/off never. A + # separate drafter (Gemma) counts as an MTP model. _mtp_canonical = _canonicalize_spec_mode(speculative_type) _mtp_effective = _mtp_canonical or "auto" _mtp_size_for_fit = _extract_model_size_b(model_identifier) @@ -3671,49 +4836,278 @@ class LlamaCppBackend: and _mtp_size_for_fit < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) ) - _mtp_will_engage = bool( + # LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras + # nor Studio emit a spec flag (mode "off", no user --spec-type), + # since _build_speculative_flags emits one for every other mode. + # Consult the env for the reserve only then, else a stale MTP env + # would over-reserve. + _spec_env: Mapping[str, str] = ( + os.environ + if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off") + else {} + ) + # Extras can run MTP even when Studio suppresses its own emission. + _user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env) + # A non-MTP model-based draft mode (draft-simple/draft-eagle3) in + # extras also loads a separate draft model that needs reserving; + # engage only when extras actually name a drafter for it. + _user_draft_via_extras = _extra_args_requests_separate_draft( + extra_args, env = _spec_env + ) and bool(_extra_args_mtp_draft_path(extra_args)) + # Mirror _build_speculative_flags: reserve only for MTP the launch + # resolver will actually emit (needs a head/drafter and a binary + # that supports --spec-type mtp). + _mtp_model_for_fit = bool( + self._nextn_predict_layers + or _is_mtp_model_name(model_identifier, model_path) + or bool(mtp_draft_path) + ) and not ( + # Drafterless Gemma falls back to ngram-mod; reserve no + # drafter VRAM for it (mirrors the launch resolver). + _is_gemma_mtp_name(model_identifier, model_path) + and not mtp_draft_path + and not self._nextn_predict_layers + ) + _mtp_binary_ok = True + _mtp_probe_raised = False + if not _user_mtp_via_extras: + try: + _mtp_binary_ok = bool( + (self.probe_server_capabilities(binary) or {}).get("mtp_token") + ) + except Exception: + _mtp_binary_ok = False + _mtp_probe_raised = True + _auto_studio_mtp = ( not _extra_args_set_spec_type(extra_args) + and _mtp_model_for_fit and ( _mtp_effective in ("mtp", "mtp+ngram") - or ( - _mtp_effective == "auto" - and ( - bool(self._nextn_predict_layers) - or _is_mtp_model_name(model_identifier, model_path) - or bool(mtp_draft_path) - ) - and not _mtp_sub_3b_for_fit - ) + or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit) + ) + and ( + _mtp_binary_ok + # Reserve on a raised (uncached) probe too: it re-probes in + # _build_speculative_flags and may still engage MTP (embedded + # head or separate drafter -- _mtp_model_for_fit covers both). + or _mtp_probe_raised ) ) + _mtp_will_engage = bool( + _user_mtp_via_extras or _user_draft_via_extras or _auto_studio_mtp + ) + # The duplicated full target-KV copy (ctx_tgt) is an MTP-only + # cost: the MTP head runs a second context over the target + # model's own KV geometry. The separate-drafter spec modes + # (draft-simple/draft-eagle3, reached via _user_draft_via_extras) + # load a small distinct drafter with its own KV and keep no such + # copy, so only charge it when the engaged mode is truly MTP. + _engaged_is_mtp = bool(_user_mtp_via_extras or _auto_studio_mtp) - # Auto-cap context to fit GPU VRAM and select GPUs. Two - # policies by whether the user set n_ctx: - # Explicit n_ctx: honor it. Try the full context with - # _select_gpus (as many GPUs as needed); cap only if it - # doesn't fit on any combination. - # Auto n_ctx=0 (native): prefer fewer GPUs with reduced - # context, since multi-GPU is slower. + # Effective draft depth: extras win (last-wins at launch), else + # the field, else the platform default (2 GPU / 3 CPU). + _extra_n_max = _extra_args_spec_draft_n_max(extra_args) + _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max + if _mtp_eff_n_max is None: + _mtp_eff_n_max = 2 if gpus else 3 + # Separate-drafter weights live on GPU (an embedded head is + # already in model_size). Size the drafter the launch loads, by + # precedence: extras --model-draft (last-wins), else Studio's + # emitted mtp_draft_path, else the env drafter. Sizing the wrong + # one would under-reserve and OOM. + _cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {}) + _studio_draft_for_budget = ( + mtp_draft_path + if ( + _mtp_will_engage + and mtp_draft_path + and not _extra_args_set_spec_type(extra_args) + ) + else None + ) + _env_draft_for_budget = _extra_args_mtp_draft_path([], env = os.environ) + _mtp_draft_for_budget = ( + _cli_draft_for_budget or _studio_draft_for_budget or _env_draft_for_budget + ) + # Drafter offloaded to CPU keeps its weights+KV off the GPU, so + # drop it from the budget (an embedded head stays in the model). + # Consult the env too: the child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT. + _draft_on_cpu = _extra_args_draft_offloaded_to_cpu(extra_args, env = os.environ) + if _draft_on_cpu: + _mtp_draft_for_budget = None + _mtp_draft_weights = 0 + if _mtp_draft_for_budget: + try: + _mtp_draft_weights = self._get_gguf_size_bytes(_mtp_draft_for_budget) + except Exception: + _mtp_draft_weights = 0 + # Draft K/V types (f16 by default; independent extras overrides). + _mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types(extra_args) + + # Byte-accurate reserve when dims allow, else None -> flat fallback. + mtp_overhead_fn: Optional[Callable[[int], int]] = None + # True when the byte reserve is the drafter weights ONLY because + # its KV couldn't be sized; the flat fraction must then stay on + # as the cushion for that unsized draft KV (it is not covered by + # the weights-only mtp_overhead_fn). + _mtp_kv_unsized = False + if _mtp_will_engage: + _probe_ctx = self._context_length or ( + effective_ctx if effective_ctx > 0 else 4096 + ) + _draft_kv_probe = self._mtp_draft_kv_bytes( + _probe_ctx, + drafter_path = _mtp_draft_for_budget, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + n_parallel = n_parallel, + ) + if ( + self._estimate_mtp_overhead_bytes( + _probe_ctx, + spec_draft_n_max = _mtp_eff_n_max, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + drafter_path = _mtp_draft_for_budget, + draft_weights_bytes = _mtp_draft_weights, + n_parallel = n_parallel, + mtp_keeps_target_ctx = _engaged_is_mtp, + ) + is not None + ): + # Reserve is weights-only when the draft KV is unsizable. + _mtp_kv_unsized = _draft_kv_probe is None + + # Closure binding this load's draft params; ctx varies. + def mtp_overhead_fn( + ctx: int, + _n: int = _mtp_eff_n_max, + _ck: Optional[str] = _mtp_draft_ck, + _cv: Optional[str] = _mtp_draft_cv, + _dp: Optional[str] = _mtp_draft_for_budget, + _w: int = _mtp_draft_weights, + _np: int = n_parallel, + _mtp: bool = _engaged_is_mtp, + ) -> int: + v = self._estimate_mtp_overhead_bytes( + ctx, + spec_draft_n_max = _n, + draft_cache_type_k = _ck, + draft_cache_type_v = _cv, + drafter_path = _dp, + draft_weights_bytes = _w, + n_parallel = _np, + mtp_keeps_target_ctx = _mtp, + ) + return v if v is not None else 0 + + def _mtp_bytes(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + + # Effective micro-batch (a user --ubatch override scales the + # compute buffer); None -> the 512 default in the estimate. + _effective_ubatch = _extra_args_n_ubatch(extra_args) + + # Layer-split compute buffer (one lump; tensor mode reserves it + # per device in _plan_tensor_parallel). Context-independent, so + # fold it into the model footprint for the branches below. Falls + # back to the flat reserve when dims are missing (returns 0), a + # safe upper bound since the tensor buffer >= the layer one. + _compute_buffer_pipeline = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = False, + ) + if _compute_buffer_pipeline <= 0: + _compute_buffer_pipeline = ( + self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + ) + model_size_fit = model_size + _compute_buffer_pipeline + + # Layer split adds a fixed per-device overhead on every GPU. The + # folded buffer covers one device; reserve the extra devices' + # share so a k-GPU split can't pin a context that OOMs a device + # (k=1 adds nothing). + _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + + # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: + # honor it, cap only if it fits no combination. Auto (native): + # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True # Per-GPU weight proportions for tensor mode (None = even). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 - # MTP draft model lives outside the main estimates; carve - # its reserve out of every fit budget and pin threshold so - # a load can't pin into the drafter's headroom. - _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0 - _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve + # Flat MTP reserve fraction: used only as the fallback when the + # byte-accurate mtp_overhead_fn can't size the draft KV (dims + # unavailable, or _mtp_kv_unsized = weights-only). A separate + # drafter on CPU uses no GPU (no reserve); an embedded head is on + # GPU regardless of draft-offload flags (keep its reserve). + _flat_mtp_engages = _mtp_will_engage and ( + mtp_overhead_fn is None or _mtp_kv_unsized + ) + _draft_cpu_no_embedded = _draft_on_cpu and not self._nextn_predict_layers + # MTP reserves GPU VRAM unless its only drafter is a separate + # CPU-offloaded one (an embedded head stays on GPU). The tensor + # path reserves like the layer path; gate both on this. + _mtp_reserves_gpu = _mtp_will_engage and not _draft_cpu_no_embedded + _flat_mtp_reserve = ( + _MTP_VRAM_RESERVE_FRAC + if (_flat_mtp_engages and not _draft_cpu_no_embedded) + else 0.0 + ) + _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve + # Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below. + _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) - # Tensor mode allocates a compute-graph buffer on every - # participating GPU, so a GPU with less free VRAM than that - # reserve can't host it and would OOM at load. Drop those - # from the tensor-parallel set up front (gpu_indices below - # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded - # from llama-server entirely, not just given zero weight). + def _restore_after_tensor_downgrade(): + # Tensor mode dropped a quantized KV and stripped the cache + # extras (it rejects quantized); layer split supports them, so + # restore the original type + extras (minus --split-mode) and + # clear the env flag so the layer launch re-emits them. + nonlocal cache_type_kv, _cache_type_from_env, extra_args + if _tensor_dropped_cache_type_kv is not None: + cache_type_kv = _tensor_dropped_cache_type_kv + _cache_type_from_env = False + extra_args = strip_split_mode_only( + _tensor_dropped_extra_args + if _tensor_dropped_extra_args is not None + else extra_args + ) + + if tensor_parallel and effective_is_vision: + logger.info( + "Tensor parallelism skipped for vision model: " + "--split-mode tensor is incompatible with --mmproj " + "in the current llama.cpp build; using layer split." + ) + tensor_parallel = False + _restore_after_tensor_downgrade() + + # Tensor mode replicates a compute buffer on every GPU, so drop + # GPUs below that reserve from the set up front (gpu_indices + # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). tp_gpus = gpus if tensor_parallel: - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - tp_gpus = [g for g in gpus if g[1] >= reserve_mib] + # Deterministic per-device compute buffer (replicated on + # every device in tensor mode); flat fallback when dims + # are unavailable. _plan_tensor_parallel uses the same. + _tp_reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = True, + ) + reserve_mib = ( + _tp_reserve_bytes // (1024 * 1024) + if _tp_reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + # Admit by usable budget (free - (1-frac)*total), not raw + # free: a partly-used big card can clear the reserve on raw + # free yet have no budget left. + tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] if tensor_parallel and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single @@ -3729,21 +5123,60 @@ class LlamaCppBackend: len(gpus), ) tensor_parallel = False - # A user --split-mode tensor in extras is appended after - # Studio's flags, so it would still reach llama-server and - # fail here; strip it so the downgrade actually applies. - extra_args = strip_split_mode_only(extra_args) + # Layer split supports a quantized KV the tensor attempt + # dropped; restore the original cache type + extras (minus + # --split-mode) so the layer launch re-emits them. + _restore_after_tensor_downgrade() if tensor_parallel and tp_gpus: - # Tensor-parallel allocation: use all usable GPUs, weight - # the split by (free - buffer), and cap context to the - # pooled VRAM after weights + per-device compute-graph - # buffers. See _plan_tensor_parallel for the policy. + # Pooled usable budget (after each device's compute buffer) + # must hold the non-shrinkable footprint: weights + the MTP + # reserve. The planner can shrink ctx/KV, not these. + _tp_weight_budget_mib = ( + sum(_gpu_usable(g) for g in tp_gpus) - len(tp_gpus) * reserve_mib + ) + _tp_flat_mtp = 2 * 1024**3 # flat reserve when dims unavailable + if not _mtp_reserves_gpu: + # No MTP, or its only drafter is CPU-offloaded (no GPU). + _tp_mtp_floor = 0 + elif mtp_overhead_fn is not None and not _mtp_kv_unsized: + _tp_mtp_floor = _mtp_bytes( + min(2048, effective_ctx) if effective_ctx > 0 else 2048 + ) + else: + # Dims unavailable / weights-only: tensor mode has no + # --fit valve, so keep the flat reserve as the unsized-KV + # cushion, never below the known byte reserve. + _tp_mtp_floor = max( + _tp_flat_mtp, + _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), + ) + _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + if _tp_weight_budget_mib <= _tp_required_mib: + logger.info( + "Tensor parallelism requested but the pooled VRAM " + "budget cannot hold the weights, MTP reserve, and " + "per-device compute buffers; falling back to layer split." + ) + tensor_parallel = False + # Restore the dropped quantized KV + original cache extras + # (minus --split-mode); layer split supports them. + _restore_after_tensor_downgrade() + + if tensor_parallel and tp_gpus: + # Tensor-parallel allocation; see _plan_tensor_parallel. target_ctx = ( effective_ctx if explicit_ctx else (self._context_length or effective_ctx) ) + # When the draft KV couldn't be sized (weights-only reserve), + # the planner's mtp_overhead_fn is non-None but covers only + # weights, so pass the flat cushion for the unsized KV (else + # the binary search spends it on context). + _tp_unsized_mtp_reserve = ( + 2 * 1024**3 if (_mtp_reserves_gpu and _mtp_kv_unsized) else 0 + ) ( effective_ctx, max_available_ctx, @@ -3755,10 +5188,14 @@ class LlamaCppBackend: target_ctx, cache_type_kv = cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + mtp_flat_reserve_bytes = _tp_unsized_mtp_reserve, # Report the UI ceiling from native ctx, not the # explicit small request. max_target_ctx = self._context_length or target_ctx, + total_by_idx = total_by_idx, + n_ubatch = _effective_ubatch, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -3767,24 +5204,38 @@ class LlamaCppBackend: # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: - ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + ranked_for_cap = sorted( + gpus, + key = lambda g: _gpu_usable( + g, _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve + ), + reverse = True, + ) best_cap = 0 + _cap_fraction = _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve for n_gpus in range(1, len(ranked_for_cap) + 1): subset = ranked_for_cap[:n_gpus] - pool_mib = sum(free for _, free in subset) + # Per-GPU-consistent pool budget (fixes mixed + # known/unknown totals); pass it as an absolute + # budget so the fit and the check below agree. + pool_budget = _pool_budget_mib(subset, _cap_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( native_ctx_for_cap, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * (_CTX_FIT_VRAM_FRACTION - _mtp_reserve): + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -3798,34 +5249,49 @@ class LlamaCppBackend: # Honor the requested context verbatim. If it fits, # pin GPUs and skip --fit; else ship -c --fit # on and let llama-server flex -ngl (CPU offload). - requested_total = model_size + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel + requested_total = ( + model_size_fit + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) + + _mtp_bytes(effective_ctx) ) gpu_indices, use_fit = self._select_gpus( - requested_total, gpus, usable_fraction = _pin_fraction + requested_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) # No silent shrink: effective_ctx stays == requested_ctx. else: # Auto context: prefer fewer GPUs, cap to fit. Same - # headroom threshold as _select_gpus (#5106). - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # headroom threshold as _select_gpus (#5106). Rank by the + # active pin fraction so the order matches the fit budget. pin_fraction = _pin_fraction + ranked = sorted( + gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True + ) for n_gpus in range(1, len(ranked) + 1): subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) + pool_budget = _pool_budget_mib(subset, pin_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( effective_ctx, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False @@ -3838,14 +5304,17 @@ class LlamaCppBackend: 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 = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel, ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = ( + _subset_model_size(n_gpus) + + kv + + _mtp_bytes(effective_ctx) + ) / (1024 * 1024) + if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) use_fit = False break @@ -3857,14 +5326,82 @@ class LlamaCppBackend: "Falling back to file-size-only GPU selection", model_size_gb = round(model_size / (1024**3), 2), ) + # Add the byte-accurate MTP reserve here too when it is + # available; otherwise _pin_fraction carries the flat + # fallback (the two are mutually exclusive by design). + _fs_total = model_size_fit + _mtp_bytes( + self._context_length or effective_ctx or 4096 + ) gpu_indices, use_fit = self._select_gpus( - model_size, gpus, usable_fraction = _pin_fraction + _fs_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 # so the slider isn't on an unusable native ctx. effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096 + elif _apple_budget_mib > 0 and effective_ctx > 0: + # No GPU on Metal: the branches above are skipped and the context + # stays at native, over-committing unified memory (#5118, #6529). + # Cap with the same fit math (--fit on stays as a backstop); only + # auto context shrinks, explicit is honored. + native_ctx_for_cap = self._context_length or effective_ctx + # Reserve the flat MTP fraction up front like the discrete + # _pin_fraction, so an unsized MTP draft (e.g. Qwen3.6-MTP, #6529) + # can't over-commit. No-op when MTP is off; exclusive with the + # byte-accurate _mtp_bytes reserve. + _apple_fit_budget_mib = int( + _apple_budget_mib * max(0.0, 1.0 - _flat_mtp_reserve) + ) + if self._can_estimate_kv(): + cap = self._fit_context_to_vram( + native_ctx_for_cap, + _apple_fit_budget_mib, + model_size_fit, + cache_type_kv, + n_parallel = n_parallel, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, + ) + _cap_footprint_mib = ( + model_size_fit + + self._estimate_kv_cache_bytes( + cap, cache_type_kv, n_parallel = n_parallel + ) + + _mtp_bytes(cap) + ) / (1024 * 1024) + # Fit returns the request unchanged when it fits OR weights + # exceed budget; only the latter over-commits, so floor to 4096. + max_available_ctx = ( + cap + if _cap_footprint_mib <= _apple_fit_budget_mib + else min(4096, native_ctx_for_cap) + ) + else: + # No KV estimate: mirror the discrete file-size-only fallback + # and floor to 4096 rather than launch at native and over-commit. + max_available_ctx = min(4096, native_ctx_for_cap) + if not explicit_ctx: + effective_ctx = max_available_ctx + + # MTP reserve at the final context, for the logs below. + _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 + if _mtp_will_engage: + _mtp_note = ( + f"MTP reserve: {_mtp_reserve_bytes / (1024**3):.2f} GB " + f"(draft KV @ {effective_ctx} + verify n_max={_mtp_eff_n_max}" + + (", flat-frac fallback" if mtp_overhead_fn is None else "") + + "), " + ) + else: + _mtp_note = "" + if effective_ctx < original_ctx: kv_est = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel @@ -3872,7 +5409,9 @@ class LlamaCppBackend: logger.info( f"Context auto-reduced: {original_ctx} -> {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_est / (1024**3):.1f} GB)" + f"est. KV cache: {kv_est / (1024**3):.1f} GB, " + f"{_mtp_note}".rstrip(", ") + + ")" ) kv_cache_bytes = self._estimate_kv_cache_bytes( @@ -3885,6 +5424,7 @@ class LlamaCppBackend: f"GGUF size: {gguf_size / (1024**3):.1f} GB, " f"{mmproj_note}" f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " + f"{_mtp_note}" f"context: {effective_ctx}, " f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" ) @@ -3894,6 +5434,17 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # Unified-memory APUs load weights into system RAM (under WSL the VM + # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an + # oversize load the OS would otherwise kill mid-flight. Base model + # only: an optional MTP drafter is dropped by the MTP-drop fallback. + if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices): + _ram_msg = self._apu_ram_shortfall_message( + model_size, self._available_system_memory_mib() + ) + if _ram_msg: + raise RuntimeError(_ram_msg) + # Audio input straight from the mmproj (clip.has_audio_encoder), # independent of token names. self._mmproj_has_audio = False @@ -3924,27 +5475,49 @@ class LlamaCppBackend: "--no-context-shift", ] + fully_gpu_offloaded = False if use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: # Fits on selected GPU(s) -- offload all layers cmd.extend(["-ngl", "-1"]) + fully_gpu_offloaded = True + server_caps = self.probe_server_capabilities(binary) + # Expose Prometheus /metrics for the engine-stats logger, only + # when the binary advertises it (older/custom binaries may not). + if server_caps.get("supports_metrics"): + cmd.append("--metrics") cmd.extend( self._ctx_integrity_flags( n_parallel, use_fit, requested_ctx, effective_ctx, - self.probe_server_capabilities(binary), + server_caps, ) ) + offload_overridden = _extra_args_set_any_flag( + extra_args, _GPU_OFFLOAD_OVERRIDE_FLAGS + ) + threads_overridden = _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS) + full_offload_tuning_active = fully_gpu_offloaded and not offload_overridden - # -1 = llama.cpp auto-detect (physical cores). Pass explicitly - # so we don't inherit llama-server's internal default, which - # has varied (hardware concurrency incl. hyperthreads on some - # builds). - cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)]) + # Thread count: an unset --threads makes llama.cpp pick physical + # cores (common_cpu_get_num_math), but an explicit --threads -1 + # resolves to hardware_concurrency() (every hyperthread), which + # contends on the memory bus and slows CPU / hybrid decode. So + # omit the flag when unset and only pin it for an explicit + # override or the Windows full-offload OpenMP cap. Pass-through + # thread flags in extra_args still win (appended last). #5692 + if ( + sys.platform == "win32" + and full_offload_tuning_active + and not threads_overridden + ): + cmd.extend(["--threads", "2"]) + elif n_threads is not None and n_threads > 0: + cmd.extend(["--threads", str(n_threads)]) # Enable Jinja chat template rendering cmd.extend(["--jinja"]) @@ -3961,7 +5534,11 @@ class LlamaCppBackend: "iq4_nl", "f32", } - if cache_type_kv and cache_type_kv in _valid_cache_types: + if ( + cache_type_kv + and cache_type_kv in _valid_cache_types + and not _cache_type_from_env + ): cmd.extend( [ "--cache-type-k", @@ -3973,6 +5550,8 @@ class LlamaCppBackend: self._cache_type_kv = cache_type_kv logger.info(f"KV cache type: {cache_type_kv}") else: + # An env-only type is left inherited (untouched) so an + # asymmetric K/V env reaches the child as set. self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor @@ -4028,6 +5607,7 @@ class LlamaCppBackend: ) self._supports_reasoning = flags["supports_reasoning"] self._reasoning_style = flags["reasoning_style"] + self._reasoning_effort_levels = flags.get("reasoning_effort_levels", []) self._reasoning_always_on = flags["reasoning_always_on"] self._supports_preserve_thinking = flags["supports_preserve_thinking"] self._supports_tools = flags["supports_tools"] @@ -4076,16 +5656,37 @@ class LlamaCppBackend: logger.info(f"Using mmproj for vision: {launch_mmproj_path}") # Option C: --api-key for direct client access when enabled - import os as _os import secrets as _secrets - if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": + if os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": self._api_key = _secrets.token_urlsafe(32) cmd.extend(["--api-key", self._api_key]) logger.info("llama-server started with --api-key for direct streaming") else: self._api_key = None + # Windows + full offload: disable KV checkpoints (WDDM/PCI-E + # overhead). CPU/partial offload keeps prompt caching. #5692. + if sys.platform == "win32" and full_offload_tuning_active: + unsupported_cache_flags: list[str] = [] + if server_caps.get("supports_cache_ram"): + cmd.extend(["--cache-ram", "0"]) + else: + unsupported_cache_flags.append("--cache-ram") + if server_caps.get("supports_ctx_checkpoints"): + cmd.extend(["--ctx-checkpoints", "0"]) + else: + unsupported_cache_flags.append("--ctx-checkpoints") + if server_caps.get("supports_no_cache_prompt"): + cmd.append("--no-cache-prompt") + else: + unsupported_cache_flags.append("--no-cache-prompt") + if unsupported_cache_flags: + logger.info( + "Skipping unsupported Windows cache flags for llama-server: %s", + ", ".join(unsupported_cache_flags), + ) + # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Studio's auto-set flags. Already # validated by the route via validate_extra_args(). @@ -4093,23 +5694,52 @@ class LlamaCppBackend: cmd.extend(str(a) for a in extra_args) logger.info(f"Appending user extra args to llama-server: {list(extra_args)}") - _log_cmd = list(cmd) - if "--api-key" in _log_cmd: - _ki = _log_cmd.index("--api-key") + 1 - if _ki < len(_log_cmd): - _log_cmd[_ki] = "" - logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") + logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") # Library paths so llama-server finds its shared libs and CUDA DLLs. - import os - import sys + env = self._llama_server_env_for_binary(binary) + # Omitting --threads relies on llama.cpp's physical-core default, so + # drop an inherited LLAMA_ARG_THREADS that would otherwise feed the + # arg handler and silently force hardware_concurrency(). #5692 + if "--threads" not in cmd: + env.pop("LLAMA_ARG_THREADS", None) - env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) + # Reconcile the inherited LLAMA_ARG_* env with Studio's final + # decision: stripping CLI extras on a tensor->layer downgrade + # can't remove env vars, so the child could run a mode/KV Studio + # didn't budget. + if not tensor_parallel: + # Layer split: clear a non-layer inherited split mode (and any + # paired tensor-split) so the child can't override the layer plan. + _inherited_sm = (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower() + if _inherited_sm and _inherited_sm != "layer": + env.pop("LLAMA_ARG_SPLIT_MODE", None) + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + else: + # Studio owns the tensor split: it emits --tensor-split when it + # picks an uneven one (CLI wins) and nothing when an even split + # is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even + # case can't be overridden by a stale env (the layer branch above + # clears it too). + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + # Tensor split aborts on a quantized KV; clear an inherited + # quantized cache type so the child uses the tensor-safe default. + for _ct_var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + _ct_raw = (env.get(_ct_var) or "").strip().lower() + if _ct_raw and _ct_raw not in self._TENSOR_PARALLEL_KV_TYPES: + env.pop(_ct_var, None) + + # Windows + full offload: PASSIVE OMP + 2 threads stop + # spin-wait burning CPU. CPU/partial offload keeps default + # OMP parallelism. #5692. + if sys.platform == "win32" and full_offload_tuning_active: + env.setdefault("OMP_WAIT_POLICY", "PASSIVE") + if not threads_overridden: + env.setdefault("OMP_NUM_THREADS", "2") # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use # shared system RAM. setdefault so a user value wins. - if self._amd_apu_wants_unified_memory(): + if self._amd_apu_wants_unified_memory(gpu_indices): env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") @@ -4121,96 +5751,6 @@ class LlamaCppBackend: f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) - if sys.platform == "win32": - # Ordering: see _build_windows_path_dirs. #5106. - path_dirs = self._build_windows_path_dirs( - binary_dir, - sys.prefix, - os.environ.get("CUDA_PATH", ""), - ) - existing_path = env.get("PATH", "") - env["PATH"] = ";".join(path_dirs) + ";" + existing_path - - # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile - # kernel files (rocblas/library/*.dat + *.hsaco); the DLL - # searches /rocblas/library/ which doesn't exist - # -> silent crash on the first GEMM. ROCBLAS_TENSILE_LIBPATH - # repoints that search at the ROCm install. - _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", "")) - if _hip_path: - _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library") - if os.path.isdir(_rocblas_lib): - env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib) - else: - # Linux: LD_LIBRARY_PATH for shared libs next to the binary - # plus CUDA runtime libs (libcudart, libcublas, etc.) - import platform - - lib_dirs = [] - # WSL: system HIP before the bundle's (which segfaults on - # /dev/dxg). Mirror install_llama_prebuilt.binary_env, which - # validates the prebuilt with this same ordering. - for _wsl_rocm in _wsl_system_rocm_lib_dirs(): - lib_dirs.append(_wsl_rocm) - if lib_dirs: - env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") - lib_dirs.append(binary_dir) - _arch = platform.machine() # x86_64, aarch64, etc. - - # Pip-installed nvidia CUDA runtime libs. The prebuilt - # binary links libcudart.so.13 / libcublas.so.13 which live - # here, not in /usr/local/cuda. - import glob as _glob - - for _nv_pattern in [ - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cu*", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cudnn", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "nvjitlink", - "lib", - ), - ]: - for _nv_dir in _glob.glob(_nv_pattern): - if os.path.isdir(_nv_dir): - lib_dirs.append(_nv_dir) - - for cuda_lib in [ - "/usr/local/cuda/lib64", - f"/usr/local/cuda/targets/{_arch}-linux/lib", - # Fallback CUDA compat paths (e.g. binary built with - # CUDA 12 where default /usr/local/cuda is CUDA 13+). - "/usr/local/cuda-12/lib64", - "/usr/local/cuda-12.8/lib64", - f"/usr/local/cuda-12/targets/{_arch}-linux/lib", - f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", - ]: - if os.path.isdir(cuda_lib): - lib_dirs.append(cuda_lib) - existing_ld = env.get("LD_LIBRARY_PATH", "") - new_ld = ":".join(lib_dirs) - env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld - # Pin to selected GPU(s). On ROCm, narrowing only # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full # set, so set HIP_VISIBLE_DEVICES too. @@ -4250,6 +5790,9 @@ class LlamaCppBackend: # retry once with --fit off before declaring the load failed. # Never retry when fit was requested (use_fit) or the caller # passed an explicit fit flag via extra args. + # Argv actually launched (post --fit off / MTP); text-only retry strips this. + _last_spawn_cmd = list(cmd) + def _spawn_and_wait(run_cmd, *, label = ""): """Start llama-server with run_cmd and wait for health. @@ -4257,6 +5800,7 @@ class LlamaCppBackend: crashes during startup and run_cmd is eligible (see _fit_off_retry_eligible). """ + nonlocal _last_spawn_cmd _fit_retry_allowed = self._fit_off_retry_eligible(run_cmd, use_fit) for _spawn_attempt in (0, 1): # Defensive kill: drop an orphan Popen a concurrent load may @@ -4291,6 +5835,7 @@ class LlamaCppBackend: # Best-effort; never block the load on logging. logger.debug(f"Could not open llama-server log file: {e}") self._llama_log_path = None + _last_spawn_cmd = list(run_cmd) self._process = subprocess.Popen( run_cmd, stdout = subprocess.PIPE, @@ -4298,7 +5843,9 @@ class LlamaCppBackend: text = True, env = env, **_windows_hidden_subprocess_kwargs(), + **_child_popen_kwargs(), ) + self._record_server_pid(self._process.pid) # Background thread to drain stdout (prevents pipe deadlock) self._stdout_thread = threading.Thread( @@ -4356,6 +5903,80 @@ class LlamaCppBackend: ) healthy = _spawn_and_wait(cmd) + # Flash-attention kernels hard-crash at startup on some ROCm/GPU + # builds (frequently inside the vision tower). Disabling FA keeps + # both vision and MTP, so retry that way before dropping either. + # Only on a hard fault with FA on; a cancel/unload stops respawn. + if not healthy and not self._cancel_event.is_set(): + _fa_rc = self._process.poll() if self._process is not None else None + _fa_cmd = ( + self._with_flash_attn_off(_last_spawn_cmd) + if self._is_signal_crash(_fa_rc) + else None + ) + if _fa_cmd is not None: + logger.warning( + "llama-server hard-crashed at startup (exit %s) with " + "flash attention on; retrying once with --flash-attn " + "off (keeps vision and MTP).", + _fa_rc, + ) + self._kill_process() + cmd = _fa_cmd + healthy = _spawn_and_wait(_fa_cmd, label = "-noflash") + + # MTP from Studio's spec flags or the user's (extra_args + # --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child + # only when neither emits a spec flag, so consult it only then. + _launch_spec_env: Mapping[str, str] = ( + os.environ + if (not _extra_args_set_spec_type(extra_args) and not spec_flags) + else {} + ) + _spec_requested_mtp = any( + "mtp" in str(t).lower() for t in spec_flags + ) or _extra_args_requests_mtp(extra_args, env = _launch_spec_env) + # Is the launched server actually running MTP+tensor? Gates the + # probe/watchdog/recovery; cleared if the MTP-drop fallback wins. + _mtp_active_for_launched_server = bool( + self._tensor_parallel and _spec_requested_mtp + ) + # MTP can pass /health then crash the flash-attn kernel on the + # first decode under tensor; probe one generation so the fallback + # catches that too. Tensor-only, so ordinary MTP stays probe-free. + if ( + healthy + and self._tensor_parallel + and _spec_requested_mtp + and not self._cancel_event.is_set() + and not self._probe_mtp_decode() + ): + # A first-decode hard fault is usually the FA kernel: retry + # FA-off (keeps MTP) before dropping speculative decoding below. + _probe_rc = self._process.poll() if self._process is not None else None + _fa_cmd = ( + self._with_flash_attn_off(_last_spawn_cmd) + if self._is_signal_crash(_probe_rc) + else None + ) + healthy = False + if _fa_cmd is not None: + logger.warning( + "MTP first-decode hard-crashed (exit %s) with flash " + "attention on; retrying with --flash-attn off.", + _probe_rc, + ) + self._kill_process() + cmd = _fa_cmd + healthy = ( + _spawn_and_wait(_fa_cmd, label = "-noflash-mtp") + and self._probe_mtp_decode() + ) + if not healthy: + logger.warning( + "MTP speculative decoding crashed on the first decode " + "under tensor parallelism; retrying without it." + ) # Any MTP request can abort the server: a separate drafter # (Gemma) on a binary that predates its arch, or an embedded # head (Qwen) the binary cannot build. Retry once with the @@ -4363,8 +5984,8 @@ class LlamaCppBackend: # loads. Gate on the spec block (not the drafter path, which # off/ngram local loads also carry) and keep # _requested_spec_mode so a duplicate /load doesn't thrash. The - # cancel check stops an /unload-killed attempt respawning. - _spec_requested_mtp = any("mtp" in str(t).lower() for t in spec_flags) + # cancel check stops an /unload-killed attempt respawning. A + # decode-probe failure above also routes here. if not healthy and _spec_requested_mtp and not self._cancel_event.is_set(): # Blame the binary only when the output shows MTP itself # failing (unknown arch / draft or context build); an @@ -4410,29 +6031,48 @@ class LlamaCppBackend: + ["--spec-default"] + cmd[_spec_start + len(spec_flags) :] ) + # User/env MTP survives in the tail; llama.cpp takes the last + # spec flag, so a trailing --spec-default overrides it too. + if _extra_args_requests_mtp(extra_args, env = _launch_spec_env): + fallback_cmd.append("--spec-default") healthy = _spawn_and_wait(fallback_cmd, label = "-retry") if healthy: self._speculative_type = "default" + _mtp_active_for_launched_server = False - # A vision GGUF launched with --mmproj can abort when the - # installed llama.cpp is too old for the model's projector - # ("Unknown projector type"); in that one case retry once - # text-only rather than failing the whole load. + # A too-old llama.cpp can reject a model's --mmproj projector + # (format message or a bare SIGSEGV); retry once text-only. if not healthy: out = "\n".join(self._stdout_lines[-50:]) + # Read the crash code before _kill_process() clears _process. + _crash_rc = self._process.poll() if self._process is not None else None self._kill_process() - if launched_with_mmproj and self._is_projector_incompatibility(out): + # Skip if a cancel/unload is pending (mirrors the MTP guard). + if ( + launched_with_mmproj + and not self._cancel_event.is_set() + and ( + self._is_projector_incompatibility(out) + or ( + self._is_signal_crash(_crash_rc) + and not self._output_has_nonprojector_diagnostic(out) + ) + ) + ): logger.warning( "llama-server could not load this model's vision " "projector (--mmproj). The installed llama.cpp build is " "likely too old for it. Loading text-only for this " "session; run 'unsloth studio update' to enable vision." ) - cmd = self._strip_mmproj_args(cmd) + cmd = self._strip_mmproj_args(_last_spawn_cmd) self._is_vision = False self._mmproj_has_audio = False self._start_llama_process(cmd, env) if not self._wait_for_health(timeout = 600.0): + # Read the exit code before _kill_process() clears it, so + # an OS-killed text-only retry still gets the OOM message. + _retry_rc = self._process.poll() if self._process is not None else None self._kill_process() raise RuntimeError( "Vision projector incompatible with this llama.cpp " @@ -4441,6 +6081,7 @@ class LlamaCppBackend: "\n".join(self._stdout_lines[-50:]), gguf_path, self._model_identifier, + _retry_rc, ) ) else: @@ -4449,6 +6090,7 @@ class LlamaCppBackend: out, gguf_path, self._model_identifier, + _crash_rc, ) ) @@ -4462,6 +6104,11 @@ class LlamaCppBackend: self._extra_args = list(extra_args) self._extra_args_source = (model_identifier, hf_variant) self._requested_n_ctx = int(n_ctx) + # Commit the known-good snapshot + whether MTP+tensor is live, then + # watch this load for a mid-generation crash. + self._last_load_kwargs = _pending_load_kwargs + self._mtp_runtime_fallback_active = _mtp_active_for_launched_server + self._start_mtp_crash_watchdog() # Catch silent CPU fallback when GPU was intended (#5106). self._gpu_offload_active = self._classify_gpu_offload( @@ -4481,6 +6128,18 @@ class LlamaCppBackend: logger.info( f"llama-server ready on port {self._port} for model '{model_identifier}'" ) + # Poll llama-server /metrics -> vLLM-style engine_stats logs + # (only when the binary exposes /metrics). + if server_caps.get("supports_metrics"): + try: + from core.inference.llama_stats import maybe_start_stats_logger + if self._stats_logger is not None: + self._stats_logger.stop() + self._stats_logger = maybe_start_stats_logger(self.base_url, logger) + except Exception as e: + logger.debug(f"engine-stats logger not started: {e}") + else: + self._stats_logger = None # Probe outside _lock (interruptible by /unload); init inside. self._is_audio = False @@ -4493,37 +6152,8 @@ class LlamaCppBackend: except Exception as exc: logger.debug("Audio probe failed: %s", exc) detected = None - if detected in ("snac", "bicodec", "dac"): - with self._lock: - if not self._healthy: - return False - try: - self.init_audio_codec(detected) - self._is_audio = True - self._audio_type = detected - except Exception as exc: - # Surface as HTTP 500 (matches pre-PR contract). - logger.warning( - "Failed to init audio codec '%s': %s", - detected, - exc, - ) - self._audio_probed = False - return False - elif detected: - # csm / whisper / audio_vlm: track type but keep _is_audio - # False -- GGUF TTS routing only fires for snac/bicodec/dac. - with self._lock: - if not self._healthy: - return False - self._audio_type = detected - - # Audio input = token probe (audio_vlm/whisper) OR mmproj encoder. - from utils.models.model_config import is_audio_input_type - - self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( - self._mmproj_has_audio - ) + if not self._apply_detected_audio(detected): + return False if not self._healthy: return False @@ -4608,6 +6238,24 @@ class LlamaCppBackend: _mtp_too_small = ( _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) ) + # Drafterless Gemma (name-only MTP, no embedded head): emitting MTP + # would abort llama-server, so every mode below falls back instead. + _mtp_drafter_missing = ( + _is_gemma_mtp_name(model_identifier, model_path) + and not mtp_draft_path + and not self._nextn_predict_layers + ) + # Embedded MTP head on an MLA model (GLM-5.2/DeepSeek/Kimi, detected by + # kv_lora_rank): llama.cpp's MLA/DSA MTP path is ~2x slower than no spec, + # so Auto drops it (override via the Settings dropdown / forced mtp, or + # UNSLOTH_MLA_MTP_ENABLED=1). Separate drafters (Gemma, mtp_draft_path) and + # non-MLA embedded heads (Qwen, no kv_lora_rank) are unaffected. + _auto_mla_embedded_mtp = ( + bool(self._nextn_predict_layers) + and self._kv_lora_rank is not None + and not bool(mtp_draft_path) + and not _mla_mtp_auto_enabled() + ) if user_owns_spec_type: # User --spec-type wins outright; suppress auto-emit to avoid a @@ -4639,6 +6287,11 @@ class LlamaCppBackend: "run `unsloth studio update`. Loading without " "speculative decoding." ) + # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins + # over env) so the child matches the binary-capability gate and + # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. + flags.append("--spec-default") + self._speculative_type = "default" self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() @@ -4649,6 +6302,8 @@ class LlamaCppBackend: if mtp_draft_path: flags.extend(["--model-draft", mtp_draft_path]) logger.info(f"Using separate MTP drafter: {mtp_draft_path}") + spec_value = mtp_token + ngram_knobs: list[str] = [] if chain_ngram: ngram_knobs = _build_ngram_mod_flags(caps) if ngram_knobs: @@ -4658,25 +6313,8 @@ class LlamaCppBackend: "llama-server lacks ngram-mod tuning " "flags; loading MTP only (no ngram chain)" ) - spec_value = mtp_token - flags.extend( - [ - "--spec-type", - spec_value, - n_max_flag, - str(draft_n_max), - ] - ) - flags.extend(ngram_knobs) - else: - flags.extend( - [ - "--spec-type", - mtp_token, - n_max_flag, - str(draft_n_max), - ] - ) + flags.extend(["--spec-type", spec_value, n_max_flag, str(draft_n_max)]) + flags.extend(ngram_knobs) self._speculative_type = "draft-mtp" chain_label = "chained ngram-mod" if chain_ngram else "MTP-only" logger.info(f"Spec decoding: {mtp_token} ({chain_label})") @@ -4697,6 +6335,20 @@ class LlamaCppBackend: logger.info("Spec decoding: ngram-mod") return True + def _fallback_drafter_not_found() -> None: + """Drafterless Gemma: use ngram-mod (or spec-default) and record why.""" + logger.warning( + "Model %s is MTP-capable but no drafter or head was found; " + "falling back. Check network or run `unsloth studio update`.", + model_identifier, + ) + if self.probe_server_capabilities(binary).get("supports_ngram_mod"): + _emit_ngram_mod() + else: + flags.append("--spec-default") + self._speculative_type = "default" + self._spec_fallback_reason = "drafter_not_found" + if effective_mode == "off": return flags # nothing to emit if effective_mode == "ngram-simple": @@ -4717,6 +6369,10 @@ class LlamaCppBackend: flags.append("--spec-default") self._speculative_type = "default" return flags + if _mtp_drafter_missing: + # Drafterless: draft-mtp would abort llama-server, so fall back. + _fallback_drafter_not_found() + return flags if _mtp_too_small: logger.warning( f"Forcing MTP on a {_mtp_size_b:.1f}B model; " @@ -4735,6 +6391,10 @@ class LlamaCppBackend: ) _emit_ngram_mod() return flags + if _mtp_drafter_missing: + # No head/drafter: keep ngram-mod, drop the draft-mtp chain. + _fallback_drafter_not_found() + return flags if _mtp_too_small: logger.warning( f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; " @@ -4746,14 +6406,42 @@ class LlamaCppBackend: # effective_mode == "auto": the promotion path. llama.cpp #22673: # MTP is compatible with mmproj, so there's no vision gate. - if is_mtp_model and not _mtp_too_small: - # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP. - _emit_mtp(chain_ngram = not gpus) + if _auto_mla_embedded_mtp: + # MLA embedded-MTP (GLM-5.2 et al.): the MTP path regresses vs spec-off + # on llama.cpp today, so Auto drops it and falls back to ngram-mod (or + # spec-off if unsupported), mirroring the sub-3B branch. Forced mtp / + # mtp+ngram (handled above) still engage; UNSLOTH_MLA_MTP_ENABLED=1 + # re-enables this promotion once upstream optimizes the path. + self._spec_fallback_reason = "mla_mtp_disabled" + _mla_caps = self.probe_server_capabilities(binary) + if _mla_caps.get("supports_ngram_mod"): + logger.info( + "Auto: MLA embedded-MTP model detected; llama.cpp's MLA/DSA " + "MTP path is slower than no speculation, so using ngram-mod " + "instead. Override via the Studio Speculative Decoding " + "dropdown or UNSLOTH_MLA_MTP_ENABLED=1." + ) + _emit_ngram_mod() + else: + logger.info( + "Auto: MLA embedded-MTP model detected; disabling speculative " + "decoding (this llama-server does not advertise ngram-mod). " + "Override via the dropdown or UNSLOTH_MLA_MTP_ENABLED=1." + ) + # spec-off: emit nothing, mirroring the sub-3B no-ngram path. + elif is_mtp_model and not _mtp_too_small: + if _mtp_drafter_missing: + # Name-only MTP, drafter did not resolve (download failed/absent). + _fallback_drafter_not_found() + else: + # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP. + _emit_mtp(chain_ngram = not gpus) elif is_mtp_model and _mtp_too_small: # Sub-3B fallback: drop the MTP draft head, keep ngram-mod when # the binary supports it. - _small_caps = self.probe_server_capabilities(binary) - if _small_caps.get("supports_ngram_mod"): + if _mtp_drafter_missing: + _fallback_drafter_not_found() + elif self.probe_server_capabilities(binary).get("supports_ngram_mod"): logger.info( f"MTP GGUF detected but model size {_mtp_size_b:.1f}B " "is below the 3B speedup threshold; using ngram-mod " @@ -4825,10 +6513,12 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False - # Reconcile a user --split-mode in extras (load_model does the same), so - # an extras-driven tensor load isn't seen as a mismatch that needlessly - # kills/reloads a healthy server. - if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + # Reconcile a user --split-mode in extras AND an inherited tensor + # LLAMA_ARG_SPLIT_MODE env, but only against a server that actually + # launched tensor: if load_model downgraded to layer split it scrubbed + # the child env, so the env must not force an endless reload of a healthy + # server. An identical request would downgrade the same way. + if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): return False # Compare on the canonical requested mode. With --spec-type in @@ -4841,6 +6531,16 @@ class LlamaCppBackend: if req_mode != backend_mode: return False + # Prior HF load fell back with drafter_not_found; a same-settings reload + # must retry the download in load_model, not dedupe to the stale fallback + # (HF loads resolve the drafter there, so gguf_path is None here). + if ( + self._spec_fallback_reason == "drafter_not_found" + and gguf_path is None + and req_mode in ("auto", "mtp", "mtp+ngram") + ): + return False + # spec_draft_n_max only matters when an MTP variant is engaged. Compare # on the resolved spec so an Auto request promoted to draft-mtp still # bounces a reload when n_max changes. @@ -4881,26 +6581,13 @@ class LlamaCppBackend: def _classify_gpu_offload( self, expected_gpu: bool, detected_gpus: list[tuple[int, int]] ) -> Optional[bool]: - """True if a GPU model buffer was allocated, False if only CPU - buffers landed despite GPU intent, None when there's no signal (no - GPU detected, no buffer-size lines, etc.).""" + """True if the model landed on a GPU, False if only CPU buffers landed + despite GPU intent, None when there's no signal. Delegates to the shared + classifier so it tracks current llama.cpp logs (offloaded-layer counts / + device_info), not just the older "model buffer size" lines.""" if not detected_gpus or not expected_gpu: return None - # llama-server logs one "model buffer size = N MiB" line per backend - # buffer; CUDA/ROCm/Metal/Vulkan/OpenCL/SYCL are GPU, CPU* are not. - gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL") - saw_buffer_line = False - saw_gpu_buffer = False - for line in self._stdout_lines: - if "model buffer size" not in line: - continue - saw_buffer_line = True - if any(marker in line for marker in gpu_markers): - saw_gpu_buffer = True - break - if not saw_buffer_line: - return None - return saw_gpu_buffer + return classify_gpu_offload_lines(self._stdout_lines) def load_cancelled(self) -> bool: """True if a load was cancelled (e.g. via unload/_cancel_event) and not @@ -4919,6 +6606,8 @@ class LlamaCppBackend: self._hf_repo = None self._mtp_draft_path = None self._spec_fallback_reason = None + self._last_load_kwargs = None + self._mtp_runtime_fallback_active = False self._hf_variant = None self._is_vision = False self._is_audio = False @@ -4936,6 +6625,7 @@ class LlamaCppBackend: self._supports_reasoning = False self._reasoning_always_on = False self._reasoning_style = "enable_thinking" + self._reasoning_effort_levels = [] self._reasoning_default = True self._supports_preserve_thinking = False self._supports_tools = False @@ -4965,7 +6655,6 @@ class LlamaCppBackend: # Clean up temp chat template file. if hasattr(self, "_chat_template_file") and self._chat_template_file: try: - import os os.unlink(self._chat_template_file.name) except Exception: pass @@ -4982,6 +6671,9 @@ class LlamaCppBackend: def _kill_process(self): """Terminate the subprocess if running.""" + # Stop the watchdog before a deliberate kill so a planned reload/unload + # isn't seen as a crash; a real crash never routes through here. + self._stop_mtp_crash_watchdog() if self._process is None: return try: @@ -4994,15 +6686,22 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"Error killing llama-server process: {e}") finally: + # getattr: teardown must tolerate a partially-built backend (failed + # __init__ or a __new__-built instance), as with _llama_log_fh below. + if getattr(self, "_stats_logger", None) is not None: + self._stats_logger.stop() + self._stats_logger = None self._process = None + self._clear_server_pid() # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). self._healthy = False # Drives _wait_for_vram_settle in the next load_model; set in finally # so both in-process and frontend Apply paths record the kill. self._last_kill_monotonic = time.monotonic() - if self._stdout_thread is not None: - self._stdout_thread.join(timeout = 2) + stdout_thread = getattr(self, "_stdout_thread", None) + if stdout_thread is not None: + stdout_thread.join(timeout = 2) self._stdout_thread = None fh = getattr(self, "_llama_log_fh", None) if fh is not None: @@ -5013,7 +6712,199 @@ class LlamaCppBackend: self._llama_log_fh = None @staticmethod - def _kill_orphaned_servers(): + def _server_pidfile_path() -> Optional[Path]: + """Pidfile recording the live llama-server PID, under the active studio root + (per-root, so concurrent Studios with distinct UNSLOTH_STUDIO_HOME stay + isolated, mirroring the reaper's custom-root isolation).""" + try: + from utils.paths.storage_roots import studio_root # noqa: WPS433 + return studio_root() / "llama-server.pid" + except Exception: + return None + + @classmethod + def _record_server_pid(cls, pid: int) -> None: + """Best-effort record of the spawned llama-server PID for orphan reaping. + + Stores ``pid:starttime`` so a later startup can reject a PID that has + since been recycled to a different process (see ``_pid_start_identity``). + A bare ``pid`` (no identity) is still accepted on read for compatibility. + """ + path = cls._server_pidfile_path() + if path is None: + return + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(f"{pid}:{cls._pid_start_identity(pid)}") + except Exception as e: + logger.debug(f"Could not write llama-server pidfile: {e}") + + @classmethod + def _clear_server_pid(cls) -> None: + """Best-effort removal of the llama-server pidfile.""" + path = cls._server_pidfile_path() + if path is None: + return + try: + path.unlink(missing_ok = True) + except Exception as e: + logger.debug(f"Could not remove llama-server pidfile: {e}") + + @staticmethod + def _pid_is_llama_server(pid: int) -> bool: + """True only if pid is a live process whose binary is a llama-server. Guards + against PID reuse before killing a recorded orphan; returns False on any + uncertainty so an unrelated process is never killed.""" + try: + import psutil + try: + proc = psutil.Process(pid) + if (proc.name() or "").lower().startswith("llama-server"): + return True + return Path(proc.exe() or "").name.lower().startswith("llama-server") + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return False + except ImportError: + pass + if sys.platform != "linux": + return False + try: + if Path(os.readlink(f"/proc/{pid}/exe")).name.lower().startswith("llama-server"): + return True + except OSError: + pass + try: + with open(f"/proc/{pid}/cmdline", "rb") as fh: + tokens = fh.read().split(b"\x00") + first = tokens[0].decode("utf-8", "replace") if tokens else "" + return Path(first).name.lower().startswith("llama-server") + except OSError: + return False + + @staticmethod + def _pid_start_identity(pid: int) -> str: + """Stable per-PID identity (process start time) guarding against PID reuse. + + Returns a token string, or "" when it cannot be determined (the caller + then falls back to the llama-server name check only).""" + try: + import psutil + try: + return str(psutil.Process(pid).create_time()) + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return "" + except ImportError: + pass + if sys.platform == "linux": + try: + with open(f"/proc/{pid}/stat", "rb") as fh: + data = fh.read() + # field 22 (starttime), counted from after the ")" that closes comm. + return data[data.rfind(b")") + 2 :].split()[19].decode() + except (OSError, IndexError): + return "" + return "" + + @staticmethod + def _pid_parent_is_alive(pid: int) -> bool: + """True if the recorded server's parent is still running, i.e. the server is + NOT orphaned. Lets the cross-session reap kill only a true orphan (parent + gone) and never a live server owned by a running Studio, regardless of which + process performs the sweep. Biased toward "alive" on uncertainty so a live + server is never mistakenly reaped.""" + try: + import psutil + + try: + ppid = psutil.Process(pid).ppid() + except psutil.NoSuchProcess: + return False # the recorded server itself is gone + except psutil.Error: + return True # cannot tell -- never risk killing a live server + if ppid <= 1: + return False # reparented to init -> orphan + return psutil.pid_exists(ppid) + except ImportError: + pass + if sys.platform == "linux": + try: + with open(f"/proc/{pid}/stat", "rb") as fh: + data = fh.read() + ppid = int(data[data.rfind(b")") + 2 :].split()[1]) + except (OSError, IndexError, ValueError): + return False + if ppid <= 1: + return False + return Path(f"/proc/{ppid}").exists() + return False + + @staticmethod + def _unlink_pidfile(path: Path) -> None: + """Best-effort removal of a resolved pidfile path.""" + try: + path.unlink(missing_ok = True) + except Exception: + pass + + @classmethod + def _reap_recorded_pid(cls) -> int: + """Kill the exact llama-server PID recorded at spawn, but only when it is a + genuine orphan -- its parent (the Studio that spawned it) is gone. This is + the cross-session backstop the parent-death reaper (Job Object / + PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Studio + (macOS, a best-effort failure, or a pre-existing orphan). Path-independent, + so it also catches an orphan the install-root match would miss. + + A live server whose parent is still running is never reaped, so constructing + a second backend in-process (the helper / advisor paths each build a + LlamaCppBackend) cannot kill the active chat server. A recorded PID that has + been recycled to a different process is rejected by the start-time identity + and the llama-server name check, so unrelated user processes are never + touched. SIGKILL falls back to SIGTERM on Windows, where os.kill maps it to + TerminateProcess and SIGKILL is undefined.""" + path = cls._server_pidfile_path() + if path is None or not path.exists(): + return 0 + + pid = -1 + identity = "" + try: + pid_str, _, identity = path.read_text().strip().partition(":") + pid = int(pid_str) + except Exception: + pid = -1 + + if pid <= 0: + cls._unlink_pidfile(path) # garbage record + return 0 + if pid == os.getpid(): + return 0 # never our own pid; leave the record alone + + if cls._pid_parent_is_alive(pid): + # Live server with a running parent -> not an orphan; keep the record so + # a later startup can still reap it if that parent later dies abnormally. + return 0 + + # Parent is gone: candidate orphan. Reject a PID recycled to something else. + if identity and cls._pid_start_identity(pid) != identity: + cls._unlink_pidfile(path) + return 0 + + killed = 0 + if cls._pid_is_llama_server(pid): + try: + os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM)) + killed = 1 + logger.info(f"Killed orphaned llama-server from pidfile (pid={pid})") + except (ProcessLookupError, PermissionError): + pass + except Exception as e: + logger.debug(f"Could not kill recorded llama-server pid {pid}: {e}") + cls._unlink_pidfile(path) + return killed + + @staticmethod + def _kill_orphaned_servers() -> int: """Kill orphaned llama-server processes started by studio. Only kills processes whose resolved binary lives under a known @@ -5025,7 +6916,15 @@ class LlamaCppBackend: Uses psutil for cross-platform support (Linux, macOS, Windows); falls back to pgrep + /proc//exe on Linux when psutil is absent. + + Returns the count of processes killed; callers arm the VRAM-settle + wait on a positive count. """ + # Cross-session backstop first: reap the exact PID we recorded at spawn, + # but only if it is a true orphan whose parent is gone (so a helper backend + # built while a chat server is live can never kill it). The root-gated + # enumeration below stays as a fallback. + killed = LlamaCppBackend._reap_recorded_pid() try: # -- Build the ownership allowlist -------------------------------- # exact_binaries -- env var overrides (exact path match). @@ -5033,20 +6932,10 @@ class LlamaCppBackend: install_roots: list[Path] = [] # Env-mode custom root (mirrors _find_llama_server_binary). - _is_custom_root = False - try: - from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 - - _resolved_sr = _sr() - _legacy_studio = Path.home() / ".unsloth" / "studio" - try: - _is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve() - except (OSError, ValueError): - _is_custom_root = _resolved_sr != _legacy_studio - if _is_custom_root: - install_roots.append(_resolved_sr / "llama.cpp") - except (ImportError, OSError, ValueError): - pass + _resolved_sr, _is_legacy = LlamaCppBackend._resolved_studio_root_and_is_legacy() + _is_custom_root = not _is_legacy + if _is_custom_root: + install_roots.append(_resolved_sr / "llama.cpp") # Primary install dir (default mode only). Env-mode skips this so a # custom-root Studio can't kill a default-install Studio's server. @@ -5117,6 +7006,7 @@ class LlamaCppBackend: continue proc.kill() + killed += 1 logger.info( f"Killed orphaned llama-server process (pid={proc.info['pid']})" ) @@ -5129,7 +7019,7 @@ class LlamaCppBackend: else: # -- Fallback: pgrep + /proc//exe (Linux only) ----------- if sys.platform != "linux": - return + return killed result = subprocess.run( ["pgrep", "-a", "-f", "llama-server"], capture_output = True, @@ -5138,7 +7028,7 @@ class LlamaCppBackend: env = child_env_without_native_path_secret(), ) if result.returncode != 0: - return + return killed for line in result.stdout.strip().splitlines(): parts = line.strip().split(None, 1) @@ -5169,6 +7059,7 @@ class LlamaCppBackend: try: os.kill(pid, signal.SIGKILL) + killed += 1 logger.info(f"Killed orphaned llama-server process (pid={pid})") except ProcessLookupError: pass @@ -5176,6 +7067,7 @@ class LlamaCppBackend: pass except Exception: logger.warning("Error during orphan server cleanup", exc_info = True) + return killed def _cleanup(self): """atexit handler to ensure llama-server is terminated.""" @@ -5198,6 +7090,137 @@ class LlamaCppBackend: return False return True + def _probe_mtp_decode(self, timeout: float = 60.0) -> bool: + """One tiny /completion to confirm MTP survives the first decode. + + MTP-draft can pass /health yet crash the flash-attn kernel only once + tokens generate (e.g. under --split-mode tensor). False on any error so + the caller can drop MTP and retry. + """ + url = f"{self.base_url}/completion" + payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False} + try: + resp = httpx.post(url, json = payload, timeout = timeout, headers = self._auth_headers) + except Exception as e: + logger.debug(f"MTP decode probe failed: {e}") + return False + if resp.status_code != 200: + logger.debug(f"MTP decode probe returned HTTP {resp.status_code}") + return False + # A crash can drop the connection or kill the process right after a reply. + if self._process is not None and self._process.poll() is not None: + return False + return True + + def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool: + """Schedule one background reload without MTP after a mid-generation death. + + MTP+tensor can crash the flash-attn kernel on a later request, after + load_model returned, past the load-time fallback and decode probe. Not a + persistent ban: a fresh load re-tries MTP. Returns True if scheduled. + """ + # Cheap async-safe gate: only our live MTP+tensor launch, not cancelled, + # with a snapshot to replay. + if self._cancel_event.is_set(): + return False + if not self._mtp_runtime_fallback_active: + return False + if not self._last_load_kwargs or self._process is None: + return False + # Single-flight: the first failure claims the reload. + with self._mtp_runtime_fallback_lock: + if self._mtp_runtime_fallback_in_progress: + return False + self._mtp_runtime_fallback_in_progress = True + snapshot = dict(self._last_load_kwargs) + proc = self._process + + def _recover(): + try: + # Confirm the process really exited (the error can arrive a beat + # early) so a transient stream error can't disable MTP. + deadline = time.monotonic() + 5.0 + while proc.poll() is None and time.monotonic() < deadline: + time.sleep(0.1) + if proc.poll() is None: + logger.debug("Generation error but llama-server is alive; keeping MTP.") + return + logger.warning( + "llama-server exited mid-generation with MTP under tensor " + "parallelism (%s); reloading without speculative decoding.", + type(exc).__name__ if exc is not None else "server exited", + ) + # Re-check under the load lock (RLock allows the nested + # load_model) so a newer load isn't clobbered by this stale replay. + requested_mode = snapshot.get("speculative_type") + with self._serial_load_lock: + if self._cancel_event.is_set(): + logger.info("MTP-crash reload skipped: load was cancelled/unloaded.") + return + if self._process is not proc: + logger.info("MTP-crash reload skipped: a newer load is already active.") + return + if self._last_load_kwargs != snapshot: + logger.info("MTP-crash reload skipped: load settings changed.") + return + snapshot["speculative_type"] = "off" + # Drop user/env MTP too: append a last-wins --spec-default. + _ea = list(snapshot.get("extra_args") or []) + if _extra_args_requests_mtp(_ea, env = os.environ): + _ea.append("--spec-default") + snapshot["extra_args"] = _ea + self.load_model(**snapshot) + # Restore the requested mode + reason load_model("off") cleared, + # so /status shows the user's mode + note (like the startup fallback). + self._requested_spec_mode = _canonicalize_spec_mode(requested_mode) + self._spec_fallback_reason = "runtime_error" + logger.info("Reloaded without MTP after the tensor-parallel crash.") + except Exception as e: + logger.error(f"Reload without MTP failed: {e}") + finally: + with self._mtp_runtime_fallback_lock: + self._mtp_runtime_fallback_in_progress = False + + threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start() + return True + + def _start_mtp_crash_watchdog(self) -> None: + """Background poll that recovers on an MTP+tensor crash even when no + request observes it (direct proxy endpoints, or nothing in flight). + + Armed only for a live MTP+tensor launch; the no-MTP reload disarms it, so + it can't loop. + """ + if not self._mtp_runtime_fallback_active: + return + proc = self._process + if proc is None: + return + # Replace any prior watchdog (loads are serialised, so at most one). + self._stop_mtp_crash_watchdog() + stop = threading.Event() + self._mtp_watchdog_stop = stop + + def _watch(): + # Exit on stop or process death. _kill_process sets stop before + # terminating, so re-check it: only a real crash (stop unset) recovers. + while not stop.wait(1.0): + if proc.poll() is not None: + if not stop.is_set(): + self._maybe_recover_from_mtp_crash() + return + + t = threading.Thread(target = _watch, daemon = True, name = "mtp-crash-watchdog") + self._mtp_watchdog_thread = t + t.start() + + def _stop_mtp_crash_watchdog(self) -> None: + """Signal the crash watchdog to exit; called before any deliberate kill.""" + stop = getattr(self, "_mtp_watchdog_stop", None) + if stop is not None: + stop.set() + self._mtp_watchdog_thread = None + def _wait_for_health( self, timeout: float = 120.0, @@ -5205,7 +7228,7 @@ class LlamaCppBackend: ) -> bool: """Poll llama-server's /health until 200; also detect early exit/crash.""" deadline = time.monotonic() + timeout - url = f"http://127.0.0.1:{self._port}/health" + url = f"{self.base_url}/health" while time.monotonic() < deadline: # Process crashed? @@ -5274,7 +7297,7 @@ class LlamaCppBackend: The memory-fit step or ``--parallel`` slot split can leave this below the requested ``-c``; requests are validated against this value. """ - url = f"http://127.0.0.1:{self._port}/props" + url = f"{self.base_url}/props" try: resp = httpx.get(url, timeout = 5.0) if resp.status_code != 200: @@ -5347,6 +7370,33 @@ class LlamaCppBackend: # ── Generation (proxy to llama-server) ──────────────────────── + @contextlib.contextmanager + def _open_stream(self, url: str, payload: dict, cancel_event): + """Open a streaming POST to llama-server, retrying through prefill, and + yield ``(response, first_token_deadline)`` once a 200 lands. Owns the + httpx.Client + auth headers for the stream's lifetime; raises + RuntimeError on a non-200. Shared scaffold for the streaming consumers, + which differ only in how they parse the SSE body.""" + stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) + with httpx.Client( + timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) + ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + with self._stream_with_retry( + client, + url, + payload, + cancel_event, + headers = self._auth_headers, + first_token_deadline = first_token_deadline, + ) as response: + if response.status_code != 200: + error_body = response.read().decode() + raise RuntimeError( + f"llama-server returned {response.status_code}: {error_body}" + ) + yield response, first_token_deadline + @staticmethod def _iter_text_cancellable( response: "httpx.Response", @@ -5495,6 +7545,38 @@ class LlamaCppBackend: finally: _cancel_closed.set() + def _respawn_if_dead(self) -> bool: + """Relaunch the llama-server if its process has exited. + + A loaded chat model can be SIGKILL'd mid-session (usually GPU/RAM pressure + from a training run on the same box), leaving a defunct process while + ``is_loaded`` still reads True. Replay the last ``load_model`` call to + recover, returning True once healthy. Serialised on ``_respawn_lock`` so + many generations hitting the dead server trigger at most one reload. + """ + with self._respawn_lock: + proc = self._process + if proc is None: + return False + if proc.poll() is None: + # Process is alive: either a concurrent caller already respawned + # it (healthy), or this connection error wasn't a dead server. + return self._healthy + kwargs = self._last_load_kwargs + if not kwargs: + return False + logger.warning( + f"llama-server for '{self._model_identifier}' exited " + f"(code {proc.returncode}); respawning to recover the session" + ) + with self._lock: + self._healthy = False + try: + return bool(self.load_model(**kwargs)) + except Exception as exc: + logger.error(f"Failed to respawn llama-server: {exc}") + return False + def generate_chat_completion( self, messages: list[dict], @@ -5512,7 +7594,8 @@ class LlamaCppBackend: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, seed: Optional[int] = None, - ) -> Generator[str | dict, None, None]: + _allow_respawn_retry: bool = True, + ) -> Generator[Union[str, dict], None, None]: """ Send a chat completion to llama-server and stream tokens back. @@ -5563,125 +7646,140 @@ class LlamaCppBackend: _metadata_finish_reason = None try: - # Prefill can use the long first-token timeout; body reads are lowered after headers. - stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client( - timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) - ) as client: - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - with self._stream_with_retry( - client, - url, - payload, + with self._open_stream(url, payload, cancel_event) as ( + response, + first_token_deadline, + ): + buffer = "" + has_content_tokens = False + reasoning_text = "" + for raw_chunk in self._iter_text_cancellable( + response, cancel_event, - headers = _auth_headers, first_token_deadline = first_token_deadline, - ) as response: - if response.status_code != 200: - error_body = response.read().decode() - raise RuntimeError( - f"llama-server returned {response.status_code}: {error_body}" - ) + ): + buffer += raw_chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() - buffer = "" - has_content_tokens = False - reasoning_text = "" - for raw_chunk in self._iter_text_cancellable( - response, - cancel_event, - first_token_deadline = first_token_deadline, - ): - buffer += raw_chunk - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) - line = line.strip() + if not line: + continue + if line == "data: [DONE]": + if in_thinking: + if has_content_tokens: + # Real thinking + content: close the tag + cumulative += "" + yield cumulative + else: + # Only reasoning_content, no content: + # model put its whole reply in reasoning + # (e.g. Qwen3 always-think). Show it as + # the main response, not a thinking block. + cumulative = reasoning_text + yield cumulative + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue - if not line: + try: + data = json.loads(line[6:]) + # Diffusion frame (per-step canvas) from the shim: forward untouched so + # the frontend renders it in place. No assistant text, so it never enters + # the cumulative content. + if data.get("type") == "diffusion_frame": + yield data continue - if line == "data: [DONE]": - if in_thinking: - if has_content_tokens: - # Real thinking + content: close the tag + # Capture server timings/usage from final chunks. + _chunk_timings = data.get("timings") + if _chunk_timings: + _metadata_timings = _chunk_timings + _chunk_usage = data.get("usage") + if _chunk_usage: + _metadata_usage = _chunk_usage + choices = data.get("choices", []) + if choices: + delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _metadata_finish_reason = _fr + + # Reasoning/thinking tokens: llama-server + # sends these as "reasoning_content"; wrap + # in tags for the frontend parser. + reasoning = delta.get("reasoning_content", "") + if reasoning: + reasoning_text += reasoning + if not in_thinking: + cumulative += "" + in_thinking = True + cumulative += reasoning + yield cumulative + + token = delta.get("content", "") + if token: + has_content_tokens = True + if in_thinking: cumulative += "" - yield cumulative - else: - # Only reasoning_content, no content: - # model put its whole reply in reasoning - # (e.g. Qwen3 always-think). Show it as - # the main response, not a thinking block. - cumulative = reasoning_text - yield cumulative - _stream_done = True - break # exit inner while - if not line.startswith("data: "): - continue + in_thinking = False + cumulative += token + yield cumulative + except json.JSONDecodeError: + logger.debug(f"Skipping malformed SSE line: {line[:100]}") + if _stream_done: + break # exit outer for + if _metadata_usage or _metadata_timings or _metadata_finish_reason: + _metadata_usage = _backfill_usage_from_timings( + _metadata_usage, _metadata_timings + ) + yield { + "type": "metadata", + # Never None: a finish-only metadata event (no usage, + # no timings) would otherwise crash consumers that do + # usage.get(...) on the non-streaming paths. + "usage": _metadata_usage or {}, + "timings": _metadata_timings, + "finish_reason": _metadata_finish_reason, + } - try: - data = json.loads(line[6:]) - # Diffusion frame (per-step canvas) from the shim: forward untouched so - # the frontend renders it in place. No assistant text, so it never enters - # the cumulative content. - if data.get("type") == "diffusion_frame": - yield data - continue - # Capture server timings/usage from final chunks. - _chunk_timings = data.get("timings") - if _chunk_timings: - _metadata_timings = _chunk_timings - _chunk_usage = data.get("usage") - if _chunk_usage: - _metadata_usage = _chunk_usage - choices = data.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - _fr = choices[0].get("finish_reason") - if _fr: - _metadata_finish_reason = _fr - - # Reasoning/thinking tokens: llama-server - # sends these as "reasoning_content"; wrap - # in tags for the frontend parser. - reasoning = delta.get("reasoning_content", "") - if reasoning: - reasoning_text += reasoning - if not in_thinking: - cumulative += "" - in_thinking = True - cumulative += reasoning - yield cumulative - - token = delta.get("content", "") - if token: - has_content_tokens = True - if in_thinking: - cumulative += "" - in_thinking = False - cumulative += token - yield cumulative - except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") - if _stream_done: - break # exit outer for - if _metadata_usage or _metadata_timings or _metadata_finish_reason: - _metadata_usage = _backfill_usage_from_timings( - _metadata_usage, _metadata_timings - ) - yield { - "type": "metadata", - # Never None: a finish-only metadata event (no usage, - # no timings) would otherwise crash consumers that do - # usage.get(...) on the non-streaming paths. - "usage": _metadata_usage or {}, - "timings": _metadata_timings, - "finish_reason": _metadata_finish_reason, - } - - except httpx.ConnectError: + except httpx.ConnectError as e: + # Server already down. If this was an MTP+tensor crash, recover by + # reloading without MTP (scheduled in the background) and fail this + # request. Otherwise the server was likely SIGKILL'd by GPU pressure + # from a concurrent training run: respawn the same config and retry the + # generation once (bounded by the private flag, no duplicate output). + if self._maybe_recover_from_mtp_crash(e): + raise RuntimeError("Lost connection to llama-server") + if _allow_respawn_retry and not cumulative and self._respawn_if_dead(): + logger.warning( + "llama-server was unreachable; respawned it and retrying the generation" + ) + yield from self.generate_chat_completion( + messages, + image_b64 = image_b64, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + max_tokens = max_tokens, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + stop = stop, + cancel_event = cancel_event, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + seed = seed, + _allow_respawn_retry = False, + ) + return raise RuntimeError("Lost connection to llama-server") except Exception as e: if cancel_event is not None and cancel_event.is_set(): return + # Died mid-generation: recover MTP, re-raise unchanged for this request. + self._maybe_recover_from_mtp_crash(e) raise # ── Tool-calling agentic loop ────────────────────────────── @@ -5710,6 +7808,7 @@ class LlamaCppBackend: seed: Optional[int] = None, disable_parallel_tool_use: bool = False, confirm_tool_calls: bool = False, + bypass_permissions: bool = False, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -5738,6 +7837,15 @@ class LlamaCppBackend: _accumulated_completion_tokens = 0 _accumulated_predicted_ms = 0.0 _accumulated_predicted_n = 0 + # GGUF buffers reasoning; emit server-side timing before answer text. + _reasoning_started_at: Optional[float] = None + _reasoning_summary_emitted = False + + def _reasoning_summary_event(started_at: float) -> dict: + return { + "type": "reasoning_summary", + "duration_ms": round((time.monotonic() - started_at) * 1000.0), + } def _strip_tool_markup( text: str, @@ -5756,6 +7864,47 @@ class LlamaCppBackend: text = pat.sub("", text) return text + def _build_metadata_event(usage, timings, finish_reason): + """Final usage+timings metadata event for the given pass, merging its + usage/timings with the running cross-iteration accumulators. None when + there is nothing to report.""" + _fu = _backfill_usage_from_timings(usage, timings) or {} + _fp = _fu.get("prompt_tokens", 0) + _tc = _fu.get("completion_tokens", 0) + _accumulated_completion_tokens + if not (usage or timings or _accumulated_completion_tokens or finish_reason): + return None + _mt = dict(timings) if timings else {} + if _accumulated_predicted_ms or _accumulated_predicted_n: + _mt["predicted_ms"] = _mt.get("predicted_ms", 0) + _accumulated_predicted_ms + _mt["predicted_n"] = _mt.get("predicted_n", 0) + _accumulated_predicted_n + if _mt["predicted_ms"] > 0: + _mt["predicted_per_second"] = _mt["predicted_n"] / ( + _mt["predicted_ms"] / 1000.0 + ) + _usage = { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + } + # Preserve KV-cache hit details (cached_tokens) so the tool path + # reports them like the standard non-tool path does, not always 0. + if _fu.get("prompt_tokens_details"): + _usage["prompt_tokens_details"] = _fu["prompt_tokens_details"] + return { + "type": "metadata", + "usage": _usage, + "timings": _mt, + "finish_reason": finish_reason, + } + + def _flush_reasoning_and_buffer(): + """Append buffered reasoning (as a block) then the held + content_buffer to the cumulative display text.""" + nonlocal cumulative_display + if reasoning_accum: + cumulative_display += "" + reasoning_accum + "" + cumulative_display += content_buffer + tool_controller = ToolLoopController( tools = tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -5794,7 +7943,7 @@ class LlamaCppBackend: if not active_tools: _append_budget_exhausted_nudge = False break - _tool_xml_signals = TOOL_XML_SIGNALS if active_tools else () + _tool_xml_signals = TOOL_XML_SIGNALS # Build payload -- stream: True so we detect tool signals # in the first 1-2 chunks without a non-streaming penalty. @@ -5827,10 +7976,6 @@ class LlamaCppBackend: payload["seed"] = seed try: - _auth_headers = ( - {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - ) - # ── Speculative buffer state machine ────────────────── # BUFFERING: accumulate content, check for tool signals # STREAMING: no tool detected, yield tokens to caller @@ -5843,6 +7988,9 @@ class LlamaCppBackend: content_buffer = "" # Raw content held during BUFFERING content_accum = "" # All content tokens (for tool parsing) reasoning_accum = "" + # Time each reasoning pass so final answers can replace tool timing. + _reasoning_started_at = None + _reasoning_summary_emitted = False cumulative_display = "" # Cumulative yielded text (with ) in_thinking = False has_content_tokens = False @@ -5853,197 +8001,258 @@ class LlamaCppBackend: _iter_finish_reason = None _stream_done = False _last_emitted = "" - provisional_render_html_tool_call_ids = set() + # Provisional tool_start cards already shown, keyed by tool_call_id. + provisional_started_tool_calls: dict[str, str] = {} + resolved_provisional_tool_call_ids: set[str] = set() _suppress_visible_output = _forced_tool_call_pending - stream_timeout = httpx.Timeout( - connect = 10, - read = 0.5, - write = 10, - pool = 10, - ) - with httpx.Client( - timeout = stream_timeout, - limits = httpx.Limits(max_keepalive_connections = 0), - ) as client: - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - with self._stream_with_retry( - client, - url, - payload, + with self._open_stream(url, payload, cancel_event) as ( + response, + first_token_deadline, + ): + raw_buf = "" + for raw_chunk in self._iter_text_cancellable( + response, cancel_event, - headers = _auth_headers, first_token_deadline = first_token_deadline, - ) as response: - if response.status_code != 200: - error_body = response.read().decode() - raise RuntimeError( - f"llama-server returned {response.status_code}: {error_body}" - ) + ): + raw_buf += raw_chunk + while "\n" in raw_buf: + line, raw_buf = raw_buf.split("\n", 1) + line = line.strip() - raw_buf = "" - for raw_chunk in self._iter_text_cancellable( - response, - cancel_event, - first_token_deadline = first_token_deadline, - ): - raw_buf += raw_chunk - while "\n" in raw_buf: - line, raw_buf = raw_buf.split("\n", 1) - line = line.strip() + if not line: + continue + if line == "data: [DONE]": + # Flush thinking state for STREAMING + if detect_state == _S_STREAMING and in_thinking: + if has_content_tokens: + cumulative_display += "" + if not _suppress_visible_output: + yield { + "type": "content", + "text": _strip_tool_markup( + cumulative_display, + final = True, + ), + } + else: + cumulative_display = reasoning_accum + if not _suppress_visible_output: + yield { + "type": "content", + "text": cumulative_display, + } + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue - if not line: + try: + chunk_data = json.loads(line[6:]) + _ct = chunk_data.get("timings") + if _ct: + _iter_timings = _ct + _cu = chunk_data.get("usage") + if _cu: + _iter_usage = _cu + + choices = chunk_data.get("choices", []) + if not choices: continue - if line == "data: [DONE]": - # Flush thinking state for STREAMING - if detect_state == _S_STREAMING and in_thinking: - if has_content_tokens: + + delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _iter_finish_reason = _fr + + # ── Structured tool_calls ── + tc_deltas = delta.get("tool_calls") + if tc_deltas: + # Preserve any visible preface before draining + # the structured tool call. + has_structured_tc = True + detect_state = _S_DRAINING + for tc_d in tc_deltas: + idx = tc_d.get("index", 0) + if idx not in tool_calls_acc: + tool_calls_acc[idx] = { + "id": tc_d.get("id", f"call_{idx}"), + "type": "function", + "function": { + "name": "", + "arguments": "", + }, + } + elif tc_d.get("id"): + # Update ID if a real one + # arrives on a later delta. + tool_calls_acc[idx]["id"] = tc_d["id"] + func = tc_d.get("function", {}) + if func.get("name"): + tool_calls_acc[idx]["function"]["name"] += func["name"] + if func.get("arguments"): + tool_calls_acc[idx]["function"]["arguments"] += func[ + "arguments" + ] + current_name = tool_calls_acc[idx]["function"].get( + "name", "" + ) + fallback_id = f"call_{idx}" + current_id = tool_calls_acc[idx].get("id", fallback_id) + already_started = ( + current_id in provisional_started_tool_calls + ) + # Empty/synthetic ids cannot reconcile with real starts. + has_real_id = bool(current_id) and current_id != fallback_id + # Show one early card per eligible streamed tool call. + _is_completed_one_shot = ( + current_name == "render_html" + and _tool_succeeded("render_html") + ) + # render_html is one-shot. + _one_shot_already_provisional = ( + current_name == "render_html" + and "render_html" + in provisional_started_tool_calls.values() + ) + # Later parallel cards only reconcile when parallel use is enabled. + _confirm_gated = ( + confirm_tool_calls and not bypass_permissions + ) + # Keep small-argument tools on the normal path. + _args_len = len( + tool_calls_acc[idx]["function"].get("arguments", "") + ) + _payload_is_large = ( + current_name == "render_html" + or _args_len >= _PROVISIONAL_ARGS_MIN_CHARS + ) + if ( + current_name + and (idx == 0 or not disable_parallel_tool_use) + and has_real_id + and not already_started + and not _is_completed_one_shot + and not _one_shot_already_provisional + and not _confirm_gated + and _payload_is_large + and any( + (tool.get("function") or {}).get("name") + == current_name + for tool in active_tools + ) + ): + provisional_started_tool_calls[current_id] = ( + current_name + ) + yield { + "type": "tool_start", + "tool_name": current_name, + "tool_call_id": current_id, + "arguments": {}, + "provenance": tool_event_provenance( + provisional = True, + ), + } + continue + + # ── Reasoning tokens ── + # Yield only in STREAMING. In BUFFERING and + # DRAINING, accumulate silently so we don't + # corrupt the consumer's prev_text tracker + # (routes/inference.py never resets it + # between tool iterations). + reasoning = delta.get("reasoning_content", "") + if reasoning: + if _reasoning_started_at is None: + _reasoning_started_at = time.monotonic() + reasoning_accum += reasoning + if detect_state == _S_STREAMING: + if not in_thinking: + cumulative_display += "" + in_thinking = True + cumulative_display += reasoning + if not _suppress_visible_output: + yield { + "type": "content", + "text": cumulative_display, + } + + # ── Content tokens ── + token = delta.get("content", "") + if token: + # First answer token ends reasoning. + if ( + _reasoning_started_at is not None + and not _reasoning_summary_emitted + ): + _reasoning_summary_emitted = True + yield _reasoning_summary_event(_reasoning_started_at) + has_content_tokens = True + content_accum += token + + if detect_state == _S_DRAINING: + pass # accumulate silently + + elif detect_state == _S_STREAMING: + if in_thinking: cumulative_display += "" + in_thinking = False + cumulative_display += token + cleaned = _strip_tool_markup_streaming(cumulative_display) + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned if not _suppress_visible_output: yield { "type": "content", - "text": _strip_tool_markup( - cumulative_display, - final = True, - ), - } - else: - cumulative_display = reasoning_accum - if not _suppress_visible_output: - yield { - "type": "content", - "text": cumulative_display, - } - _stream_done = True - break # exit inner while - if not line.startswith("data: "): - continue - - try: - chunk_data = json.loads(line[6:]) - _ct = chunk_data.get("timings") - if _ct: - _iter_timings = _ct - _cu = chunk_data.get("usage") - if _cu: - _iter_usage = _cu - - choices = chunk_data.get("choices", []) - if not choices: - continue - - delta = choices[0].get("delta", {}) - _fr = choices[0].get("finish_reason") - if _fr: - _iter_finish_reason = _fr - - # ── Structured tool_calls ── - tc_deltas = delta.get("tool_calls") - if tc_deltas: - # llama-server can emit visible assistant - # preface content before native structured - # tool_calls. Preserve content_accum as - # the assistant pre-tool text and still - # drain/execute the structured call. - has_structured_tc = True - detect_state = _S_DRAINING - for tc_d in tc_deltas: - idx = tc_d.get("index", 0) - if idx not in tool_calls_acc: - tool_calls_acc[idx] = { - "id": tc_d.get("id", f"call_{idx}"), - "type": "function", - "function": { - "name": "", - "arguments": "", - }, - } - elif tc_d.get("id"): - # Update ID if a real one - # arrives on a later delta. - tool_calls_acc[idx]["id"] = tc_d["id"] - func = tc_d.get("function", {}) - if func.get("name"): - tool_calls_acc[idx]["function"]["name"] += func[ - "name" - ] - if func.get("arguments"): - tool_calls_acc[idx]["function"]["arguments"] += ( - func["arguments"] - ) - current_name = tool_calls_acc[idx]["function"].get( - "name", "" - ) - fallback_id = f"call_{idx}" - current_id = tool_calls_acc[idx].get("id", fallback_id) - already_started = ( - current_id in provisional_render_html_tool_call_ids - ) - has_real_id = current_id != fallback_id - if ( - current_name == "render_html" - and not _tool_succeeded("render_html") - and any( - ( - (tool.get("function") or {}).get("name") - == "render_html" - ) - for tool in active_tools - ) - and not already_started - and not provisional_render_html_tool_call_ids - and has_real_id - ): - provisional_render_html_tool_call_ids.add( - current_id - ) - yield { - "type": "tool_start", - "tool_name": "render_html", - "tool_call_id": current_id, - "arguments": {}, - "provenance": tool_event_provenance( - provisional = True, - ), - } - continue - - # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). - reasoning = delta.get("reasoning_content", "") - if reasoning: - reasoning_accum += reasoning - if detect_state == _S_STREAMING: - if not in_thinking: - cumulative_display += "" - in_thinking = True - cumulative_display += reasoning - if not _suppress_visible_output: - yield { - "type": "content", - "text": cumulative_display, + "text": cleaned, } - # ── Content tokens ── - token = delta.get("content", "") - if token: - has_content_tokens = True - content_accum += token + elif detect_state == _S_BUFFERING: + content_buffer += token + stripped_buf = content_buffer.lstrip() + if not stripped_buf: + continue - if detect_state == _S_DRAINING: - pass # accumulate silently + # Check tool signal prefixes. + is_prefix = False + is_match = False + for sig in _tool_xml_signals: + if stripped_buf.startswith(sig): + is_match = True + break + if sig.startswith(stripped_buf): + is_prefix = True + break - elif detect_state == _S_STREAMING: - if in_thinking: - cumulative_display += "" - in_thinking = False - cumulative_display += token + if is_match: + # Tool signal -- flush any visible + # prefix before DRAINING so the + # route sends it before tool_start. + _flush_reasoning_and_buffer() cleaned = _strip_tool_markup_streaming( - cumulative_display + cumulative_display, + force = True, + ) + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned + if not _suppress_visible_output: + yield { + "type": "content", + "text": cleaned, + } + detect_state = _S_DRAINING + elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS: + pass # keep buffering + else: + # Not a tool -- flush buffer + detect_state = _S_STREAMING + # Flush reasoning accumulated + # during BUFFERING. + _flush_reasoning_and_buffer() + cleaned = _strip_tool_markup( + cumulative_display, ) if len(cleaned) > len(_last_emitted): _last_emitted = cleaned @@ -6053,73 +8262,10 @@ class LlamaCppBackend: "text": cleaned, } - elif detect_state == _S_BUFFERING: - content_buffer += token - stripped_buf = content_buffer.lstrip() - if not stripped_buf: - continue - - # Check tool signal prefixes. - is_prefix = False - is_match = False - for sig in _tool_xml_signals: - if stripped_buf.startswith(sig): - is_match = True - break - if sig.startswith(stripped_buf): - is_prefix = True - break - - if is_match: - # Tool signal -- flush any visible - # prefix before DRAINING so the - # route sends it before tool_start. - if reasoning_accum: - cumulative_display += "" - cumulative_display += reasoning_accum - cumulative_display += "" - cumulative_display += content_buffer - cleaned = _strip_tool_markup_streaming( - cumulative_display, - force = True, - ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned - if not _suppress_visible_output: - yield { - "type": "content", - "text": cleaned, - } - detect_state = _S_DRAINING - elif ( - is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS - ): - pass # keep buffering - else: - # Not a tool -- flush buffer - detect_state = _S_STREAMING - # Flush reasoning accumulated - # during BUFFERING. - if reasoning_accum: - cumulative_display += "" - cumulative_display += reasoning_accum - cumulative_display += "" - cumulative_display += content_buffer - cleaned = _strip_tool_markup( - cumulative_display, - ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned - if not _suppress_visible_output: - yield { - "type": "content", - "text": cleaned, - } - - except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") - if _stream_done: - break # exit outer for + except json.JSONDecodeError: + logger.debug(f"Skipping malformed SSE line: {line[:100]}") + if _stream_done: + break # exit outer for # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: @@ -6130,11 +8276,7 @@ class LlamaCppBackend: detect_state = _S_STREAMING if content_buffer: # Flush reasoning first. - if reasoning_accum: - cumulative_display += "" - cumulative_display += reasoning_accum - cumulative_display += "" - cumulative_display += content_buffer + _flush_reasoning_and_buffer() if not _suppress_visible_output: yield { "type": "content", @@ -6144,9 +8286,10 @@ class LlamaCppBackend: ), } elif reasoning_accum and not has_content_tokens: - # Reasoning-only response: show reasoning as plain - # text, matching the final streaming pass for - # models that put everything in reasoning. + # Reasoning-only reply: show it as plain text. + if _reasoning_started_at is not None and not _reasoning_summary_emitted: + _reasoning_summary_emitted = True + yield _reasoning_summary_event(_reasoning_started_at) cumulative_display = reasoning_accum if not _suppress_visible_output: yield { @@ -6255,31 +8398,11 @@ class LlamaCppBackend: # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} - _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} - _fc = _fu.get("completion_tokens", 0) - _fp = _fu.get("prompt_tokens", 0) - _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: - _mt = dict(_iter_timings) if _iter_timings else {} - if _accumulated_predicted_ms or _accumulated_predicted_n: - _mt["predicted_ms"] = ( - _mt.get("predicted_ms", 0) + _accumulated_predicted_ms - ) - _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n - _mt["predicted_n"] = _tn - _tms = _mt["predicted_ms"] - if _tms > 0: - _mt["predicted_per_second"] = _tn / (_tms / 1000.0) - yield { - "type": "metadata", - "usage": { - "prompt_tokens": _fp, - "completion_tokens": _tc, - "total_tokens": _fp + _tc, - }, - "timings": _mt, - "finish_reason": _iter_finish_reason, - } + _meta = _build_metadata_event( + _iter_usage, _iter_timings, _iter_finish_reason + ) + if _meta is not None: + yield _meta return # Safety net caught tool XML -- treat as tool call. @@ -6330,31 +8453,11 @@ class LlamaCppBackend: content_accum = _strip_tool_markup(content_accum, final = True) if content_accum: yield {"type": "content", "text": content_accum} - _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} - _fc = _fu.get("completion_tokens", 0) - _fp = _fu.get("prompt_tokens", 0) - _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: - _mt = dict(_iter_timings) if _iter_timings else {} - if _accumulated_predicted_ms or _accumulated_predicted_n: - _mt["predicted_ms"] = ( - _mt.get("predicted_ms", 0) + _accumulated_predicted_ms - ) - _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n - _mt["predicted_n"] = _tn - _tms = _mt["predicted_ms"] - if _tms > 0: - _mt["predicted_per_second"] = _tn / (_tms / 1000.0) - yield { - "type": "metadata", - "usage": { - "prompt_tokens": _fp, - "completion_tokens": _tc, - "total_tokens": _fp + _tc, - }, - "timings": _mt, - "finish_reason": _iter_finish_reason, - } + _meta = _build_metadata_event( + _iter_usage, _iter_timings, _iter_finish_reason + ) + if _meta is not None: + yield _meta return # ── Execute tool calls ── @@ -6377,20 +8480,30 @@ class LlamaCppBackend: for tc in tool_calls or []: func = tc.get("function", {}) tool_name = func.get("name", "") - provisional_render_html_match = ( - tool_name == "render_html" - and tc.get("id") in provisional_render_html_tool_call_ids - ) + provisional_match = tc.get("id") in provisional_started_tool_calls decision = tool_controller.prepare_call( tc, forced = _forced_tool_call_pending, - provisional = provisional_render_html_match, + provisional = provisional_match, ) if not decision.should_execute: if content_text and not assistant_appended: conversation.append(assistant_msg) assistant_appended = True + if provisional_match: + # A provisional tool card is already on screen for this + # id; close it so it never dangles when the controller + # turns the call into an internal no-op (duplicate / + # disabled / render_html_repeat). + resolved_provisional_tool_call_ids.add(decision.tool_call_id) + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": "", + "provenance": decision.provenance, + } completion = tool_controller.record_noop(decision) conversation.append(completion.model_message()) if _forced_tool_call_pending: @@ -6410,7 +8523,9 @@ class LlamaCppBackend: decision.as_assistant_tool_call() ) - needs_confirm = bool(confirm_tool_calls) + # Bypass wins over the confirm gate at the loop level too, + # so a direct internal caller with both flags never prompts. + needs_confirm = bool(confirm_tool_calls) and not bypass_permissions approval_id = new_approval_id() if needs_confirm else "" decision_slot = ( begin_tool_decision(session_id, approval_id) if needs_confirm else None @@ -6433,6 +8548,7 @@ class LlamaCppBackend: == "deny" ): decision_slot = None + resolved_provisional_tool_call_ids.add(decision.tool_call_id) yield { "type": "tool_end", "tool_name": decision.tool_name, @@ -6471,16 +8587,30 @@ class LlamaCppBackend: timeout = _effective_timeout, session_id = session_id, rag_scope = rag_scope, + disable_sandbox = bypass_permissions, ) if decision.tool_name == "search_knowledge_base": _kb_search_count += 1 completion = tool_controller.record_result(decision, result) + resolved_provisional_tool_call_ids.add(decision.tool_call_id) yield completion.tool_end_event() conversation.append(completion.tool_message()) if _forced_tool_call_pending: _forced_tool_call_pending = False + # Close provisional cards not resolved by execution/no-op handling. + for _pid, _pname in provisional_started_tool_calls.items(): + if _pid not in resolved_provisional_tool_call_ids: + resolved_provisional_tool_call_ids.add(_pid) + yield { + "type": "tool_end", + "tool_name": _pname, + "tool_call_id": _pid, + "result": "", + "provenance": tool_event_provenance(provisional = True), + } + # Clear tool status badge before next generation/final pass. yield {"type": "status", "text": ""} if tool_controller.force_final_answer or not tool_controller.active_tools(): @@ -6489,10 +8619,32 @@ class LlamaCppBackend: continue except httpx.ConnectError: + # Mark unresolved provisional cards as failed before raising. + for _pid, _pname in provisional_started_tool_calls.items(): + if _pid not in resolved_provisional_tool_call_ids: + resolved_provisional_tool_call_ids.add(_pid) + yield { + "type": "tool_end", + "tool_name": _pname, + "tool_call_id": _pid, + "result": "Error: lost connection to llama-server before the tool call completed.", + "provenance": tool_event_provenance(provisional = True), + } raise RuntimeError("Lost connection to llama-server") except Exception as e: if cancel_event is not None and cancel_event.is_set(): return + # Same cleanup for other mid-iteration failures. + for _pid, _pname in provisional_started_tool_calls.items(): + if _pid not in resolved_provisional_tool_call_ids: + resolved_provisional_tool_call_ids.add(_pid) + yield { + "type": "tool_end", + "tool_name": _pname, + "tool_call_id": _pid, + "result": "Error: the tool call was interrupted before it completed.", + "provenance": tool_event_provenance(provisional = True), + } raise # ── Tool iteration cap reached -- synthesize final answer ── @@ -6546,131 +8698,107 @@ class LlamaCppBackend: in_thinking = False has_content_tokens = False reasoning_text = "" + _final_reasoning_started_at: Optional[float] = None + _final_reasoning_summary_emitted = False _metadata_usage = None _metadata_timings = None _metadata_finish_reason = None _stream_done = False try: - stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client( - timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) - ) as client: - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - with self._stream_with_retry( - client, - url, - stream_payload, + with self._open_stream(url, stream_payload, cancel_event) as ( + response, + first_token_deadline, + ): + buffer = "" + for raw_chunk in self._iter_text_cancellable( + response, cancel_event, - headers = _auth_headers, first_token_deadline = first_token_deadline, - ) as response: - if response.status_code != 200: - error_body = response.read().decode() - raise RuntimeError( - f"llama-server returned {response.status_code}: {error_body}" - ) + ): + buffer += raw_chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() - buffer = "" - for raw_chunk in self._iter_text_cancellable( - response, - cancel_event, - first_token_deadline = first_token_deadline, - ): - buffer += raw_chunk - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) - line = line.strip() + if not line: + continue + if line == "data: [DONE]": + if in_thinking: + if ( + _final_reasoning_started_at is not None + and not _final_reasoning_summary_emitted + ): + _final_reasoning_summary_emitted = True + yield _reasoning_summary_event(_final_reasoning_started_at) + if has_content_tokens: + cumulative += "" + yield { + "type": "content", + "text": _strip_tool_markup(cumulative, final = True), + } + else: + cumulative = reasoning_text + yield {"type": "content", "text": cumulative} + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue - if not line: - continue - if line == "data: [DONE]": - if in_thinking: - if has_content_tokens: + try: + chunk_data = json.loads(line[6:]) + # Capture server timings/usage from final chunks. + _chunk_timings = chunk_data.get("timings") + if _chunk_timings: + _metadata_timings = _chunk_timings + _chunk_usage = chunk_data.get("usage") + if _chunk_usage: + _metadata_usage = _chunk_usage + choices = chunk_data.get("choices", []) + if choices: + delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _metadata_finish_reason = _fr + + reasoning = delta.get("reasoning_content", "") + if reasoning: + if _final_reasoning_started_at is None: + _final_reasoning_started_at = time.monotonic() + reasoning_text += reasoning + if not in_thinking: + cumulative += "" + in_thinking = True + cumulative += reasoning + yield {"type": "content", "text": cumulative} + + token = delta.get("content", "") + if token: + if ( + _final_reasoning_started_at is not None + and not _final_reasoning_summary_emitted + ): + _final_reasoning_summary_emitted = True + yield _reasoning_summary_event(_final_reasoning_started_at) + has_content_tokens = True + if in_thinking: cumulative += "" - yield { - "type": "content", - "text": _strip_tool_markup(cumulative, final = True), - } - else: - cumulative = reasoning_text - yield {"type": "content", "text": cumulative} - _stream_done = True - break # exit inner while - if not line.startswith("data: "): - continue - - try: - chunk_data = json.loads(line[6:]) - # Capture server timings/usage from final chunks. - _chunk_timings = chunk_data.get("timings") - if _chunk_timings: - _metadata_timings = _chunk_timings - _chunk_usage = chunk_data.get("usage") - if _chunk_usage: - _metadata_usage = _chunk_usage - choices = chunk_data.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - _fr = choices[0].get("finish_reason") - if _fr: - _metadata_finish_reason = _fr - - reasoning = delta.get("reasoning_content", "") - if reasoning: - reasoning_text += reasoning - if not in_thinking: - cumulative += "" - in_thinking = True - cumulative += reasoning - yield {"type": "content", "text": cumulative} - - token = delta.get("content", "") - if token: - has_content_tokens = True - if in_thinking: - cumulative += "" - in_thinking = False - cumulative += token - cleaned = _strip_tool_markup(cumulative) - # Emit only when cleaned text grows (monotonic). - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned - yield {"type": "content", "text": cleaned} - except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") - if _stream_done: - break # exit outer for - _final_usage = _metadata_usage or {} - _final_completion = _final_usage.get("completion_tokens", 0) - _final_prompt = _final_usage.get("prompt_tokens", 0) - _total_completion = _final_completion + _accumulated_completion_tokens - if _metadata_usage or _metadata_timings or _metadata_finish_reason: - _merged_timings = dict(_metadata_timings) if _metadata_timings else {} - if _accumulated_predicted_ms or _accumulated_predicted_n: - _merged_timings["predicted_ms"] = ( - _merged_timings.get("predicted_ms", 0) + _accumulated_predicted_ms - ) - _total_predicted_n = ( - _merged_timings.get("predicted_n", 0) + _accumulated_predicted_n - ) - _merged_timings["predicted_n"] = _total_predicted_n - _total_predicted_ms = _merged_timings["predicted_ms"] - if _total_predicted_ms > 0: - _merged_timings["predicted_per_second"] = _total_predicted_n / ( - _total_predicted_ms / 1000.0 - ) - yield { - "type": "metadata", - "usage": { - "prompt_tokens": _final_prompt, - "completion_tokens": _total_completion, - "total_tokens": _final_prompt + _total_completion, - }, - "timings": _merged_timings, - "finish_reason": _metadata_finish_reason, - } + in_thinking = False + cumulative += token + cleaned = _strip_tool_markup(cumulative) + # Emit only when cleaned text grows (monotonic). + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned + yield {"type": "content", "text": cleaned} + except json.JSONDecodeError: + logger.debug(f"Skipping malformed SSE line: {line[:100]}") + if _stream_done: + break # exit outer for + _meta = _build_metadata_event( + _metadata_usage, _metadata_timings, _metadata_finish_reason + ) + if _meta is not None: + yield _meta except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") @@ -6708,8 +8836,6 @@ class LlamaCppBackend: continue if not isinstance(block, dict): return True - if block.get("type") == "text" and isinstance(block.get("text"), str): - continue if isinstance(block.get("text"), str): continue return True @@ -6730,9 +8856,7 @@ class LlamaCppBackend: parts = [] for block in content: if isinstance(block, dict): - if block.get("type") == "text" and isinstance(block.get("text"), str): - parts.append(block["text"]) - elif isinstance(block.get("text"), str): + if isinstance(block.get("text"), str): parts.append(block["text"]) elif isinstance(block, str): parts.append(block) @@ -6747,8 +8871,7 @@ class LlamaCppBackend: system_text = _block_text(system) try: - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client(timeout = 10, headers = _auth_headers) as client: + with httpx.Client(timeout = 10, headers = self._auth_headers) as client: def _tokenize(text: str) -> int: r = client.post( @@ -6776,7 +8899,7 @@ class LlamaCppBackend: try: # llama-server's /apply-template renders tool declarations # into the prompt when ``tools`` is supplied, so pass them - # through — otherwise tool-schema tokens go uncounted. + # through, otherwise tool-schema tokens go uncounted. template_body = {"messages": template_messages} if tools: template_body["tools"] = tools @@ -6827,12 +8950,44 @@ class LlamaCppBackend: logger.debug(f"Audio type detection failed: {e}") return None + def _apply_detected_audio(self, detected: Optional[str]) -> bool: + """Apply a probed audio codec under self._lock. Returns True to continue + the load (codec inited OK, or nothing to init), False to abort (server + unhealthy or codec init failed). Shared by the fast-path retry and the + main load path.""" + if detected in ("snac", "bicodec", "dac"): + with self._lock: + if not self._healthy: + return False + try: + self.init_audio_codec(detected) + self._is_audio = True + self._audio_type = detected + except Exception as exc: + # Surface as HTTP 500 (matches pre-PR contract). + logger.warning("Failed to init audio codec '%s': %s", detected, exc) + self._audio_probed = False + return False + elif detected: + # csm / whisper / audio_vlm: track type but keep _is_audio False -- + # GGUF TTS routing only fires for snac/bicodec/dac. + with self._lock: + if not self._healthy: + return False + self._audio_type = detected + # Audio input = token probe (audio_vlm/whisper) OR mmproj encoder. + from utils.models.model_config import is_audio_input_type + + self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( + self._mmproj_has_audio + ) + return True + def _detect_audio_type_strict(self) -> Optional[str]: """Codec name on match, None on non-audio, raises on transport/JSON errors.""" if not self.is_loaded: return None - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client(timeout = 10, headers = _auth_headers) as client: + with httpx.Client(timeout = 10, headers = self._auth_headers) as client: def _detok(tid: int) -> str: # Non-200 means "marker not in vocab" -- keep probing. @@ -6946,8 +9101,9 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client: + with httpx.Client( + timeout = httpx.Timeout(300, connect = 10), headers = self._auth_headers + ) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}") diff --git a/studio/backend/core/inference/llama_http.py b/studio/backend/core/inference/llama_http.py new file mode 100644 index 0000000000..b554949c3e --- /dev/null +++ b/studio/backend/core/inference/llama_http.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared pooled httpx.AsyncClient for NON-streaming calls to the local llama-server. + +Streaming generation must NOT use this. It relies on ``Connection: close`` and +``max_keepalive_connections=0`` so a client disconnect tears down the upstream +socket and stops GPU decode (PR #5749). This pooled client is only for short +request/response proxy calls (non-streaming completions, embeddings) where +reusing a connection removes per-request setup cost. Per-request ``timeout`` is +still passed at each call site. +""" + +from __future__ import annotations + +import asyncio +import weakref + +import httpx + +_LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32) + + +def _new_client() -> httpx.AsyncClient: + try: + return httpx.AsyncClient(limits = _LIMITS) + except Exception: + # Mirror external_provider: an unsupported env proxy scheme can raise. + return httpx.AsyncClient(limits = _LIMITS, trust_env = False) + + +# One client per running event loop: an httpx client binds its transport to the +# loop it first runs on, so a single global instance breaks across a lifespan +# restart or a second test loop. Weak keys let a finished loop drop its client. +_clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]" = ( + weakref.WeakKeyDictionary() +) + + +def nonstreaming_client() -> httpx.AsyncClient: + loop = asyncio.get_running_loop() + client = _clients.get(loop) + if client is None or client.is_closed: + client = _new_client() + _clients[loop] = client + return client + + +async def aclose() -> None: + clients = list(_clients.values()) + _clients.clear() + for client in clients: + try: + await client.aclose() + except Exception: + pass diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 69a86fa3ba..b42be5ee0d 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -13,7 +13,8 @@ Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md from __future__ import annotations -from typing import Iterable, Optional +import os +from typing import Iterable, Mapping, Optional # Each group = every alias (short + long) of one hard-denied flag. # Extend the matching group when llama.cpp adds a new alias. @@ -124,7 +125,9 @@ def is_managed_flag(flag: str) -> bool: # from inherited extras so they can't last-wins-override an Apply that # re-sets the same field. _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) -_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}) +_CACHE_TYPE_K_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k"}) +_CACHE_TYPE_V_FLAGS: frozenset[str] = frozenset({"-ctv", "--cache-type-v"}) +_CACHE_FLAGS: frozenset[str] = _CACHE_TYPE_K_FLAGS | _CACHE_TYPE_V_FLAGS _SPEC_FLAGS: frozenset[str] = frozenset( { "--spec-default", @@ -133,13 +136,22 @@ _SPEC_FLAGS: frozenset[str] = frozenset( "--spec-ngram-size", "--draft-min", "--draft-max", - # MTP path (llama.cpp #22673). --model-draft and aliases are - # Studio-managed since the separate-drafter support (Gemma 4): an - # inherited copy must not last-wins-override the auto-detected - # drafter. Explicit extras for the current load are never stripped. + # MTP path (llama.cpp #22673). The drafter selectors (local --model-draft + # and HF --spec-draft-hf aliases) are Studio-managed since the separate- + # drafter support (Gemma 4): an inherited copy must not last-wins-override + # the auto-detected drafter. Explicit extras for the current load are never + # stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld, + # --spec-draft-device) are deliberately NOT stripped: the VRAM budget reads + # them via the same parsers the child honors, so they stay consistent on + # inherit, and stripping them would silently move a CPU-offloaded drafter + # back onto the GPU. "--model-draft", "-md", "--spec-draft-model", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", "--spec-draft-n-max", "--spec-draft-n-min", "--spec-draft-p-min", @@ -274,6 +286,20 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return _last_flag_value(args, _CACHE_FLAGS) +def parse_cache_override_per_axis( + args: Optional[Iterable[str]], +) -> tuple[Optional[str], Optional[str]]: + """Last-wins --cache-type-k / --cache-type-v values kept apart, as (k, v). + + parse_cache_override collapses both axes to one last-wins value; this keeps + them separate so an asymmetric K/V can be budgeted by its heavier axis. + """ + return ( + _last_flag_value(args, _CACHE_TYPE_K_FLAGS), + _last_flag_value(args, _CACHE_TYPE_V_FLAGS), + ) + + def resolve_cache_type_kv( args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str] ) -> Optional[str]: @@ -309,6 +335,60 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral return override.strip().lower() == "tensor" +def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool: + """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio + emits --split-mode only on its tensor branch, so a tensor env on the layer + path would run the child tensor-parallel unbudgeted; this flips the budget + to tensor. Only tensor is heavier, so other modes are ignored.""" + raw = (os.environ if env is None else env).get("LLAMA_ARG_SPLIT_MODE") + return bool(raw) and raw.strip().lower() == "tensor" + + +def _effective_tensor_parallel( + extra_args: Optional[Iterable[str]], + tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Tensor-parallel decision including the inherited LLAMA_ARG_SPLIT_MODE env. + + resolve_tensor_parallel (extras + toggle), flipped on when extras set no split + mode but the child inherits a tensor split env. Shared by load_model (which + budgets and launches it) and the tensor-fallback wrapper (so an env-only + tensor crash still retries layer split).""" + resolved = resolve_tensor_parallel(extra_args, tensor_parallel) + if ( + not resolved + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + return True + return resolved + + +def _tensor_parallel_matches_loaded( + extra_args: Optional[Iterable[str]], + requested_tensor_parallel: bool, + loaded_tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Whether a duplicate load request matches a loaded server's tensor state. + + Env-only tensor mode is a launch hint load_model may downgrade to layer split + (capacity/buffer), scrubbing the child env. So only let an inherited tensor env + raise a match against a server that *actually* launched tensor; on a downgraded + (layer) server the env is ignored, and an identical request would downgrade the + same way -- avoiding an endless reload of a healthy server.""" + requested = resolve_tensor_parallel(extra_args, requested_tensor_parallel) + if ( + loaded_tensor_parallel + and not requested + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + requested = True + return requested == loaded_tensor_parallel + + _MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) _MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) diff --git a/studio/backend/core/inference/llama_stats.py b/studio/backend/core/inference/llama_stats.py new file mode 100644 index 0000000000..6047aedbc0 --- /dev/null +++ b/studio/backend/core/inference/llama_stats.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Translate llama-server's Prometheus /metrics into a periodic, vLLM-style +engine-stats log line (generation/prompt throughput, requests in flight). + +llama-server already computes these (it needs `--metrics`); this lifts them +into Studio's structured log so the terminal shows serving health, not just +per-request access lines. Emitted only while there is activity. +""" + +import os +import re +import threading +import time +import urllib.request + +# Prometheus body lines: "llamacpp:[{labels}] " (skip "#" HELP/TYPE). +_METRIC_RE = re.compile(r"^llamacpp:(\w+)(?:\{[^}]*\})?\s+([0-9.eE+-]+)", re.MULTILINE) +_OFF = {"0", "false", "no", "off"} + + +class LlamaServerStatsLogger: + """Daemon poller that logs vLLM-style engine stats from llama-server. + + Keeps retrying through transient scrape failures; the backend stops it via + stop() on unload/reload, so a brief /metrics stall does not silence stats. + """ + + def __init__( + self, + base_url, + logger, + interval_s = 10.0, + ): + self._url = f"{base_url.rstrip('/')}/metrics" + self._log = logger + self._interval = max(1.0, float(interval_s)) + self._stop = threading.Event() + self._thread = None + + def start(self): + if self._thread is None: + self._thread = threading.Thread(target = self._run, name = "llama-stats", daemon = True) + self._thread.start() + + def stop(self): + self._stop.set() + + def _scrape(self): + try: + with urllib.request.urlopen(self._url, timeout = 3) as r: + if r.status != 200: + return None + body = r.read().decode("utf-8", "replace") + except Exception: + return None + out = {} + for k, v in _METRIC_RE.findall(body): + try: # a malformed value must not kill the daemon thread + out[k] = float(v) + except ValueError: + continue + return out + + def _run(self): + misses = 0 + prev = None # (monotonic_t, tokens_predicted_total, prompt_tokens_total) + while not self._stop.wait(self._interval): + m = self._scrape() + if not m: + misses += 1 + if misses == 3: # transient stall (load/GC); keep polling. + self._log.debug("engine_stats: /metrics scrape failing, still retrying") + continue # real shutdown is driven by stop() from _kill_process + misses = 0 + # Generation tokens come from tokens_predicted_total (counter) and + # predicted_tokens_seconds (gauge); n_decode_total counts + # llama_decode() calls, not tokens, so it must not feed tok/s. + now = time.monotonic() + predicted = m.get("tokens_predicted_total", 0.0) + prompt = m.get("prompt_tokens_total", 0.0) + gen_delta = prompt_delta = 0.0 + if prev is not None and now > prev[0]: + dt = now - prev[0] + gen_delta = max(0.0, (predicted - prev[1]) / dt) + prompt_delta = max(0.0, (prompt - prev[2]) / dt) + prev = (now, predicted, prompt) + # Prefer llama.cpp's own throughput gauges; fall back to the counter + # delta for binaries that expose only the counters. + gen_tps = m.get("predicted_tokens_seconds") or gen_delta + prompt_tps = m.get("prompt_tokens_seconds") or prompt_delta + running, waiting = ( + int(m.get("requests_processing", 0)), + int(m.get("requests_deferred", 0)), + ) + # Gate on real activity this tick so a stale gauge never logs at idle. + if running or waiting or gen_delta or prompt_delta: + self._log.info( + "engine_stats", + gen_tok_s = round(float(gen_tps), 1), + prompt_tok_s = round(float(prompt_tps), 1), + running = running, + waiting = waiting, + ) + + +def maybe_start_stats_logger(base_url, logger): + """Start a stats logger unless UNSLOTH_STUDIO_ENGINE_STATS disables it.""" + if (os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "").strip().lower() in _OFF: + return None + try: + interval = float(os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS_INTERVAL_S", "10")) + except ValueError: + interval = 10.0 + sl = LlamaServerStatsLogger(base_url, logger, interval) + sl.start() + return sl diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 5a36d90c5d..7bd4a7d6e9 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -117,10 +117,25 @@ def join_stdio_command(parts: list[str]) -> str: def stdio_mcp_enabled() -> bool: """stdio MCP servers spawn local processes as the backend user (bypassing the - sandbox), so allowed only when the host is the user's own machine. The Tauri - app sets UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1; localhost/self-hosted users can opt - in with the same var. Off for Colab and any network (0.0.0.0) bind.""" - return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1" + sandbox), so allowed only when the host is the user's own machine. On startup + a loopback bind defaults UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see + utils.host_policy.apply_stdio_mcp_loopback_default, called from run.py); the + Tauri app does the same. Off for Colab and any network (0.0.0.0) bind unless + an operator sets the var out-of-band; set it to 0 to force-disable. + + When stdio is on only because of that loopback auto-default, an explicit + `unsloth studio run --disable-tools` turns it back off (a local stdio command + is server-side code execution). An explicit operator opt-in via the env var + still wins -- including the documented `=1` network opt-in, where the process + tool policy is False merely by the external-host default, not by choice.""" + if os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") != "1": + return False + from state.tool_policy import get_tool_policy + from utils.host_policy import loopback_default_active + + if loopback_default_active() and get_tool_policy() is False: + return False + return True # Probe timeouts for discovering a server's tool list. OAuth needs minutes for diff --git a/studio/backend/core/inference/message_content.py b/studio/backend/core/inference/message_content.py new file mode 100644 index 0000000000..b7c499a087 --- /dev/null +++ b/studio/backend/core/inference/message_content.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Normalize chat-message `content` (string or OpenAI multimodal list) to text. + +String-only formatting paths called string ops directly on `content` and broke +on the list form (#4383). `content_to_text` collapses either shape to a string, +dropping non-text parts. No heavy imports, so it is unit-testable alone. +""" + +from __future__ import annotations + +from typing import Any + + +def content_to_text(content: Any) -> str: + """Plain text of a `content`: str unchanged, list/tuple text parts newline-joined + (non-text dropped), None to "", else str(content).""" + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, (list, tuple)): + parts = [] + for item in content: + if isinstance(item, str): + if item: + parts.append(item) + elif isinstance(item, dict): + # Skip non-text parts (image_url, input_audio, ...). + part_type = item.get("type") + if part_type is not None and part_type != "text": + continue + text = item.get("text") + if isinstance(text, str) and text: + parts.append(text) + return "\n".join(parts) + return str(content) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4ccac2912e..c980dbde2d 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -18,7 +18,6 @@ import atexit import base64 import os import signal -import structlog from loggers import get_logger import multiprocessing as mp import queue @@ -30,15 +29,15 @@ from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union from utils.hardware import prepare_gpu_selection +# Re-exported from the shared helper so GGUF, training, and inference share one +# type; kept importable here for backwards compatibility. +from utils.hf_xet_fallback import DownloadStallError + logger = get_logger(__name__) _CTX = mp.get_context("spawn") -class DownloadStallError(RuntimeError): - """Raised when the worker reports no download progress for too long.""" - - # Dispatcher timeout constants (seconds) _DISPATCH_READ_TIMEOUT = 30.0 _DISPATCH_POLL_INTERVAL = 0.5 @@ -61,7 +60,6 @@ class InferenceOrchestrator: self._cmd_queue: Any = None self._resp_queue: Any = None self._cancel_event: Any = None # mp.Event — set to cancel generation - self._lock = threading.Lock() self._gen_lock = threading.Lock() # Serializes generation # Dispatcher state for compare mode (adapter-controlled requests): @@ -76,7 +74,6 @@ class InferenceOrchestrator: self.active_model_name: Optional[str] = None self.models: dict = {} self.loading_models: set = set() - self.loaded_local_models: list = [] from core.inference.defaults import get_default_models self._static_models = get_default_models() @@ -84,8 +81,6 @@ class InferenceOrchestrator: self._top_hub_cache: Optional[list[str]] = None self._top_models_ready = threading.Event() - self._current_transformers_major: Optional[str] = None # "4" or "5" - atexit.register(self._cleanup) logger.info("InferenceOrchestrator initialized (subprocess mode)") @@ -177,6 +172,9 @@ class InferenceOrchestrator: daemon = True, ) self._proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Inference subprocess started (pid=%s)", self._proc.pid) def _cancel_generation(self) -> None: @@ -389,6 +387,108 @@ class InferenceOrchestrator: return logger.warning("Timed out waiting for gen_done after cancel") + # ------------------------------------------------------------------ + # Generation command + token-stream helpers (shared by all paths) + # ------------------------------------------------------------------ + + def _build_generate_cmd( + self, + request_id: str, + image_b64: Optional[str], + *, + messages: list = None, + system_prompt: str = "", + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 256, + repetition_penalty: float = 1.0, + use_adapter = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + ) -> dict: + """Build the 'generate' command shared by the locked and dispatched paths.""" + cmd = { + "type": "generate", + "request_id": request_id, + "messages": messages or [], + "system_prompt": system_prompt, + "image_base64": image_b64, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + # Only forward template kwargs the caller set, for older worker compat. + if use_adapter is not None: + cmd["use_adapter"] = use_adapter + if tools is not None: + cmd["tools"] = tools + if enable_thinking is not None: + cmd["enable_thinking"] = enable_thinking + if reasoning_effort is not None: + cmd["reasoning_effort"] = reasoning_effort + if preserve_thinking is not None: + cmd["preserve_thinking"] = preserve_thinking + return cmd + + def _consume_token_stream( + self, + read_one, + drain_on_cancel, + *, + crash_context: str, + cancel_event = None, + stats_holder: Optional[dict] = None, + read_timeout: float = 30.0, + ) -> Generator[str, None, None]: + """Yield tokens from a response stream until gen_done/gen_error. + + ``read_one(timeout)`` returns the next response (or None on timeout) and + owns the queue choice — the shared resp_queue under _gen_lock, or a + per-request mailbox on the dispatcher path — so this loop stays agnostic + of which queue is read. On cancel, ``drain_on_cancel()`` consumes the + cancel ack from that same source so stale events don't leak into the + next request. + """ + while True: + resp = read_one(read_timeout) + if resp is None: + # Check subprocess health + if not self._ensure_subprocess_alive(): + yield f"Error: {self._subprocess_crash_message(crash_context)}" + return + continue + + rtype = resp.get("type", "") + if rtype == "status": + continue + # Subprocess-level error (no request_id); request-scoped failures + # arrive as gen_error below. + if rtype == "error" and not resp.get("request_id"): + yield f"Error: {resp.get('error', 'Unknown error')}" + return + + if rtype == "token": + # Cancel from route (e.g. SSE connection closed). + if cancel_event is not None and cancel_event.is_set(): + self._cancel_generation() + drain_on_cancel() + return + yield resp.get("text", "") + elif rtype == "gen_done": + if stats_holder is not None: + stats_holder["stats"] = resp.get("stats") + return + elif rtype == "gen_error": + yield f"Error: {resp.get('error', 'Unknown error')}" + return + # ------------------------------------------------------------------ # Dispatcher — per-request mailbox routing for compare mode # ------------------------------------------------------------------ @@ -451,13 +551,12 @@ class InferenceOrchestrator: continue # No matching mailbox (a _gen_lock reader or orphaned). Can't - # un-get from mp.Queue, so just log. - if rtype not in ("status",): - logger.debug( - "Dispatcher: no mailbox for request_id=%s type=%s, dropping", - rid, - rtype, - ) + # un-get from mp.Queue, so just log. (status was handled above.) + logger.debug( + "Dispatcher: no mailbox for request_id=%s type=%s, dropping", + rid, + rtype, + ) def _generate_dispatched( self, @@ -502,30 +601,23 @@ class InferenceOrchestrator: if image is not None: image_b64 = self._pil_to_base64(image) - cmd = { - "type": "generate", - "request_id": request_id, - "messages": messages or [], - "system_prompt": system_prompt, - "image_base64": image_b64, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "max_new_tokens": max_new_tokens, - "repetition_penalty": repetition_penalty, - } - - if use_adapter is not None: - cmd["use_adapter"] = use_adapter - if tools is not None: - cmd["tools"] = tools - if enable_thinking is not None: - cmd["enable_thinking"] = enable_thinking - if reasoning_effort is not None: - cmd["reasoning_effort"] = reasoning_effort - if preserve_thinking is not None: - cmd["preserve_thinking"] = preserve_thinking + cmd = self._build_generate_cmd( + request_id, + image_b64, + messages = messages, + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + use_adapter = use_adapter, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) # Create mailbox BEFORE sending command mailbox: queue.Queue = queue.Queue() @@ -540,36 +632,22 @@ class InferenceOrchestrator: yield f"Error: {exc}" return - # Read tokens from our private mailbox + def read_mailbox(timeout): + try: + return mailbox.get(timeout = timeout) + except queue.Empty: + return None + + # Read tokens from our private mailbox (the dispatcher owns resp_queue). try: - while True: - try: - resp = mailbox.get(timeout = _DISPATCH_READ_TIMEOUT) - except queue.Empty: - # Timeout — check subprocess health - if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message('generation')}" - return - continue - - rtype = resp.get("type", "") - - if rtype == "token": - # Cancel from route (e.g. SSE connection closed) - if cancel_event is not None and cancel_event.is_set(): - self._cancel_generation() - self._drain_mailbox(mailbox, timeout = 5.0) - return - yield resp.get("text", "") - - elif rtype == "gen_done": - if stats_holder is not None: - stats_holder["stats"] = resp.get("stats") - return - - elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" - return + yield from self._consume_token_stream( + read_mailbox, + lambda: self._drain_mailbox(mailbox, timeout = 5.0), + crash_context = "generation", + cancel_event = cancel_event, + stats_holder = stats_holder, + read_timeout = _DISPATCH_READ_TIMEOUT, + ) finally: with self._mailbox_lock: self._mailboxes.pop(request_id, None) @@ -636,7 +714,9 @@ class InferenceOrchestrator: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, + subject: Optional[str] = None, ) -> bool: """Load a model for inference. @@ -659,6 +739,8 @@ class InferenceOrchestrator: "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, + "subject": subject, "gpu_ids": gpu_ids, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( @@ -716,7 +798,6 @@ class InferenceOrchestrator: ) if resp.get("success"): - self._current_transformers_major = needed_major model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) self.models[self.active_model_name] = { @@ -738,7 +819,8 @@ class InferenceOrchestrator: logger.info("Model '%s' loaded successfully in subprocess", model_name) return True else: - error = resp.get("error", "Failed to load model") + # Worker reports failures (consent gate included) under "message". + error = resp.get("message") or resp.get("error") or "Failed to load model" self.loading_models.discard(model_name) self.active_model_name = None self.models.clear() @@ -862,6 +944,7 @@ class InferenceOrchestrator: session_id: Optional[str] = None, rag_scope: Optional[dict] = None, confirm_tool_calls: bool = False, + bypass_permissions: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, **_unused, @@ -924,6 +1007,7 @@ class InferenceOrchestrator: session_id = session_id, rag_scope = rag_scope, confirm_tool_calls = confirm_tool_calls, + bypass_permissions = bypass_permissions, ) def generate_with_adapter_control( @@ -982,127 +1066,42 @@ class InferenceOrchestrator: self._wait_dispatcher_idle() # Serialize generation: two concurrent readers on resp_queue would - # consume and drop each other's token events. + # consume and drop each other's token events. Hold _gen_lock across the + # cmd build + send + whole stream so we stay the sole resp_queue reader. with self._gen_lock: - yield from self._generate_locked( + request_id = str(uuid.uuid4()) + image_b64 = self._pil_to_base64(image) if image is not None else None + cmd = self._build_generate_cmd( + request_id, + image_b64, messages = messages, system_prompt = system_prompt, - image = image, temperature = temperature, top_p = top_p, top_k = top_k, min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, - cancel_event = cancel_event, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, - stats_holder = stats_holder, ) - def _generate_locked( - self, - messages: list = None, - system_prompt: str = "", - image = None, - temperature: float = 0.7, - top_p: float = 0.9, - top_k: int = 40, - min_p: float = 0.0, - max_new_tokens: int = 256, - repetition_penalty: float = 1.0, - cancel_event = None, - use_adapter = None, - tools: Optional[list] = None, - enable_thinking: Optional[bool] = None, - reasoning_effort: Optional[str] = None, - preserve_thinking: Optional[bool] = None, - stats_holder: Optional[dict] = None, - ) -> Generator[str, None, None]: - """Actual generation logic — must be called under _gen_lock.""" - request_id = str(uuid.uuid4()) - - # Convert PIL Image to base64 if needed - image_b64 = None - if image is not None: - image_b64 = self._pil_to_base64(image) - - cmd = { - "type": "generate", - "request_id": request_id, - "messages": messages or [], - "system_prompt": system_prompt, - "image_base64": image_b64, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "max_new_tokens": max_new_tokens, - "repetition_penalty": repetition_penalty, - } - - if use_adapter is not None: - cmd["use_adapter"] = use_adapter - # Only forward template kwargs the caller set, for older worker compat. - if tools is not None: - cmd["tools"] = tools - if enable_thinking is not None: - cmd["enable_thinking"] = enable_thinking - if reasoning_effort is not None: - cmd["reasoning_effort"] = reasoning_effort - if preserve_thinking is not None: - cmd["preserve_thinking"] = preserve_thinking - - try: - self._send_cmd(cmd) - except RuntimeError as exc: - yield f"Error: {exc}" - return - - # We are the only resp_queue reader (under _gen_lock). - while True: - resp = self._read_resp(timeout = 30.0) - - if resp is None: - # Check subprocess health - if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message('generation')}" - return - continue - - rtype = resp.get("type", "") - - # Status messages — skip - if rtype == "status": - continue - - # Error without request_id = subprocess-level error - resp_rid = resp.get("request_id") - if rtype == "error" and not resp_rid: - yield f"Error: {resp.get('error', 'Unknown error')}" + try: + self._send_cmd(cmd) + except RuntimeError as exc: + yield f"Error: {exc}" return - if rtype == "token": - # Cancel from route (e.g. SSE connection closed) - if cancel_event is not None and cancel_event.is_set(): - self._cancel_generation() - # Wait for the cancel ack so stale events don't leak into - # the next request. - self._drain_until_gen_done(timeout = 5.0) - return - yield resp.get("text", "") - - elif rtype == "gen_done": - if stats_holder is not None: - stats_holder["stats"] = resp.get("stats") - return - - elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" - return + yield from self._consume_token_stream( + self._read_resp, + lambda: self._drain_until_gen_done(timeout = 5.0), + crash_context = "generation", + cancel_event = cancel_event, + stats_holder = stats_holder, + ) def reset_generation_state(self): """Cancel any ongoing generation and reset state.""" @@ -1138,8 +1137,6 @@ class InferenceOrchestrator: if not self.active_model_name: raise RuntimeError("No active model") - import uuid - request_id = str(uuid.uuid4()) cmd = { @@ -1252,8 +1249,6 @@ class InferenceOrchestrator: return with self._gen_lock: - import uuid - request_id = str(uuid.uuid4()) # numpy array -> list for mp.Queue serialization @@ -1282,38 +1277,12 @@ class InferenceOrchestrator: yield f"Error: {exc}" return - # Yield tokens — same pattern as _generate_locked - while True: - resp = self._read_resp(timeout = 30.0) - - if resp is None: - if not self._ensure_subprocess_alive(): - yield ("Error: " + self._subprocess_crash_message("audio input generation")) - return - continue - - rtype = resp.get("type", "") - - if rtype == "status": - continue - - if rtype == "error" and not resp.get("request_id"): - yield f"Error: {resp.get('error', 'Unknown error')}" - return - - if rtype == "token": - if cancel_event is not None and cancel_event.is_set(): - self._cancel_generation() - self._drain_until_gen_done(timeout = 5.0) - return - yield resp.get("text", "") - - elif rtype == "gen_done": - return - - elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" - return + yield from self._consume_token_stream( + self._read_resp, + lambda: self._drain_until_gen_done(timeout = 5.0), + crash_context = "audio input generation", + cancel_event = cancel_event, + ) # ------------------------------------------------------------------ # Local helpers (no subprocess needed) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 3b6a393f3d..0c96378d6c 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -154,6 +154,7 @@ def run_safetensors_tool_loop( session_id: Optional[str] = None, rag_scope: Optional[dict] = None, confirm_tool_calls: bool = False, + bypass_permissions: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -235,12 +236,39 @@ def run_safetensors_tool_loop( cumulative_display = "" last_emitted = "" provisional_render_html_started = False + provisional_resolved = False provisional_render_html_id = f"call_{next_call_id}" + # When a human confirmation gate is active the real tool_start is keyed + # by an approval id and carries awaiting_confirmation, so an early + # provisional card (keyed by tool_call_id, no approval) would show the + # tool as "running" before the user has approved it. Suppress the early + # card in that case and let the gated tool_start be the first signal. + _provisional_confirm_gated = bool(confirm_tool_calls) and not bypass_permissions gen = _call_single_turn(single_turn, conversation, active_tools) prev_cumulative = "" - for cumulative in gen: + _gen_iter = iter(gen) + while True: + try: + cumulative = next(_gen_iter) + except StopIteration: + break + except Exception: + # The model pipeline raised mid-stream. If a provisional + # render_html card was already surfaced, close it as errored so + # the UI never leaves a tool card spinning after the turn fails. + if provisional_render_html_started and not provisional_resolved: + provisional_resolved = True + yield { + "type": "tool_end", + "tool_name": "render_html", + "tool_call_id": provisional_render_html_id, + "result": "Error: generation was interrupted before the tool call completed.", + "provenance": _tool_event_provenance(provisional = True), + } + raise + if cancel_event is not None and cancel_event.is_set(): return @@ -256,6 +284,7 @@ def run_safetensors_tool_loop( if detect_state == _state_draining: if ( not _tool_succeeded("render_html") + and not _provisional_confirm_gated and any( ((tool.get("function") or {}).get("name") == "render_html") for tool in active_tools @@ -294,6 +323,7 @@ def run_safetensors_tool_loop( detect_state = _state_draining if ( not _tool_succeeded("render_html") + and not _provisional_confirm_gated and any( ((tool.get("function") or {}).get("name") == "render_html") for tool in active_tools @@ -352,6 +382,7 @@ def run_safetensors_tool_loop( detect_state = _state_draining if ( not _tool_succeeded("render_html") + and not _provisional_confirm_gated and any( ((tool.get("function") or {}).get("name") == "render_html") for tool in active_tools @@ -459,7 +490,8 @@ def run_safetensors_tool_loop( tool_protocol_active = False, ), } - if provisional_render_html_started: + if provisional_render_html_started and not provisional_resolved: + provisional_resolved = True yield { "type": "tool_end", "tool_name": "render_html", @@ -502,6 +534,18 @@ def run_safetensors_tool_loop( if content_text and not assistant_appended: conversation.append(assistant_msg) assistant_appended = True + if provisional_match and not provisional_resolved: + # A provisional render_html card is already on screen for + # this id; close it so it never dangles when the controller + # turns the call into an internal no-op (duplicate / repeat). + provisional_resolved = True + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": "", + "provenance": decision.provenance, + } completion = tool_controller.record_noop(decision) conversation.append(completion.model_message()) logger.info( @@ -517,7 +561,9 @@ def run_safetensors_tool_loop( else: assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call()) - needs_confirm = bool(confirm_tool_calls) + # Bypass wins over the confirm gate at the loop level too, so a + # direct internal caller passing both flags never prompts. + needs_confirm = bool(confirm_tool_calls) and not bypass_permissions approval_id = new_approval_id() if needs_confirm else "" decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None start_event = decision.tool_start_event() @@ -538,6 +584,8 @@ def run_safetensors_tool_loop( == "deny" ): decision_slot = None + if provisional_match: + provisional_resolved = True yield { "type": "tool_end", "tool_name": decision.tool_name, @@ -575,6 +623,7 @@ def run_safetensors_tool_loop( timeout = eff_timeout, session_id = session_id, rag_scope = rag_scope, + disable_sandbox = bypass_permissions, ) except Exception as exc: logger.exception("Tool %s raised: %s", decision.tool_name, exc) @@ -583,6 +632,8 @@ def run_safetensors_tool_loop( kb_search_count += 1 completion = tool_controller.record_result(decision, result) + if provisional_match: + provisional_resolved = True yield completion.tool_end_event() conversation.append(completion.tool_message()) diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py index 73687165b8..3ceb1a268a 100644 --- a/studio/backend/core/inference/tensor_fallback.py +++ b/studio/backend/core/inference/tensor_fallback.py @@ -13,7 +13,7 @@ import logging from typing import Awaitable, Callable, Optional from core.inference.llama_server_args import ( - resolve_tensor_parallel, + _effective_tensor_parallel, strip_split_mode_only, ) @@ -34,18 +34,20 @@ async def load_with_tensor_fallback( True on success; it *raises* on a hard crash (llama-server aborts on some archs / older builds), which is treated the same as a False return. - Tensor mode can be requested by the toggle or by a ``--split-mode tensor`` - in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether - tensor mode is actually engaged, and it strips ``--split-mode`` from the - extras so the layer retry can't relaunch the same failing tensor load. A - non-tensor load keeps its original contract and propagates exceptions. + Tensor mode can be requested by the toggle, by a ``--split-mode tensor`` in + ``extra_args`` (an allowed shadow flag), or by an inherited + ``LLAMA_ARG_SPLIT_MODE=tensor`` env (load_model engages it the same way), so + the retry is keyed on whether tensor mode is actually engaged, and it forces + ``--split-mode layer`` on the retry so neither leftover extras nor the + inherited tensor env can relaunch the same failing tensor load. A non-tensor + load keeps its original contract and propagates exceptions. ``cancelled()`` distinguishes a real tensor-start failure from a user cancellation: ``attempt_load`` also returns False when the load was cancelled, so without this the helper would restart a load the user just cancelled. """ - tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor) + tensor_requested = _effective_tensor_parallel(extra_args, requested_tensor) try: success = await attempt_load(requested_tensor, extra_args) except Exception as exc: @@ -67,4 +69,8 @@ async def load_with_tensor_fallback( "(this model may not support tensor parallelism)", label, ) - return await attempt_load(False, strip_split_mode_only(extra_args)) + # Force --split-mode layer (CLI wins over env) so neither leftover extras nor + # an inherited LLAMA_ARG_SPLIT_MODE=tensor can re-engage tensor and re-crash + # the retry; load_model and the child both honor the explicit layer override. + layer_extras = strip_split_mode_only(extra_args) or [] + return await attempt_load(False, [*layer_extras, "--split-mode", "layer"]) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 8d5d45269e..ca3d1e4cbc 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -7,25 +7,31 @@ Tolerates missing closing tags in either ``{json}`` or ``v...`` shape. """ -import json -import re +from core import tool_healing as _tool_healing -# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing unclosed -# runs so truncated tails don't leak markup. The [\w-] name set matches OpenAI's -# so hyphenated MCP tool names (mcp__srv__list-issues) parse like built-ins. -_TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r".*?", re.DOTALL), -] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r".*$", re.DOTALL), - re.compile(r".*$", re.DOTALL), -] +_TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS + + +def parse_tool_calls_from_text( + content: str, + *, + id_offset: int = 0, + allow_incomplete: bool = True, +) -> list[dict]: + return _tool_healing.parse_tool_calls_from_text( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + ) + + +def strip_tool_markup(text: str, *, final: bool = False) -> str: + return _tool_healing.strip_tool_call_markup(text, final = final) # Prefixes the streaming buffer watches for to gate in-progress text. -TOOL_XML_SIGNALS = ("", "", "<|tool_call>", "\s*\{") -_TC_FUNC_START_RE = re.compile(r"\s*") -_TC_END_TAG_RE = re.compile(r"") -_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -# [\w-] so hyphenated MCP param names (issue-number) aren't dropped. -_TC_PARAM_START_RE = re.compile(r"\s*") -_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") -_PARAM_CLOSE_TAG = "" -_FUNC_CLOSE_TAG = "" - - -def _inside_open_parameter(content: str, pos: int) -> bool: - """Return True when ``pos`` falls inside an unclosed parameter value.""" - last_param_start = -1 - for match in _TC_PARAM_START_RE.finditer(content, 0, pos): - last_param_start = match.start() - if last_param_start < 0: - return False - last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos) - last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos) - return last_param_start > max(last_param_close, last_func_close) - - -def strip_tool_markup(text: str, *, final: bool = False) -> str: - """Strip tool-call XML from streamed text. - - ``final=False`` only removes closed pairs (used during streaming so - in-progress XML stays buffered). ``final=True`` also removes a - trailing unclosed run and trims the result. - """ - pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in pats: - text = pat.sub("", text) - return text.strip() if final else text - - -def parse_tool_calls_from_text( - content: str, - *, - id_offset: int = 0, - allow_incomplete: bool = True, -) -> list[dict]: - """Parse OpenAI-format ``tool_calls`` from model text. - - Returns a list of ``{"id", "type", "function": {"name", "arguments"}}`` - dicts. ``arguments`` is always a JSON string so callers can hand it - straight back into an OpenAI-style response. - - Handles two shapes: - - - JSON inside ```` tags: - ``{"name":"web_search","arguments":{"query":"..."}}`` - - XML-style function blocks: - ``v`` - - ``allow_incomplete=True`` keeps the historical healing behavior for - missing closing tags. ``allow_incomplete=False`` accepts only - well-formed wrappers so disabled Auto-Heal can still parse valid - local tool protocol without repairing truncated output. - """ - tool_calls: list[dict] = [] - - # Pattern 1: {json}. Balanced-brace scan, skipping braces in - # JSON strings. - for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 # opening { - depth, i = 0, brace_start - in_string = False - while i < len(content): - ch = content[i] - if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 - continue - if ch == '"': - in_string = False - elif ch == '"': - in_string = True - elif ch == "{": - depth += 1 - elif ch == "}": - depth -= 1 - if depth == 0: - break - i += 1 - if depth != 0: - continue - if not allow_incomplete: - tail_after_json = content[i + 1 :].lstrip() - if _TC_END_TAG_RE.match(tail_after_json) is None: - continue - json_str = content[brace_start : i + 1] - try: - obj = json.loads(json_str) - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass - - # Pattern 2: v... -- closing tags optional; - # isn't a body boundary since code values can contain it. - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - if not allow_incomplete: - # Bound the body at the closing tag rather than - # the end of the response, so a complete call followed by - # trailing prose is still accepted (matching the JSON-style - # path, which already tolerates trailing text). - # rfind picks the last , so a literal - # inside a code parameter value stays in the body. - close_idx = body.rfind(_FUNC_CLOSE_TAG) - if close_idx < 0: - continue - body = body[:close_idx] - else: - body = _TC_FUNC_CLOSE_RE.sub("", body) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - # Single param: take everything to body end so an embedded - # in code strings is preserved. - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) - - return tool_calls - - def has_tool_signal(text: str) -> bool: """Return True if ``text`` contains any tool-call XML signal.""" return any(s in text for s in TOOL_XML_SIGNALS) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 43c9610282..a5c193ff39 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -316,6 +316,196 @@ def _build_safe_env(workdir: str) -> dict[str, str]: return env +# Credential env vars dropped even in bypass mode so tool code cannot read the +# operator's keys. Over-strips on purpose (a benign var is harmless to lose). +_BYPASS_ENV_SECRET_NAMES = frozenset( + { + "HF_TOKEN", + "HF_HUB_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "HUGGINGFACE_TOKEN", + "HUGGINGFACEHUB_API_TOKEN", + "WANDB_API_KEY", + "GH_TOKEN", + "GITHUB_TOKEN", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GEMINI_API_KEY", + "GOOGLE_API_KEY", + "GROQ_API_KEY", + "OPENROUTER_API_KEY", + "REPLICATE_API_TOKEN", + "COHERE_API_KEY", + "MISTRAL_API_KEY", + "NGC_API_KEY", + "KAGGLE_KEY", + "MYSQL_PWD", # exact name: markers use PASSWD, not PWD (PWD is the cwd var) + "LD_PRELOAD", + # Auth brokers / capability handles: not secrets by value, but they + # hand the child the operator's live agent (ssh/gpg), kube config, or + # docker daemon. Names are listed because there is no value signal to + # key off. URL config vars (HTTP_PROXY, PIP_INDEX_URL, DATABASE_URL, + # ...) are intentionally NOT name-listed: a benign proxy/index without + # credentials must keep working in bypass mode, while a credentialed + # value is dropped by _is_secret_env_value() regardless of its name. + "SSH_AUTH_SOCK", + "SSH_AGENT_PID", + "GPG_AGENT_INFO", + "GNUPGHOME", + "KUBECONFIG", + "DOCKER_HOST", + } +) +_BYPASS_ENV_SECRET_PREFIXES = ("AWS_", "AZURE_", "GOOGLE_", "GCP_", "GCLOUD_", "DYLD_") +_BYPASS_ENV_SECRET_MARKERS = ( + "TOKEN", + "API_KEY", + "APIKEY", + "SECRET", + "PASSWORD", + "PASSWD", + "CREDENTIAL", + "PRIVATE_KEY", + "AUTH", # e.g. NPM_CONFIG__AUTH (npm _auth), REDISCLI_AUTH + # Azure App Service connection strings: SQLCONNSTR_/CUSTOMCONNSTR_/... and + # WEBSITE_CONTENTAZUREFILECONNECTIONSTRING carry DB/storage credentials. + "CONNSTR", + "CONNECTIONSTRING", +) +# Non-secret hardening flags that match a secret prefix/marker but must be KEPT +# so bypass mode does not silently undo an operator's opt-out. AWS_EC2_METADATA_ +# DISABLED tells the AWS SDK/CLI not to pull instance-role creds from IMDS; +# dropping it would re-open that path for a bypassed tool. +_BYPASS_ENV_KEEP_NAMES = frozenset( + { + "AWS_EC2_METADATA_DISABLED", + "AWS_EC2_METADATA_V1_DISABLED", + } +) +# Matches a URL that embeds userinfo before the host, covering both +# "scheme://user:pass@host" and token-only "scheme://token@host" (and +# percent-encoded variants). The userinfo must precede the first '/', so an '@' +# in a path or query does not false-positive. Used to scrub credential-bearing +# URL values regardless of the variable's name. +_URL_USERINFO_RE = re.compile(r"://[^/\s@]+@") +# Connection-string credential fields (ADO.NET / Azure storage / Service Bus): +# "...;Password=...", "...;AccountKey=...", "...;SharedAccessKey=...". Catches +# credential-bearing values whose names dodge the name classifier. "accesskey" +# also covers Shared/Secret AccessKey via substring; the Name fields (e.g. +# SharedAccessKeyName=) do not match since "=" must follow the keyword. +_SECRET_VALUE_RE = re.compile(r"(?i)(?:password|pwd|accountkey|accesskey)\s*=\s*[^\s;]") + +# Names that hold no secret value but point SDKs at the operator's real +# home/cache/config (cached tokens, cred files), defeating the HOME repoint. +# Startup always sets HF_HOME (-> $HF_HOME/token), so this is the live leak. +# Dropped in bypass mode so tools fall back to the empty repointed HOME. +_BYPASS_ENV_CRED_LOCATION_NAMES = frozenset( + { + # HF cache roots (token lives under $HF_HOME/token) + "HF_HOME", + "HF_HUB_CACHE", + "HUGGINGFACE_HUB_CACHE", + "HF_XET_CACHE", + "TRANSFORMERS_CACHE", + "HF_DATASETS_CACHE", + "HF_ASSETS_CACHE", + # XDG base dirs (resolved before $HOME) + "XDG_CONFIG_HOME", + "XDG_CACHE_HOME", + "XDG_DATA_HOME", + # explicit cred/config file pointers honoured before $HOME + "NETRC", + "PGPASSFILE", + "BOTO_CONFIG", + "PIP_CONFIG_FILE", + "CLOUDSDK_CONFIG", + "KAGGLE_CONFIG_DIR", + "DOCKER_CONFIG", + "WANDB_DIR", + "WANDB_CONFIG_DIR", + "WANDB_CACHE_DIR", + # package-manager / git / cloud config pointers to real cred files + "NPM_CONFIG_USERCONFIG", + "NPM_CONFIG_GLOBALCONFIG", + "YARN_RC_FILENAME", + "GIT_CONFIG_GLOBAL", + "GIT_CONFIG_SYSTEM", + "CARGO_HOME", + "RCLONE_CONFIG", + # auth-helper scripts that hand creds to git/ssh + "GIT_ASKPASS", + "SSH_ASKPASS", + # shell startup hook: bash -c sources $BASH_ENV (can re-export secrets) + "BASH_ENV", + # Windows: HOMEDRIVE+HOMEPATH compose a home that bypasses HOME + "HOMEDRIVE", + "HOMEPATH", + } +) +# Windows profile dirs SDKs read creds under; repointed (not dropped) since +# callers expect them present. +_BYPASS_ENV_WINDOWS_PROFILE_VARS = ("USERPROFILE", "APPDATA", "LOCALAPPDATA") + + +def _is_secret_env_name(name: str) -> bool: + """True if an env var name looks like it carries a credential.""" + upper = name.upper() + if upper in _BYPASS_ENV_KEEP_NAMES: + return False # non-secret hardening flag; keep it + if upper in _BYPASS_ENV_SECRET_NAMES: + return True + if any(upper.startswith(p) for p in _BYPASS_ENV_SECRET_PREFIXES): + return True + return any(marker in upper for marker in _BYPASS_ENV_SECRET_MARKERS) + + +def _is_cred_location_env_name(name: str) -> bool: + """True for vars that point SDKs at the real home/cache/config (cached creds).""" + return name.upper() in _BYPASS_ENV_CRED_LOCATION_NAMES + + +def _is_secret_env_value(value: str) -> bool: + """True if a value embeds credentials regardless of its name. + + Catches URL userinfo (``scheme://user:token@host`` in DATABASE_URL / + PIP_INDEX_URL / HTTP_PROXY) and connection-string credential fields + (``...;Password=...`` / ``...;AccountKey=...``) whose names dodge the name + classifier. + """ + if not value: + return False + return _URL_USERINFO_RE.search(value) is not None or _SECRET_VALUE_RE.search(value) is not None + + +def _build_bypass_env(workdir: str) -> dict[str, str]: + """Env for bypass exec: full host env (unrestricted) minus credential vars, + with HOME/TMPDIR repointed at the workdir so SDKs cannot read cached creds. + + Note: stripping the child env is necessary but not sufficient on its own - + a same-UID child can still read the parent's environment via procfs, so + callers also harden the parent (see _harden_parent_against_proc_env_leak). + """ + env = { + k: v + for k, v in os.environ.items() + if not _is_secret_env_name(k) + and not _is_secret_env_value(v) + and not _is_cred_location_env_name(k) + } + env["HOME"] = workdir + env["TMPDIR"] = workdir + # Windows tempfile / SDKs honour TEMP/TMP, not TMPDIR; repoint all three so + # the bypassed tool writes under the per-session sandbox dir on every OS. + env["TEMP"] = workdir + env["TMP"] = workdir + # Windows SDKs read creds under the profile dirs, not $HOME; repoint set + # ones to the workdir (HOMEDRIVE/HOMEPATH are dropped above). + for var in _BYPASS_ENV_WINDOWS_PROFILE_VARS: + if var in os.environ: + env[var] = workdir + return env + + def _sandbox_preexec(): """Best-effort sandbox setup for sandboxed subprocesses (modules are resolved at import time so the forked child runs no imports).""" @@ -377,6 +567,65 @@ def _sandbox_preexec(): pass +def _bypass_preexec(): + """Minimal pre-exec for bypass exec: os.setsid() only. + + Required, not a restriction: _kill_process_tree does killpg(getpgid(child)), + so without a new session a timeout/cancel would kill the Studio server too. + """ + try: + os.setsid() + except OSError: + pass + + +# Hardening the Studio parent is done once (PR_SET_DUMPABLE is process-global +# and sticky); guarded so repeated bypass calls do not re-issue the prctl. +_parent_proc_hardened = False + + +def _harden_parent_against_proc_env_leak() -> bool: + """Make the Studio process's /proc//environ unreadable to its children. + + Stripping the child env is not enough on Linux: a bypassed same-UID child + runs unsandboxed and can read /proc//environ to recover the + tool-executing process's *unfiltered* secrets (HF_TOKEN, cloud keys, ...). + Clearing the dumpable flag (PR_SET_DUMPABLE=0) reparents this process's + /proc entries to root, so a same-UID child can no longer read its environ. + + Returns True when the process is hardened or hardening is unnecessary (no + /proc leak off Linux), and False when it is needed but could not be applied + (e.g. prctl denied by a seccomp policy). Callers must fail closed - refuse + the unsandboxed exec - when this returns False, rather than running with the + parent environ still readable. + + Scope: this closes the direct parent read (the demonstrated leak). It is a + mitigation, not a full boundary - a bypassed tool is unsandboxed by design, + so it can still walk /proc to a same-UID *ancestor* (e.g. the launching + shell) or read on-disk credentials by absolute path. Complete isolation + needs a separate uid / PID+mount namespace, which is out of scope here; the + UI already warns the mode is dangerous. Applied lazily on first bypass exec + so non-bypass operation is unchanged. + """ + global _parent_proc_hardened + if _parent_proc_hardened: + return True + if sys.platform != "linux": + return True # no /proc//environ same-UID leak to close + if _libc is None: + return False # on Linux but cannot issue prctl -> cannot harden + try: + # prctl(PR_SET_DUMPABLE=4, SUID_DUMP_DISABLE=0). ctypes returns the + # syscall result (-1 on failure) and does NOT raise, so check it. + ret = _libc.prctl(4, 0, 0, 0, 0) + except (OSError, AttributeError): + return False + if ret != 0: + return False + _parent_proc_hardened = True + return True + + def _get_shell_cmd(command: str) -> list[str]: """Return the platform-appropriate shell invocation for a command string.""" if sys.platform == "win32": @@ -524,10 +773,10 @@ RENDER_HTML_TOOL = { "function": { "name": "render_html", "description": ( - "Render a self-contained HTML/CSS/JavaScript artifact for the user. " + "Render a self-contained HTML/CSS/JavaScript canvas for the user. " "Call this at most once per assistant response unless the user " "explicitly asks for changes in that response. Future user requests " - "for new artifacts may call render_html once. Put the entire document " + "for new canvases may call render_html once. Put the entire document " "in code, including any CSS in