diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index d430d2c172..f4189a159e 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -376,7 +376,7 @@ case "$MODE" in hermes) patch_hermes_tools none invoke_via_connect "$OUT" -z "$PROMPT" ;; openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ + CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; *) invoke_via_connect "$OUT" "$PROMPT" ;; esac @@ -449,7 +449,7 @@ case "$MODE" in fi ;; opencode) invoke_via_connect "$out" run "$prompt" ;; hermes) invoke_via_connect "$out" -z "$prompt" ;; - openclaw) invoke_via_connect "$out" agent --local --agent ci \ + openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;; *) invoke_via_connect "$out" "$prompt" ;; esac @@ -527,6 +527,154 @@ case "$MODE" in echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; + # ── resume: does a launched agent's session survive exit and resume? ──── + # Unlike the other modes, this drives the real LAUNCH path (`unsloth start + # ...`, the interactive default), not the --no-launch recipe. That + # path relocates each agent's home to a throwaway temp dir wiped on exit, so + # a session cannot be resumed -- unless --persist routes it to the stable + # Unsloth agents dir instead. We run one headless turn per pass and check + # whether the turn left a session in a persistent store (deterministic, no + # reliance on the model recalling anything), for a baseline pass and a + # --persist pass, and assert the expected split for this agent. + resume) + CODEWORD="PLATYPUS7" + T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK." + T2="What codeword did I ask you to remember? Reply with just that word." + WORK="$WORKDIR_BASE/${AGENT}-resume" + + # STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to. + # Read it from a --no-launch probe (which also writes the agent's config + # there). codex/pi relocate their whole home/HOME here; opencode/claude keep + # their session data in a fixed user dir, so STABLE_HOME stays empty for them. + parse_connect + case "$AGENT" in + codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;; + pi) STABLE_HOME="$(raw_env HOME)" ;; + *) STABLE_HOME="" ;; + esac + + # The persistent stores a session would land in if it were NOT wiped. We + # count files here before/after each turn; a positive delta means the + # session persisted (is resumable), zero means it went to a wiped temp dir. + resume_tracked_dirs() { + case "$AGENT" in + codex) printf '%s\n' "$HOME/.codex" ;; + opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;; + claude) printf '%s\n' "$HOME/.claude" ;; + pi) printf '%s\n' "$HOME/.pi" ;; + *) : ;; + esac + [ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME" + } + count_session_files() { + local total=0 d n + while IFS= read -r d; do + [ -n "$d" ] && [ -d "$d" ] || continue + n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n)) + done < <(resume_tracked_dirs) + echo "$total" + } + + # The headless first-turn subcommand per agent (mirrors file-edit's map), + # forwarded verbatim through the launch path as passthrough args. + set_t1_cmd() { + case "$AGENT" in + claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;; + codex) T1_CMD=(exec "$T1") ;; + opencode) T1_CMD=(run "$T1") ;; + pi) T1_CMD=(-p "$T1") ;; + *) guide_fail "resume mode does not cover agent '$AGENT'" ;; + esac + } + + # Run one headless turn through the launch path. $1=outfile, $2="" or + # "--persist", rest = the agent subcommand. --yolo auto-approves so no tool + # prompt can hang; --api-key attaches to the already-served CI model. + launch_turn() { + local out="$1" rflag="$2"; shift 2 + local flag=(); [ -n "$rflag" ] && flag=("$rflag") + run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \ + --api-key "$UNSLOTH_API_KEY" "$@" + local rc=$? + redact "$out" + return "$rc" + } + + # One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED + # from the session-store delta. Runs in the main shell (not a command + # substitution) so a hang's guide_fail actually fails the job and the + # progress lines reach the CI log. $1 = "" (baseline) or "--persist". + RESULT="" + run_pass() { + local rflag="$1" label="baseline" + [ -n "$rflag" ] && label="resume" + rm -rf "$WORK"; mkdir -p "$WORK" + set_t1_cmd + local out="$LOGS_DIR/${AGENT}-resume-${label}.txt" + local before after rc + before="$(count_session_files)" + pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK" + launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$? + popd >/dev/null || true + after="$(count_session_files)" + echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})" + # The turn must succeed for the delta to mean anything: an agent that writes a + # session file then errors would otherwise be misread as PERSISTED. Mirror the + # file-edit mode and fail the pass on a non-zero launch (the flagship codex recall + # below stays WARN-only, driven by its own launch_turn calls). + [ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \ + guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; } + if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi + } + + run_pass ""; BASELINE="$RESULT" + # Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix. + # opencode/claude persist either way, so the baseline already proves it and a + # second full CPU turn only risks a timeout; skip it for them. + case "$AGENT" in + codex|pi) run_pass "--persist"; RESUME="$RESULT" ;; + *) RESUME="n/a (persists either way)" ;; + esac + + # Expected: codex/pi relocate their whole home to the temp dir, so a plain + # launch is WIPED and only --persist PERSISTS. opencode/claude keep their + # session data in a fixed user dir, so the baseline already PERSISTS. + case "$AGENT" in + codex|pi) EXPECT_BASELINE="WIPED" ;; + opencode|claude) EXPECT_BASELINE="PERSISTED" ;; + esac + + echo "──────────────────────────────────────────────" + echo "[$AGENT] RESUME EXPERIMENT" + echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})" + echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}" + echo "──────────────────────────────────────────────" + + [ "$BASELINE" = "$EXPECT_BASELINE" ] \ + || guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}" + case "$AGENT" in + codex|pi) + [ "$RESUME" = "PERSISTED" ] \ + || guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;; + esac + + # Flagship behavioral proof (codex only, WARN-only): after a --persist plant, + # resume the session and check the model actually recalls the codeword. A + # miss is not a failure (the CI model is small); the mechanism gate above is + # the real assertion. + if [ "$AGENT" = "codex" ]; then + rm -rf "$WORK"; mkdir -p "$WORK" + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true + if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then + echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}" + else + echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed" + fi + fi + echo "[$AGENT] resume OK" + ;; + *) echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2 exit 2 diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index ae4b386589..1bb4c2bb58 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,7 +364,9 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ + tests/test_gemma_2b_mapper_key.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 47f75dc1ba..25796bd5cf 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -471,6 +471,176 @@ jobs: redacted-configs/ retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ + # Job: resume + # Does a conversation started with `unsloth start ` survive exit + # and resume? This drives the REAL launch path (not the --no-launch + # recipe the other jobs use). A plain launch relocates the agent home to + # a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the + # session to the stable Unsloth agents dir so it persists. opencode/claude + # keep their session data in a fixed user dir, so they persist either way. + # Dispatch-only: it is an end-to-end experiment, not a PR gate. + # ═════════════════════════════════════════════════════════════════════ + resume: + name: resume (${{ matrix.agent }}) + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + # codex/pi relocate their whole home (resume broken without --persist); + # opencode/claude keep session data in a fixed dir (resume already works). + # One agent from each class proves the split end to end; openclaw/hermes + # share codex's relocation mechanism and are covered by the unit tests. + agent: [codex, opencode, claude, pi] + env: + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18904' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: Resume experiment (launch path) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: resume-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ # Job 3: prompt-cache # (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0 diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 0ef2ad1e9d..1275d12216 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -2,8 +2,8 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # Multi-language supply-chain audit. Triggers: -# - PRs touching any dependency manifest (Python / npm / Cargo) or -# this workflow file, +# - PRs touching any dependency manifest (Python / npm / Cargo), a +# scanner or its allowlist baseline, or this workflow file, # - push to main / pip, # - nightly @ 04:13 UTC so newly-published advisories surface even # when no PR opens, @@ -57,7 +57,9 @@ on: - 'studio/src-tauri/Cargo.lock' - 'pyproject.toml' - 'scripts/scan_packages.py' + - 'scripts/scan_packages_baseline.json' - 'scripts/scan_npm_packages.py' + - 'scripts/scan_npm_packages_baseline.json' - '.github/workflows/security-audit.yml' push: branches: [main, pip] diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index aebf90380a..f540c11da4 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -444,6 +444,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -464,8 +466,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -938,6 +956,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -956,8 +976,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index d562294d42..03c0a8580d 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -430,6 +430,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -450,8 +452,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -825,6 +843,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -848,8 +868,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 0bc216d65a..0453c9212a 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -634,6 +634,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -656,8 +658,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): body = {**body, "stream": True} @@ -1063,6 +1081,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -1082,8 +1102,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── status, data = post("/v1/chat/completions", { @@ -1334,42 +1370,75 @@ jobs: try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { } } - - name: Hide Visual Studio + CMake (simulate a host with no build tools) + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # A Program Files dir can hold a transient handle (Defender / MSBuild node) - # so Rename-Item intermittently fails with "Access is denied"; retry to ride it out. - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { + [void] $blocked.Add( + [Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\')) + } + } } } - # 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-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff') - Write-Host "Hid VS: $d" - } + # Normalized comparison so registry spellings (trailing slash, + # unexpanded %VAR%) still match. + function Test-Blocked([string]$p) { + $n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\') + return $blocked.Contains($n) } - # 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-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off') - $hidden += $c.Source - Write-Host "Hid cmake: $($c.Source)" - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not (Test-Blocked $_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + # install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment + # rebuild the session Path from these scopes mid-install, so filter + # them too. Originals are saved for the cleanup step. + foreach ($scope in @('Machine', 'User')) { + $orig = [Environment]::GetEnvironmentVariable('Path', $scope) + if (-not $orig) { continue } + Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline + $kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';' + [Environment]::SetEnvironmentVariable('Path', $kept, $scope) + Write-Host "Filtered $scope Path scope." + } + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH<&1 | Tee-Object -FilePath logs/install.log @@ -1480,19 +1553,19 @@ jobs: [ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; } echo "Inference OK without Visual Studio: $CONTENT" - - name: Restore Visual Studio + CMake + - name: Clean no-build-tools simulation 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) } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + foreach ($scope in @('Machine', 'User')) { + $saved = Join-Path $root "orig-path-$scope.txt" + if (Test-Path -LiteralPath $saved) { + [Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope) + Write-Host "Restored $scope Path scope." } } + Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue - name: Stop Studio if: always() @@ -1540,21 +1613,34 @@ jobs: with: python-version: '3.12' - - name: Hide Visual Studio + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # Retry the rename: a Program Files dir can hold a transient handle that - # makes Rename-Item intermittently fail with "Access is denied". - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { [void] $blocked.Add($dir) } + } } } - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" } - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not $blocked.Contains($_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH< /tmp/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." + if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 } + python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json + if ($LASTEXITCODE -ne 0) { + Write-Host "::error::resolver exited non-zero" + if (Test-Path resolve.json) { Get-Content resolve.json } + exit 1 + } + Get-Content resolve.json + Write-Host "Prebuilt resolver ran with no Visual Studio present." - - name: Restore Visual Studio + - name: Clean no-build-tools simulation 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" } - } + Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue # ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ── pester: diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 888b3d70a3..5b92f1a3e0 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -6,9 +6,9 @@ # windows-latest runner: # # 1. install.ps1 --local --no-torch installs Studio AND auto-fetches -# the prebuilt llama.cpp Windows binary (llama-bNNNN-bin-win-cpu- -# x64 from ggml-org/llama.cpp). Hitting the source-build fallback -# is treated as an Unsloth bug -- Studio must always pick the +# the prebuilt llama.cpp Windows binary (app--windows-x64-cpu +# from unslothai/llama.cpp). Hitting the source-build fallback is +# treated as an Unsloth bug -- Studio must always pick the # prebuilt on Windows. # 2. unsloth studio update --local is idempotent. Two consecutive # runs both report "prebuilt up to date and validated", no diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index e492d21e99..6becccc90a 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -285,6 +285,92 @@ jobs: tests/vllm_compat/test_extended_module_imports.py \ -v --tb=short + # Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike + # the static symbol/source greps above, this drives unsloth's actual + # source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only + # runner under the tests/conftest.py spoof harness -- no GPU, no training. + # Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple + # per-token-logps return, restructured PEFT ref-adapter block) by asserting + # the generated Unsloth trainer still satisfies the transform contracts. + grpo-fake-run: + name: GRPO fake-run (latest + main TRL, CPU spoof) + runs-on: ubuntu-latest + timeout-minutes: 18 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - name: Clone unsloth-zoo @ main + run: | + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install CPU torch + ecosystem + TRL latest + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' + # Ecosystem floors unsloth needs; TRL itself is installed last so it + # can pull the transformers/peft it requires. + pip install \ + 'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \ + 'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \ + 'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow + pip install --upgrade trl + pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo" + pip install --no-deps -e ./unsloth + - name: Fake-run vs TRL latest + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + # Disable dynamo/inductor at the process level, before conftest.py's early + # `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner + # (defense in depth; the CPU fake-train also flips this at runtime). + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge + # TRL break does not red every PR. github.event_name is valid in a step if. + - name: Fake-run vs TRL main (scheduled / dispatch only) + if: ${{ github.event_name != 'pull_request' }} + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + pip install --upgrade "git+https://github.com/huggingface/trl" + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + # Daily-only: same suites but with --strict on importable upstream # tags. Schedule-only so PR jobs stay fast; cron tolerates a flake. daily-fresh-fetch: diff --git a/README.md b/README.md index e3fd4e6980..849ee2e87b 100644 --- a/README.md +++ b/README.md @@ -212,7 +212,7 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i ```bash 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. +- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. This also starts a public Cloudflare quick tunnel by default, which publishes an internet-reachable `https://*.trycloudflare.com` URL even behind a firewall. Both the raw port and the tunnel expose Studio beyond this machine, so only use this on a network you trust; pass `--no-cloudflare` to drop the public link while keeping the network bind. ```bash unsloth studio -H 0.0.0.0 -p 8888 ``` diff --git a/install.ps1 b/install.ps1 index 9114f80af9..0797cd3868 100644 --- a/install.ps1 +++ b/install.ps1 @@ -469,6 +469,17 @@ function Install-UnslothStudio { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when the command pins an index, clear every uv index env var so + # it wins, then restore in finally. Other installs keep the user's mirror. + $savedUvIndex = $null + if ($Command.ToString() -match '--default-index') { + $savedUvIndex = @{} + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue + } + } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -488,6 +499,7 @@ function Install-UnslothStudio { return [int]$LASTEXITCODE } finally { $ErrorActionPreference = $prevEap + if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } } } @@ -2155,7 +2167,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $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.7.1" "unsloth-zoo>=2026.7.1" } + $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.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2169,7 +2181,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2200,7 +2212,7 @@ exit 0 # 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 } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base so the install # still completes; Studio setup retries ROCm afterwards. @@ -2209,7 +2221,7 @@ exit 0 # 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 } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $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) @@ -2223,7 +2235,7 @@ exit 0 } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { 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 --default-index $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) @@ -2235,7 +2247,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2247,7 +2259,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2275,7 +2287,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) @@ -2306,7 +2318,7 @@ exit 0 # 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 + # is a PEP 503 index uv resolves via --default-index, 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 @@ -2322,7 +2334,7 @@ exit 0 $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 } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $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) @@ -2331,7 +2343,7 @@ exit 0 } 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 } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $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) diff --git a/install.sh b/install.sh index 796d80e401..3f4ea92387 100755 --- a/install.sh +++ b/install.sh @@ -159,6 +159,12 @@ run_maybe_quiet() { run_install_cmd() { _label="$1" shift + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when we pass --default-index, neutralize every uv index env var so + # the pinned index wins. Other installs keep the user's mirror. + case " $* " in + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + esac if _is_verbose; then "$@" && return 0 _rc=$? @@ -2190,9 +2196,9 @@ _expected_torch_flavor_tag() { esac } -# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# Whether index ($1) supports a plain --default-index 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 +# resolves (torch + every transitive dep) via --default-index -- 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() { @@ -2706,7 +2712,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. 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.7.1" "unsloth-zoo>=2026.7.1" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # 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. @@ -2721,7 +2727,7 @@ if [ "$_MIGRATED" = true ]; then # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2744,7 +2750,7 @@ if [ "$_MIGRATED" = true ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2870,7 +2876,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; 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_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." # Pass explicit wheel URLs so the matched trio is @@ -2893,18 +2899,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$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_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm @@ -2925,7 +2931,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --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.7.1" "unsloth-zoo>=2026.7.1" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2943,7 +2949,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" + --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" 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..." @@ -2964,7 +2970,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2975,7 +2981,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2999,14 +3005,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; 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 + # Repair when flavor is wrong AND the index is plain --default-index 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" \ + --default-index "$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="" @@ -3017,7 +3023,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; 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" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi diff --git a/pyproject.toml b/pyproject.toml index 506f1e4c75..c6d5d8c6d3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "wheel>=0.42.0", "packaging", "numpy", @@ -94,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "torchvision", "unsloth[triton]", ] @@ -629,7 +629,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "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_packages_baseline.json b/scripts/scan_packages_baseline.json index 1d34cfb66d..27fa801a2a 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -39,7 +39,7 @@ "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", + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3719: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", "evidence_hash": "2d691bc373ab872aad23c744104596ba6d0d9f3b35aa101c7edbff4429b174c1" }, { @@ -55,23 +55,23 @@ "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", - "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" + "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", + "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" }, { "package": "datasets", "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", - "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" + "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", + "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" }, { "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__)", + "evidence": "L1052: return importlib.import_module(\".\" + module_name, self.__name__)", "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { @@ -79,7 +79,7 @@ "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, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", + "evidence": "Env: L236: value = os.environ[key]\nNetwork: L691: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L712: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L731: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", "evidence_hash": "671190a6106c6ee9674e5e5942dc0940e1d2f8c78d5faf674413c2345b783fd9" }, { @@ -90,12 +90,20 @@ "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()", "evidence_hash": "894862e547cf91b90cd6e4b495db3fb05b7490ef0d63de7e795a7e3d9447d850" }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", + "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + }, { "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 client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence": "Archive: L1353: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", "evidence_hash": "73a7a72013e9f800627ea07e6dbc3beeb8c905a6a5480c8fd896f0063173d25c" }, { @@ -103,8 +111,8 @@ "file": "fastmcp/cli/apps_dev.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L624: history.replaceState(null, \"\", url); sha256:fd8dbfa8af4dea2ce43f4d441f3f81239de341b76a2eb0a33c446f6757ce5f43\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: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", - "evidence_hash": "6ada4a9111213bdee5ea24c70a72ec4acdc8ffe0de4a01fd9835bc261ccab8f8" + "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" }, { "package": "fonttools", @@ -132,19 +140,35 @@ }, { "package": "huggingface-hub", - "file": "huggingface_hub/hf_api.py", + "file": "huggingface_hub/_sandbox.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", - "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" + "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", + "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", + "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", - "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" + "evidence": "L4677: while True: sha256:04afb38843e4125d1476f3f04bdad0edf1f63f8d75ad49a713b13e4bc68612fb", + "evidence_hash": "18877a2502c862b46a5d7e33fa7c39ab4ef32da7e1b07f596fd455f4376770c6" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", + "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" }, { "package": "huggingface-hub", @@ -159,8 +183,8 @@ "file": "huggingface_hub/utils/_http.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:0ab4fed32d3af10f361963371f681923481377508a405b5d8770cef75f859168", - "evidence_hash": "1484f6b92f41c427ba8cbc7c4695a94975fea683dfa83aa510b4b0e982be4721" + "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", + "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" }, { "package": "huggingface-hub", @@ -218,6 +242,22 @@ "evidence": "L5: import socket sha256:915068303029fa5806199f256fb74504c65f253f9aee8ea23d8e384bb772b1c7", "evidence_hash": "30be130f165f418dfd37b144c5ae333de184b95f828ab8bd4010a67b84a5f814" }, + { + "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) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", + "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", + "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" + }, { "package": "numba", "file": "numba/pycc/decorators.py", @@ -231,7 +271,7 @@ "file": "numba/tests/support.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)", + "evidence": "L1016: os.dup2(w, fd) | L1021: os.dup2(save, fd)", "evidence_hash": "fea7aa03d48bf0f4386302fa444984c4f5dfc772cfec3f1df199fd33a52eec10" }, { @@ -298,13 +338,21 @@ "evidence": "Env: L105: token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L150: environment_token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L415: http_client: httpx.Client | None = None, | L531: http_client: httpx.Client | None = None, | L649: http_client: httpx.AsyncClient | None = None, | L767: http_client: httpx.AsyncClient | None = None,", "evidence_hash": "92dbec8ccd79c1e0bc41e93cdd0bdbb091220616c6a1352873196e9dda6bd85c" }, + { + "package": "openai", + "file": "openai/resources/beta/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", + "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" + }, { "package": "openai", "file": "openai/resources/beta/threads/runs/runs.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1074: while True: sha256:ef6d59a4a10b73a5af491f10af2885b7a309fda9468eb0f9572d19558d3ceb9f", - "evidence_hash": "43c03b55fedcbc980e5e6649c3c4493729128d280cc868349ab9590908ea5f99" + "evidence": "L1053: while True: sha256:973bb1aeca2e17e022872dc343a1bf5d8fe33bfa59fe01e2f8fe875522db5bce", + "evidence_hash": "24626e4aa53047a515ead563b42c07c43f73a2c5b82978fa59f58ffc2859e19b" }, { "package": "openai", @@ -319,7 +367,7 @@ "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3803: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", + "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" }, { @@ -495,16 +543,16 @@ "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", - "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" + "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", + "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" }, { "package": "scikit-learn", "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", - "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" + "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", + "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" }, { "package": "scikit-learn", @@ -642,6 +690,14 @@ "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run(\nL2693: cmd.split(), capture_output=True, text=True, check=True\nL2694: ) | L2995: cmd_output = subprocess.run(\nL2996: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL2997: ) | L3707: out = subprocess.check_output(\nL3708: [\"ldd\", os.path.join(search, file)]\nL3709: ) | L3791: jobs.append(functools.partial(subprocess.check_call, cmd)) | L3876: subprocess.check_call(\nL3877: shlex.split(halide_cmd_gen.get_command_line())\nL3878: ) | L4336: subprocess.check_output(\nL4337: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4338: ) | L4591: output = subprocess.check_output(\nL4592: cmd_parts,\nL4593: stderr=subprocess.STDOUT,\nL4594: text=True,\nL4595: env=os.environ,\nL4596: )", "evidence_hash": "c09774087b702a6c5d6e2e85d9239c7c241ec938fbe9c0153e8f0b5c0710389b" }, + { + "package": "torch", + "file": "torch/_inductor/codecache.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1727: content = base64.b64decode(data)\nSubprocess: L3270: subprocess.run(\nL3271: cmd, capture_output=True, text=True, check=True\nL3272: ) | L3583: cmd_output = subprocess.run(\nL3584: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL3585: ) | L4338: out = subprocess.check_output(\nL4339: [\"ldd\", os.path.join(search, file)]\nL4340: ) | L4422: jobs.append(functools.partial(subprocess.check_call, cmd)) | L4507: subprocess.check_call(\nL4508: shlex.split(halide_cmd_gen.get_command_line())\nL4509: ) | L4992: subprocess.check_output(\nL4993: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4994: ) | L5247: output = subprocess.check_output(\nL5248: cmd_parts,\nL5249: stderr=subprocess.STDOUT,\nL5250: text=True,\nL5251: env=os.environ,\nL5252: )", + "evidence_hash": "87f77b5f51cb84fe9950fdeeb90fe8710e1b863100e90b5e2cfb228a725bee06" + }, { "package": "torch", "file": "torch/ao/__init__.py", @@ -695,7 +751,7 @@ "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:", + "evidence": "Env: L4900: env = os.environ.copy()\nNetwork: L4962: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4980: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", "evidence_hash": "704a851b9d68c9b885b9e15538bd7e96f03875503b618fe6f126c4438edd7386" }, { @@ -706,6 +762,14 @@ "evidence": "L32: import socket sha256:89faaaa8bc908e02dad73fd59b2b481fa91189c84b39b556c2766e71d2783bf3", "evidence_hash": "3d23d77ace91812a07cb9508cf352185d154176e8e8c8b9b28fa92cdbcfe0d53" }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L32: import socket sha256:ba439cbf568b194872f1d974c02b0487e51f677b67e379400522d0992600bd2d", + "evidence_hash": "88e98b227573997f86eedea8e885a407b0dd549d46d4a3f0b840ec5aafe66865" + }, { "package": "torchvision", "file": "torchvision/datasets/utils.py", @@ -743,8 +807,8 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1663: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", - "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" + "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", + "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" }, { "package": "transformers", @@ -759,15 +823,15 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", - "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" + "evidence": "L1699: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", + "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L284: value = os.environ[key] | L300: value = os.environ[key] | L2129: env = os.environ.copy() | L2251: for k in list(os.environ.keys()):\nNetwork: L2561: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", + "evidence": "Env: L288: value = os.environ[key] | L304: value = os.environ[key] | L2165: env = os.environ.copy() | L2287: for k in list(os.environ.keys()):\nNetwork: L2597: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", "evidence_hash": "73ff16aee09cf163fb3a7a04dfa2cf610595bde2f19460a579397695f728e3f4" }, { @@ -799,16 +863,16 @@ "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", - "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" + "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", + "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" }, { "package": "trl", "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", - "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" + "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", + "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" }, { "package": "trl", @@ -866,6 +930,14 @@ "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 resp:", "evidence_hash": "278ff15b0b702d37d7f0b30a1e55a31bf2b11883685718a47478fbb5ce7f5212" }, + { + "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)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", + "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" + }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", @@ -919,16 +991,16 @@ "file": "tests/test_mlx_save_export_regressions.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L164: temporary_location=\"/tmp/ignored\", sha256:78837e80d48e872ef191aaacfe5e1c621a98a20df486a70a41d1a932d074a5b3", - "evidence_hash": "dd11376e664d0d7e7f4cc4baf57eacd4b7ae7b03222dce3912ce68b63dbfca1e" + "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:9f8502377b19666288b28399633dfc6740a64d0cb70ad1615e38b1269f94bf37", + "evidence_hash": "b7262d6e58f2ebad961dd3e64ca6c32bba356b5044d7a642d7dbd36a58cb6c81" }, { "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\", sha256:32532cadc357beee1009f4e86481bdbe60a0b7bf47f6bb022b05ec1b8e15aed0", - "evidence_hash": "49f5b67379de17178f21a9bc93b79d6b94a70ecbdd16de86574934aac30a071d" + "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:06789b55e8f31426c233f37ff7d3729cc9e1f61c0829abd2c00c39216c63c7ad", + "evidence_hash": "ad4913d9099eb9b70e09d6860b242eb5f48c67e46d9bf4ae35c1c38a267d753b" }, { "package": "unsloth-zoo", @@ -951,7 +1023,7 @@ "file": "unsloth_zoo/llama_cpp.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "b9f3b1652349fa8ef9ac2d1715978aca1e1632165851a00a2698dd47189e410c" }, { @@ -959,7 +1031,7 @@ "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()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "9cd0b1bb59c7eb1d814d7636dfd167c34f265eb7c4521a9d88b2bdcfd535b926" }, { @@ -1002,6 +1074,14 @@ "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, + { + "package": "cffi", + "file": "cffi/_cffi_gen_src.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", + "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" + }, { "package": "cffi", "file": "cffi/setuptools_ext.py", @@ -1127,7 +1207,7 @@ "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)", + "evidence": "Obfusc: L874: __import__(modname)\nExec: L808: eval(co, globs, ns)", "evidence_hash": "649a7d750f903478243b0bcb9e8020521b505fc7fedc5b696ec01f4efc096109" }, { @@ -1159,7 +1239,7 @@ "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)", + "evidence": "Obfusc: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)", "evidence_hash": "9e81164131d16056fb56ad3cd11b8d129d1ff4f5855031e8b501e0335d5c14ed" }, { @@ -1175,16 +1255,16 @@ "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)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", - "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" + "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", + "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", - "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" + "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", + "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" }, { "package": "numpy", @@ -1199,7 +1279,7 @@ "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:", + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3776: def eval(image: Image, *args: Callable[[int], float]) -> Image:", "evidence_hash": "c2c1e7ae44e15862caf8de549d09db7b35e93282450f07ef61aaf5450a408c13" }, { @@ -1255,7 +1335,7 @@ "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):", + "evidence": "Obfusc: L1287: __import__(module_name)\nExec: L1114: if lib_type not in eval(expected):", "evidence_hash": "368651e9818ed2d1bb009027d3bcfbf94ae30639c0882a6c2bddde97b8c4f1e5" }, { @@ -1271,7 +1351,7 @@ "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')\",", + "evidence": "Obfusc: L387: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", "evidence_hash": "067d41014f72a61d8b4adf25f3659d1f66a0e909f732223f48837aa7684df4e6" }, { @@ -1279,7 +1359,7 @@ "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)", + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L447: exec(finder, loc, loc)", "evidence_hash": "a78d7f5af7eb4ba92656cda258c195b92f6337c585c97d0823e47a9d4a2eb15d" }, { @@ -1322,12 +1402,20 @@ "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace) | L909: exec(ln, {}, namespace) | L920: exec(c, namespace, funclocals)", "evidence_hash": "ab4f5819576a70038301668b8f3e4a781c4b757b146117d5d93eab1896a5a6cd" }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", + "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" + }, { "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)\nExec: L683: result = eval(arg, frame_globals, eval_locals) | L708: result = eval(cmd, frame_globals, eval_locals) | L716: exec(cmd, frame_globals, eval_locals)", + "evidence": "Anti: L1052: self._old_trace = sys.gettrace() | L1053: sys.settrace(self._settrace_callback) | L1113: sys.settrace(self._old_trace)\nExec: L684: result = eval(arg, frame_globals, eval_locals) | L709: result = eval(cmd, frame_globals, eval_locals) | L717: exec(cmd, frame_globals, eval_locals)", "evidence_hash": "dc2afd1769d357c15b69802bd2799fafa059c0b1dcdd4937528fb5b601962f1b" }, { @@ -1343,7 +1431,7 @@ "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)", + "evidence": "Obfusc: L44: code = compile(dest_ast, \"\", \"exec\")\nExec: L47: exec(code, globals_dict)", "evidence_hash": "76374f96feed416eec390458843621f33524cfb8d93ef0f3eb4cb1b47d0ad748" }, { @@ -1359,7 +1447,7 @@ "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)", + "evidence": "Obfusc: L599: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", "evidence_hash": "c7c0650f0c74a086d224112f77ee76634b8f47afc047ce27fee8c7fc45560512" }, { @@ -1391,7 +1479,7 @@ "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):", + "evidence": "Obfusc: L1158: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L1136: def eval(self):", "evidence_hash": "c409327ef6420cc0c7224506fcb82b11bbc9838a6f2f97c9c2cfc00a40c4cdbf" }, { @@ -1407,7 +1495,7 @@ "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\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3385: exec(f\"import {model_location}\", globals()) | L3388: modeling_file = eval(model_location) | L3401: exec(\nL3402: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3403: ) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3407: globals(),\nL3408: locals(),\nL3409: ) | L3560: source = eval(f\"modeling_file.{module}\") | L3574: source = eval(f\"modeling_file.{module}\") | L3675: source = eval(f\"modeling_file.{module}\") | L3713: source = eval(f\"{model_location}.{module}\") | L3784: source = eval(f\"{model_location}.{module}\") | L3832: source = eval(f\"{model_location}.{module}\") | L4054: source = eval(f\"{model_location}.{module}\") | L4065: exec(\nL4066: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4067: globals(),\nL4068: ) | L4131: source = eval(f\"{model_location}.{module}\") | L4172: module_cls = eval(f\"{model_location}.{module}\") | L4209: module_cls = eval(f\"{model_location}.{module}\") | L4276: exec(\nL4277: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4278: globals(),\nL4279: ) | L4341: exec(inner_training_loop, globals()) | L4349: function = eval(f\"{model_location}.{module}\") | L4427: function = eval(f\"{model_location}.{module}\") | L4562: source = eval(f\"{model_location}.torch\") | L4569: function = eval(f\"source.nn.{module}\") | L4628: exec(\nL4629: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4630: globals(),\nL4631: locals(),\nL4632: ) | L4634: exec(\nL4635: f\"{model_location}.nn.{module}.forward = forward\",\nL4636: globals(),\nL4637: locals(),\nL4638: ) | L4642: exec(\nL4643: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4644: globals(),\nL4645: locals(),\nL4646: ) | L4648: exec(\nL4649: f\"combined_module.nn.{module}.forward = forward\",\nL4650: globals(),\nL4651: locals(),\nL4652: ) | L4669: exec(\nL4670: f\"{model_location}.{module} = combined_module.{module}\",\nL4671: globals(),\nL4672: locals(),\nL4673: ) | L4683: check_dicts = dir(eval(f\"{model_location}\")) | L4685: item = eval(f\"{model_location}.{check}\") | L4695: exec(\nL4696: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4697: globals(),\nL4698: locals(),\nL4699: )", + "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3389: exec(f\"import {model_location}\", globals()) | L3392: modeling_file = eval(model_location) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3407: ) | L3409: exec(\nL3410: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3411: globals(),\nL3412: locals(),\nL3413: ) | L3564: source = eval(f\"modeling_file.{module}\") | L3578: source = eval(f\"modeling_file.{module}\") | L3679: source = eval(f\"modeling_file.{module}\") | L3717: source = eval(f\"{model_location}.{module}\") | L3788: source = eval(f\"{model_location}.{module}\") | L3836: source = eval(f\"{model_location}.{module}\") | L4058: source = eval(f\"{model_location}.{module}\") | L4069: exec(\nL4070: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4071: globals(),\nL4072: ) | L4135: source = eval(f\"{model_location}.{module}\") | L4176: module_cls = eval(f\"{model_location}.{module}\") | L4213: module_cls = eval(f\"{model_location}.{module}\") | L4280: exec(\nL4281: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4282: globals(),\nL4283: ) | L4345: exec(inner_training_loop, globals()) | L4353: function = eval(f\"{model_location}.{module}\") | L4431: function = eval(f\"{model_location}.{module}\") | L4566: source = eval(f\"{model_location}.torch\") | L4573: function = eval(f\"source.nn.{module}\") | L4632: exec(\nL4633: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4634: globals(),\nL4635: locals(),\nL4636: ) | L4638: exec(\nL4639: f\"{model_location}.nn.{module}.forward = forward\",\nL4640: globals(),\nL4641: locals(),\nL4642: ) | L4646: exec(\nL4647: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4648: globals(),\nL4649: locals(),\nL4650: ) | L4652: exec(\nL4653: f\"combined_module.nn.{module}.forward = forward\",\nL4654: globals(),\nL4655: locals(),\nL4656: ) | L4673: exec(\nL4674: f\"{model_location}.{module} = combined_module.{module}\",\nL4675: globals(),\nL4676: locals(),\nL4677: ) | L4687: check_dicts = dir(eval(f\"{model_location}\")) | L4689: item = eval(f\"{model_location}.{check}\") | L4699: exec(\nL4700: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4701: globals(),\nL4702: locals(),\nL4703: )", "evidence_hash": "ec1875fd32d00fe885e566ebda75163e46e838ca31020abb57e0991892c2bdf7" }, { @@ -1423,8 +1511,8 @@ "file": "unsloth_zoo/mlx/loader.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L2218: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L140: mx.eval(model.parameters()) | L176: mx.eval(model.parameters()) | L2022: model.eval() | L2605: mx.eval(model.parameters()) | L2721: mx.eval(module.weight) | L4030: mx.eval(model.parameters()) | L4058: mx.eval(model.parameters()) | L4178: mx.eval(model.parameters())", - "evidence_hash": "9b29dade82912216c8b4808aa293b79749aa80ef1d2be35edd93bec7632810f1" + "evidence": "Obfusc: L2869: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L148: mx.eval(model.parameters()) | L180: mx.eval(model.parameters()) | L732: mx.eval(model.parameters()) | L733: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L799: mx.eval(model.parameters()) | L802: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L2673: model.eval() | L3256: mx.eval(model.parameters()) | L3372: mx.eval(module.weight) | L5666: mx.eval(model.parameters()) | L5716: mx.eval(model.parameters()) | L5859: mx.eval(model.parameters())", + "evidence_hash": "7b44760032c5df6d379ccfdd0bff3d23f857f64e08210fa0fba8d2881d457634" }, { "package": "unsloth-zoo", @@ -1439,7 +1527,7 @@ "file": "unsloth_zoo/saving_utils.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L3241: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3123: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3169: exec(save_pretrained, globals(), functions)", + "evidence": "Obfusc: L4015: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3897: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3943: exec(save_pretrained, globals(), functions)", "evidence_hash": "530b2383acd9fe8330aa65cd0bf86164aaacd47770e7c8d0752195bee36396ec" }, { @@ -1449,118 +1537,6 @@ "severity": "HIGH", "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)", "evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d" - }, - { - "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) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", - "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" - }, - { - "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)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", - "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", - "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L586: while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe", - "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb" - }, - { - "package": "tensorboard", - "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", - "check": "Python wheel ships large JS bundle (uncommon; manually review)", - "severity": "HIGH", - "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", - "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", - "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" - }, - { - "package": "fastmcp-slim", - "file": "fastmcp/cli/apps_dev.py", - "check": "Enumerates filesystem AND makes network calls", - "severity": "CRITICAL", - "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", - "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/_sandbox.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", - "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/_sandbox.py", - "check": "Writes to /tmp and executes (staged dropper)", - "severity": "CRITICAL", - "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", - "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/hf_api.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", - "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/utils/_http.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", - "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" - }, - { - "package": "cffi", - "file": "cffi/_cffi_gen_src.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", - "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" - }, - { - "package": "multiprocess", - "file": "multiprocess/forkserver.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L6: import socket sha256:6c707119169286c9a798e2c8d13a48614e481d8a503950916fd4ffb4c94d3182", - "evidence_hash": "50fec0f0522a8e4e636bf348b752002d7935d8455af31fb78c6f11e2eba19f6d" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3569: os.dup2(conn.fileno(), i) | L3601: \"test needs os.dup2()\") | L3619: os.dup2(fd, newfd) | L20: import socket sha256:c824dc0f409f242420c3fbb324790c53cb3078d2c8b07ee8f2a05694b01c2946", - "evidence_hash": "3878a2b430c175dbc5877a95195bfe52f9588ff73fb74e2261ed5e33087915ad" } ] } diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py new file mode 100644 index 0000000000..706346daad --- /dev/null +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Standalone free-VRAM probe for the bundled ggml Vulkan backend. + +Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the +Vulkan instance never lives in the long-running backend process. Loads the +bundled ggml Vulkan backend from ```` and prints one +``\\t\\t\\t`` line per device to stdout. +Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi +order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU +sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses +it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm +fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM. + +Uses only the standard library so it stays runnable as a bare script. +""" + +import ctypes +import os +import sys + +# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ... +_GGML_BACKEND_DEVICE_TYPE_IGPU = 2 + + +def _igpu_flags(base, lib, count: int) -> list[bool]: + """Per-device integrated-GPU flags via ggml's backend registry. + + The Vulkan reg enumerates devices in the same order as + ``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device = + i``), so reg index == device ordinal. Returns all-False on any failure so + the reader never over-caps a discrete card. + """ + flags = [False] * count + try: + lib.ggml_backend_vk_reg.restype = ctypes.c_void_p + lib.ggml_backend_vk_reg.argtypes = [] + base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t + base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p] + base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p + base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + base.ggml_backend_dev_type.restype = ctypes.c_int + base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p] + + reg = lib.ggml_backend_vk_reg() + if not reg: + return flags + dev_count = base.ggml_backend_reg_dev_count(reg) + for i in range(min(count, dev_count)): + dev = base.ggml_backend_reg_dev_get(reg, i) + if dev: + flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + except Exception: + # Best-effort: any failure degrades to "discrete" so the memory + # readings still get through instead of crashing the probe. + pass + return flags + + +def main() -> int: + if len(sys.argv) < 2: + return 0 + bindir = sys.argv[1] + + # Hold add_dll_directory's handle for the rest of main() (the documented + # idiom) so bindir stays on the search path while the sibling ggml DLLs + # resolve below. + _dll_dir = None + if sys.platform == "win32": + base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll" + try: + _dll_dir = os.add_dll_directory(bindir) + except Exception: + pass + else: + base_name, vk_name = "libggml-base.so", "libggml-vulkan.so" + + # RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr + # falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode). + _rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0) + try: + base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global) + lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global) + except OSError as e: + print(f"ggml-vulkan load failed: {e}", file = sys.stderr) + return 1 + + lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int + lib.ggml_backend_vk_get_device_count.argtypes = [] + lib.ggml_backend_vk_get_device_memory.restype = None + lib.ggml_backend_vk_get_device_memory.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ] + + count = lib.ggml_backend_vk_get_device_count() + igpu = _igpu_flags(base, lib, count) + rows = [] + for i in range(count): + free, total = ctypes.c_size_t(0), ctypes.c_size_t(0) + lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total)) + rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value)) + sys.stdout.write("\n".join(rows)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index dfd4c1c0bc..897db8262d 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -12,6 +12,43 @@ import json import logging from typing import Optional +_THINK_OPEN = "" +_THINK_CLOSE = "" + + +def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: + """Return the trailing open ```` prefill of a rendered prompt. + + Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation + prompt with ``\\n`` so the model starts reasoning immediately. + Because that opening tag is part of the *prompt*, skip_prompt streaming + never emits it, and the frontend's ````/```` parser shows + the reasoning as plain text instead of a thinking block. (The GGUF path + is unaffected: llama-server's reasoning parser returns + ``reasoning_content``, which gets re-wrapped in think tags.) + + Returns the exact prompt tail to re-emit at the start of the generated + stream (e.g. ``"\\n"``), or ``""`` when the prompt does not end + with an open think block, including the ``enable_thinking=False`` case + where templates prefill an already-closed ``\\n\\n``. + + ``special_tokens`` is the tokenizer's special-token list. If ```` + is one, the streamer's skip_special_tokens strips the model's closing tag, + so re-emitting the open would leave an unclosed block that swallows the + answer. In that case return ``""`` and fall back to plain text. + """ + if not prompt: + return "" + open_idx = prompt.rfind(_THINK_OPEN) + if open_idx == -1: + return "" + tail = prompt[open_idx:] + if _THINK_CLOSE in tail or tail.strip() != _THINK_OPEN: + return "" + if special_tokens and _THINK_CLOSE in set(special_tokens): + return "" + return tail + logger = logging.getLogger(__name__) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 167706f701..7e69e05124 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -1178,13 +1178,22 @@ class InferenceBackend: add_special_tokens = False, return_tensors = "pt", ).to(model.device) + prompt_text = input_text else: # Text-only path for a vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device) + prompt_text = formatted_prompt # Stream with TextIteratorStreamer + background thread try: + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill swallowed by skip_prompt (see + # generate_stream). + think_prefix = detect_think_prefill( + prompt_text, getattr(raw_tokenizer, "all_special_tokens", None) + ) from transformers import TextIteratorStreamer import threading @@ -1233,7 +1242,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False @@ -1467,6 +1480,16 @@ class InferenceBackend: from transformers import TextIteratorStreamer import threading + from core.inference.chat_template_helpers import detect_think_prefill + + # skip_prompt swallows an open prefilled by the template; + # re-emit it so the frontend can render the thinking block. + # gpt-oss emits its own tags via HarmonyTextStreamer. + think_prefix = ( + "" + if self._is_gpt_oss_model() + else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) + ) # gpt-oss models: HarmonyTextStreamer parses the multi-channel # harmony protocol into tags @@ -1550,7 +1573,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py new file mode 100644 index 0000000000..b6a939c87b --- /dev/null +++ b/studio/backend/core/inference/llama_admission.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Admission control for local llama-server generation requests. + +The helpers in this module deliberately know nothing about FastAPI, SSE, or the +OpenAI-compatible route shape. They only coordinate how many upstream generation +requests may be active for one llama-server backend and provide a cancellable +FIFO queue for excess requests. +""" + +from __future__ import annotations + +import asyncio +import os +import threading +from collections import deque +from dataclasses import dataclass +from typing import Deque, Optional + + +ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL" +ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT" +ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL" +ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE" + +DEFAULT_ADMISSION_ENABLED = True +DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None +DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0 +DEFAULT_ADMISSION_MAX_QUEUE = 64 + + +@dataclass(frozen = True) +class LlamaAdmissionConfig: + enabled: bool = DEFAULT_ADMISSION_ENABLED + queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE + + +@dataclass(frozen = True) +class LlamaAdmissionSnapshot: + key: str + capacity: int + active: int + queued: int + + +class LlamaAdmissionError(Exception): + def __init__( + self, + message: str, + *, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + super().__init__(message) + self.snapshot = snapshot + + +class LlamaAdmissionQueueFull(LlamaAdmissionError): + pass + + +class LlamaAdmissionTimeout(LlamaAdmissionError): + pass + + +class LlamaAdmissionCancelled(LlamaAdmissionError): + pass + + +def _bool_env(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + value = value.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + return default + + +def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def _positive_float_env(name: str, default: float) -> float: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else default + + +def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = int(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def llama_admission_config_from_env() -> LlamaAdmissionConfig: + return LlamaAdmissionConfig( + enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED), + queue_timeout_s = _optional_positive_float_env( + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + ), + keepalive_interval_s = _positive_float_env( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + ), + max_queue = _optional_positive_int_env( + ADMISSION_MAX_QUEUE_ENV, + DEFAULT_ADMISSION_MAX_QUEUE, + ), + ) + + +@dataclass +class _Waiter: + loop: asyncio.AbstractEventLoop + future: asyncio.Future + cancelled: bool = False + granted_lease: Optional["LlamaAdmissionLease"] = None + + +class LlamaAdmissionLease: + def __init__(self, queue: Optional["LlamaAdmissionQueue"]): + self._queue = queue + self._released = False + self._release_lock = threading.Lock() + + def release(self) -> None: + queue = None + with self._release_lock: + if self._released: + return + self._released = True + queue = self._queue + if queue is not None: + queue.release() + + async def __aenter__(self) -> "LlamaAdmissionLease": + return self + + async def __aexit__(self, *_args) -> None: + self.release() + + +class LlamaAdmissionReservation: + def __init__( + self, + *, + queue: Optional["LlamaAdmissionQueue"], + lease: Optional[LlamaAdmissionLease] = None, + waiter: Optional[_Waiter] = None, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + self._queue = queue + self._lease = lease + self._waiter = waiter + self.snapshot = snapshot + + @property + def is_cancelled(self) -> bool: + return self._lease is None and self._waiter is None + + def lease_nowait(self) -> Optional[LlamaAdmissionLease]: + if self._lease is not None: + return self._lease + if self._waiter is None or not self._waiter.future.done(): + return None + if self._waiter.future.cancelled(): + self._waiter.cancelled = True + self._waiter = None + return None + self._lease = self._waiter.future.result() + self._waiter = None + return self._lease + + async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]: + lease = self.lease_nowait() + if lease is not None: + return lease + if self._waiter is None: + return None + waiter = self._waiter + try: + await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s) + except asyncio.CancelledError: + if waiter.future.cancelled(): + waiter.cancelled = True + if self._waiter is waiter: + self._waiter = None + return None + raise + return self.lease_nowait() + + def cancel(self) -> None: + lease = self.lease_nowait() + if lease is not None: + lease.release() + self._lease = None + return + if self._queue is not None and self._waiter is not None: + self._queue.cancel(self._waiter) + self._waiter = None + + def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]: + if self._queue is None: + return self.snapshot + return self._queue.snapshot() + + +class LlamaAdmissionQueue: + def __init__(self, key: str): + self.key = key + self._lock = threading.Lock() + self._active = 0 + self._capacity = 1 + self._waiters: Deque[_Waiter] = deque() + + def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation: + capacity = max(1, int(capacity or 1)) + if not config.enabled: + return LlamaAdmissionReservation( + queue = None, + lease = LlamaAdmissionLease(None), + snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0), + ) + + loop = asyncio.get_running_loop() + with self._lock: + self._capacity = capacity + self._prune_waiters_locked() + self._grant_waiters_locked() + if self._active < self._capacity and not self._waiters: + self._active += 1 + return LlamaAdmissionReservation( + queue = self, + lease = LlamaAdmissionLease(self), + snapshot = self._snapshot_locked(), + ) + if config.max_queue is not None and len(self._waiters) >= config.max_queue: + raise LlamaAdmissionQueueFull( + "llama-server generation queue is full", + snapshot = self._snapshot_locked(), + ) + waiter = _Waiter( + loop = loop, + future = loop.create_future(), + ) + self._waiters.append(waiter) + return LlamaAdmissionReservation( + queue = self, + waiter = waiter, + snapshot = self._snapshot_locked(), + ) + + def release(self) -> None: + with self._lock: + if self._active > 0: + self._active -= 1 + self._grant_waiters_locked() + + def cancel(self, waiter: _Waiter) -> None: + lease_to_release = None + with self._lock: + waiter.cancelled = True + try: + self._waiters.remove(waiter) + except ValueError: + pass + if waiter.granted_lease is not None: + lease_to_release = waiter.granted_lease + waiter.granted_lease = None + if not waiter.future.done(): + waiter.loop.call_soon_threadsafe(waiter.future.cancel) + if lease_to_release is not None: + lease_to_release.release() + + def snapshot(self) -> LlamaAdmissionSnapshot: + with self._lock: + self._prune_waiters_locked() + return self._snapshot_locked() + + def is_idle(self) -> bool: + with self._lock: + self._prune_waiters_locked() + return self._active == 0 and not self._waiters + + def _grant_waiters_locked(self) -> None: + self._prune_waiters_locked() + while self._waiters and self._active < self._capacity: + waiter = self._waiters.popleft() + if waiter.cancelled or waiter.future.done(): + continue + self._active += 1 + lease = LlamaAdmissionLease(self) + waiter.granted_lease = lease + waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + + def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None: + if waiter.cancelled or waiter.future.done(): + waiter.granted_lease = None + if not waiter.future.done(): + waiter.future.cancel() + lease.release() + return + try: + waiter.future.set_result(lease) + waiter.granted_lease = None + except asyncio.InvalidStateError: + waiter.granted_lease = None + lease.release() + + def _prune_waiters_locked(self) -> None: + self._waiters = deque( + waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() + ) + + def _snapshot_locked(self) -> LlamaAdmissionSnapshot: + return LlamaAdmissionSnapshot( + key = self.key, + capacity = self._capacity, + active = self._active, + queued = len(self._waiters), + ) + + +_QUEUES_LOCK = threading.Lock() +_QUEUES: dict[str, LlamaAdmissionQueue] = {} + + +def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue: + with _QUEUES_LOCK: + queue = _QUEUES.get(key) + if queue is None: + queue = LlamaAdmissionQueue(key) + _QUEUES[key] = queue + # base_url carries a fresh ephemeral port on every model load, so + # each load registers a new key. Drop the now-idle queues from prior + # loads so the registry can't grow without bound on a long-running + # server. Queues with in-flight requests are kept until they drain. + for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]: + del _QUEUES[stale_key] + return queue + + +def reset_llama_admission_queues() -> None: + with _QUEUES_LOCK: + _QUEUES.clear() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 757467a008..b06c6eb5cb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -67,7 +67,6 @@ from core.tool_healing import ( _strip_bracket_tag_calls, apply_tool_strip_patterns, strip_outside_think, - 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 @@ -101,6 +100,10 @@ class LlamaServerNotFoundError(RuntimeError): Subclasses RuntimeError so existing handlers still catch it.""" +class _LlamaStreamCancelled(Exception): + """Internal signal for an expected client/request cancellation.""" + + # 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 " @@ -378,6 +381,19 @@ def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool: return True if result[0] is None else result[0] +def _hf_env_offline() -> bool: + """True when an HF offline env var is set to any truthy value (1/true/yes/on). + + Mirrors utils.models.model_config._env_offline so a user-set HF_HUB_OFFLINE=true + (not just "1") still routes through the local-cache reuse path below. + """ + try: + from utils.models.model_config import _env_offline + return _env_offline() + except Exception: + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + + @contextlib.contextmanager def _hf_offline_if_dns_dead(): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; @@ -838,6 +854,112 @@ def _gguf_snapshot_files(snapshot: Path) -> list[str]: ] +def _cached_hf_snapshot_file( + repo_id: str, + filename: str, + *, + expected_size: Optional[int] = None, +) -> Optional[str]: + """Return a cached snapshot file even when HF's current-ref probe misses it.""" + if not filename: + return None + parts = [part for part in filename.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + candidate = snap.joinpath(*parts) + if not candidate.is_file(): + continue + if expected_size: + try: + if candidate.stat().st_size < expected_size: + continue + except OSError: + continue + return str(candidate) + except Exception as e: + logger.debug("Snapshot cache lookup failed for %s/%s: %s", repo_id, filename, e) + return None + + +def _snapshot_has_all_shards( + main_path: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> bool: + """True when every shard sits beside ``main_path`` in the same cache snapshot. + + llama.cpp loads a split GGUF by resolving its siblings from the main shard's + directory, so a cached main shard is only safe to reuse when the rest of the + set is co-located; otherwise the caller must fetch the whole set together. + """ + root = Path(main_path) + for _ in [part for part in main_filename.replace("\\", "/").split("/") if part]: + root = root.parent + for shard in shards: + parts = [part for part in shard.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return False + sibling = root.joinpath(*parts) + try: + if not sibling.is_file(): + return False + expected = expected_sizes.get(shard) + if expected and sibling.stat().st_size < expected: + return False + except OSError: + return False + return True + + +def _resolve_repo_id_casing(hf_repo: str) -> str: + """Map a requested repo id to its cached canonical casing, or return it unchanged. + + A case-variant request (for example a lowercased id) resolves to the + canonical-cased cache directory so the main GGUF and its companions + (mmproj / MTP drafter) all read the same cache entry. Returns ``hf_repo`` + unchanged when resolution is unavailable or errors. + """ + try: + from utils.paths import resolve_cached_repo_id_case + return resolve_cached_repo_id_case(hf_repo) + except Exception: + return hf_repo + + +def _cached_colocated_split_main( + repo_id: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> Optional[str]: + """Main-shard path from a cache snapshot that also holds every sibling shard. + + A newer snapshot may hold only the first shard while an older snapshot has the + complete split set. ``_cached_hf_snapshot_file`` would return that newer partial + main and the co-location check would then force a refetch, so scan snapshots for + one where the whole set is present and return that main path instead. None when + no snapshot holds the full set. + """ + main_parts = [part for part in main_filename.replace("\\", "/").split("/") if part] + if not main_parts or any(part in (".", "..") for part in main_parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + main_path = snap.joinpath(*main_parts) + if not main_path.is_file(): + continue + expected_main = expected_sizes.get(main_filename) + try: + if expected_main and main_path.stat().st_size < expected_main: + continue + except OSError: + continue + if _snapshot_has_all_shards(str(main_path), main_filename, shards, expected_sizes): + return str(main_path) + except Exception as e: + logger.debug("Co-located split snapshot lookup failed for %s: %s", repo_id, e) + return None + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -1318,6 +1440,50 @@ def _backfill_usage_from_timings(usage, timings): return out +def _vulkan_lib_filename() -> str: + return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so" + + +# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit +# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared +# system RAM, so hold back the same margin rather than inventing a larger one. +_IGPU_HOST_RESERVE_MIB = 1024 + + +def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int: + """Reserve host headroom on an integrated (shared-memory) Vulkan GPU. + + An iGPU's reported free "VRAM" is really free system RAM, so sizing + context/offload against all of it would push the host into swap or the OOM + killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from + ggml's device type, so a discrete card is never touched; only ever reduces. + """ + if not is_igpu: + return free_mib + return max(0, free_mib - _IGPU_HOST_RESERVE_MIB) + + +def _llama_lib_dir(binary: str) -> Path: + # The installer exposes llama-server as a top-level entrypoint into build/bin/, + # where the ggml backend libs live, so callers looking for sibling libs (Vulkan + # detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a + # symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to + # a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it + # cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's + # exec target too, so a wrapper-based install still finds build/bin. + resolved = Path(binary).resolve() + try: + with open(resolved, "rb") as _f: + _head = _f.read(256) + if _head.startswith(b"#!"): + _m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore")) + if _m: + return (resolved.parent / _m.group(1)).resolve().parent + except OSError: + pass + return resolved.parent + + def _is_external_link(path: Path) -> bool: """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink or a Windows directory junction / reparse point. Such a link resolves into @@ -1375,6 +1541,7 @@ class LlamaCppBackend: self._context_length: Optional[int] = None self._effective_context_length: Optional[int] = None self._max_context_length: Optional[int] = None + self._effective_parallel_slots: int = 1 self._chat_template: Optional[str] = None self._chat_template_override: Optional[str] = None self._supports_reasoning: bool = False @@ -1560,6 +1727,15 @@ class LlamaCppBackend: """Return the effective context length the server is running at.""" return self._effective_context_length or self._context_length + @property + def effective_parallel_slots(self) -> int: + """Return the serving-slot count the active llama-server actually uses.""" + try: + slots = int(getattr(self, "_effective_parallel_slots", 1)) + except (TypeError, ValueError): + slots = 1 + return max(1, slots) + @property def max_context_length(self) -> Optional[int]: """Return the largest context that fits on this hardware at load time. @@ -1576,6 +1752,16 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + def _commit_effective_parallel_slots(self, n_parallel: int) -> None: + try: + slots = int(n_parallel) + except (TypeError, ValueError): + slots = 1 + self._effective_parallel_slots = max(1, slots) + + def _reset_effective_parallel_slots(self) -> None: + self._effective_parallel_slots = 1 + @staticmethod def _read_rss_bytes(pid: int) -> Optional[int]: """Resident set size of ``pid`` in bytes, from /proc//status (Linux). @@ -1781,6 +1967,13 @@ class LlamaCppBackend: return False return self._supports_tools + @property + def supports_tool_passthrough(self) -> bool: + # supports_tools is forced off for DiffusionGemma (its agentic loop drops the + # per-step canvas frames), but client passthrough skips that loop, so it uses + # the real _supports_tools. + return self._supports_tools + @property def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv @@ -2153,6 +2346,30 @@ class LlamaCppBackend: return total + @staticmethod + def _is_vulkan_backend(binary: Optional[str] = None) -> bool: + """True if the installed llama.cpp build is Vulkan-only. + + The official prebuilts are single-backend, so the Vulkan ggml lib next + to llama-server identifies a Vulkan build. Keeps the free-memory probe + and GPU pin in ggml's Vulkan device-index space. For a custom + multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer + to that backend (torch-usable, better-understood probe/pin). + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return False + lib_dir = _llama_lib_dir(binary) + if not (lib_dir / _vulkan_lib_filename()).is_file(): + return False + for _backend in ("cuda", "hip"): + sibling = ( + f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so" + ) + if (lib_dir / sibling).is_file(): + return False + return True + @staticmethod def _resolve_visible_physical_ids() -> Optional[list[int]]: """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on @@ -2315,11 +2532,42 @@ class LlamaCppBackend: return True @staticmethod - def _get_gpu_free_memory() -> list[tuple[int, int]]: + def _visible_devices_mask(env_name: str) -> Optional[set[int]]: + """Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset. + + ``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask + ("") yields an empty set (all devices hidden), distinct from an unset + var (None, no mask). Used by the nvidia-smi probe. + """ + raw = os.environ.get(env_name) + if raw is None: + return None + try: + return set(int(x.strip()) for x in raw.split(",") if x.strip()) + except ValueError: + return None + + @staticmethod + def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]: + """``--device Vulkan,...`` to pin a Vulkan launch to selected GPUs. + + The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory + reports and the registry names ``Vulkan``). Pin by that name, NOT via + GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw + vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices + and deduplicating ICDs), so a compact ordinal there could select a + different physical device or the CPU rasterizer. + """ + if not gpu_indices: + return [] + return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)] + + @staticmethod + def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]: """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()] + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)] @staticmethod def _apple_metal_memory_budget_bytes() -> int: @@ -2350,7 +2598,7 @@ class LlamaCppBackend: return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION) @staticmethod - def _get_gpu_memory() -> list[tuple[int, int, int]]: + def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]: """Query free AND total memory per GPU. Order: @@ -2362,9 +2610,18 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. + On a Vulkan build the ggml Vulkan probe is authoritative, so the indices + are ggml's compact Vulkan ordinals (the space the pin selects via + ``--device Vulkan``). It reports ``total`` for discrete cards and 0 + for an iGPU (shared RAM) so the fit falls back to free*frac there. + Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no - supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. + supported GPU is reachable. """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if LlamaCppBackend._is_vulkan_backend(binary): + return LlamaCppBackend._get_gpu_free_memory_vulkan(binary) # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( @@ -2380,16 +2637,7 @@ class LlamaCppBackend: **_windows_hidden_subprocess_kwargs(), ) if result.returncode == 0: - allowed: Optional[set[int]] = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - # `if x.strip()` filters trailing-comma masks ("0,1,"). - # Empty mask (CVD="") yields an empty set -> all GPUs - # filtered out, per codebase convention. - allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) - except ValueError: - pass + allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES") gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): parts = [p.strip() for p in line.split(",")] @@ -2454,6 +2702,91 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: + """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. + + Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance + in this process) and returns (device_index, free_mib, total_mib) sorted + by index. The index is ggml's compact Vulkan ordinal -- the one the + registry names ``Vulkan`` and load_model pins with ``--device``, + NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set + ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the + list already reflects it. iGPUs leave a host-RAM margin (see + ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass + their real total through. [] when no Vulkan build or device is reachable. + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return [] + binary_dir = _llama_lib_dir(binary) + if not (binary_dir / _vulkan_lib_filename()).is_file(): + return [] + + env = child_env_without_native_path_secret() + # Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so + # the probe enumerates the same device list the launch will, named + # Vulkan0..N in the compact order reported here and pinned by that name + # via --device -- probe, mask, and pin stay in one index space. Do NOT + # filter the mask in Python: ggml parses the env var in raw + # vkEnumeratePhysicalDevices space while this probe reports the compact + # post-filter ordinal, so a Python filter would compare mismatched spaces. + if sys.platform != "win32": + # Let the loader resolve sibling ggml libs next to the binary. + existing_ld = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ( + f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir) + ) + probe_script = Path(__file__).with_name("_vulkan_probe.py") + try: + result = subprocess.run( + [sys.executable, str(probe_script), str(binary_dir)], + capture_output = True, + text = True, + timeout = 15, + env = env, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + logger.debug( + f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}" + ) + return [] + except Exception as e: + logger.debug(f"vulkan GPU probe failed: {e}") + return [] + + gpus: list[tuple[int, int, int]] = [] + for line in result.stdout.strip().splitlines(): + parts = line.split("\t") + if len(parts) != 4: + continue + try: + idx = int(parts[0]) + free_mib = int(parts[1]) // (1024 * 1024) + is_igpu = parts[2] == "1" + # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the + # fit stays on free*frac (the host reserve below is its + # headroom); a discrete card passes its real total through. + total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024) + except ValueError: + continue + capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu) + if capped < free_mib: + logger.info( + f"Vulkan device VK{idx} is an integrated GPU sharing system " + f"RAM; reserving {free_mib - capped}MiB host headroom " + f"({free_mib}->{capped}MiB usable)" + ) + gpus.append((idx, capped, total_mib)) + gpus.sort(key = lambda g: g[0]) + if gpus: + logger.info( + "Vulkan GPU memory detected: " + + ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus) + ) + return gpus + @staticmethod def _available_system_memory_mib() -> Optional[int]: """Available system RAM in MiB (psutil, then /proc/meminfo), or None if @@ -2682,7 +3015,8 @@ class LlamaCppBackend: 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) + # _llama_lib_dir resolves the llama-server symlink to the real build/bin. + binary_dir = str(_llama_lib_dir(binary)) if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. @@ -3846,7 +4180,11 @@ class LlamaCppBackend: # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM # (capped at the training context). An explicit in-range n_ctx overrides it. maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 - gpu = os.environ.get("DG_GPU", "0") + # No visible CUDA GPU: a genuine CPU host, or a GPU host masked with + # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child + # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. + cpu_only = self._effective_gpu_count() == 0 + gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") cmd = list(shim_cmd) + [ "--gguf", @@ -3866,6 +4204,12 @@ class LlamaCppBackend: # refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import # unsloth`). The shim never imports unsloth, so set it here as unsloth does. env["UNSLOTH_IS_PRESENT"] = "1" + # The shim's `import unsloth_zoo` aborts in get_device_type() ("needs a GPU") + # when no accelerator is visible, even though it only drives the CPU + # visual-server binary and does no torch GPU work. Allow the CPU device so the + # runner starts; the visual server still runs on the CPU llama.cpp build. + if cpu_only: + env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. @@ -3986,6 +4330,15 @@ class LlamaCppBackend: "Install it with: pip install huggingface_hub" ) + resolved_hf_repo = _resolve_repo_id_casing(hf_repo) + if resolved_hf_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved_hf_repo, + hf_repo, + ) + hf_repo = resolved_hf_repo + # Resolve the filename from the variant gguf_filename = None gguf_extra_shards: list[str] = [] @@ -4031,10 +4384,12 @@ class LlamaCppBackend: # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards + expected_sizes: dict[str, int] = {} try: from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) + expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -4043,7 +4398,26 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - if not force: + # Cross-snapshot / case-variant cache reuse is offline-only (see the download + # path below); online, hf_hub_download fetches the current revision and + # resumes partials, so an old snapshot must not be counted as cached here or + # the preflight would under-count the download and skip the disk fallback. + offline = _hf_env_offline() + # A split GGUF whose shards are not co-located in a single snapshot is + # refetched as a whole set later, so it must not be counted as cached here. + split_needs_refetch = False + if offline and not force and gguf_extra_shards: + # Scan all snapshots for one that holds the whole set co-located, so a + # newer snapshot with only the first shard does not mask an older + # complete one and needlessly trip the disk fallback. + if ( + _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + is None + ): + split_needs_refetch = True + if not force and not split_needs_refetch: for p in path_infos: if not p.size: continue @@ -4051,6 +4425,15 @@ class LlamaCppBackend: cached_path = try_to_load_from_cache(hf_repo, p.path) except Exception: cached_path = None + if ( + not (isinstance(cached_path, str) and os.path.exists(cached_path)) + and offline + ): + cached_path = _cached_hf_snapshot_file( + hf_repo, + p.path, + expected_size = p.size, + ) if isinstance(cached_path, str) and os.path.exists(cached_path): try: on_disk = os.path.getsize(cached_path) @@ -4113,6 +4496,13 @@ class LlamaCppBackend: ) else: gguf_extra_shards = [] + # Record the fallback's size so the later cache-reuse probe can + # size-verify it; only for a single-file fallback, since + # _find_smallest_fitting_variant returns the whole-variant size + # and using that as the first shard's expected size would reject + # a valid cached first shard of a split fallback. + if not gguf_extra_shards: + expected_sizes[fallback_file] = fallback_size else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4132,25 +4522,45 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # 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 = cancel_event, - on_status = lambda m: logger.info(m), - force_download = force, - ) - for shard in gguf_extra_shards: - if cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( + local_path = None + # Reuse a cached copy from another snapshot / case-variant repo dir only when + # offline. Online, fall through to hf_hub_download so its revision/etag check + # fetches the current file (and resumes a partial) instead of serving a stale + # same-name blob from an older revision. + if not force and _hf_env_offline(): + if gguf_extra_shards: + # A split GGUF must load every shard from one snapshot; reuse only a + # snapshot that holds the whole set co-located, scanning past a newer + # snapshot that has just the first shard while an older one is complete. + local_path = _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + else: + local_path = _cached_hf_snapshot_file( + hf_repo, + gguf_filename, + expected_size = expected_sizes.get(gguf_filename), + ) + if local_path is None: + local_path = hf_hub_download_with_xet_fallback( hf_repo, - shard, + gguf_filename, hf_token, cancel_event = cancel_event, + on_status = lambda m: logger.info(m), force_download = force, ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = cancel_event, + force_download = force, + ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4228,6 +4638,17 @@ class LlamaCppBackend: if target is None or cancel_event.is_set(): return None + # Offline, resolve the companion straight from the cache snapshot that + # holds it. resolve_cached_repo_id_case can return a partial lower-case + # spelling when any dir exists under the requested casing, so calling + # hf_hub_download with hf_repo would miss the canonical file and silently + # drop the companion. _cached_hf_snapshot_file scans every case variant. + if _hf_env_offline(): + cached = _cached_hf_snapshot_file(hf_repo, target) + if cached: + logger.info("Resolved %s from local HF cache: %s", label, cached) + return cached + try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). @@ -4276,6 +4697,29 @@ class LlamaCppBackend: cancel_event = cancel_event, ) + def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: + """A drafter already in this repo's local HF cache, reused offline when a + fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all + cached snapshots; else an existing ``MTP/`` copy (any precision -- the + target verifies every drafted token). None if none is cached.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + + roots: list[Path] = [] + subdirs: list[Path] = [] + for snap in _iter_hf_cache_snapshots(hf_repo): # newest first + for f in sorted(_gguf_snapshot_files(snap)): + if _is_companion_gguf_path(f) and "mmproj" not in f.lower(): + (roots if "/" not in f else subdirs).append(snap / f) + # Keep snapshot order (newest first), root before any MTP/ copy, so a + # newer main GGUF pairs with the newest cached drafter, not a stale one. + for cand in roots + subdirs: + if cand.is_file(): + return str(cand) + except Exception as e: + logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e) + return None + def _download_mtp( self, *, @@ -4292,11 +4736,25 @@ class LlamaCppBackend: are intentionally skipped. Returns the local path, or None. """ + # Offline, reuse any drafter already on disk (a fresh copy can't be + # fetched). Online, _download_companion_gguf/hf_hub_download reuse the + # current cached file and refetch a changed one, so skip the probe here + # rather than pair new weights with a stale draft. + if _hf_env_offline(): + cached = self._cached_repo_mtp_drafter(hf_repo) + if cached: + logger.info(f"Reusing cached MTP drafter (offline): {cached}") + return cached + def _pick_mtp(candidates: list[str]) -> Optional[str]: + # Root-level only: MTP/ subdir copies now share the mtp- prefix but + # are explicit-selection, not auto-fetch (they'd sort ahead of root). mtp_files = sorted( f for f in candidates - if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-") + if f.lower().endswith(".gguf") + and "/" not in f + and Path(f).name.lower().startswith("mtp-") ) return mtp_files[0] if mtp_files else None @@ -4998,6 +5456,7 @@ class LlamaCppBackend: # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() + is_vulkan_backend = self._is_vulkan_backend(binary) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -5006,6 +5465,19 @@ class LlamaCppBackend: # dead; cleanup runs even on exception so a transient hiccup # can't quarantine future loads. if hf_repo: + # Resolve the requested repo id to its cached canonical casing once, + # up front, so the main GGUF and its companions (mmproj / MTP drafter) + # all resolve from the same cache entry. Otherwise a case-variant + # request resolves the main file from the canonical cache dir while the + # companions keep the requested casing and miss the cached files. + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): model_path = self._download_gguf( hf_repo = hf_repo, @@ -5224,7 +5696,8 @@ class LlamaCppBackend: model_size = gguf_size + mmproj_size # 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() + # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. + _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} @@ -5997,7 +6470,12 @@ class LlamaCppBackend: # 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): + # CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals. + if ( + model_size is not None + and not is_vulkan_backend + and self._amd_apu_wants_unified_memory(gpu_indices) + ): _ram_msg = self._apu_ram_shortfall_message( model_size, self._available_system_memory_mib() ) @@ -6260,6 +6738,12 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) + # Vulkan pins via --device (a cmd arg, unlike the env-based + # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's + # last-wins parsing lets a user --device override Studio's pick. + if is_vulkan_backend and gpu_indices is not None: + cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # 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(). @@ -6311,23 +6795,25 @@ class LlamaCppBackend: 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(gpu_indices): + # shared system RAM. setdefault so a user value wins. Not on Vulkan + # (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids. + if not is_vulkan_backend and 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") # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. - if self._apply_datacenter_env(env, gpu_indices): + if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices): multi_gpu = self._effective_gpu_count(gpu_indices) > 1 logger.info( f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) # 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. - if gpu_indices is not None: + # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so + # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device + # (above), not here. + if gpu_indices is not None and not is_vulkan_backend: pinned = ",".join(str(i) for i in gpu_indices) env["CUDA_VISIBLE_DEVICES"] = pinned try: @@ -6721,6 +7207,7 @@ class LlamaCppBackend: ) self._healthy = True + self._commit_effective_parallel_slots(n_parallel) # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] @@ -7258,6 +7745,7 @@ class LlamaCppBackend: self._context_length = None self._effective_context_length = None self._max_context_length = None + self._reset_effective_parallel_slots() self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -7313,6 +7801,7 @@ class LlamaCppBackend: # 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() + self._reset_effective_parallel_slots() if self._process is None: return try: @@ -8154,7 +8643,7 @@ class LlamaCppBackend: ): """Open one streaming POST and let cancel interrupt prefill or reads.""" if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled _cancel_closed = threading.Event() _response_ref: list = [None] @@ -8199,13 +8688,13 @@ class LlamaCppBackend: ) as response: _response_ref[0] = response if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled yield response return except (httpx.RequestError, RuntimeError): # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled raise finally: _cancel_closed.set() @@ -8408,6 +8897,8 @@ class LlamaCppBackend: "finish_reason": _metadata_finish_reason, } + except _LlamaStreamCancelled: + return 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 @@ -9532,6 +10023,8 @@ class LlamaCppBackend: break continue + except _LlamaStreamCancelled: + return except httpx.ConnectError: # Mark unresolved provisional cards as failed before raising. for _pid, _pname in provisional_started_tool_calls.items(): @@ -9714,6 +10207,8 @@ class LlamaCppBackend: if _meta is not None: yield _meta + except _LlamaStreamCancelled: + return except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") except Exception as e: diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 62d268e15f..6287b184a6 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -5,6 +5,7 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm instead of torch/transformers for model loading and generation. """ +import os import threading from typing import Optional, Generator from core.inference.runtime_context import runtime_context_length @@ -41,6 +42,48 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): } +def _mlx_distributed_rank_size(group = None): + """Return ``(rank, world_size)`` for an optional MLX distributed group.""" + if group is None: + return 0, 1 + rank = int(group.rank()) + world_size = int(group.size()) + if world_size < 1: + raise ValueError(f"Invalid MLX distributed world_size={world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.") + return rank, world_size + + +def _mlx_distributed_backend_from_env(): + if os.environ.get("MLX_JACCL_COORDINATOR") and os.environ.get("MLX_IBV_DEVICES"): + return "jaccl" + return None + + +def _init_mlx_distributed(): + """Initialize MLX distributed state, falling back to singleton metadata.""" + import mlx.core as mx + + group = None + rank = 0 + world_size = 1 + distributed = getattr(mx, "distributed", None) + init = getattr(distributed, "init", None) if distributed is not None else None + if callable(init): + backend = _mlx_distributed_backend_from_env() + if backend is None: + group = init() + else: + try: + group = init(backend = backend) + except TypeError: + group = init() + if group is not None: + rank, world_size = _mlx_distributed_rank_size(group) + return group, rank, world_size + + def _make_mlx_presence_penalty_processor(penalty: float): """Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path. @@ -52,7 +95,7 @@ def _make_mlx_presence_penalty_processor(penalty: float): def _processor(tokens, logits): if state["prompt_len"] is None: - # First call = prompt only; latch its length. + # First call is prompt-only; latch its length. state["prompt_len"] = int(tokens.shape[0]) return logits generated = tokens[state["prompt_len"] :] @@ -61,22 +104,17 @@ def _make_mlx_presence_penalty_processor(penalty: float): import mlx.core as mx vocab = logits.shape[-1] - # Bound generated ids to the valid range [0, vocab) before they index - # logits. MLX does no bounds checking and out-of-bounds indexing is - # documented undefined behavior (crash / memory corruption), unlike the - # torch path's harmless negative wrap -- so this bound is load-bearing - # here and matches the torch filter seen[(seen >= 0) & (seen < vocab)]. - # MLX has no boolean-mask filtering (data-dependent output shape is - # unsupported), so instead of compacting the id list we route every - # out-of-range or negative id to a scratch slot at index ``vocab`` that - # is dropped before the subtract. That scratch slot can never collide - # with a real token, so real ids (including id 0) are penalized exactly - # once and stray ids are ignored. + # Bound ids to [0, vocab) before indexing logits: MLX does no bounds + # checking and out-of-bounds indexing is undefined behavior (crash / + # corruption), unlike torch's harmless negative wrap. MLX also lacks + # boolean-mask filtering, so out-of-range/negative ids route to a + # scratch slot at index vocab (dropped before the subtract) that never + # collides with a real token: real ids (including 0) are penalized + # once, strays ignored. valid = (generated >= 0) & (generated < vocab) safe = mx.where(valid, generated, vocab).astype(mx.int32) - # Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate - # ids are idempotent, so presence applies once per distinct token; the - # scratch column is discarded and the full-width subtract stays on-device. + # Scatter penalty into a (vocab + 1)-wide mask: duplicate ids are + # idempotent (presence applies once per token); scratch column dropped. mask = mx.zeros((vocab + 1,), dtype = logits.dtype) mask[safe] = penalty logits = logits - mask[:vocab] @@ -93,7 +131,7 @@ class MLXInferenceBackend: self.loaded_local_models = [] self.device = "mlx" self._generation_lock = threading.Lock() - # usage/timings of the latest generation; shipped on gen_done. + # usage/timings of the latest generation, shipped on gen_done. self.last_generation_stats = None self._model = None @@ -101,6 +139,9 @@ class MLXInferenceBackend: self._processor = None self._is_vlm = False self._config = {} + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 # Recorded for unload to release pinned memory back to the OS. self._memory_limits_applied = {} @@ -145,19 +186,26 @@ class MLXInferenceBackend: trust_remote_code = False, gpu_ids = None, dtype = None, + parallel_mode = None, + distributed_group = None, ) -> bool: import mlx.core as mx - # Keep the token so the native-template fallback can fetch a - # gated model's repo template later during generation. + # Keep the token so the native-template fallback can fetch a gated + # model's repo template during generation. self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) + distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group) + is_distributed = distributed_group is not None and distributed_size > 1 + self._distributed_group = distributed_group + self._distributed_rank = distributed_rank + self._distributed_world_size = distributed_size - # GGUF guard. GGUF models are served by llama-server in the parent - # process, not mlx-lm here. Reaching this with is_gguf=True means the - # route's first detection flaked (transient HF Hub) but the subprocess - # re-detected GGUF; raise loudly instead of a cryptic mlx_lm error. + # GGUF guard: GGUF is served by llama-server in the parent process, + # not mlx-lm. Reaching here with is_gguf=True means the route's + # detection flaked but the subprocess re-detected GGUF; raise loudly + # instead of a cryptic mlx_lm error. if getattr(config, "is_gguf", False): raise RuntimeError( f"MLXInferenceBackend cannot load GGUF model '{model_name}': " @@ -176,11 +224,26 @@ class MLXInferenceBackend: is_lora = getattr(config, "is_lora", False) logger.info( - "Loading %s via %s (is_lora=%s)", + "Loading %s via %s (is_lora=%s, distributed=%s, rank=%s/%s, mode=%s)", model_name, "mlx-vlm" if is_vision else "mlx-lm", is_lora, + is_distributed, + distributed_rank, + distributed_size, + parallel_mode, ) + if is_distributed and parallel_mode not in ("pipeline", "tensor"): + raise ValueError( + "Unsloth: distributed MLX inference requires parallel_mode='pipeline' " + "or parallel_mode='tensor'." + ) + if is_distributed and is_lora: + raise ValueError( + "Unsloth: distributed MLX inference for LoRA adapter repos " + "is not supported yet. Merge/export the adapter into an MLX model " + "before distributed inference." + ) try: from unsloth_zoo.mlx.loader import FastMLXModel @@ -190,14 +253,23 @@ class MLXInferenceBackend: "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e + load_kwargs = { + "max_seq_length": max_seq_length, + "dtype": dtype, + "load_in_4bit": load_in_4bit, + "token": hf_token, + "trust_remote_code": trust_remote_code, + "text_only": False if is_vision else True, + } + if is_distributed: + if parallel_mode == "pipeline": + load_kwargs["pipeline_group"] = distributed_group + else: + load_kwargs["tensor_group"] = distributed_group + model, tokenizer_or_processor = FastMLXModel.from_pretrained( model_name, - max_seq_length = max_seq_length, - dtype = dtype, - load_in_4bit = load_in_4bit, - token = hf_token, - trust_remote_code = trust_remote_code, - text_only = False if is_vision else True, + **load_kwargs, ) if is_vision: @@ -217,8 +289,7 @@ class MLXInferenceBackend: self.models[model_name] = { # Per-model token for the native-template fallback (matches transformers). "hf_token": hf_token, - # Per-model consent for the native-template reload: re-use the exact - # trust_remote_code this model was loaded with (matches transformers). + # Per-model trust_remote_code reused by the native-template reload (matches transformers). "trust_remote_code": trust_remote_code, "model": self._model, "tokenizer": self._tokenizer, @@ -234,8 +305,7 @@ class MLXInferenceBackend: "has_audio_input": False, "context_length": runtime_context_length(self._model, max_seq_length), } - # Capture chat_template_info so the worker IPC reply ships it back and - # the route layer classifies capabilities like the other paths. + # Capture chat_template_info for the worker IPC reply and route capability classification. self._populate_chat_template_info(model_name) logger.info("Model %s loaded successfully", model_name) @@ -293,6 +363,9 @@ class MLXInferenceBackend: self._model = None self._tokenizer = None self._processor = None + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 if self.active_model_name == model_name: self.active_model_name = None gc.collect() @@ -320,8 +393,7 @@ class MLXInferenceBackend: max_new_tokens = 256, repetition_penalty = 1.0, cancel_event = None, - # Reasoning / tool kwargs forwarded by the route + worker; rendered via - # apply_chat_template_for_generation like the transformers path. + # Reasoning / tool kwargs, rendered via apply_chat_template_for_generation (transformers parity). tools = None, enable_thinking = None, reasoning_effort = None, @@ -334,7 +406,6 @@ class MLXInferenceBackend: # Reset so a failed run cannot surface stale stats. self.last_generation_stats = None - # Build messages with system prompt full_messages = [] if system_prompt: full_messages.append({"role": "system", "content": system_prompt}) @@ -351,7 +422,6 @@ class MLXInferenceBackend: {"type": "text", "text": content}, ] elif isinstance(content, list): - # Prepend image if not already present has_image = any( p.get("type") == "image" for p in content if isinstance(p, dict) ) @@ -415,6 +485,7 @@ class MLXInferenceBackend: from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, + detect_think_prefill, render_with_native_template_fallback, ) @@ -429,11 +500,11 @@ class MLXInferenceBackend: if prompt is None: raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") - # Same parity fix as the transformers backend: if the template dropped the - # requested tools, fall back to the native template so MLX text models keep - # advertising them. ``self._tokenizer`` is this entry's model_info tokenizer, - # so probe and native render share a renderer. (The VLM path renders via the - # processor for image tokens and is intentionally not wired here.) + # Parity with the transformers backend: if the template dropped the + # requested tools, fall back to the native template so MLX text models + # keep advertising them. self._tokenizer is this entry's tokenizer, so + # probe and native render share a renderer. (VLM renders via the + # processor for image tokens and is not wired here.) model_info = self.models.get(self.active_model_name, {}) prompt = render_with_native_template_fallback( formatted_prompt = prompt, @@ -448,6 +519,15 @@ class MLXInferenceBackend: hf_token = model_info.get("hf_token"), ) + # An open prefilled by the template lives in the prompt, not + # the generated tokens; re-emit it so the frontend renders the block. + think_prefix = detect_think_prefill( + prompt, getattr(self._tokenizer, "all_special_tokens", None) + ) + # Emit it before the first token so the block renders during prefill. + if think_prefix: + yield think_prefix + sampler = make_sampler( temp = temperature, top_p = top_p, @@ -455,7 +535,7 @@ class MLXInferenceBackend: min_p = float(min_p or 0.0), min_tokens_to_keep = 1, ) - # Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths). + # Repetition and/or presence penalty processors (GGUF/safetensors parity). logits_processors = [] if repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, @@ -496,12 +576,11 @@ class MLXInferenceBackend: ): final_response = response token_ids.append(response.token) - # Decode full sequence with skip_special_tokens cumulative = self._tokenizer.decode( token_ids, skip_special_tokens = True, ) - yield cumulative + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break @@ -544,8 +623,7 @@ class MLXInferenceBackend: ) # Pick the chat-template-aware caller: processors with their own - # apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it - # directly; else fall back to the nested tokenizer. + # apply_chat_template + chat_template (e.g. Qwen2.5-VL), else the nested tokenizer. chat_target = self._processor if ( getattr(self._processor, "apply_chat_template", None) is None @@ -566,16 +644,21 @@ class MLXInferenceBackend: # mlx_vlm's stream_generate handles pixel_values (None for text-only) images = [image] if image is not None else None - cumulative = "" + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill from the prompt (see _generate_text). + cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) + # Emit it before the first token so the block renders during prefill. + if cumulative: + yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), image is not None, ) - # mlx_vlm.stream_generate forwards **kwargs into generate_step, which - # builds the sampler + logits_processors internally. - # GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=`` - # silently falls into **kwargs and is ignored, stuck at greedy 0.0. + # stream_generate forwards **kwargs into generate_step (builds the + # sampler + logits_processors internally). GOTCHA: generate_step expects + # temperature= (long form); temp= is silently ignored, stuck at greedy 0.0. vlm_kwargs = dict( max_tokens = max_new_tokens, temperature = temperature, @@ -589,7 +672,7 @@ class MLXInferenceBackend: ) if presence_penalty: # Presence needs a custom processor: pass the full list (repetition + - # presence) instead of the repetition_penalty shortcut so both apply once. + # presence) instead of the repetition_penalty shortcut so both apply. from mlx_lm.sample_utils import make_logits_processors _vlm_processors = [] @@ -634,7 +717,7 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported — generate normally + # MLX LoRA adapter toggling not yet supported; generate normally yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) def reset_generation_state(self): diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index cf5d24c367..4fa0d3ed26 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -50,6 +50,18 @@ _DISPATCH_DRAIN_TIMEOUT = 5.0 _UNLOAD_GEN_LOCK_TIMEOUT = 15.0 +class GenStreamError(str): + """A stream chunk carrying a real backend/generation error, not model text. + + Subclasses str so existing display/logging consumers are unaffected, while + callers that must abort a distributed run on error (raise_on_streamed_error) + can distinguish a real error from model output whose visible text starts with + "Error:" by checking isinstance(chunk, GenStreamError). + """ + + __slots__ = () + + class InferenceOrchestrator: """ Inference backend orchestrator — subprocess-based. @@ -482,13 +494,13 @@ class InferenceOrchestrator: initial_resp_queue = self._resp_queue while True: if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: - yield f"Error: {self._subprocess_crash_message(crash_context)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return 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)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return continue @@ -498,7 +510,7 @@ class InferenceOrchestrator: # 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')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return if rtype == "token": @@ -513,7 +525,7 @@ class InferenceOrchestrator: stats_holder["stats"] = resp.get("stats") return elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return # ------------------------------------------------------------------ @@ -640,11 +652,11 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return # Latch the target model so the recheck below can detect a switch that completed # between _start_dispatcher and mailbox registration (mirrors the locked path's @@ -655,7 +667,7 @@ class InferenceOrchestrator: # so without this early-out a compare request would enqueue a generate on the # outgoing model and delay the switch. if self._unload_pending: - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under @@ -727,7 +739,7 @@ class InferenceOrchestrator: # _stop_dispatcher joins the dispatcher, which itself takes that lock. if orphaned_dispatcher: self._stop_dispatcher() - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return try: @@ -735,7 +747,7 @@ class InferenceOrchestrator: except RuntimeError as exc: with self._mailbox_lock: self._mailboxes.pop(request_id, None) - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return def read_mailbox(timeout): @@ -813,6 +825,59 @@ class InferenceOrchestrator: self._stop_dispatcher() return True + def share_distributed_object( + self, + obj, + timeout: Optional[float] = 300.0, + ): + """Share a small object through the worker's MLX distributed group.""" + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess is not running") + + self._wait_dispatcher_idle() + with self._mailbox_lock: + if self._mailboxes: + raise RuntimeError( + "Cannot share distributed objects while compare requests are active" + ) + request_id = str(uuid.uuid4()) + cmd = { + "type": "share_object", + "request_id": request_id, + "object": obj, + } + + with self._gen_lock: + self._send_cmd(cmd) + deadline = None if timeout is None else time.monotonic() + timeout + while deadline is None or time.monotonic() < deadline: + remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) + continue + + rtype = resp.get("type", "") + rid = resp.get("request_id") + if rid and rid != request_id: + logger.debug( + "Skipping response for request_id=%s while sharing request_id=%s", + rid, + request_id, + ) + continue + if rtype == "shared": + return resp.get("object") + if rtype == "share_error": + raise RuntimeError(resp.get("error", "Failed to share object")) + if rtype == "error": + raise RuntimeError(resp.get("error", "Subprocess error")) + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for distributed object share") + # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend # ------------------------------------------------------------------ @@ -828,6 +893,8 @@ class InferenceOrchestrator: approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, subject: Optional[str] = None, + tensor_parallel: bool = False, + mlx_distributed: bool = False, ) -> bool: """Load a model for inference. @@ -853,6 +920,11 @@ class InferenceOrchestrator: "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "subject": subject, "gpu_ids": gpu_ids, + "tensor_parallel": bool(tensor_parallel), + "mlx_distributed": bool(mlx_distributed), + "mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline") + if mlx_distributed + else None, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( gpu_ids, @@ -1338,11 +1410,11 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return expected_model = self.active_model_name @@ -1359,7 +1431,7 @@ class InferenceOrchestrator: # so we never generate on the wrong one. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None @@ -1385,7 +1457,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( @@ -1544,10 +1616,10 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return expected_model = self.active_model_name @@ -1556,7 +1628,7 @@ class InferenceOrchestrator: # cleared or swapped the model while we waited. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return request_id = str(uuid.uuid4()) @@ -1583,7 +1655,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index d4b102e422..05dee39283 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker. from __future__ import annotations import base64 +import json from loggers import get_logger import os import queue as _queue @@ -26,6 +27,9 @@ from typing import Any logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +_SHARE_OBJECT_MAX_BYTES = 1 << 20 +_SHARE_OBJECT_ERROR_SIZE = -1 + # studio/backend root, prepended to sys.path so the spawned subprocess can # import the utils/core packages. _BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent) @@ -75,6 +79,17 @@ def _send_response(resp_queue: Any, response: dict) -> None: logger.error("Failed to send response: %s", exc) +def _encode_share_object(obj: Any) -> bytes: + data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8") + if len(data) > _SHARE_OBJECT_MAX_BYTES: + raise ValueError("Distributed object share payload is too large") + return data + + +def _decode_share_object(data: Any) -> Any: + return json.loads(bytes(data.tolist()).decode("utf-8")) + + def _clean_token(value: str | None) -> str | None: """Normalize an HF token: blank or whitespace-only becomes None.""" return value if value and value.strip() else None @@ -329,14 +344,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) try: - success = backend.load_model( - config = mc, - max_seq_length = config.get("max_seq_length", 2048), - load_in_4bit = load_in_4bit, - hf_token = hf_token, - trust_remote_code = trust_remote_code, - gpu_ids = config.get("resolved_gpu_ids"), - ) + load_kwargs = { + "config": mc, + "max_seq_length": config.get("max_seq_length", 2048), + "load_in_4bit": load_in_4bit, + "hf_token": hf_token, + "trust_remote_code": trust_remote_code, + "gpu_ids": config.get("resolved_gpu_ids"), + } + if getattr(backend, "device", None) == "mlx": + load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode") + load_kwargs["distributed_group"] = config.get("_mlx_distributed_group") + success = backend.load_model(**load_kwargs) finally: heartbeat_stop.set() @@ -521,6 +540,67 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: ) +def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None: + """Share a small Python object across MLX distributed ranks.""" + request_id = cmd.get("request_id", "") + group = getattr(backend, "_distributed_group", None) + rank = int(getattr(backend, "_distributed_rank", 0) or 0) + world_size = int(getattr(backend, "_distributed_world_size", 1) or 1) + obj = cmd.get("object") + + try: + if group is None or world_size <= 1: + shared = obj + else: + import mlx.core as mx + if rank == 0: + if obj is None: + mx.eval(mx.distributed.all_sum(mx.array(0), group = group)) + shared = None + else: + try: + data = mx.array(_encode_share_object(obj), dtype = mx.uint8) + except Exception: + mx.eval( + mx.distributed.all_sum( + mx.array(_SHARE_OBJECT_ERROR_SIZE), + group = group, + ) + ) + raise + mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group)) + mx.eval(mx.distributed.all_sum(data, group = group)) + shared = obj + else: + size = int(mx.distributed.all_sum(mx.array(0), group = group).item()) + if size == _SHARE_OBJECT_ERROR_SIZE: + raise RuntimeError("Failed to share distributed object") + if size == 0: + shared = None + else: + data = mx.zeros(size, dtype = mx.uint8) + data = mx.distributed.all_sum(data, group = group) + shared = _decode_share_object(data) + _send_response( + resp_queue, + { + "type": "shared", + "request_id": request_id, + "object": shared, + }, + ) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "share_error", + "request_id": request_id, + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + }, + ) + + def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None: """Handle TTS audio generation — returns WAV bytes + sample_rate.""" request_id = cmd.get("request_id", "") @@ -720,9 +800,29 @@ def run_inference_process( exc, ) try: - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed backend = MLXInferenceBackend() + if config.get("mlx_distributed"): + group, rank, size = _init_mlx_distributed() + config["_mlx_distributed_group"] = group + if size <= 1: + # A singleton group (MLX built without distributed support, + # or an invalid launch env/hostfile) would leave nonzero ranks + # looping forever on share_distributed_object. Fail the load + # instead of silently continuing without sharding. + raise RuntimeError( + "MLX distributed launch requested but initialized a singleton " + "group (size 1). Ensure the installed MLX has distributed " + "support and the launch environment/hostfile is valid, or run " + "without distributed." + ) + logger.info( + "MLX distributed initialized in worker: rank=%s size=%s mode=%s", + rank, + size, + config.get("mlx_parallel_mode"), + ) _send_response( resp_queue, {"type": "status", "message": "Loading model..."}, @@ -764,6 +864,8 @@ def run_inference_process( if _drain_skip_generate(cmd, resp_queue, drain_event): continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) @@ -977,6 +1079,9 @@ def run_inference_process( continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) + elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 20b2305a5a..96d1b90b16 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -62,7 +62,6 @@ from loggers import get_logger import time from pathlib import Path from typing import Any, Dict, List, Optional, Callable -from dataclasses import dataclass import pandas as pd from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -71,7 +70,7 @@ from core.inference.llama_cpp import _hf_offline_if_dns_dead from utils.models import is_vision_model, detect_audio_type from utils.models.model_config import _env_offline from utils.datasets import format_and_template_dataset -from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER +from utils.datasets.completion_masking import apply_completion_masking from utils.datasets.iterable import is_streaming_dataset as detect_streaming_dataset from utils.datasets.raw_text import prepare_raw_text_dataset, resolve_column_names from utils.paths import ( @@ -86,6 +85,11 @@ from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from .training import ( + TrainingProgress, + create_mlx_trainer_adapter, + should_use_mlx_training_backend, +) logger = get_logger(__name__) @@ -104,31 +108,16 @@ def _build_report_targets(training_args) -> list[str] | str: return report_to or "none" -@dataclass -class TrainingProgress: - """Training progress tracking""" - - epoch: float = 0 - step: int = 0 - total_steps: int = 0 - loss: Optional[float] = None - learning_rate: Optional[float] = None - is_training: bool = False - is_completed: bool = False - error: Optional[str] = None - status_message: str = "Ready to train" # Current stage - elapsed_seconds: Optional[float] = None - eta_seconds: Optional[float] = None - grad_norm: Optional[float] = None - num_tokens: Optional[int] = None - eval_loss: Optional[float] = None - - class UnslothTrainer: """ Unsloth Training Backend """ + def __new__(cls, *args, **kwargs): + if cls is UnslothTrainer and should_use_mlx_training_backend(): + return create_mlx_trainer_adapter(*args, **kwargs) + return super().__new__(cls) + def __init__(self): self.model = None self.tokenizer = None @@ -3466,8 +3455,6 @@ class UnslothTrainer: # ========== TRAIN ON RESPONSES ONLY ========== # Raw-text datasets always train on all tokens. - instruction_part = None - response_part = None is_cpt = training_args.get("is_cpt", False) train_on_responses_enabled = ( False @@ -3484,113 +3471,93 @@ class UnslothTrainer: # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip + # Markers auto-detected from the chat template first, manual table + # as fallback; gpt-oss stays on its manual markers. See + # apply_completion_masking. if ( train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset_final_format == "alpaca") ): - try: - logger.info("Configuring train on responses only...\n") + from unsloth.chat_templates import train_on_responses_only - # Template mapping for this model - model_name_lower = self.model_name.lower() + logger.info("Configuring train on responses only...\n") - if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Detected template: {template_name}\n") + def _notify(level, message): + if level == "warning": + logger.warning(message) + else: + logger.info(f"{message}\n") - if template_name in TEMPLATE_TO_RESPONSES_MAPPER: - instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ - "instruction" - ] - response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"] + # No try/except: the helper handles detection failures and + # double misses itself, so an exception here is a real masking + # failure that must fail the run, not silently train on full + # sequences. + self.trainer, masking_applied = apply_completion_masking( + self.trainer, + self.model_name, + train_on_responses_only, + num_proc = config_args["dataset_num_proc"], + notify = _notify, + ) - logger.info(f"Instruction marker: {instruction_part[:50]}...\n") - logger.info(f"Response marker: {response_part[:50]}...\n") + if not masking_applied: + train_on_responses_enabled = False + + if masking_applied: + try: + # ── Safety net: check if all samples were filtered out ── + # train_on_responses_only masks non-response tokens with -100; a + # row becomes all -100 (Unsloth drops it) when the response + # template is not found in the formatted text. Usually a + # dataset/template mismatch (already-formatted data, or 'Train on + # completions' on data that doesn't match the model's chat + # template); only sometimes max_seq_length truncating the response + # away. Skip this len()-based check for streaming. + if detect_streaming_dataset(self.trainer.train_dataset): + logger.info("Skipping post-filter length check for streaming dataset\n") else: - logger.info( - f"No response mapping found for template: {template_name}\n" + filtered_len = len(self.trainer.train_dataset) + original_dataset_obj = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset ) - train_on_responses_enabled = False - else: - logger.info(f"No template mapping found for model: {self.model_name}\n") - train_on_responses_enabled = False - - except Exception as e: - logger.warning(f"Could not configure train on responses: {e}") - train_on_responses_enabled = False - - # Apply train on responses only if we have valid parts - if ( - train_on_responses_enabled - and instruction_part - and response_part - and not self.is_audio_vlm - and not self.is_audio - and not (is_deepseek_ocr or dataset_final_format == "alpaca") - ): - try: - from unsloth.chat_templates import train_on_responses_only - - self.trainer = train_on_responses_only( - self.trainer, - instruction_part = instruction_part, - response_part = response_part, - num_proc = config_args["dataset_num_proc"], - ) - logger.info("Train on responses only configured successfully\n") - - # ── Safety net: check if all samples were filtered out ── - # train_on_responses_only masks non-response tokens with -100; - # a row becomes all -100 (and Unsloth drops it) when the response - # template is not found in the formatted text. That is usually a - # dataset/template mismatch (already-formatted data, or 'Train on - # completions' applied to data that doesn't match the model's chat - # template), and only sometimes max_seq_length truncating the - # response away. Skip this len()-based check for streaming. - if detect_streaming_dataset(self.trainer.train_dataset): - logger.info("Skipping post-filter length check for streaming dataset\n") - else: - filtered_len = len(self.trainer.train_dataset) - original_dataset_obj = ( - dataset["dataset"] if isinstance(dataset, dict) else dataset - ) - original_len = len(original_dataset_obj) - dropped = original_len - filtered_len - drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0 - - if filtered_len == 0 or drop_pct > 30: - max_seq = training_args.get("max_seq_length", 2048) - error_msg = ( - f"{dropped}/{original_len} samples ({drop_pct}%) were " - f"dropped after applying 'Train on completions': after " - f"masking, those rows had no trainable response tokens " - f"left. The usual cause is that this model's response " - f"template was not found in the formatted samples, so " - f"every token was masked out. That typically means the " - f"dataset is already formatted, or its structure does " - f"not match the model's chat template, so 'Train on " - f"completions' should be turned off for this dataset. " - f"Less commonly, a max_seq_length ({max_seq}) shorter " - f"than the prompt can truncate the response away; only " - f"raise it if your samples are actually longer than that." + original_len = len(original_dataset_obj) + dropped = original_len - filtered_len + drop_pct = ( + round(100 * dropped / original_len, 1) if original_len > 0 else 0 ) - logger.error(error_msg) - self._update_progress(error = error_msg, is_training = False) - return - if dropped > 0: - logger.info( - f"⚠️ {dropped}/{original_len} samples " - f"({drop_pct}%) were dropped (all labels " - f"masked). {filtered_len} samples remain.\n" - ) - logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + if filtered_len == 0 or drop_pct > 30: + max_seq = training_args.get("max_seq_length", 2048) + error_msg = ( + f"{dropped}/{original_len} samples ({drop_pct}%) were " + f"dropped after applying 'Train on completions': after " + f"masking, those rows had no trainable response tokens " + f"left. The usual cause is that this model's response " + f"template was not found in the formatted samples, so " + f"every token was masked out. That typically means the " + f"dataset is already formatted, or its structure does " + f"not match the model's chat template, so 'Train on " + f"completions' should be turned off for this dataset. " + f"Less commonly, a max_seq_length ({max_seq}) shorter " + f"than the prompt can truncate the response away; only " + f"raise it if your samples are actually longer than that." + ) + logger.error(error_msg) + self._update_progress(error = error_msg, is_training = False) + return - except Exception as e: - logger.warning(f"Failed to apply train on responses only: {e}") - train_on_responses_enabled = False + if dropped > 0: + logger.info( + f"⚠️ {dropped}/{original_len} samples " + f"({drop_pct}%) were dropped (all labels " + f"masked). {filtered_len} samples remain.\n" + ) + logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + + except Exception as e: + logger.warning(f"Post-masking dataset size check failed: {e}") else: if train_on_responses_enabled and is_deepseek_ocr: logger.info("Train on responses handled by DeepSeek OCR collator\n") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index f4233fcf04..2ddda19951 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -14,17 +14,19 @@ import json as _json import math import multiprocessing as mp import os +import platform import queue import re import shutil import threading import time +import traceback import structlog from datetime import datetime, timezone from loggers import get_logger -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Optional, Tuple, Any, TYPE_CHECKING +from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt @@ -98,6 +100,107 @@ def _coerce_optional_nonneg_float(name: str, value): return coerced +def is_apple_silicon_training_platform() -> bool: + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def is_mlx_training_device(device: Any) -> bool: + return ( + str(device).lower() == "mlx" + or str(device).lower().endswith(".mlx") + or getattr(device, "name", "").lower() == "mlx" + ) + + +def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool: + if device is not None: + return is_mlx_training_device(device) + return is_apple_silicon_training_platform() + + +def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: + """Build the normalized worker config shared by Studio and the CLI adapter.""" + config = { + "model_name": values["model_name"], + "project_name": values.get("project_name"), + "training_type": values.get("training_type", "LoRA/QLoRA"), + "hf_token": values.get("hf_token", ""), + "load_in_4bit": values.get("load_in_4bit", True), + "max_seq_length": values.get("max_seq_length", 2048), + "vision_image_size": values.get("vision_image_size"), + "hf_dataset": values.get("hf_dataset", ""), + "local_datasets": values.get("local_datasets"), + "local_eval_datasets": values.get("local_eval_datasets"), + "format_type": values.get("format_type", ""), + "subset": values.get("subset"), + "train_split": values.get("train_split", "train"), + "eval_split": values.get("eval_split"), + "eval_steps": values.get("eval_steps", 0.00), + "dataset_streaming": values.get("dataset_streaming", False), + "dataset_slice_start": values.get("dataset_slice_start"), + "dataset_slice_end": values.get("dataset_slice_end"), + "custom_format_mapping": values.get("custom_format_mapping"), + "is_dataset_image": values.get("is_dataset_image", False), + "is_dataset_audio": values.get("is_dataset_audio", False), + "is_embedding": values.get("is_embedding", False), + "num_epochs": values.get("num_epochs", 3), + "learning_rate": values.get("learning_rate", "2e-4"), + "embedding_learning_rate": values.get("embedding_learning_rate"), + "batch_size": values.get("batch_size", 2), + "gradient_accumulation_steps": values.get("gradient_accumulation_steps", 4), + "warmup_steps": values.get("warmup_steps"), + "warmup_ratio": values.get("warmup_ratio"), + "max_steps": values.get("max_steps", 0), + "save_steps": values.get("save_steps", 0), + "weight_decay": values.get("weight_decay", 0.001), + "max_grad_norm": values.get("max_grad_norm", 0.0), + "max_grad_value": _coerce_optional_nonneg_float( + "max_grad_value", values.get("max_grad_value") + ), + "max_grad_leaf_norm": _coerce_optional_nonneg_float( + "max_grad_leaf_norm", values.get("max_grad_leaf_norm") + ), + "cast_norm_output_to_input_dtype": _coerce_optional_bool( + values.get("cast_norm_output_to_input_dtype"), True + ), + "random_seed": _coerce_seed(values.get("random_seed")), + "packing": values.get("packing", False), + "optim": values.get("optim", "adamw_8bit"), + "lr_scheduler_type": values.get("lr_scheduler_type", "linear"), + "use_lora": values.get("use_lora", True), + "lora_r": values.get("lora_r", 16), + "lora_alpha": values.get("lora_alpha", 16), + "lora_dropout": values.get("lora_dropout", 0.0), + "target_modules": values.get("target_modules"), + "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), + "use_rslora": values.get("use_rslora", False), + "use_loftq": values.get("use_loftq", False), + "train_on_completions": values.get("train_on_completions", False), + "finetune_vision_layers": values.get("finetune_vision_layers", True), + "finetune_language_layers": values.get("finetune_language_layers", True), + "finetune_attention_modules": values.get("finetune_attention_modules", True), + "finetune_mlp_modules": values.get("finetune_mlp_modules", True), + "enable_wandb": values.get("enable_wandb", False), + "wandb_token": values.get("wandb_token"), + "wandb_project": values.get("wandb_project", "unsloth-training"), + "enable_tensorboard": values.get("enable_tensorboard", False), + "tensorboard_dir": values.get("tensorboard_dir", "runs"), + "resume_from_checkpoint": values.get("resume_from_checkpoint"), + "trust_remote_code": values.get("trust_remote_code", False), + "approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"), + "subject": values.get("subject"), + "gpu_ids": values.get("gpu_ids"), + "s3_config": values.get("s3_config"), + "disable_xet": values.get("disable_xet", False), + } + for key in ("output_dir", "allow_external_output_dir"): + if key in values: + config[key] = values.get(key) + if config["training_type"] == "Full Finetuning": + config["load_in_4bit"] = False + return config + + _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") @@ -133,7 +236,7 @@ def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" -def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: +def _cleanup_cancelled_checkpoints(output_dir: Union[str, os.PathLike]) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. Completed ``checkpoint-/`` dirs survive. Symlinked output_dir / children @@ -183,7 +286,7 @@ PLOT_HEIGHT = 3.5 @dataclass class TrainingProgress: - """Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules.""" + """Shared training progress payload for Studio and backend-aware trainers.""" epoch: float = 0 step: int = 0 @@ -200,6 +303,423 @@ class TrainingProgress: num_tokens: Optional[int] = None eval_loss: Optional[float] = None peak_memory_gb: Optional[float] = None + output_dir: Optional[str] = None + + +class _MLXTrainerAdapter: + """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path.""" + + def __init__(self): + self.model = None + self.tokenizer = None + self.trainer = None + self.training_thread = None + self.training_progress = TrainingProgress() + self.progress_callbacks: list[Callable[[TrainingProgress], None]] = [] + self.is_training = False + self.should_stop = False + self.save_on_stop = True + self.load_in_4bit = True + self.output_dir = None + + self.is_cpt = False + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model_name = None + self.max_seq_length = None + + self._model_config: dict[str, Any] = {} + self._peft_config: dict[str, Any] = {} + self._dataset_config: dict[str, Any] = {} + self._event_queue: Optional[queue.Queue] = None + self._stop_queue: Optional[queue.Queue] = None + self._pump_thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + + def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None: + try: + from utils.transformers_version import activate_transformers_for_subprocess + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc)) + + def add_progress_callback(self, callback: Callable[[TrainingProgress], None]): + self.progress_callbacks.append(callback) + + def _update_progress(self, **kwargs): + with self._lock: + for key, value in kwargs.items(): + if hasattr(self.training_progress, key): + setattr(self.training_progress, key, value) + progress = self.training_progress + for callback in self.progress_callbacks: + try: + callback(progress) + except Exception: + pass + + def load_model( + self, + model_name: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + is_dataset_image: bool = False, + is_dataset_audio: bool = False, + trust_remote_code: bool = False, + full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, + ) -> bool: + self.model_name = model_name + self.max_seq_length = max_seq_length + self.load_in_4bit = load_in_4bit + self._audio_type = None + self._activate_transformers_for_model(model_name, hf_token) + try: + from utils.models import detect_audio_type, is_vision_model + + self._audio_type = detect_audio_type(model_name, hf_token) + if self._audio_type == "audio_vlm": + self.is_audio = False + self.is_audio_vlm = bool(is_dataset_audio) + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False + vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and bool(is_dataset_image) + except Exception as exc: + logger.warning("MLX trainer adapter model type detection failed", error = str(exc)) + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model = object() + self.tokenizer = object() + self._model_config = { + "model_name": model_name, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "hf_token": hf_token or "", + "is_dataset_image": bool(is_dataset_image), + "is_dataset_audio": bool(is_dataset_audio), + "trust_remote_code": bool(trust_remote_code), + "gpu_ids": gpu_ids, + } + self._update_progress( + is_training = False, + is_completed = False, + error = None, + step = 0, + loss = 0.0, + epoch = 0, + status_message = f"Queued MLX model load: {model_name}", + ) + return True + + def prepare_model_for_training( + self, + use_lora: bool = True, + finetune_vision_layers: bool = True, + finetune_language_layers: bool = True, + finetune_attention_modules: bool = True, + finetune_mlp_modules: bool = True, + target_modules: Optional[Union[list, str]] = None, + lora_r: int = 16, + lora_alpha: int = 16, + lora_dropout: float = 0.0, + use_gradient_checkpointing: Union[str, bool] = "unsloth", + use_rslora: bool = False, + use_loftq: bool = False, + ) -> bool: + self._peft_config = { + "use_lora": bool(use_lora), + "lora_r": lora_r, + "lora_alpha": lora_alpha, + "lora_dropout": lora_dropout, + "target_modules": target_modules, + "gradient_checkpointing": use_gradient_checkpointing, + "use_rslora": bool(use_rslora), + "use_loftq": bool(use_loftq), + "finetune_vision_layers": bool(finetune_vision_layers), + "finetune_language_layers": bool(finetune_language_layers), + "finetune_attention_modules": bool(finetune_attention_modules), + "finetune_mlp_modules": bool(finetune_mlp_modules), + } + self._update_progress(status_message = "Queued MLX training setup") + return True + + def load_and_format_dataset( + self, + dataset_source: Optional[str], + format_type: str = "auto", + local_datasets: Optional[list[str]] = None, + local_eval_datasets: Optional[list[str]] = None, + custom_format_mapping: Optional[dict[str, Any]] = None, + subset: Optional[str] = None, + train_split: str = "train", + eval_split: Optional[str] = None, + dataset_streaming: bool = False, + eval_steps: float = 0.00, + dataset_slice_start: Optional[int] = None, + dataset_slice_end: Optional[int] = None, + is_cpt: bool = False, + s3_config: dict = None, + ) -> Optional[tuple]: + self._dataset_config = { + "hf_dataset": dataset_source or "", + "local_datasets": local_datasets, + "local_eval_datasets": local_eval_datasets, + "format_type": format_type or "", + "custom_format_mapping": custom_format_mapping, + "subset": subset, + "train_split": train_split or "train", + "eval_split": eval_split, + "dataset_streaming": bool(dataset_streaming), + "eval_steps": eval_steps or 0.0, + "dataset_slice_start": dataset_slice_start, + "dataset_slice_end": dataset_slice_end, + "s3_config": s3_config, + } + self.is_cpt = bool(is_cpt) + self._update_progress(status_message = "Queued MLX dataset load") + return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None) + + def start_training( + self, + dataset = None, + eval_dataset = None, + **training_args, + ) -> bool: + if self.is_training and self.training_thread and self.training_thread.is_alive(): + return False + if self._pump_thread and self._pump_thread.is_alive(): + self._pump_thread.join(timeout = 2.0) + if self._pump_thread.is_alive(): + self._update_progress(error = "Previous training event pump is still finalizing") + return False + if not self._model_config: + self._update_progress(error = "Model not loaded") + return False + if not self._dataset_config: + self._update_progress(error = "Dataset not loaded") + return False + if self.is_cpt: + self._update_progress( + error = "Continued Pretraining is not supported for MLX training yet.", + is_training = False, + is_completed = False, + ) + return False + + config = self._build_worker_config(training_args) + event_queue = queue.Queue() + stop_queue = queue.Queue() + self._event_queue = event_queue + self._stop_queue = stop_queue + self.should_stop = False + self.is_training = True + self.training_progress = TrainingProgress( + is_training = True, + status_message = "Initializing MLX training...", + ) + + self.training_thread = threading.Thread( + target = self._run_training_thread, + args = (config, event_queue, stop_queue), + daemon = True, + ) + self._pump_thread = threading.Thread( + target = self._pump_events, + args = (event_queue, self.training_thread), + daemon = True, + ) + self.training_thread.start() + self._pump_thread.start() + return True + + def _build_worker_config(self, training_args: dict[str, Any]) -> dict[str, Any]: + peft = { + "use_lora": True, + "lora_r": 16, + "lora_alpha": 16, + "lora_dropout": 0.0, + "target_modules": None, + "gradient_checkpointing": "unsloth", + "use_rslora": False, + "use_loftq": False, + "finetune_vision_layers": True, + "finetune_language_layers": True, + "finetune_attention_modules": True, + "finetune_mlp_modules": True, + **self._peft_config, + } + output_dir = training_args.get("output_dir") + if output_dir: + output_dir = os.path.abspath(os.path.expanduser(str(output_dir))) + values = { + **self._model_config, + **self._dataset_config, + **training_args, + "training_type": ( + "Continued Pretraining" + if self.is_cpt + else "LoRA/QLoRA" + if peft["use_lora"] + else "Full Finetuning" + ), + **peft, + "output_dir": output_dir, + "allow_external_output_dir": bool(output_dir), + } + config = _build_training_worker_config(values) + config["resolved_gpu_ids"] = None + config["gpu_selection"] = None + return config + + def _run_training_thread( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + try: + self._run_mlx_worker(config, event_queue, stop_queue) + except Exception as exc: + if event_queue is not None: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def _run_mlx_worker( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + from .worker import run_mlx_training_process + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + ) + + def _pump_events(self, event_queue: queue.Queue, training_thread: threading.Thread): + while True: + event = None + try: + event = event_queue.get(timeout = 0.25) + except queue.Empty: + pass + if event is not None: + self._handle_event(event) + continue + if not training_thread.is_alive(): + self._drain_events(event_queue) + with self._lock: + if self.training_progress.is_training: + self.training_progress.is_training = False + if self.should_stop: + self.training_progress.status_message = "Training stopped." + elif ( + not self.training_progress.error + and not self.training_progress.is_completed + ): + self.training_progress.error = "Training process exited unexpectedly" + self.is_training = False + self._event_queue = None + self._stop_queue = None + return + + def _drain_events(self, event_queue: Optional[queue.Queue] = None): + event_queue = event_queue or self._event_queue + if event_queue is None: + return + while True: + try: + self._handle_event(event_queue.get_nowait()) + except queue.Empty: + return + + def _handle_event(self, event: dict[str, Any]): + etype = event.get("type") + if etype == "status": + self._update_progress( + status_message = event.get("status_message") or event.get("message") or "" + ) + return + if etype == "progress": + self._update_progress( + step = event.get("step", self.training_progress.step), + epoch = event.get("epoch", self.training_progress.epoch), + loss = event.get("loss", self.training_progress.loss), + learning_rate = event.get("learning_rate", self.training_progress.learning_rate), + total_steps = event.get("total_steps", self.training_progress.total_steps), + elapsed_seconds = event.get( + "elapsed_seconds", + self.training_progress.elapsed_seconds, + ), + eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds), + grad_norm = event.get("grad_norm", self.training_progress.grad_norm), + num_tokens = event.get("num_tokens", self.training_progress.num_tokens), + eval_loss = event.get("eval_loss", self.training_progress.eval_loss), + peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb), + ) + return + if etype == "complete": + status_message = event.get("status_message") or "Training completed" + output_dir = event.get("output_dir") + was_cancelled = self.should_stop or status_message.strip().lower() in { + "training cancelled", + "training stopped", + } + self.output_dir = output_dir + self._update_progress( + is_training = False, + is_completed = not was_cancelled, + error = None, + status_message = status_message, + output_dir = output_dir, + ) + self.is_training = False + return + if etype == "error": + self._update_progress( + is_training = False, + is_completed = False, + error = event.get("error") or event.get("message") or "Training failed", + ) + self.is_training = False + return + + def stop_training(self, save: bool = True): + self.should_stop = True + self.save_on_stop = bool(save) + if self._stop_queue is not None: + self._stop_queue.put({"type": "stop", "save": save}) + status_message = ( + "Stopping training and saving checkpoint..." if save else "Cancelling training..." + ) + self._update_progress(status_message = status_message) + return True + + def get_training_progress(self) -> TrainingProgress: + pump_thread = self._pump_thread + training_thread = self.training_thread + if ( + pump_thread is not None + and pump_thread.is_alive() + and (training_thread is None or not training_thread.is_alive()) + and threading.current_thread() is not pump_thread + ): + pump_thread.join(timeout = 5.0) + if pump_thread is None or not pump_thread.is_alive(): + self._drain_events() + with self._lock: + return replace(self.training_progress) + + +def create_mlx_trainer_adapter(*args, **kwargs): + return _MLXTrainerAdapter(*args, **kwargs) class TrainingBackend: @@ -296,86 +816,7 @@ class TrainingBackend: # treat this fresh setup as a recoverable death. self._pump_running = False - # Build config dict for the subprocess - config = { - "model_name": kwargs["model_name"], - "project_name": kwargs.get("project_name"), - "training_type": kwargs.get("training_type", "LoRA/QLoRA"), - "hf_token": kwargs.get("hf_token", ""), - "load_in_4bit": kwargs.get("load_in_4bit", True), - "max_seq_length": kwargs.get("max_seq_length", 2048), - "vision_image_size": kwargs.get("vision_image_size"), - "hf_dataset": kwargs.get("hf_dataset", ""), - "local_datasets": kwargs.get("local_datasets"), - "local_eval_datasets": kwargs.get("local_eval_datasets"), - "format_type": kwargs.get("format_type", ""), - "subset": kwargs.get("subset"), - "train_split": kwargs.get("train_split", "train"), - "eval_split": kwargs.get("eval_split"), - "eval_steps": kwargs.get("eval_steps", 0.00), - "dataset_streaming": kwargs.get("dataset_streaming", False), - "dataset_slice_start": kwargs.get("dataset_slice_start"), - "dataset_slice_end": kwargs.get("dataset_slice_end"), - "custom_format_mapping": kwargs.get("custom_format_mapping"), - "is_dataset_image": kwargs.get("is_dataset_image", False), - "is_dataset_audio": kwargs.get("is_dataset_audio", False), - "is_embedding": kwargs.get("is_embedding", False), - "num_epochs": kwargs.get("num_epochs", 3), - "learning_rate": kwargs.get("learning_rate", "2e-4"), - "embedding_learning_rate": kwargs.get("embedding_learning_rate"), - "batch_size": kwargs.get("batch_size", 2), - "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), - "warmup_steps": kwargs.get("warmup_steps"), - "warmup_ratio": kwargs.get("warmup_ratio"), - "max_steps": kwargs.get("max_steps", 0), - "save_steps": kwargs.get("save_steps", 0), - "weight_decay": kwargs.get("weight_decay", 0.001), - "max_grad_norm": kwargs.get("max_grad_norm", 0.0), - "max_grad_value": _coerce_optional_nonneg_float( - "max_grad_value", kwargs.get("max_grad_value") - ), - "max_grad_leaf_norm": _coerce_optional_nonneg_float( - "max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm") - ), - "cast_norm_output_to_input_dtype": _coerce_optional_bool( - kwargs.get("cast_norm_output_to_input_dtype"), True - ), - # MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises). - "random_seed": _coerce_seed(kwargs.get("random_seed")), - "packing": kwargs.get("packing", False), - "optim": kwargs.get("optim", "adamw_8bit"), - "lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"), - "use_lora": kwargs.get("use_lora", True), - "lora_r": kwargs.get("lora_r", 16), - "lora_alpha": kwargs.get("lora_alpha", 16), - "lora_dropout": kwargs.get("lora_dropout", 0.0), - "target_modules": kwargs.get("target_modules"), - "gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"), - "use_rslora": kwargs.get("use_rslora", False), - "use_loftq": kwargs.get("use_loftq", False), - "train_on_completions": kwargs.get("train_on_completions", False), - "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), - "finetune_language_layers": kwargs.get("finetune_language_layers", True), - "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), - "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), - "enable_wandb": kwargs.get("enable_wandb", False), - "wandb_token": kwargs.get("wandb_token"), - "wandb_project": kwargs.get("wandb_project", "unsloth-training"), - "enable_tensorboard": kwargs.get("enable_tensorboard", False), - "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), - "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), - "trust_remote_code": kwargs.get("trust_remote_code", False), - "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), - "subject": kwargs.get("subject"), - "gpu_ids": kwargs.get("gpu_ids"), - "s3_config": kwargs.get("s3_config"), - # Flipped to True only by the HTTP-fallback respawn after a stall. - "disable_xet": kwargs.get("disable_xet", False), - } - - # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. - if config["training_type"] == "Full Finetuning": - config["load_in_4bit"] = False + config = _build_training_worker_config(kwargs) # Split GPU validation from placement around the VRAM hook: # * Explicit gpu_ids are validated here (raises -> the route returns 400 @@ -401,7 +842,7 @@ class TrainingBackend: ) defer_auto_selection = False - if _hw.DEVICE == _hw.DeviceType.MLX: + if should_use_mlx_training_backend(device = _hw.DEVICE): config["resolved_gpu_ids"] = None config["gpu_selection"] = None elif gpu_ids: @@ -1022,17 +1463,22 @@ class TrainingBackend: self._progress.is_training = True elif etype == "complete": - self._progress.is_training = False - self._progress.is_completed = True - self._output_dir = event.get("output_dir") msg = event.get("status_message", "Training completed") + stopped = self._should_stop or msg.strip().lower() in { + "training cancelled", + "training stopped", + } + self._progress.is_training = False + self._progress.is_completed = not stopped + self._output_dir = event.get("output_dir") + self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: db_action = "create_and_finalize" else: db_action = "finalize" db_action_kwargs = { - "status": "stopped" if self._should_stop else "completed", + "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, } diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 17dc1299ca..5fb8fc2bb6 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1309,14 +1309,18 @@ def _normalize_mlx_studio_scheduler(value): def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: - """Resolve Studio local dataset uploads without importing the GPU trainer.""" + """Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer.""" from utils.paths import resolve_dataset_path all_files: list[str] = [] for dataset_file in file_paths or []: - file_path = ( - dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file)) - ) + dataset_path = Path(os.path.expanduser(str(dataset_file))) + if dataset_path.is_absolute(): + file_path = str(dataset_path) + elif dataset_path.exists(): + file_path = str(dataset_path.resolve()) + else: + file_path = str(resolve_dataset_path(str(dataset_file))) file_path_obj = Path(file_path) if file_path_obj.is_dir(): @@ -1355,6 +1359,58 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: raise ValueError(f"Unsupported dataset format: {files[0]}") +_MLX_WORKER_COMPLETE = "_mlx_worker_complete" + + +def _start_mlx_stop_poller(stop_queue): + import queue as _queue + import threading + + stop_save = [True] + stop_requested = [False] + trainer_ref = [None] + + def is_stop_requested(): + return stop_requested[0] + + def poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 0.25) + if msg and msg.get("type") == _MLX_WORKER_COMPLETE: + return + if msg and msg.get("type") == "stop": + stop_save[0] = msg.get("save", True) + stop_requested[0] = True + trainer = trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = poll_stop, daemon = True) + stop_thread.start() + return stop_save, stop_requested, trainer_ref, is_stop_requested, stop_thread + + +def _resolve_mlx_output_dir(config, model_name): + from utils.paths import resolve_output_dir, default_run_dir_name + + output_dir = config.get("output_dir", "") + if not output_dir: + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + return str(resolve_output_dir(output_dir)) + if config.get("allow_external_output_dir"): + output_path = Path(output_dir).expanduser() + if not output_path.is_absolute(): + output_path = Path.cwd() / output_path + return str(output_path.resolve()) + return str(resolve_output_dir(output_dir)) + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -1363,8 +1419,6 @@ def _run_mlx_training(event_queue, stop_queue, config): """ import time import math - import threading - import queue as _queue from pathlib import Path def _send(event_type, **kwargs): @@ -1374,31 +1428,9 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) - _stop_save = [True] - _stop_requested = [False] - _trainer_ref = [None] - - def _is_stop_requested(): - return _stop_requested[0] - - def _poll_stop(): - while True: - try: - msg = stop_queue.get(timeout = 1.0) - if msg and msg.get("type") == "stop": - _stop_save[0] = msg.get("save", True) - _stop_requested[0] = True - trainer = _trainer_ref[0] - if trainer is not None: - trainer.stop_requested = True - return - except _queue.Empty: - continue - except (EOFError, OSError): - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() + _stop_save, _stop_requested, _trainer_ref, _is_stop_requested, _stop_thread = ( + _start_mlx_stop_poller(stop_queue) + ) _send("status", status_message = "Loading MLX libraries...") @@ -1699,6 +1731,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column). format_type = config.get("format_type", "") custom_format_mapping = config.get("custom_format_mapping") + dataset_final_format = "" try: from utils.datasets import format_and_template_dataset def _fmt_progress(status_message = "", **_kw): @@ -1764,6 +1797,7 @@ def _run_mlx_training(event_queue, stop_queue, config): ) if info.get("success", True): dataset = info.get("dataset", dataset) + dataset_final_format = str(info.get("final_format", "") or "").lower() if eval_dataset is not None: ev = format_and_template_dataset( eval_dataset, @@ -1804,21 +1838,14 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 5. Build output dir ── # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it - from utils.paths import resolve_output_dir, ensure_dir + from utils.paths import ensure_dir - output_dir = config.get("output_dir", "") - if not output_dir: - output_dir = build_default_output_dir_name( - model_name, - config.get("project_name"), - ) - output_dir = str(resolve_output_dir(output_dir)) + output_dir = _resolve_mlx_output_dir(config, model_name) ensure_dir(Path(output_dir)) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1: - # Studio sometimes sends fraction-of-total-steps eval_steps_val = max(1, int(eval_steps_val * max_steps)) else: eval_steps_val = int(eval_steps_val) @@ -1869,6 +1896,9 @@ def _run_mlx_training(event_queue, stop_queue, config): eval_steps = eval_steps_val, ) + # Also gates the masking skip below, so defined outside the feature-detect block. + raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" + # Feature-detect optional fields so this PR works without the paired zoo bump. _supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {}) if "cast_norm_output_to_input_dtype" in _supported_fields: @@ -1882,7 +1912,6 @@ def _run_mlx_training(event_queue, stop_queue, config): if "max_grad_leaf_norm" in _supported_fields: mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm if "append_eos" in _supported_fields: - raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" # Studio SFT formatting owns rendered examples; raw/CPT text still # needs MLX to append EOS like the CUDA raw-text path. mlx_config_kwargs["append_eos"] = bool(raw_text_mode) @@ -1903,29 +1932,27 @@ def _run_mlx_training(event_queue, stop_queue, config): _send("eval_configured") # ── 7. Apply train_on_responses_only if requested ── - if config.get("train_on_completions", False): + # Auto-detect markers from the chat template first, manual table as + # fallback. Mirror the CUDA skips: raw/CPT text has no chat turns and + # Alpaca-rendered text lacks the chat markers. Also check the resolved + # format, since format_type="auto" can land on alpaca or raw text. + if ( + config.get("train_on_completions", False) + and not raw_text_mode + and format_type != "alpaca" + and dataset_final_format not in ("alpaca", "raw_text") + ): _send("status", status_message = "Configuring response-only training...") - try: - from utils.datasets import ( - MODEL_TO_TEMPLATE_MAPPER, - TEMPLATE_TO_RESPONSES_MAPPER, - ) - - template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower()) - markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None - if markers: - trainer = train_on_responses_only( - trainer, - instruction_part = markers["instruction"], - response_part = markers["response"], - ) - else: - _send( - "status", - status_message = f"train_on_completions skipped (no template for {model_name})", - ) - except Exception as e: - _send("status", status_message = f"train_on_completions failed: {e}") + # No catch: the helper handles detection failures and double misses, so + # an exception here is a real masking failure that must fail the run, + # not silently train on full sequences. + from utils.datasets.completion_masking import apply_completion_masking + trainer, _masking_applied = apply_completion_masking( + trainer, + model_name, + train_on_responses_only, + notify = lambda level, message: _send("status", status_message = message), + ) # ── 8. Setup wandb / tensorboard ── wandb_run = None @@ -2043,12 +2070,27 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint = resume_from_checkpoint) + _save_model = trainer.save_model + + def _skip_internal_final_save(*args, **kwargs): + raise ValueError("worker owns final save") + + trainer.save_model = _skip_internal_final_save + try: + trainer.train(resume_from_checkpoint = resume_from_checkpoint) + finally: + trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested and not _stop_save[0]: - # User clicked "Cancel" (save=False) — skip saving - _send("complete", output_dir = None, status_message = "Training cancelled") + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: _send("status", status_message = "Saving model...") mx.synchronize() @@ -2067,6 +2109,79 @@ def _run_mlx_training(event_queue, stop_queue, config): pass +def _is_current_process_apple_silicon() -> bool: + import platform + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def run_mlx_training_process( + *, + event_queue: Any, + stop_queue: Any, + config: dict, + transformers_activated: bool = False, +) -> None: + """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter.""" + model_name = config["model_name"] + + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + + if not transformers_activated: + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE != _hw.DeviceType.MLX: + event_queue.put( + { + "type": "error", + "error": "MLX training requires Apple Silicon with the MLX backend available.", + "stack": "", + "ts": time.time(), + } + ) + return + + if config.get("is_dataset_audio"): + event_queue.put( + { + "type": "error", + "error": "Audio dataset training is not yet supported on Apple Silicon.", + "stack": "", + "ts": time.time(), + } + ) + return + + try: + try: + _run_mlx_training(event_queue, stop_queue, config) + finally: + try: + stop_queue.put({"type": _MLX_WORKER_COMPLETE}) + except (EOFError, OSError, ValueError): + pass + except Exception as exc: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None: """Subprocess entrypoint. Fresh Python — no stale module state. @@ -2141,36 +2256,26 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) + from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend + + mlx_backend_requested = is_apple_silicon_training_platform() + + mlx_transformers_activated = False + if mlx_backend_requested and _is_current_process_apple_silicon(): + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + mlx_transformers_activated = True + from utils.hardware import hardware as _hw _hw.detect_hardware() - if _hw.DEVICE == _hw.DeviceType.MLX: - if config.get("is_dataset_audio"): - event_queue.put( - { - "type": "error", - "error": "Audio dataset training is not yet supported on Apple Silicon.", - "stack": "", - "ts": time.time(), - } - ) - return - # Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.) - # Must happen before any transformers/mlx-lm imports in _run_mlx_training. - # Non-fatal: fall through with whatever version is installed, but log - # the failure instead of swallowing it (issue #6103). - _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) - try: - _run_mlx_training(event_queue, stop_queue, config) - except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } - ) + if mlx_backend_requested or should_use_mlx_training_backend(device = _hw.DEVICE): + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + transformers_activated = mlx_transformers_activated, + ) return # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -2693,7 +2798,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer, TrainingProgress + from core.training.training import TrainingProgress + from core.training.trainer import UnslothTrainer from utils.paths import ( ensure_dir, resolve_output_dir, diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index f05d8359ec..44701c0b64 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -105,6 +105,10 @@ def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id): assert paths.is_valid_repo_id(repo_id) +def test_repo_id_validation_accepts_max_length_namespaced_repo(): + assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}") + + @pytest.mark.parametrize( "repo_id", [ @@ -121,6 +125,48 @@ def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id): assert not paths.is_valid_repo_id(repo_id) +def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + + path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M") + + assert path is not None + assert path.name == "models--owner--repo--variant--q4_k_m.json" + + +@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64]) +def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + repo_id = f"{'a' * 96}/{'b' * 96}" + + assert paths.is_valid_repo_id(repo_id) + assert download_manifest.write_cancel_marker("model", repo_id, variant, "http") + assert download_manifest.write_manifest( + "model", + repo_id, + variant, + [download_manifest.ExpectedFile(path = "model.gguf", size = 1)], + "http", + ) + + marker_path = state_dir.marker_path("model", repo_id, variant) + manifest_path = state_dir.manifest_path("model", repo_id, variant) + + assert marker_path is not None + assert manifest_path is not None + assert "--sha256-" in marker_path.name + assert len(marker_path.name.encode("utf-8")) <= 255 + assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255 + assert download_manifest.has_cancel_marker("model", repo_id, variant) + assert download_manifest.read_manifest("model", repo_id, variant) is not None + assert list(download_manifest.iter_variant_markers("model", repo_id)) == [ + (variant, marker_path) + ] + assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [ + (variant, manifest_path) + ] + + class _RecordingLogger: def __init__(self): self.warnings = [] diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 2abdb0fb79..18daa4f84e 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]: def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: """The separate MTP drafter to fetch with every variant: the repo-root ``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery - (Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename - prefix, first in sort order) so download and load resolve the same file; - the higher-precision ``MTP/`` subdir copies are for explicit selection and - are not auto-fetched. None for repos with the head baked into the main - GGUF (Qwen).""" + (Gemma 4). Same pick as the loader's drafter resolution (root-level + ``mtp-`` prefix, first in sort order) so download and load resolve the same + file; the higher-precision ``MTP/`` subdir copies are for explicit + selection and are not auto-fetched. None for repos with the head baked into + the main GGUF (Qwen).""" + # Root-level only: the MTP/ subdir copies now share the mtp- prefix too. candidates = sorted( ( s for s in siblings - if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-") + if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-") ), key = lambda s: getattr(s, "rfilename"), ) diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py index afcb0b41dc..5435202565 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -181,15 +181,20 @@ def is_valid_repo_id(repo_id: str) -> bool: """Validate Hugging Face ``repo_name`` or ``namespace/repo_name`` IDs.""" if not repo_id or repo_id != repo_id.strip(): return False - if len(repo_id) > _MAX_REPO_ID_LENGTH or repo_id.endswith(".git"): + if repo_id.endswith(".git"): return False if "--" in repo_id or ".." in repo_id: return False segments = repo_id.split("/") if len(segments) not in (1, 2): return False + # Match huggingface_hub.validate_repo_id: the 96-char limit applies per + # segment (repo name / namespace), not to the whole "namespace/repo_name" + # string, so long-but-valid repo names are not falsely rejected. return all( - segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None + segment not in ("", ".", "..") + and len(segment) <= _MAX_REPO_ID_LENGTH + and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None for segment in segments ) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index a304477a3d..183e934724 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -11,8 +11,9 @@ cache lifecycle. Two subdirectories: manifests/ .json per-download expected-files manifest cancelled/ .json per-download cancel marker -The ```` mirrors HF's cache dir naming so a state file can be -eyeballed next to the on-disk repo it describes: +The ```` mirrors HF's cache dir naming while the resulting manifest, +cancel-marker, and atomic-write temp filenames fit common filesystem basename +limits. Very long repo IDs use a stable hash in the state key: models---- full snapshot models------variant-- GGUF variant @@ -49,6 +50,11 @@ _MANIFESTS_SUBDIR = "manifests" _CANCELLED_SUBDIR = "cancelled" _WORKERS_SUBDIR = "workers" _SAFE_VARIANT_FRAGMENT = re.compile(r"^[a-z0-9._-]{1,64}$") +_MAX_STATE_BASENAME_BYTES = 255 +_STATE_EXTENSION = ".json" +# _atomic_write_json writes "..tmp-<8hex>" beside the final file. +_ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8 +_MAX_VARIANT_FRAGMENT_LENGTH = 64 def state_root() -> Optional[Path]: @@ -84,16 +90,35 @@ def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str: return f"{repo_type}s--{repo_id.replace('/', '--')}".lower() +def _filename_bytes(name: str) -> int: + return len(name.encode("utf-8")) + + +def _state_filename_fits(entry_key: str) -> bool: + filename = f"{entry_key}{_STATE_EXTENSION}" + return _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD <= _MAX_STATE_BASENAME_BYTES + + +def _state_repo_key(repo_type: RepoType, repo_id: str) -> str: + base = repo_cache_basename(repo_type, repo_id) + variant_prefix = f"{base}--variant--" + longest_variant_key = f"{variant_prefix}{'x' * _MAX_VARIANT_FRAGMENT_LENGTH}" + if _state_filename_fits(longest_variant_key): + return base + digest = hashlib.sha256(base.encode("utf-8")).hexdigest()[:32] + return f"{repo_type}s--sha256-{digest}" + + def variant_filename_prefix(repo_type: RepoType, repo_id: str) -> str: """Lowercased prefix every variant-keyed state file for this repo shares. The single source the download_manifest enumerators match against, so the scheme in :func:`_entry_key` cannot drift from them silently.""" - return f"{repo_cache_basename(repo_type, repo_id)}--variant--" + return f"{_state_repo_key(repo_type, repo_id)}--variant--" def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str: - base = repo_cache_basename(repo_type, repo_id) + base = _state_repo_key(repo_type, repo_id) if not variant: return base normalized_variant = variant.strip().lower() diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0f27b695fe..53b0f14b09 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1533,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel): tool_use_id: str content: Union[str, list] = "" + @field_validator("content", mode = "before") + @classmethod + def _coerce_null_content(cls, v): + # Some clients send null content for an empty tool result; the str|list + # union would 400 on it, so treat null as "". + return "" if v is None else v + + +# Block types the converter translates explicitly. Anything else (thinking / +# redacted_thinking, a provider block a resumed session replays, or a future type) +# is accepted as an unknown block and dropped by the converter, rather than 400-ing +# the whole request on strict validation. +_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"}) + + +class AnthropicUnknownBlock(BaseModel): + type: str + model_config = {"extra": "allow"} + + @field_validator("type") + @classmethod + def _only_unknown_types(cls, v): + # Known types parse as their typed models above (so a malformed known block + # still fails cleanly); this fallback only catches the rest. + if v in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError("known block type handled by its typed model") + return v + AnthropicContentBlock = Union[ AnthropicTextBlock, AnthropicImageBlock, AnthropicToolUseBlock, AnthropicToolResultBlock, + AnthropicUnknownBlock, ] @@ -1583,6 +1612,40 @@ class AnthropicMessage(BaseModel): role: Literal["user", "assistant"] content: Union[str, list[AnthropicContentBlock]] + @model_validator(mode = "before") + @classmethod + def _normalize_content(cls, data): + # Role-aware leniency that never silently drops real user input: + # - assistant: a resumed tool-only turn's null content -> "" (str|list would + # 400 on null; "" keeps the converter's `for block in content` safe). + # Unknown blocks (thinking / future types) validate via + # AnthropicUnknownBlock and are dropped by the converter. + # - user: keep strict. Null user content stays None so str|list rejects it + # (400) rather than forwarding an empty prompt; and reject block types the + # converter cannot translate, since it silently skips unknown user blocks + # -- a user turn made only of them would validate yet send no content + # (silent data loss). + if not isinstance(data, dict): + return data + content = data.get("content") + if data.get("role") == "assistant": + # Coerce only an explicit null (resumed tool-only turn). A missing + # content key stays malformed so the required-field check still 400s. + if "content" in data and content is None: + return {**data, "content": ""} + return data + if isinstance(content, list): + for block in content: + btype = ( + block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + ) + # Guard the value: a non-string type is unsupported too, and a + # membership test on an unhashable value would raise TypeError + # (escaping as a 500 instead of a clean 400). + if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError(f"unsupported content block type {btype!r} in a user message") + return data + class AnthropicTool(BaseModel): # Client tools have input_schema; server tools may only have type/name. diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 57a291291e..a5b75b7335 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -10,6 +10,7 @@ import binascii import json import os import re +import shutil from itertools import islice from pathlib import Path from typing import Any @@ -59,6 +60,9 @@ UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"} SEED_UPLOAD_DIR = seed_uploads_root() UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root() _SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$") +# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...) +# never match: those directories can be shared by several recipes. +_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$") def _validate_safe_id(value: str, label: str) -> str: @@ -580,6 +584,39 @@ async def remove_unstructured_file(block_id: str, file_id: str): return {"status": "ok"} +@router.delete("/seed/unstructured-block/{block_id}") +async def remove_unstructured_block(block_id: str): + """Delete a block's upload directory; files on disk still count toward its quota. + + Only uid-namespaced directories may be bulk-deleted: they have exactly one + owning block. Legacy node-id directories (n1, ...) can be shared by other + recipes, so they are managed file-by-file instead. + """ + _validate_safe_id(block_id, "block_id") + if not _UPLOAD_UID_RE.match(block_id): + raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted") + + block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve() + if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()): + raise HTTPException(400, "Invalid block_id: outside upload root") + if not block_dir.exists(): + return {"status": "ok", "deleted": False} + + try: + shutil.rmtree(block_dir) + except OSError as exc: + raise log_and_http_error( + exc, + 500, + "failed to delete uploaded files", + event = "data_recipe.seed.unstructured_block_delete_failed", + log = logger, + ) from exc + if block_dir.exists(): + raise HTTPException(500, "failed to delete uploaded files") + return {"status": "ok", "deleted": True} + + @router.post("/seed/inspect-upload", response_model = SeedInspectResponse) def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: if payload.file_ids is not None: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5332037e0d..ec5d309810 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -13,7 +13,7 @@ from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response from starlette.requests import ClientDisconnect -from typing import Any, List, Optional, Union +from typing import Any, Callable, List, Optional, Union import json import httpx from loggers import get_logger @@ -28,6 +28,16 @@ import re as _re from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +from core.inference.llama_admission import ( + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + LlamaAdmissionLease, + LlamaAdmissionQueueFull, + LlamaAdmissionReservation, + LlamaAdmissionTimeout, + get_llama_admission_queue, + llama_admission_config_from_env, +) def _positive_int_or_none(value: Any) -> Optional[int]: @@ -40,6 +50,43 @@ def _positive_int_or_none(value: Any) -> Optional[int]: return value_int if value_int > 0 else None +def _nonnegative_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int >= 0 else None + + +_MLX_MPI_DISTRIBUTED_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) + + +def _mlx_distributed_launch_detected() -> bool: + if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None: + world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE")) + if world_size is not None and world_size > 1: + return True + return bool( + os.environ.get("MLX_HOSTFILE") + or os.environ.get("MLX_IBV_DEVICES") + or os.environ.get("MLX_JACCL_COORDINATOR") + or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT")) + ) + return any( + _nonnegative_int_or_none(os.environ.get(rank_env)) is not None + and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1 + for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS + ) + + def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. @@ -231,10 +278,87 @@ def _effective_max_tokens(payload): ) +_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT" + + +def _positive_float_env(env_name: str, default): + """Parse a positive float from an env var. A parseable non-positive value + returns ``None`` (0 disables the guarded feature); only unparseable or unset + values fall back to ``default``.""" + raw_value = os.environ.get(env_name) + if raw_value is None or not raw_value.strip(): + return default + try: + value = float(raw_value.strip()) + except ValueError: + return default + return value if value > 0 else None + + +def _effective_openai_max_tokens_from_values(max_tokens, max_completion_tokens = None): + """Resolve the OpenAI-compatible generation cap from raw request values. + + Prefers ``max_completion_tokens`` over the deprecated ``max_tokens``, and + returns ``None`` when both are omitted so callers keep their context-window + default (OpenAI treats an omitted cap as bounded only by the context + window). Explicit client caps pass through unchanged. + """ + + def _validate_explicit(value, param: str): + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be an integer.", + status = 400, + code = "invalid_type", + param = param, + ), + ) + # The legacy completions spec declares ``minimum: 0`` for max_tokens, + # so 0 is a valid (if degenerate) cap and only negatives are rejected. + # The chat fields never reach here with 0 (pydantic enforces ge=1). + if value < 0: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be at least 0.", + status = 400, + code = "invalid_value", + param = param, + ), + ) + return value + + max_tokens = _validate_explicit(max_tokens, "max_tokens") + max_completion_tokens = _validate_explicit(max_completion_tokens, "max_completion_tokens") + return max_completion_tokens if max_completion_tokens is not None else max_tokens + + +def _effective_openai_max_tokens(payload): + return _effective_openai_max_tokens_from_values( + getattr(payload, "max_tokens", None), + getattr(payload, "max_completion_tokens", None), + ) + + def _wants_multiple_choices(payload) -> bool: return (payload.n or 1) > 1 +def _has_openai_tool_history(messages) -> bool: + for message in messages or []: + if isinstance(message, dict): + if message.get("role") == "tool" or message.get("tool_calls"): + return True + continue + if getattr(message, "role", None) == "tool" or getattr(message, "tool_calls", None): + return True + return False + + def _raise_unsupported_openai_parameter(param: str, message: str) -> None: raise HTTPException( status_code = 400, @@ -294,6 +418,14 @@ def _openai_stream_error_chunk(exc) -> dict: return openai_error_body(_friendly_error(exc), status = 500) +def _openai_stream_error_sse(error: dict) -> str: + return f"data: {json.dumps(error)}\n\ndata: [DONE]\n\n" + + +def _openai_stream_error_sse_bytes(error: dict) -> bytes: + return _openai_stream_error_sse(error).encode("utf-8") + + def _openai_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 @@ -508,20 +640,62 @@ def _drop_parallel_tool_call_deltas(chunk) -> bool: return changed -def _cap_parallel_tool_calls_sse_line(raw_line: str) -> str: - """Drop tool_call deltas whose index >= 1 from one streamed OpenAI SSE - ``data:`` line so only the first tool call survives (parallel_tool_calls=false, - best-effort). Non-tool / unparseable payloads are returned byte-for-byte.""" - payload = raw_line[len("data: ") :] +def _add_empty_content_to_reasoning_deltas(chunk: dict) -> bool: + """Make reasoning-only deltas palatable to strict OpenAI adapters. + + Some clients built on OpenAI-compatible streams ignore or reject chunks whose + delta only contains non-standard ``reasoning_content``. Preserve that field, + but add an empty standard ``content`` member so the chunk is still a valid + text-delta shape and downstream parsers keep the stream alive. + """ + changed = False + choices = chunk.get("choices") + if not isinstance(choices, list): + return False + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") + if not isinstance(delta, dict): + continue + if "reasoning_content" in delta and "content" not in delta: + delta["content"] = "" + changed = True + return changed + + +def _normalize_openai_passthrough_sse_line( + raw_line: str, *, cap_parallel_tool_calls: bool = False +) -> str: + """Normalize one passthrough OpenAI SSE ``data:`` line before relaying. + + The function is intentionally narrow: it leaves comments, blank events, + ``[DONE]``, and unparseable upstream bytes untouched; parsed chunks are + re-serialized only when a compatibility mutation is actually required. + """ + if not raw_line.startswith("data:"): + return raw_line + # Both mutations key off JSON object keys, so a line without either quoted + # key can never change; skip the parse on the per-token common case. + if '"reasoning_content"' not in raw_line and not ( + cap_parallel_tool_calls and '"tool_calls"' in raw_line + ): + return raw_line + payload = raw_line[len("data:") :].lstrip() if payload.strip() in ("", "[DONE]"): return raw_line try: obj = json.loads(payload) except Exception: return raw_line - if not _drop_parallel_tool_call_deltas(obj): + if not isinstance(obj, dict): return raw_line - return "data: " + json.dumps(obj, separators = (",", ":")) + changed = _add_empty_content_to_reasoning_deltas(obj) + if cap_parallel_tool_calls and _drop_parallel_tool_call_deltas(obj): + changed = True + if not changed: + return raw_line + return "data: " + json.dumps(obj, separators = (",", ":"), ensure_ascii = False) def _prompt_tokens_details(upstream): @@ -538,6 +712,49 @@ def _wants_stream_usage(payload) -> bool: return bool((payload.stream_options or {}).get("include_usage")) +_OPENAI_PASSTHROUGH_TERMINAL_GRACE_S = 2.0 +_SSE_DONE_LINE = "data: [DONE]" + + +def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]: + """Classify OpenAI-compatible chat stream terminal markers. + + Some llama-server builds can emit the logical final chunk (``finish_reason``) + and optional usage chunk, then keep the HTTP stream open without sending the + OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close + the client stream promptly while preserving an optional trailing usage chunk. + """ + if not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + if data_str == "[DONE]": + return "done" + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + return _openai_passthrough_terminal_state_from_data(data) + + +def _openai_passthrough_terminal_state_from_data(data) -> Optional[str]: + """Dict-level core of ``_openai_passthrough_sse_line_terminal_state`` for + callers that already parsed the chunk (avoids a re-parse per relayed line).""" + if not isinstance(data, dict): + return None + if _monitor_openai_error_message(data): + return "error" + choices = data.get("choices") + if isinstance(choices, list): + if not choices and isinstance(data.get("usage"), dict): + return "usage" + for choice in choices: + if isinstance(choice, dict) and choice.get("finish_reason") is not None: + return "finish" + elif isinstance(data.get("usage"), dict): + return "usage" + return None + + def _openai_stream_usage_chunk( payload, completion_id, created, model_name, stream_usage, stream_timings ): @@ -604,12 +821,16 @@ def _chat_content_chunk(completion_id, created, model_name, text) -> str: def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str: - """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block).""" + """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block). + + Carries ``content: ""`` alongside, like the GGUF and passthrough paths, so + strict OpenAI adapters don't drop the reasoning-only delta. + """ return _chat_chunk_sse( completion_id, created, model_name, - delta = ChoiceDelta(reasoning_content = text), + delta = ChoiceDelta(content = "", reasoning_content = text), finish_reason = None, ) @@ -844,8 +1065,11 @@ def _llama_streaming_generation_timeout() -> httpx.Timeout: def _set_stream_response_read_timeout( - response: httpx.Response, read_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S + response: httpx.Response, read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S ) -> None: + # ``read_timeout_s = None`` clears httpx's read timeout (wait indefinitely), + # used when the stall guard is disabled so a stale first-token deadline + # can't keep timing out post-first-chunk gaps. try: timeout_ext = response.request.extensions.get("timeout") if isinstance(timeout_ext, dict): @@ -856,6 +1080,295 @@ def _set_stream_response_read_timeout( _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 +_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0 +_OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n" +_OPENAI_LLAMA_ADMISSION_POLL_S = 0.25 + + +def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int: + """Serving slots available for one local llama-server backend. + + The loaded backend is the source of truth because it may have reduced + ``--parallel`` at load time to keep the model on GPU. The app state is a + launch-intent fallback for tests and for the short window before a backend + reports its committed runtime slots. + """ + slots = _positive_int_or_none(getattr(llama_backend, "effective_parallel_slots", None)) + if slots is not None: + return slots + try: + slots = getattr(request.app.state, "llama_parallel_slots", None) + except Exception: + slots = None + return _positive_int_or_none(slots) or 1 + + +def _openai_llama_admission_reserve( + *, request: Optional[Request], llama_backend +) -> tuple[LlamaAdmissionReservation, LlamaAdmissionConfig]: + config = llama_admission_config_from_env() + capacity = _openai_llama_admission_capacity(request, llama_backend) + key = str(getattr(llama_backend, "base_url", "llama-server")) + reservation = get_llama_admission_queue(key).reserve( + capacity = capacity, + config = config, + ) + return reservation, config + + +def _openai_admission_request_path(request: Optional[Request]) -> Optional[str]: + try: + return str(request.url.path) if request is not None else None + except Exception: + return None + + +def _openai_admission_log( + event: str, + reservation: Optional[LlamaAdmissionReservation] = None, + *, + snapshot = None, + request: Optional[Request], + mode: str, + wait_started_at: Optional[float] = None, + completion_id: Optional[str] = None, + level: str = "debug", +) -> None: + if snapshot is None and reservation is not None: + snapshot = reservation.snapshot_now() + wait_ms = None + if wait_started_at is not None: + wait_ms = int(max(0.0, time.monotonic() - wait_started_at) * 1000) + log = getattr(logger, level, logger.debug) + log( + "openai admission %s: mode=%s path=%s completion_id=%s capacity=%s active=%s queued=%s wait_ms=%s", + event, + mode, + _openai_admission_request_path(request), + completion_id, + getattr(snapshot, "capacity", None), + getattr(snapshot, "active", None), + getattr(snapshot, "queued", None), + wait_ms, + ) + + +def _openai_admission_error_body(exc: Exception, *, status_code: int) -> dict: + snapshot = getattr(exc, "snapshot", None) + message = str(exc) + if snapshot is not None: + message = ( + f"{message} " + f"(active={snapshot.active}, queued={snapshot.queued}, capacity={snapshot.capacity})" + ) + return openai_error_body(message, status = status_code) + + +def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: + return HTTPException( + status_code = status_code, + detail = _openai_admission_error_body(exc, status_code = status_code), + ) + + +def _openai_admission_timeout_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionTimeout: + return LlamaAdmissionTimeout( + "Timed out waiting for an available local llama-server generation slot", + snapshot = reservation.snapshot_now(), + ) + + +def _openai_admission_cancelled_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionCancelled: + return LlamaAdmissionCancelled( + "Client disconnected before an upstream llama-server generation slot was available", + snapshot = reservation.snapshot_now(), + ) + + +async def _raise_if_openai_admission_cancelled( + reservation: LlamaAdmissionReservation, *, request: Optional[Request], cancel_event +) -> None: + if reservation.is_cancelled: + raise _openai_admission_cancelled_error(reservation) + if await _preheader_cancelled(cancel_event, request): + reservation.cancel() + raise _openai_admission_cancelled_error(reservation) + + +async def _wait_for_openai_admission_non_streaming( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +) -> LlamaAdmissionLease: + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + wait_s = _OPENAI_LLAMA_ADMISSION_POLL_S + if deadline is not None: + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + continue + if lease is not None: + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _openai_admission_wait_stream_chunks( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +): + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + keepalive_interval_s = max(0.001, config.keepalive_interval_s) + next_keepalive_at = time.monotonic() + keepalive_interval_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + now = time.monotonic() + wait_s = min(_OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001)) + if deadline is not None: + remaining_s = deadline - now + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + lease = None + if lease is not None: + yield lease + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + now = time.monotonic() + if now >= next_keepalive_at: + next_keepalive_at = now + keepalive_interval_s + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _close_openai_admitted_stream_iterator(iterator, *, cancelled: bool) -> None: + if iterator is None: + return + if cancelled: + athrow = getattr(iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + return + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + + +def _openai_compat_stream_stall_timeout(): + """Max silent gap after an OpenAI passthrough stream has produced data. + + If the socket goes silent after valid SSE data, this bounds how long the + client is kept open. Defaults to the backend-wide stall timeout so this + path stalls out like every sibling stream; set the env var to tighten it + for local serving, or to 0 to disable the guard. + """ + return _positive_float_env( + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, + _DEFAULT_STREAM_STALL_TIMEOUT_S, + ) + + +def _openai_passthrough_upstream_headers(*, llama_backend = None) -> dict: + headers = {} + auth_headers = getattr(llama_backend, "_auth_headers", None) + if isinstance(auth_headers, dict): + headers.update(auth_headers) + headers["Connection"] = "close" + return headers class _CompatSameTaskTimeout: @@ -983,7 +1496,8 @@ async def _aclose_stream_resources( """Tear down an httpx streaming generator's resources in the required order: cancel + await each watcher task, then aclose() the byte/line iterator, the response, and the client. Each step swallows its own exceptions so teardown - always completes. See _anthropic_passthrough_stream for the ordering rationale.""" + always completes; a close-time CancelledError is re-raised only after every + step has run. See _anthropic_passthrough_stream for the ordering rationale.""" for watcher in watchers: if watcher is not None: watcher.cancel() @@ -991,21 +1505,30 @@ async def _aclose_stream_resources( await watcher except (asyncio.CancelledError, Exception): pass + close_cancelled = False if iterator is not None: try: await iterator.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if resp is not None: try: await resp.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if client is not None: try: await client.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass + if close_cancelled: + raise asyncio.CancelledError() async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: @@ -1028,6 +1551,7 @@ async def _send_stream_with_preheader_cancel( req: httpx.Request, cancel_event = None, request: Optional[Request] = None, + mark_cancel_on_cancel: bool = True, ) -> Optional[httpx.Response]: if cancel_event is None and request is None: return await client.send(req, stream = True) @@ -1059,7 +1583,7 @@ async def _send_stream_with_preheader_cancel( await _stop_send_task() return None except asyncio.CancelledError: - if cancel_event is not None: + if mark_cancel_on_cancel and cancel_event is not None: cancel_event.set() await _stop_send_task() raise @@ -1078,11 +1602,19 @@ async def _aiter_llama_stream_items( request: Optional[Request] = None, first_token_deadline: Optional[float] = None, response: Optional[httpx.Response] = None, - post_first_item_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, + post_first_item_read_timeout_s: Optional[ + Union[float, Callable[[], Optional[float]]] + ] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ): if first_token_deadline is None: first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S last_item_at: Optional[float] = None + + def _post_first_timeout_s() -> Optional[float]: + if callable(post_first_item_read_timeout_s): + return post_first_item_read_timeout_s() + return post_first_item_read_timeout_s + while True: if cancel_event is not None and cancel_event.is_set(): return @@ -1103,15 +1635,14 @@ async def _aiter_llama_stream_items( async with _same_task_timeout(remaining_s): item = await async_iter.__anext__() else: + timeout_s = _post_first_timeout_s() if ( request is not None and response is not None - and post_first_item_read_timeout_s is not None + and timeout_s is not None and last_item_at is not None ): - stall_remaining_s = post_first_item_read_timeout_s - ( - time.monotonic() - last_item_at - ) + stall_remaining_s = timeout_s - (time.monotonic() - last_item_at) if stall_remaining_s <= 0: raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") _set_stream_response_read_timeout(response, stall_remaining_s) @@ -1128,19 +1659,16 @@ async def _aiter_llama_stream_items( if now >= first_token_deadline: raise continue - if ( - request is not None - and post_first_item_read_timeout_s is not None - and now - last_item_at < post_first_item_read_timeout_s - ): + timeout_s = _post_first_timeout_s() + if request is not None and timeout_s is not None and now - last_item_at < timeout_s: continue raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") - if ( - last_item_at is None - and response is not None - and post_first_item_read_timeout_s is not None - ): - _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) + if last_item_at is None and response is not None: + # The first-token read deadline no longer applies once a chunk has + # arrived: switch to the stall timeout, or clear the read timeout + # entirely when the stall guard is disabled (callable returns None) + # so a long gap can't trip the stale first-token deadline. + _set_stream_response_read_timeout(response, _post_first_timeout_s()) last_item_at = time.monotonic() yield item @@ -1519,6 +2047,20 @@ def _effective_enable_tools(payload) -> Optional[bool]: return policy if policy is not None else payload.enable_tools +def _explicit_studio_tool_loop_requested(payload) -> bool: + """True when the request itself asks Studio to execute local tools. + + Process-wide CLI policy can default Studio's tool loop on for ordinary chat, + but it must not steal OpenAI-compatible client tools or response_format + requests from the llama-server passthrough path. A policy of ``False`` + (--disable-tools) vetoes even an explicit ``enable_tools: true`` ask. + """ + from state.tool_policy import get_tool_policy + + policy = get_tool_policy() + return policy is not False and (payload.enable_tools is True or bool(payload.mcp_enabled)) + + # Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts so # is_disconnected() never fires. POST /inference/cancel looks up in-flight # cancel_events here by cancel_id (per-run) or session_id / completion_id @@ -1659,6 +2201,39 @@ async def _await_disconnect_then_cancel(request, cancel_event) -> None: return +def _cancelable_nonstreaming_client() -> httpx.AsyncClient: + return httpx.AsyncClient( + limits = httpx.Limits(max_connections = 1, max_keepalive_connections = 0), + trust_env = False, + ) + + +async def _await_cancel_or_disconnect_then_close_client( + *, cancel_event, request: Optional[Request], client: httpx.AsyncClient +) -> None: + """Close a dedicated non-streaming upstream client on cancel/disconnect. + + The shared ``nonstreaming_client()`` is pooled, so cancelable generation calls + use a per-request client. Closing it interrupts a blocked llama-server + request without affecting unrelated pooled non-streaming calls. + """ + try: + while True: + if cancel_event is not None and cancel_event.is_set(): + break + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + break + await asyncio.sleep(0.1) + try: + await client.aclose() + except Exception: + pass + except asyncio.CancelledError: + return + + async def _stop_local_disconnect_cancel_watcher(watcher) -> None: watcher.cancel() try: @@ -2769,6 +3344,21 @@ def _automatic_model_load_may_run() -> bool: return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0 +def _no_model_loaded_detail(base: str) -> str: + """Append a pointer to the opt-in auto-switch toggle to a "no model loaded" + error, but only when it's off. Auto-switch (default off) cold-loads a + requested downloaded GGUF, so an off toggle is the usual reason a request + naming a listed model still 400/503s; surface the fix. With it on the name + simply didn't resolve to a local GGUF, so the hint would mislead and is omitted.""" + from utils.openai_auto_switch_settings import get_openai_auto_switch_enabled + + if get_openai_auto_switch_enabled(): + return base + return base + ( + " Or enable Model auto-switch (Settings > API) to load a requested model automatically." + ) + + async def _maybe_auto_switch_model( requested_model: Optional[str], fastapi_request: Request, @@ -3038,10 +3628,14 @@ def _remote_gguf_companion_bytes( info = model_info(repo, token = hf_token, files_metadata = True) total = 0 for sibling in info.siblings or []: - base = Path(sibling.rfilename or "").name.lower() + name = sibling.rfilename or "" + base = Path(name).name.lower() if not base.endswith(".gguf"): continue - if base.startswith("mtp-") or (include_mmproj and "mmproj" in base): + # Root-level mtp- only: -hf auto-fetches the repo-root drafter, not + # the MTP/ subdir copies (which now share the mtp- prefix too). + is_root_mtp = "/" not in name and base.startswith("mtp-") + if is_root_mtp or (include_mmproj and "mmproj" in base): total += getattr(sibling, "size", 0) or 0 return total except Exception as e: @@ -3426,6 +4020,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre status_code = 400, detail = "gpu_ids is not supported for GGUF models yet.", ) + if not config.is_gguf and _mlx_distributed_launch_detected(): + raise HTTPException( + status_code = 400, + detail = ( + "Studio does not support distributed MLX inference under " + "mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio " + "without the distributed launcher." + ), + ) # Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it. effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) @@ -5863,7 +6466,7 @@ async def openai_chat_completions( if not backend.active_model_name: raise HTTPException( status_code = 400, - detail = "No model loaded. Call POST /inference/load first.", + detail = _no_model_loaded_detail("No model loaded. Call POST /inference/load first."), ) # Clean public id so the response never echoes a local path; the audio # branch below receives this sanitized label too. @@ -6040,11 +6643,11 @@ async def openai_chat_completions( # unaware of `role="tool"` messages and assistant messages that only # carry `tool_calls` (content=None) — both of which are valid in # multi-turn client-side tool loops. - effective_max_tokens = _effective_max_tokens(payload) + effective_max_tokens = _effective_openai_max_tokens(payload) normalized_stop = _normalize_stop_sequences(payload.stop) - _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) + _has_tool_messages = _has_openai_tool_history(payload.messages) # Route guided-decoding requests through the verbatim passthrough so # ``response_format`` (JSON schema) reaches llama-server and the model's # GBNF-constrained output comes back unmodified. The non-passthrough GGUF @@ -6053,12 +6656,41 @@ async def openai_chat_completions( # free-form sampling. Guided decoding does not require ``supports_tools`` -- # the grammar machinery is independent of tool-call parsing. _has_response_format = _extract_response_format(payload) is not None - _tools_passthrough = llama_backend.supports_tools and ( - (payload.tools and len(payload.tools) > 0) or _has_tool_messages + _has_tool_catalog = bool(payload.tools and len(payload.tools) > 0) + _has_active_tool_catalog = _has_tool_catalog and payload.tool_choice != "none" + _has_client_tool_contract = _has_active_tool_catalog or _has_tool_messages + # The Studio tool loop needs a tool-capable backend, so a request that asks + # for it on a backend that can't run it (DiffusionGemma forces supports_tools + # off) must not steal client tools from the passthrough (#6851). + _studio_tool_loop_requested = ( + _explicit_studio_tool_loop_requested(payload) and llama_backend.supports_tools ) + _client_disabled_tool_calls = payload.tool_choice == "none" and not _studio_tool_loop_requested + _supports_tool_passthrough = getattr( + llama_backend, "supports_tool_passthrough", llama_backend.supports_tools + ) + _tools_passthrough = _supports_tool_passthrough and _has_client_tool_contract if ( using_gguf - and not _effective_enable_tools(payload) + and not _studio_tool_loop_requested + and _has_client_tool_contract + and not _supports_tool_passthrough + ): + raise _reject( + 400, + openai_error_body( + ( + "Client-supplied tools or tool-call history require a GGUF chat template " + "with tool-call support; the current model/template does not advertise tools." + ), + status = 400, + code = "unsupported_parameter", + param = "tools" if payload.tools else "messages", + ), + ) + if ( + using_gguf + and not _studio_tool_loop_requested and (_tools_passthrough or _has_response_format) ): if _wants_multiple_choices(payload): @@ -6104,12 +6736,20 @@ async def openai_chat_completions( completion_id, monitor_id = monitor_id, ) - return await _openai_passthrough_non_streaming( - llama_backend, - payload, - model_name, - monitor_id = monitor_id, - ) + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + try: + return await _openai_passthrough_non_streaming( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + finally: + _tracker.__exit__(None, None, None) # ── Parse messages (handles multimodal content parts) ───── # Reuse the pre-hook parse when auto-switch did it, else parse now. @@ -6169,6 +6809,8 @@ async def openai_chat_completions( ) def _gguf_chat_delta_line(delta: ChoiceDelta, finish_reason = None) -> str: + if delta.reasoning_content is not None and delta.content is None: + delta = delta.model_copy(update = {"content": ""}) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -6192,8 +6834,12 @@ async def openai_chat_completions( from state.tool_policy import get_tool_policy as _get_tool_policy_g _cli_policy = _get_tool_policy_g() - _tools_on = _effective_enable_tools(payload) - _mcp_allowed = bool(payload.mcp_enabled) and _cli_policy is not False + _tools_on = False if _client_disabled_tool_calls else _effective_enable_tools(payload) + _mcp_allowed = ( + not _client_disabled_tool_calls + and bool(payload.mcp_enabled) + and _cli_policy is not False + ) use_tools = (_tools_on or _mcp_allowed) and llama_backend.supports_tools if use_tools: @@ -6292,6 +6938,24 @@ async def openai_chat_completions( bypass_permissions = bool(payload.bypass_permissions), ) + _tool_admission_mode = "chat_tool_stream" if payload.stream else "chat_tool_nonstream" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + _tool_sentinel = object() _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) @@ -6300,6 +6964,8 @@ async def openai_chat_completions( async def gguf_tool_stream(): gen = None + next_task = None + stream_completed = False disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -6337,7 +7003,14 @@ async def openai_chat_completions( api_monitor.finish(monitor_id, "cancelled") return - event = await asyncio.to_thread(next, gen, _tool_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _tool_sentinel) + ) + try: + event = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if event is _tool_sentinel: break @@ -6436,6 +7109,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -6448,20 +7122,157 @@ async def openai_chat_completions( # Recover if an MTP+tensor crash killed the server mid-stream. get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: - await _stop_local_disconnect_cancel_watcher(disconnect_watcher) - if gen is not None: - try: - gen.close() - except (RuntimeError, ValueError): - pass - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF tool stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) if payload.stream: + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_tool_stream(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_tool_stream() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_tool_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( - gguf_tool_stream(), - unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + admitted_gguf_tool_stream(), + unstarted_cleanup = _gguf_tool_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6507,10 +7318,61 @@ async def openai_chat_completions( except (RuntimeError, ValueError): pass + drain_task = None + + async def _drain_cancelled_gguf_tool_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass + + admission_lease = None + admission_wait_started_at = None try: - full_text, completion_usage, completion_finish = await asyncio.to_thread( - _drain_gguf_tool_loop + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_tool_loop)) + full_text, completion_usage, completion_finish = await asyncio.shield(drain_task) reasoning_text, visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _responses_should_parse_think_markers( @@ -6556,6 +7418,48 @@ async def openai_chat_completions( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_tool_task() + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) except Exception as e: logger.error(f"Error during GGUF tool completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -6576,6 +7480,8 @@ async def openai_chat_completions( ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) finally: + if admission_lease is not None: + admission_lease.release() _tracker.__exit__(None, None, None) # ── Standard GGUF path (no tools) ───────────────────── @@ -6610,11 +7516,31 @@ async def openai_chat_completions( _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) async def gguf_stream_chunks(): disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) + gen = None + next_task = None + stream_completed = False try: yield _chat_role_chunk(completion_id, created, model_name) @@ -6633,7 +7559,14 @@ async def openai_chat_completions( cancel_event.set() api_monitor.finish(monitor_id, "cancelled") return - cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _gguf_sentinel) + ) + try: + cumulative = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if cumulative is _gguf_sentinel: break # Capture server metadata for the final usage chunk @@ -6702,6 +7635,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -6712,14 +7646,156 @@ async def openai_chat_completions( logger.error(f"Error during GGUF streaming: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: - await _stop_local_disconnect_cancel_watcher(disconnect_watcher) - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) + + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_stream_chunks(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_stream_chunks() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( - gguf_stream_chunks(), - unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + admitted_gguf_stream_chunks(), + unstarted_cleanup = _gguf_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6728,57 +7804,189 @@ async def openai_chat_completions( }, ) else: + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + admission_lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + try: # ``n`` requests several independent completions; the single # decode slot yields one at a time, so loop sequentially. - _n = payload.n or 1 + drain_task = None - _choices = [] - _monitor_replies = [] - _prompt_tokens = 0 - _sum_completion = 0 - _prompt_details = None - for _idx in range(_n): - # Stop spawning the remaining choices once cancelled. - if cancel_event.is_set(): - break - full_text = "" - completion_usage = None - completion_finish = None - for token in gguf_generate(_idx): - if isinstance(token, dict): - if token.get("type") == "metadata": - completion_usage = token.get("usage") - completion_finish = token.get("finish_reason") + async def _drain_cancelled_gguf_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() continue - full_text = token + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass - reasoning_text, visible_text = _extract_responses_reasoning( - full_text, - parse_think_markers = _responses_should_parse_think_markers( - payload, - llama_backend, - ), - ) - message_kwargs = {"content": visible_text} - if reasoning_text: - message_kwargs["reasoning_content"] = reasoning_text - _choices.append( - CompletionChoice( - index = _idx, - message = CompletionMessage(**message_kwargs), - finish_reason = _clamp_finish_reason(completion_finish), + def _drain_gguf_choices(): + _n = payload.n or 1 + _choices = [] + _monitor_replies = [] + _prompt_tokens = 0 + _sum_completion = 0 + _prompt_details = None + for _idx in range(_n): + # Stop spawning the remaining choices once cancelled. + if cancel_event.is_set(): + break + full_text = "" + completion_usage = None + completion_finish = None + for token in gguf_generate(_idx): + if isinstance(token, dict): + if token.get("type") == "metadata": + completion_usage = token.get("usage") + completion_finish = token.get("finish_reason") + continue + full_text = token + + reasoning_text, visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _responses_should_parse_think_markers( + payload, + llama_backend, + ), ) + message_kwargs = {"content": visible_text} + if reasoning_text: + message_kwargs["reasoning_content"] = reasoning_text + _choices.append( + CompletionChoice( + index = _idx, + message = CompletionMessage(**message_kwargs), + finish_reason = _clamp_finish_reason(completion_finish), + ) + ) + _monitor_replies.append(visible_text) + if completion_usage: + # The prompt is shared across all n choices, so count its + # tokens ONCE (OpenAI bills only generated tokens for each + # extra choice). Only completion_tokens accumulates. + _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens + _sum_completion += completion_usage.get("completion_tokens") or 0 + if _prompt_details is None: + _prompt_details = completion_usage.get("prompt_tokens_details") + return ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, ) - _monitor_replies.append(visible_text) - if completion_usage: - # The prompt is shared across all n choices, so count its - # tokens ONCE (OpenAI bills only generated tokens for each - # extra choice). Only completion_tokens accumulates. - _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens - _sum_completion += completion_usage.get("completion_tokens") or 0 - if _prompt_details is None: - _prompt_details = completion_usage.get("prompt_tokens_details") + + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_choices)) + ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, + ) = await asyncio.shield(drain_task) response = ChatCompletion( id = completion_id, @@ -6810,6 +8018,11 @@ async def openai_chat_completions( api_monitor.finish(monitor_id) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_task() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -6829,7 +8042,10 @@ async def openai_chat_completions( ), ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) - + finally: + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) # ── Standard Unsloth path ───────────────────────────────── # Decode image (from content parts OR legacy field) @@ -7149,7 +8365,7 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if gen is not None: @@ -7494,7 +8710,7 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) @@ -7983,7 +9199,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); @@ -7992,8 +9208,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if not isinstance(body, dict): raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - if body.get("max_tokens") is None: - body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR + _resolved_max_tokens = _effective_openai_max_tokens_from_values(body.get("max_tokens")) + body["max_tokens"] = ( + _resolved_max_tokens + if _resolved_max_tokens is not None + else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR) + ) target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) prompt_text = _flatten_monitor_prompt(body.get("prompt", "")) @@ -8086,7 +9306,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return api_monitor.finish(monitor_id, "cancelled") return @@ -8101,7 +9321,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return finally: await _aclose_stream_resources( @@ -8195,7 +9415,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); @@ -8935,7 +10155,7 @@ async def _responses_stream( # so the client sees a useful error instead of a dangling stream. raise HTTPException( status_code = 400, - detail = ( + detail = _no_model_loaded_detail( "Streaming /v1/responses requires a GGUF model loaded via " "llama-server. Use non-streaming /v1/responses, " "/v1/chat/completions, or load a GGUF model." @@ -8957,6 +10177,52 @@ async def _responses_stream( ) body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + def _responses_admission_failed_sse(exc: Exception, *, status_code: int) -> str: + return ( + "event: response.failed\n" + "data: " + + json.dumps( + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": _llama_public_model_id(llama_backend, payload.model) + or payload.model, + "output": [], + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + "error": { + "code": status_code, + "message": str(exc), + }, + }, + } + ) + + "\n\n" + ) async def event_generator(): # Clean public id for every response envelope. Prefer the loaded model's @@ -9766,14 +11032,111 @@ async def _responses_stream( api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) + async def admitted_event_generator(): + lease = reservation.lease_nowait() + admission_wait_started_at = None + stream_started = False + stream_cancelled = False + iterator = None + try: + if lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "debug", + ) + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = None, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = None, + ) + iterator = event_generator() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _responses_admission_failed_sse(exc, status_code = 503) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + + async def _responses_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + return _SameTaskStreamingResponse( - event_generator(), + admitted_event_generator(), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, + unstarted_cleanup = _responses_admission_unstarted_cleanup, ) @@ -10031,7 +11394,7 @@ async def anthropic_count_tokens( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # Same Anthropic → OpenAI translation as anthropic_messages: system is @@ -10043,8 +11406,11 @@ async def anthropic_count_tokens( # Apply the same sanitization /messages does before generation, so the count # matches the prompt the real request would build (otherwise empty-assistant # sentinels / synthetic tool history inflate the count or hit the fallback). - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) openai_tools = anthropic_tools_to_openai(payload.tools or []) or None @@ -10101,7 +11467,7 @@ async def anthropic_messages( if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # max_tokens is a required field on the Anthropic Messages API; real Anthropic @@ -10152,7 +11518,7 @@ async def anthropic_messages( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # Advertised repo id after an auto-switch load, else a clean public id, never @@ -10172,8 +11538,11 @@ async def anthropic_messages( # builders apply the same strip; without it an Anthropic /v1/messages caller # replaying a prior provider-side tool_use forwards fake builtin tool # history to a backend with no matching function declarations. - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) # Enforce vision guard + re-encode embedded images to PNG so the Anthropic @@ -10222,7 +11591,9 @@ async def anthropic_messages( and not _has_image ) client_tools = ( - not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools + not server_tools + and len(openai_client_tools) > 0 + and getattr(llama_backend, "supports_tool_passthrough", llama_backend.supports_tools) ) # Anthropic tool_choice.disable_parallel_tool_use caps the response to a @@ -10840,13 +12211,15 @@ def _build_passthrough_payload( ): body = { "messages": openai_messages, - "tools": openai_tools, - "tool_choice": tool_choice, "temperature": temperature, "top_p": top_p, "top_k": top_k, "stream": stream, } + if openai_tools: + body["tools"] = openai_tools + if tool_choice is not None: + body["tool_choice"] = tool_choice if seed is not None: body["seed"] = seed if stream and stream_options is not None: @@ -11053,13 +12426,17 @@ async def _anthropic_passthrough_stream( yield event return finally: - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # Same shape as the OpenAI passthrough: a close-time CancelledError + # re-raised by _aclose_stream_resources must not skip the tracker exit. + try: + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + _tracker.__exit__(None, None, None) for line in emitter.finish(): yield line @@ -11556,6 +12933,9 @@ def _build_openai_passthrough_body( system_prompt, _, _ = _extract_content_parts(payload.messages) messages = _set_or_prepend_system_message(messages, system_prompt) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" + tools = payload.tools + if payload.tool_choice == "none" and not _has_openai_tool_history(payload.messages): + tools = None # Forward per-request reasoning fields (enable_thinking / reasoning_effort / # preserve_thinking) via chat_template_kwargs so the Jinja template renders # in the caller's mode, gated on the active template's capabilities exactly @@ -11571,12 +12951,12 @@ def _build_openai_passthrough_body( ) return _build_passthrough_payload( messages, - payload.tools, + tools, payload.temperature, payload.top_p, payload.top_k, # Honor max_completion_tokens on the tools/response_format passthrough too. - _effective_max_tokens(payload), + _effective_openai_max_tokens(payload), payload.stream, stop = payload.stop, min_p = payload.min_p, @@ -11600,33 +12980,219 @@ async def _openai_passthrough_stream( completion_id, monitor_id: Optional[str] = None, ): - """Streaming client-side pass-through for /v1/chat/completions. + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionCancelled as exc: + lease.release() + _tracker.__exit__(None, None, None) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + return await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = lease, + tracker = _tracker, + ) + + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "debug", + ) + + async def _queued_stream(): + admitted_started = False + admitted_body_owns_cleanup = False + admitted_response = None + admitted_body_cancelled = False + try: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + admitted_response = await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = wait_item, + tracker = _tracker, + ) + admitted_started = True + iterator = admitted_response.body_iterator + admitted_body_owns_cleanup = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + admitted_body_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = admitted_body_cancelled, + ) + if not admitted_body_owns_cleanup: + cleanup = getattr(admitted_response, "_unstarted_cleanup", None) + if cleanup is not None: + await cleanup() + return + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse(_openai_admission_error_body(exc, status_code = 503)) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if not admitted_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _queued_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + return _SameTaskStreamingResponse( + _queued_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + unstarted_cleanup = _queued_unstarted_cleanup, + ) + + +async def _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id: Optional[str] = None, + *, + admission_lease: LlamaAdmissionLease, + tracker, +): + """Streaming client-side pass-through after Studio granted an upstream slot. Forwards the client's OpenAI function-calling request to llama-server and - relays the SSE stream back verbatim, preserving llama-server's native - response ``id``, ``finish_reason`` (including ``"tool_calls"``), - ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so - the client sees a standard OpenAI response. + relays the SSE stream back with minimal normalization (reasoning-only + deltas gain ``content: ""``; errors and missing terminal markers get a + closing ``[DONE]``), preserving llama-server's native response ``id``, + ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and + any client-requested trailing ``usage`` chunk so the client sees a + standard OpenAI response. Reasoning/tool-call splitting is delegated to llama-server (``--jinja --reasoning-format auto``), so ``delta.content`` carries no raw markup and is deliberately not re-parsed locally, unlike the ``/completion`` paths. """ + _tracker = tracker target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body( - payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend - ) - # Text-form tool calls from small models get promoted to structured calls on - # the way back (declared client tools only); requests without tools or with - # auto_heal_tool_calls=false keep the verbatim relay. tool_choice constrains - # the allowlist ("none" disables, a forced function narrows to it). - _allowed_tools = heal_gate( - payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") - ) + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) - _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) - _tracker = _TrackedCancel(cancel_event, *_cancel_keys) - _tracker.__enter__() client = None resp = None send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None @@ -11648,6 +13214,17 @@ async def _openai_passthrough_stream( # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + ) + # Text-form tool calls from small models get promoted to structured calls on + # the way back (declared client tools only); requests without tools or with + # auto_heal_tool_calls=false keep the unhealed relay. tool_choice constrains + # the allowlist ("none" disables, a forced function narrows to it). + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) + # Keep the pre-header window short so accepted SSE clients receive # immediate headers in the common timeout-reduced stall. client = httpx.AsyncClient( @@ -11661,12 +13238,16 @@ async def _openai_passthrough_stream( while True: try: - req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} - ) + req = client.build_request("POST", target_url, json = body, headers = upstream_headers) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S send_task = asyncio.create_task( - _send_stream_with_preheader_cancel(client, req, cancel_event, request = request) + _send_stream_with_preheader_cancel( + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, + ) ) done, _ = await asyncio.wait( {send_task}, @@ -11692,13 +13273,17 @@ async def _openai_passthrough_stream( if resp is None and send_task is not None and not send_task.done(): break if resp is None: + if cancel_event is not None: + cancel_event.set() api_monitor.finish(monitor_id, "cancelled") - await _aclose_send_task(send_task) try: - await client.aclose() - except Exception: - pass - _tracker.__exit__(None, None, None) + await _aclose_send_task(send_task) + await _aclose_stream_resources(client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", @@ -11723,6 +13308,7 @@ async def _openai_passthrough_stream( await resp.aclose() except Exception: pass + resp = None # Opt-in overflow policy: shrink and retry instead of a fatal 400. if ( _truncate_budget > 0 @@ -11751,11 +13337,14 @@ async def _openai_passthrough_stream( disconnect_watcher = None nonlocal resp, send_task, first_token_deadline, _truncate_budget + nonlocal client monitor_done = False saw_finish_reason = False saw_done = False saw_stream_error = False + saw_stream_item = False saw_tool_call_delta = False + terminal_seen = False last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) @@ -11816,6 +13405,13 @@ async def _openai_passthrough_stream( lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) return lines + stall_timeout_s = _openai_compat_stream_stall_timeout() + + def _terminal_read_timeout_s() -> Optional[float]: + if terminal_seen: + return _OPENAI_PASSTHROUGH_TERMINAL_GRACE_S + return stall_timeout_s + def _heal_transform(chunk_data: dict, raw_line: str) -> list: """SSE lines to emit in place of one upstream line (healing on).""" choices = chunk_data.get("choices") @@ -11884,24 +13480,47 @@ async def _openai_passthrough_stream( try: while True: - if send_task is not None and not send_task.done(): - try: - resp = await send_task - except httpx.RequestError as e: - logger.error("openai passthrough stream: upstream unreachable: %s", e) - api_monitor.fail(monitor_id, _friendly_error(e)) - yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" - return - send_task = None - elif send_task is not None: - try: - resp = send_task.result() - except httpx.RequestError as e: - logger.error("openai passthrough stream: upstream unreachable: %s", e) - api_monitor.fail(monitor_id, _friendly_error(e)) - yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" - return - send_task = None + if send_task is not None: + last_keepalive_at = time.monotonic() + while not send_task.done(): + # Wake often enough that _preheader_cancelled keeps + # cancel/disconnect latency sub-second during prefill; + # keepalives still pace off last_keepalive_at. + wait_timeout = min( + _STREAM_DISCONNECT_POLL_TIMEOUT_S, + _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S, + ) + done, _ = await asyncio.wait( + {send_task}, + timeout = wait_timeout, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + break + if await _preheader_cancelled(cancel_event, request): + api_monitor.finish(monitor_id, "cancelled") + return + # The downstream SSE response is already committed; + # keep strict clients and proxies from treating a long + # llama-server prefill/header wait as a dead stream. + now = time.monotonic() + if ( + now - last_keepalive_at + >= _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S + ): + last_keepalive_at = now + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + if resp is None: + try: + resp = send_task.result() + except httpx.RequestError as e: + logger.error( + "openai passthrough stream: upstream unreachable: %s", e + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield _openai_stream_error_sse(_openai_stream_error_chunk(e)) + return + send_task = None if resp is None: api_monitor.finish(monitor_id, "cancelled") @@ -11929,12 +13548,16 @@ async def _openai_passthrough_stream( ): _truncate_budget -= 1 req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} + "POST", target_url, json = body, headers = upstream_headers ) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S send_task = asyncio.create_task( _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, ) ) continue @@ -11949,7 +13572,7 @@ async def _openai_passthrough_stream( ) ) api_monitor.fail(monitor_id, err_text[:500]) - yield f"data: {json.dumps(error_payload)}\n\n" + yield _openai_stream_error_sse(error_payload) return cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) @@ -11963,12 +13586,14 @@ async def _openai_passthrough_stream( request = request, first_token_deadline = first_token_deadline, response = resp, + post_first_item_read_timeout_s = _terminal_read_timeout_s, ): if not raw_line: continue - if not raw_line.startswith("data: "): + if not raw_line.startswith("data:"): continue - data_text = raw_line[6:].strip() + saw_stream_item = True + data_text = raw_line[5:].strip() if data_text == "[DONE]": saw_done = True # Upstream ended without a finish chunk: heal the residue @@ -12000,13 +13625,11 @@ async def _openai_passthrough_stream( yield raw_line + "\n\n" monitor_done = True break - # Honor parallel_tool_calls=false (best-effort): drop tool_call - # deltas with index>=1 so only the first call streams. Only - # lines carrying tool_calls are reparsed; everything else is - # relayed byte-for-byte. - if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: - raw_line = _cap_parallel_tool_calls_sse_line(raw_line) - data_text = raw_line[6:].strip() + raw_line = _normalize_openai_passthrough_sse_line( + raw_line, + cap_parallel_tool_calls = payload.parallel_tool_calls is False, + ) + data_text = raw_line[5:].strip() try: chunk_data = json.loads(data_text) except json.JSONDecodeError: @@ -12033,8 +13656,9 @@ async def _openai_passthrough_stream( if _monitor_openai_error_message(chunk_data): saw_stream_error = True # With healing active, a content-bearing line may be replaced by - # held/promoted chunks; otherwise the single upstream line - # relays verbatim (monitored exactly as emitted either way). + # held/promoted chunks; otherwise the single (already + # normalized) line relays unchanged (monitored exactly as + # emitted either way). if ( healer is not None and not healer.dormant @@ -12084,6 +13708,27 @@ async def _openai_passthrough_stream( yield out_line + "\n\n" if monitor_event == "done": monitor_done = True + break + terminal_state = ( + _openai_passthrough_terminal_state_from_data(chunk_data) + if out_line is raw_line + else _openai_passthrough_sse_line_terminal_state(out_line) + ) + if terminal_state == "usage" or ( + terminal_state == "finish" and not _wants_stream_usage(payload) + ): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + saw_done = True + monitor_done = True + break + if terminal_state == "finish": + terminal_seen = True if monitor_done: break if not saw_done and not saw_stream_error and not cancel_event.is_set(): @@ -12105,7 +13750,7 @@ async def _openai_passthrough_stream( llama_backend.context_length, ) yield finish_line + "\n\n" - done_line = "data: [DONE]" + done_line = _SSE_DONE_LINE _monitor_openai_sse_line( monitor_id, done_line, @@ -12121,6 +13766,29 @@ async def _openai_passthrough_stream( except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise + except httpx.ReadTimeout as e: + if terminal_seen and not saw_stream_error and not cancel_event.is_set(): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + api_monitor.finish(monitor_id) + return + if cancel_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") + return + logger.error( + "openai passthrough stream %s: %s", + "stalled mid-response" if saw_stream_item else "timeout", + e, + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) + err = _openai_stream_error_chunk(e) + yield _openai_stream_error_sse(err) except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: # Watcher closed resp on cancel. Emit nothing extra; the client # initiated the cancel or already disconnected. @@ -12129,6 +13797,16 @@ async def _openai_passthrough_stream( get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) raise api_monitor.finish(monitor_id, "cancelled") + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error_payload = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error_payload) except Exception as e: if cancel_event.is_set(): api_monitor.finish(monitor_id, "cancelled") @@ -12138,25 +13816,37 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, _friendly_error(e)) get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) err = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(err)}\n\n" + yield _openai_stream_error_sse(err) finally: - await _aclose_send_task(send_task) - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # _aclose_stream_resources re-raises a close-time CancelledError + # only after finishing teardown, and the tracker exits either way. + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) async def _unstarted_cleanup() -> None: # Client disconnected before the body stream started, so _stream()'s # finally never ran. Release the eagerly-opened upstream resp/client # and the cancel-registry entry here; the watchers and line iterator # are created inside _stream(), so there is nothing else to close. - await _aclose_send_task(send_task) - await _aclose_stream_resources(resp = resp, client = client) - _tracker.__exit__(None, None, None) + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( _stream(), @@ -12168,10 +13858,22 @@ async def _openai_passthrough_stream( }, unstarted_cleanup = _unstarted_cleanup, ) - except BaseException: - await _aclose_send_task(send_task) - await _aclose_stream_resources(resp = resp, client = client) - _tracker.__exit__(None, None, None) + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + if cancel_event is not None: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + else: + detail = exc.detail if isinstance(exc, HTTPException) else _friendly_error(exc) + api_monitor.fail(monitor_id, str(detail)) + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) raise @@ -12180,6 +13882,109 @@ async def _openai_passthrough_non_streaming( payload, model_name, monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, +): + """Non-streaming pass-through guarded by local llama-server admission.""" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_nonstream", + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + level = "debug", + ) + lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + return await _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + raise + finally: + if lease is not None: + lease.release() + + +async def _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, ): """Non-streaming client-side pass-through for /v1/chat/completions. @@ -12188,20 +13993,67 @@ async def _openai_passthrough_non_streaming( ``tool_calls``, and accurate ``usage`` token counts. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) body = _build_openai_passthrough_body( payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + body["stream"] = False + body.pop("stream_options", None) _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) - while True: - try: - resp = await nonstreaming_client().post( + + async def _post(body_to_send): + if cancel_event is None and request is None: + return await nonstreaming_client().post( target_url, - json = body, + json = body_to_send, + headers = upstream_headers, timeout = _llama_non_streaming_generation_timeout(), ) + + if cancel_event is None: + cancel = threading.Event() + else: + cancel = cancel_event + client = _cancelable_nonstreaming_client() + watcher = asyncio.create_task( + _await_cancel_or_disconnect_then_close_client( + cancel_event = cancel, + request = request, + client = client, + ) + ) + try: + try: + response = await client.post( + target_url, + json = body_to_send, + headers = upstream_headers, + timeout = _llama_non_streaming_generation_timeout(), + ) + except httpx.RequestError: + if cancel.is_set(): + raise asyncio.CancelledError() + raise + if cancel.is_set(): + raise asyncio.CancelledError() + return response + finally: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + try: + await client.aclose() + except Exception: + pass + + while True: + try: + resp = await _post(body) except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise @@ -12269,15 +14121,14 @@ async def _openai_passthrough_non_streaming( "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], } try: - retry_resp = await nonstreaming_client().post( - target_url, - json = retry_body, - timeout = _llama_non_streaming_generation_timeout(), - ) + retry_resp = await _post(retry_body) if retry_resp.status_code == 200: retry_data = retry_resp.json() if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): resp, data = retry_resp, retry_data + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except (httpx.RequestError, ValueError) as exc: logger.warning("tool-call nudge retry failed; keeping original: %s", exc) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 914699f540..bbee374334 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -32,6 +32,7 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents from utils.openai_auto_switch_settings import ( DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, @@ -174,6 +175,19 @@ def update_helper_precache( return _helper_precache_response(enabled) +class CodingAgentsResponse(BaseModel): + # All agents `unsloth start` supports, in the CLI's declared order. + agents: tuple[str, ...] = CODING_AGENTS + # Subset of `agents` whose CLI binary was found on PATH; the frontend uses + # this to default the API-keys panel to a command the user can run as-is. + detected: list[str] + + +@router.get("/coding-agents", response_model = CodingAgentsResponse) +def get_coding_agents(current_subject: str = Depends(get_current_subject)) -> CodingAgentsResponse: + return CodingAgentsResponse(detected = detect_installed_coding_agents()) + + @router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) def get_openai_auto_switch( current_subject: str = Depends(get_current_subject), diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 0c6550a3bb..54454f6563 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1770,3 +1770,187 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "plain" + + +def test_resumed_session_thinking_and_null_content_do_not_400(): + # A resumed session replays assistant turns with `thinking` (and sometimes null) + # content. Those must be accepted (thinking dropped by the converter), not 400ed. + from pydantic import ValidationError + + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "secret reasoning", "signature": "s"}, + {"type": "text", "text": "the answer"}, + {"type": "tool_use", "id": "t1", "name": "f", "input": {}}, + ], + }, + {"role": "assistant", "content": None}, # tool-only turn serialized as null + ], + ) + # Known blocks still parse as their typed models; only the unknown one is loose. + assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock" + assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock" + assert req.messages[2].content == "" # null coerced + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content")) + assert assistant["content"] == "the answer" + assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded + + # A malformed KNOWN block still fails cleanly instead of being swallowed. + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}], + ) + + +def test_user_null_content_rejected(): + # The null->"" leniency is assistant-only; a null user content must be rejected + # at the boundary, not coerced into an empty prompt and forwarded to the model. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": None}], + ) + + +def test_user_unknown_block_rejected_not_silently_dropped(): + # The converter skips user blocks it cannot translate, so a user turn whose only + # block is unknown would validate yet forward no content. Reject at the boundary + # to avoid that silent data loss (the assistant fallback is unaffected). + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "document", "source": {}}]}, + ], + ) + + +def test_user_translatable_blocks_still_accepted(): + # text / image / tool_result are translatable, so a real user message built from + # them must still pass; the unknown-block guard only trips on other types. + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "AA"}, + }, + {"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, + ], + } + ], + ) + assert [type(b).__name__ for b in req.messages[0].content] == [ + "AnthropicTextBlock", + "AnthropicImageBlock", + "AnthropicToolResultBlock", + ] + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai) + + +def test_user_malformed_known_block_still_rejected(): + # The guard only allow-lists a user block's *type*; the union still validates its + # shape, so a known-but-malformed block (tool_result without tool_use_id) fails. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "tool_result", "content": "x"}]}, + ], + ) + + +def test_user_content_block_non_string_type_rejected_cleanly(): + # A user block whose `type` is a non-string (unhashable list / dict, or a stray + # int) must fail as a clean validation error, not raise TypeError from the + # frozenset membership test and escape as a 500. + from pydantic import ValidationError + for bad_type in ([], {}, 5): + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": [{"type": bad_type}]}], + ) + + +def test_assistant_missing_content_key_still_rejected(): + # The null -> "" leniency is only for an EXPLICIT null. An assistant message that + # omits content entirely stays malformed and must fail required-field validation. + from pydantic import ValidationError + + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant"}], + ) + # An explicit null is still accepted and coerced (regression guard). + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": None}, + ], + ) + assert req.messages[1].content == "" + + +def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch): + # user -> assistant(null) -> user is now accepted: the null assistant turn coerces + # to "" and is dropped. The route must then coalesce the two remaining user turns + # so a strict GGUF chat template does not 400 on non-alternating roles. + backend = _mock_backend(monkeypatch, context_length = 2048) + + class _Req: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/messages") + method = "POST" + + async def is_disconnected(self): + return False + + payload = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "please continue"}, + ], + ) + + response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t")) + assert response.status_code == 200 + + [(_path, kwargs)] = backend.calls + user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"] + assert len(user_turns) == 1 # the two user turns were merged, not left adjacent + merged = user_turns[0]["content"] + if isinstance(merged, list): + merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict)) + assert "first question" in merged and "please continue" in merged diff --git a/studio/backend/tests/test_coding_agents.py b/studio/backend/tests/test_coding_agents.py new file mode 100644 index 0000000000..b19da1dded --- /dev/null +++ b/studio/backend/tests/test_coding_agents.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for coding-agent CLI detection used by the API-keys settings panel.""" + +from unittest.mock import patch + +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents + + +def test_matches_unsloth_start_subcommands(): + # Each entry must be an actual `unsloth start ` subcommand name + # (unsloth_cli/commands/start.py). Spelled out here rather than imported + # from that module, which pulls in the CLI's heavier dependencies. + assert CODING_AGENTS == ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def test_detects_only_agents_present_on_path(): + installed = {"claude", "opencode"} + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: f"/usr/bin/{name}" if name in installed else None, + ): + assert detect_installed_coding_agents() == ["claude", "opencode"] + + +def test_returns_empty_list_when_nothing_is_installed(): + with patch("utils.coding_agents.shutil.which", return_value = None): + assert detect_installed_coding_agents() == [] + + +def test_preserves_declared_order_regardless_of_path_lookup_order(): + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: name if name in ("pi", "claude", "hermes") else None, + ): + assert detect_installed_coding_agents() == ["claude", "hermes", "pi"] + + +def test_treats_a_path_lookup_error_as_not_installed(): + # An advisory check: shutil.which raising for one entry (e.g. a permission + # error walking a PATH directory) should not take down the whole endpoint, + # and should not stop the remaining agents from being checked. + def flaky_which(name: str): + if name == "codex": + raise OSError("permission denied") + return name if name == "claude" else None + + with patch("utils.coding_agents.shutil.which", side_effect = flaky_which): + assert detect_installed_coding_agents() == ["claude"] diff --git a/studio/backend/tests/test_completion_masking.py b/studio/backend/tests/test_completion_masking.py new file mode 100644 index 0000000000..be0d8a69bd --- /dev/null +++ b/studio/backend/tests/test_completion_masking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy: auto-detect first, manual table fallback. + +Covers utils.datasets.completion_masking.apply_completion_masking, shared by +the CUDA trainer (core/training/trainer.py) and the MLX worker +(core/training/worker.py): + - unmapped models use chat template auto-detection (previously masking was + silently disabled), + - gpt-oss goes auto-first too (its quantized checkpoints ship a template + the manual markers cannot match), + - an auto-detection failure falls back to the template table markers, + - a table miss after an auto failure warns and leaves the trainer unchanged. +""" + +from __future__ import annotations + +import pytest + +from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers +from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER + + +class _Trainer: + """Sentinel trainer; train_fn wraps it in a new object when applied.""" + + +class _Recorder: + """Fake train_on_responses_only that records calls.""" + + def __init__(self): + self.calls = [] + + def __call__(self, trainer, **kwargs): + self.calls.append(kwargs) + wrapped = _Trainer() + wrapped.wrapped_from = trainer + return wrapped + + +def _detect_ok(processor): + return "", "" + + +def _detect_fail(processor): + raise ValueError( + "Unsloth: Could not reliably auto-detect response_part - " + "pass instruction_part and response_part." + ) + + +_AUTO = {"instruction_part": "", "response_part": ""} + + +class _Notes: + def __init__(self): + self.messages = [] + + def __call__(self, level, message): + self.messages.append((level, message)) + + def warnings(self): + return [m for level, m in self.messages if level == "warning"] + + +def test_unmapped_model_uses_auto_detection(): + # Unmapped model: the auto path applies masking (was silently disabled). + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, notify = notes, detect_fn = _detect_ok + ) + + assert applied is True + assert result.wrapped_from is trainer + assert train_fn.calls == [dict(_AUTO)] # applied with the detected markers + assert notes.warnings() == [] + + +def test_mapped_model_prefers_auto_detection(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_uses_auto_detection_first(): + # The quantized gpt-oss checkpoints ship a template without the + # <|channel|>final header, where the manual markers match nothing; auto + # derives markers from the template the checkpoint actually ships. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_detection_failure_falls_back_to_manual_markers(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_fail + ) + + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +def test_auto_failure_falls_back_to_template_table(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is True + assert result.wrapped_from is trainer + expected = TEMPLATE_TO_RESPONSES_MAPPER["qwen3"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + }, + ] + assert any("falling back to the template table" in m for m in notes.warnings()) + + +def test_application_failure_propagates_not_fallback(): + # Detection succeeds; a failure while APPLYING the masking must propagate, + # never silently fall back to full-sequence training. + def train_fn(trainer, **kwargs): + raise RuntimeError("dataset map worker crashed") + + with pytest.raises(RuntimeError, match = "dataset map worker crashed"): + apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok) + + +def test_preset_tokenizer_markers_used_directly(): + # Preset unsloth marker attrs skip detection; zoo reuses them on a bare call. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.processing_class = _Tok() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_table_miss_warns_and_disables_without_crashing(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is False + assert result is trainer # unchanged: full sequence training + assert train_fn.calls == [] # detection failed; nothing applied + assert any("could not be applied" in m for m in notes.warnings()) + assert any("full sequences" in m for m in notes.warnings()) + + +def test_num_proc_forwarded_only_when_given(): + # CUDA path passes num_proc; the MLX path omits it. + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_ok + ) + assert train_fn.calls == [dict(_AUTO, num_proc = 4)] + + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_fail + ) + assert train_fn.calls[0]["num_proc"] == 4 + + train_fn = _Recorder() + apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok) + assert train_fn.calls == [dict(_AUTO)] + + +def test_manual_fallback_failure_propagates_to_caller(): + # Errors while applying the manual fallback must propagate to the caller. + def train_fn(trainer, **kwargs): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match = "boom"): + apply_completion_masking(_Trainer(), "unsloth/gpt-oss-20b", train_fn) + + +def test_notify_is_optional(): + train_fn = _Recorder() + _, applied = apply_completion_masking( + _Trainer(), "some-org/not-in-any-mapper", train_fn, detect_fn = _detect_fail + ) + assert applied is False + + +def test_lookup_manual_markers(): + template, instruction, response = lookup_manual_markers("unsloth/Qwen3-0.6B") + assert template == "qwen3" + assert instruction == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["instruction"] + assert response == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["response"] + + template, instruction, response = lookup_manual_markers("some-org/unknown") + assert (template, instruction, response) == (None, None, None) + + template, instruction, response = lookup_manual_markers(None) + assert (template, instruction, response) == (None, None, None) + + +def test_renamed_gpt_oss_gets_template_markers(): + # Name-detected as gpt-oss but not in the exact-name table: must use the + # gpt-oss markers, not fall through to full-sequence training. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "some-org/gpt-oss-20b-sft", train_fn, detect_fn = _detect_fail + ) + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +class _FakeTokenizerWrapper: + """mlx-lm TokenizerWrapper semantics: plain reads delegate to the wrapped + tokenizer, underscore attrs do not (so preset markers are hidden).""" + + def __init__(self, tokenizer): + object.__setattr__(self, "_tokenizer", tokenizer) + + def __getattr__(self, attr): + if attr.startswith("_"): + return object.__getattribute__(self, attr) + return getattr(object.__getattribute__(self, "_tokenizer"), attr) + + +_FakeTokenizerWrapper.__name__ = "TokenizerWrapper" + + +def test_mlx_tokenizer_wrapper_unwrapped_for_preset_markers(): + # Markers live on the inner HF tokenizer that the wrapper hides; the helper + # must unwrap so the preset bare-call path still fires on MLX. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(_Tok()) + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_mlx_tokenizer_wrapper_unwrapped_for_detection(): + # Detection must see the real tokenizer, not the wrapper, so it does not + # depend on the loader's __call__ patch. + class _Tok: + pass + + inner = _Tok() + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(inner) + train_fn = _Recorder() + seen = [] + + def detect(processor): + seen.append(processor) + return "", "" + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect + ) + assert applied is True + assert seen == [inner] diff --git a/studio/backend/tests/test_data_recipe_seed.py b/studio/backend/tests/test_data_recipe_seed.py index 09e22116ed..58bbd24061 100644 --- a/studio/backend/tests/test_data_recipe_seed.py +++ b/studio/backend/tests/test_data_recipe_seed.py @@ -124,3 +124,100 @@ def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, e assert result.status == "error" assert result.error == "Text extraction failed." assert _block_files(seed_route) == [] + + +_TEST_UPLOAD_UID = "0f" * 16 + + +def test_remove_unstructured_block_deletes_directory(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = _TEST_UPLOAD_UID) + assert _block_files(seed_route, _TEST_UPLOAD_UID) != [] + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": True} + assert not (seed_route.UNSTRUCTURED_UPLOAD_ROOT / _TEST_UPLOAD_UID).exists() + + +def test_remove_unstructured_block_missing_directory_is_ok(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": False} + + +def test_remove_unstructured_block_rejects_unsafe_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("../escape")) + + assert exc.value.status_code == 400 + + +def test_remove_unstructured_block_rejects_legacy_node_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = "n1") + assert _block_files(seed_route, "n1") != [] + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("n1")) + + assert exc.value.status_code == 400 + assert _block_files(seed_route, "n1") != [] + + +def test_remove_unstructured_block_rejects_symlink_escape(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + outside = tmp_path / "outside" + outside.mkdir() + (outside / "victim.txt").write_text("keep me") + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + root.mkdir(parents = True) + (root / _TEST_UPLOAD_UID).symlink_to(outside) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert exc.value.status_code == 400 + assert (outside / "victim.txt").exists() + + +def test_remove_unstructured_block_fails_if_directory_remains(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + block_dir = root / _TEST_UPLOAD_UID + block_dir.mkdir(parents = True) + (block_dir / "victim.txt").write_text("keep me") + + calls = [] + + def noop_rmtree(path, *args, **kwargs): + calls.append((path, args, kwargs)) + + monkeypatch.setattr(seed_route.shutil, "rmtree", noop_rmtree) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert calls + assert exc.value.status_code == 500 + assert block_dir.exists() + + +def test_total_upload_quota_is_scoped_per_block(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES", 10) + + first = _run_upload(seed_route, "a.txt", b"123456789") + assert first.status == "ok" + + with pytest.raises(seed_route.HTTPException) as exc: + _run_upload(seed_route, "b.txt", b"123") + assert exc.value.status_code == 413 + + # Another block starts with its own untouched budget. + other = _run_upload(seed_route, "c.txt", b"123", block_id = "other") + assert other.status == "ok" diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 4d73213d15..2fff744b64 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -287,7 +287,9 @@ def test_degrades_when_shared_helper_import_raises_importerror(): def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim - retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails.""" + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails. + The backend loads lazily (first use of a heavy helper), so this triggers the load explicitly + before asserting the retry/degrade behavior.""" import importlib import os @@ -321,11 +323,15 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.meta_path.insert(0, finder) try: degraded = importlib.import_module("utils.hf_xet_fallback") - # First attempt without the light env, then a retry with it set. + # Import is light (lazy backend); unsloth_zoo not loaded yet. + assert seen_env == [], seen_env + # First use of a heavy helper triggers the load (attempt without the light env, then a retry + # with it set); accessing DownloadStallError drives it via __getattr__. + stall_error = degraded.DownloadStallError assert seen_env == [None, "1"], seen_env # Both attempts raised -> Studio still boots in degraded mode. - assert issubclass(degraded.DownloadStallError, RuntimeError) - # The env override must not leak past the import. + assert issubclass(stall_error, RuntimeError) + # The env override must not leak past the load. assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None finally: sys.meta_path.remove(finder) @@ -333,3 +339,31 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.modules.update(saved) if saved_shim is not None: sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_importing_child_should_disable_xet_stays_light(monkeypatch): + """Regression guard for the stale-transformers-sidecar bug: importing the shim (and + ``child_should_disable_xet``) must NOT pull in ``transformers``/``unsloth_zoo``. The worker calls + this at startup to decide the Xet env flip BEFORE activating the sidecar; an eager import here + would cache the default transformers 4.57.x in sys.modules, defeating the sidecar sys.path prepend + and breaking 5.x models (Qwen3.5/GLM/gemma-4).""" + import importlib + + for name in [ + m + for m in list(sys.modules) + if m == "transformers" + or m.startswith("transformers.") + or m == "unsloth_zoo" + or m.startswith("unsloth_zoo.") + or m == "utils.hf_xet_fallback" + ]: + monkeypatch.delitem(sys.modules, name, raising = False) + + mod = importlib.import_module("utils.hf_xet_fallback") + # The lightweight decision works without the heavy backend. + assert mod.child_should_disable_xet({"disable_xet": True}) is True + assert mod.child_should_disable_xet({}) is False + # And nothing heavy was imported as a side effect. + assert "transformers" not in sys.modules, "importing the shim must not import transformers" + assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo" diff --git a/studio/backend/tests/test_inference_dispatcher_resilience.py b/studio/backend/tests/test_inference_dispatcher_resilience.py index 60d7503485..6184496d78 100644 --- a/studio/backend/tests/test_inference_dispatcher_resilience.py +++ b/studio/backend/tests/test_inference_dispatcher_resilience.py @@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env(): continue calls.append(node) - assert len(calls) == 4 + assert len(calls) == 5 for call in calls: assert any( kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index e9941d9e62..e97ca47717 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. +"""install_llama_prebuilt.py: the --resolve-prebuilt probe (plans against the fork +by default; --published-repo overrides). These back the in-app update for source-build (markerless) installs: the backend asks the installer whether an official prebuilt exists for this host without @@ -24,9 +25,7 @@ if str(_studio) not in sys.path: ilp = importlib.import_module("install_llama_prebuilt") -if not hasattr(ilp, "published_repo_for_host") or not hasattr( - ilp, "resolve_simple_install_release_plans" -): +if not hasattr(ilp, "resolve_simple_install_release_plans"): pytest.skip("PR symbols not present - check branch", allow_module_level = True) FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp @@ -56,71 +55,25 @@ def _host(**kw): return ilp.HostInfo(**base) -def test_published_repo_for_host(): - # CPU-only Linux (x64 and arm64) -> ggml-org upstream. - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) - == UPSTREAM - ) - # GPU Linux -> fork. - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) - == FORK - ) - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK - # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). - assert ( - ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) - == UPSTREAM - ) - # GPU Windows -> fork. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) - ) - == FORK - ) - # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). - assert ( - ilp.published_repo_for_host( - _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True - ) - == FORK - ) - # The tooling hint is Linux-only: Windows CPU stays on ggml-org. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True), - linux_amd_tooling_present = True, - ) - == UPSTREAM - ) - - -def test_macos_intel_and_arm_both_route_to_fork(): - # macOS uses the unslothai fork's own Mac prebuilts for BOTH arm64 and Intel; - # there is no longer any upstream-on-macOS default path, so the obsolete - # pre-macOS-26 pin (b9415) is gone. - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_x86_64 = True, machine = "x86_64") - ) - == FORK +def test_force_cpu_clears_all_gpu_attributes_including_intel(): + # --cpu-fallback is the "select the CPU prebuilt even when a GPU is present" + # escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or + # the planner still prepends the Vulkan asset on an Intel-GPU host. + host = _host( + is_linux = True, + is_x86_64 = True, + has_usable_nvidia = True, + has_physical_nvidia = True, + has_rocm = True, + rocm_gfx_target = "gfx1100", + has_intel_gpu = True, ) + forced = ilp._apply_host_overrides(host, force_cpu = True) + assert forced.has_usable_nvidia is False + assert forced.has_physical_nvidia is False + assert forced.has_rocm is False + assert forced.rocm_gfx_target is None + assert forced.has_intel_gpu is False def test_macos_upstream_pin_only_for_explicit_pre26_upstream(): @@ -188,15 +141,13 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys): assert out["repo"] == FORK -def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): - # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the - # fork so a HIP source build is not offered an upstream CPU prebuilt. - monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) - monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") +def _run_resolve_capture_host(monkeypatch, capsys): + """Drive --resolve-prebuilt and return the host the resolver was handed.""" seen = {} def _resolver(tag, host, repo, published_release_tag): seen["repo"] = repo + seen["host"] = host raise ilp.PrebuiltFallback("no asset") monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) @@ -207,10 +158,33 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): ) assert ilp.main() == ilp.EXIT_SUCCESS out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + return seen, out + + +def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys): + # CPU-only Linux host (no GPU): the dispatch routes to the fork, which now + # ships the CPU prebuilt -- it no longer falls back to ggml-org upstream. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == FORK assert out["repo"] == FORK +def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, capsys): + # A CPU-only host that merely has ROCm/HIP SDK tools on PATH (no AMD GPU, so + # detect_host leaves has_rocm False) is a valid CPU-prebuilt target. The probe + # must NOT reclassify it as ROCm from tool presence alone and suppress the CPU + # bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source + # build. The host is left CPU-only and resolves against the fork. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr( + ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == FORK + assert seen["host"].has_rocm is False + + # Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer # sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection. @@ -360,3 +334,386 @@ def test_sm103_host_drops_cuda128_windows_build(): ) kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129]) assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name] + + +def _upstream_release(tag, asset_names): + return { + "tag_name": tag, + "assets": [ + {"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names + ], + } + + +def test_direct_upstream_arm64_intel_prefers_vulkan(): + # Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU + # second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset). + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-arm64" in kinds + assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz" + + +def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only(): + # A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical + # True, usable False) + an Intel iGPU must NOT get the Vulkan archive even + # when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES + # and could grab the reserved card. It falls through to the CPU asset. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-cpu"] + + +def test_direct_upstream_arm64_without_intel_is_cpu_only(): + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64") + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-arm64"] + + +def test_direct_upstream_x86_intel_prefers_vulkan(): + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-cpu" in kinds + + +def test_linux_vulkan_health_glob_matches_bare_cpu_lib(): + # The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU + # libs so a valid Vulkan install is not re-flagged unhealthy every check. + choice = ilp.AssetChoice( + repo = UPSTREAM, + tag = "b9925", + name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", + url = "https://example/x", + source_label = "upstream", + install_kind = "linux-vulkan", + ) + groups = ilp.runtime_payload_health_groups(choice) + assert ["libggml-cpu*.so*"] in groups + assert ["libggml-cpu-*.so*"] not in groups + + +def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin(): + # Routing fork -> upstream also drops the fork release pin, which is in a + # different tag namespace and would make the upstream resolver miss. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False) + assert repo == UPSTREAM + assert tag == "" + assert routed.has_intel_gpu is True + + +def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin(): + # A pin set WITH an explicit upstream repo is already on upstream -> kept. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + _routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False) + assert repo == UPSTREAM + assert tag == "b9596" + + +def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): + # --cpu-fallback suppresses Vulkan routing even for an Intel host. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True) + assert repo == FORK + assert tag == "b9596-mix-abc" + assert routed is host + + +def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted(): + # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1): + # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or + # Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted(): + # An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_non_intel_unchanged(): + host = _host(is_linux = True, is_x86_64 = True) + routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + assert routed is host + + +def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): + # The --resolve-prebuilt probe must agree with the install path: an + # auto-detected Intel host resolves against upstream (Vulkan), not the fork. + monkeypatch.setattr( + ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == UPSTREAM + assert out["repo"] == UPSTREAM + + +# --------------------------------------------------------------------------- +# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake +# winreg module stands in for the real registry so the walk runs anywhere. +# --------------------------------------------------------------------------- + + +class _FakeRegKey: + def __init__( + self, + subkeys = None, + values = None, + denied = False, + ): + self.subkeys = subkeys or {} + self.values = values or {} + self.denied = denied + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +class _FakeWinreg: + HKEY_LOCAL_MACHINE = object() + + def __init__(self, root_key): + self._root_key = root_key + + def OpenKey(self, parent, name): + if parent is self.HKEY_LOCAL_MACHINE: + # Pin the production constant: a typo'd class GUID must fail here, + # not silently return the fake tree. + if name != ilp._WINDOWS_DISPLAY_CLASS_KEY: + raise FileNotFoundError(name) + if self._root_key is None: + raise FileNotFoundError(name) + return self._root_key + key = parent.subkeys.get(name) + if key is None: + # Real winreg raises OSError, never KeyError, for a missing key. + raise FileNotFoundError(name) + if key.denied: + raise PermissionError(name) + return key + + def QueryInfoKey(self, key): + return (len(key.subkeys), len(key.values), 0) + + def EnumKey(self, key, index): + return list(key.subkeys)[index] + + def QueryValueEx(self, key, value_name): + if value_name not in key.values: + raise FileNotFoundError(value_name) + return (key.values[value_name], 1) + + +def _probe_with_display_class(monkeypatch, adapters): + # The helper lazily does `import winreg`; plant the fake in sys.modules the + # same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path. + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters))) + return ilp.windows_intel_gpu_in_registry() + + +def test_windows_intel_registry_matches_vendor_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678", + "DriverDesc": "Intel(R) Arc(TM) A770 Graphics", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}), + }, + ) + is True + ) + + +def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684", + "DriverDesc": "NVIDIA GeForce RTX 4090", + } + ), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_1002&DEV_744C", + "DriverDesc": "AMD Radeon RX 7900 XTX", + } + ), + }, + ) + is False + ) + + +def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch): + # The real class key carries an ACL-restricted "Properties" subkey and can + # deny access to individual adapter keys; neither may abort the walk. + assert ( + _probe_with_display_class( + monkeypatch, + { + "Properties": _FakeRegKey(denied = True), + "0000": _FakeRegKey(denied = True), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_missing_class_key_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None)) + assert ilp.windows_intel_gpu_in_registry() is False + + +def _detect_windows_host( + monkeypatch, + winreg_fake, + powershell_stdout = "", +): + """Drive the real detect_host() as a GPU-less Windows host with a fake + registry, recording every run_capture invocation. Pins the wiring the + unit tests above cannot see: registry-first, CIM only on a registry miss.""" + monkeypatch.setitem(sys.modules, "winreg", winreg_fake) + monkeypatch.setattr(ilp.platform, "system", lambda: "Windows") + monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64") + for _env in ( + "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "HIP_PATH", + "ROCM_PATH", + ): + monkeypatch.delenv(_env, raising = False) + monkeypatch.setattr( + ilp.shutil, + "which", + lambda name: "powershell" if name in ("powershell", "pwsh") else None, + ) + captured = [] + + def _fake_run_capture(command, **kwargs): + captured.append(command[0]) + if command[0] == "powershell": + return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "") + return SimpleNamespace(returncode = 1, stdout = "", stderr = "") + + monkeypatch.setattr(ilp, "run_capture", _fake_run_capture) + return ilp.detect_host(), captured + + +def test_detect_host_registry_intel_skips_cim_probe(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}), + } + ) + ) + host, captured = _detect_windows_host(monkeypatch, winreg) + assert host.has_intel_gpu is True + assert "powershell" not in captured + + +def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}), + } + ) + ) + host, captured = _detect_windows_host( + monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured + + +def test_windows_intel_registry_unexpected_error_is_false(monkeypatch): + # The probe is advisory: even a non-OSError bug in the walk must return + # False (deferring to the CIM fallback), never crash detect_host. + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg()) + assert ilp.windows_intel_gpu_in_registry() is False + + +def test_detect_host_cim_rescues_exploding_registry(monkeypatch): + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + host, captured = _detect_windows_host( + monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py new file mode 100644 index 0000000000..2f04e81926 --- /dev/null +++ b/studio/backend/tests/test_llama_admission.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import asyncio +import os +import sys +import threading + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_admission +from core.inference.llama_admission import ( + ADMISSION_CONTROL_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + DEFAULT_ADMISSION_MAX_QUEUE, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + LlamaAdmissionConfig, + LlamaAdmissionQueueFull, + get_llama_admission_queue, + llama_admission_config_from_env, + reset_llama_admission_queues, +) + + +@pytest.fixture(autouse = True) +def _reset_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_admission_config_defaults(monkeypatch): + for name in ( + ADMISSION_CONTROL_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ): + monkeypatch.delenv(name, raising = False) + + config = llama_admission_config_from_env() + + assert config.enabled is True + assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE + + +def test_admission_config_env_overrides(monkeypatch): + monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off") + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25") + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0") + + config = llama_admission_config_from_env() + + assert config.enabled is False + assert config.queue_timeout_s is None + assert config.keepalive_interval_s == 0.25 + assert config.max_queue is None + + +def test_admission_config_positive_queue_timeout_env(monkeypatch): + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600") + + config = llama_admission_config_from_env() + + assert config.queue_timeout_s == 600.0 + + +def test_fifo_capacity_one_grants_next_waiter_on_release(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + third = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert third.lease_nowait() is None + assert queue.snapshot().queued == 2 + + first_lease.release() + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + second_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + third_lease.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_queue_full_rejects_excess_waiter(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(max_queue = 1) + + first = queue.reserve(capacity = 1, config = config) + queued = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert queued.lease_nowait() is None + with pytest.raises(LlamaAdmissionQueueFull): + queue.reserve(capacity = 1, config = config) + + asyncio.run(_run()) + + +def test_disabled_admission_bypasses_active_slot_limit(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(enabled = False) + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert second.lease_nowait() is not None + assert queue.snapshot().active == 0 + assert queue.snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + await asyncio.sleep(0) + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_before_delivery_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_external_waiter_future_cancel_invalidates_reservation(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + second._waiter.future.cancel() + + assert second.lease_nowait() is None + assert second.is_cancelled is True + assert await second.wait(0.01) is None + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_wait_returns_none_when_waiter_future_cancelled_during_wait(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + wait_task = asyncio.create_task(second.wait(1.0)) + await asyncio.sleep(0) + second._waiter.future.cancel() + + assert await asyncio.wait_for(wait_task, timeout = 0.1) is None + assert second.is_cancelled is True + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_capacity_increase_promotes_existing_waiter_fifo(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert queue.snapshot().active == 1 + assert queue.snapshot().queued == 1 + + third = queue.reserve(capacity = 2, config = config) + + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + snapshot = queue.snapshot() + assert snapshot.capacity == 2 + assert snapshot.active == 2 + assert snapshot.queued == 1 + + first_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + + second_lease.release() + third_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_lease_release_is_idempotent_under_concurrent_calls(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + reservation = queue.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + threads = [threading.Thread(target = lease.release) for _ in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_new_key_evicts_idle_prior_load_queues(): + # Each model load carries a fresh ephemeral port, so a new base_url key must + # not leave the drained queues from earlier loads accumulating forever. + get_llama_admission_queue("http://127.0.0.1:1001") + get_llama_admission_queue("http://127.0.0.1:1002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"} + + get_llama_admission_queue("http://127.0.0.1:1003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"} + + +def test_new_key_retains_in_flight_prior_load_queue(): + config = LlamaAdmissionConfig() + busy = get_llama_admission_queue("http://127.0.0.1:2001") + + async def _run(): + reservation = busy.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + # A new load must not drop a queue that still has an in-flight request. + get_llama_admission_queue("http://127.0.0.1:2002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"} + + # Once it drains, the next load reclaims it. + lease.release() + get_llama_admission_queue("http://127.0.0.1:2003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"} + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py new file mode 100644 index 0000000000..5525bc3ea9 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import os +import sys + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend + + +@pytest.fixture +def backend(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0) + monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None) + return LlamaCppBackend() + + +def test_effective_parallel_slots_initial_value_is_one(backend): + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend): + backend._commit_effective_parallel_slots(3) + + assert backend.effective_parallel_slots == 3 + + +@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"]) +def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value): + backend._commit_effective_parallel_slots(value) + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_reset_returns_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend._reset_effective_parallel_slots() + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_unload_resets_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend.unload_model() + + assert backend.effective_parallel_slots == 1 diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index 2a2e113585..08e1334ac9 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -137,9 +137,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path): @pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"]) def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo): - # The freshness check queries whichever release repo the marker records, - # so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right - # "latest" tag. + # The freshness check queries whichever release repo the marker records: + # new installs record the fork, legacy CPU/macOS markers still say ggml-org, + # and both must get the right "latest" tag. install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9000", published_repo = repo) bin_path = _fake_binary(install_dir, layout = "cmake") diff --git a/studio/backend/tests/test_llama_cpp_stream_cancel.py b/studio/backend/tests/test_llama_cpp_stream_cancel.py new file mode 100644 index 0000000000..1c73f4d17c --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_stream_cancel.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import contextlib +import os +import sys +import threading + +import httpx +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled + + +def _backend_stub() -> LlamaCppBackend: + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = object() + backend._healthy = True + backend._port = 48848 + backend._effective_context_length = 4096 + backend._supports_reasoning = False + backend._reasoning_always_on = False + backend._reasoning_style = "enable_thinking" + backend._supports_preserve_thinking = False + return backend + + +def test_stream_cancel_uses_internal_exception_not_generator_exit(): + class FakeResponse: + status_code = 200 + + def close(self): + pass + + class FakeStream: + def __enter__(self): + return FakeResponse() + + def __exit__(self, *_args): + return False + + class FakeClient: + def stream(self, *_args, **_kwargs): + return FakeStream() + + cancel_event = threading.Event() + + with pytest.raises(Exception) as exc_info: + with LlamaCppBackend._stream_with_retry( + FakeClient(), + "http://llama.test/v1/chat/completions", + {}, + cancel_event, + ): + cancel_event.set() + raise httpx.ReadError("client closed") + + assert exc_info.type is _LlamaStreamCancelled + assert not issubclass(exc_info.type, GeneratorExit) + + +def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch): + backend = _backend_stub() + + @contextlib.contextmanager + def fake_open_stream(*_args, **_kwargs): + raise _LlamaStreamCancelled + + monkeypatch.setattr(backend, "_open_stream", fake_open_stream) + + chunks = list( + backend.generate_chat_completion( + [{"role": "user", "content": "hi"}], + cancel_event = threading.Event(), + ) + ) + + assert chunks == [] diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 4ceffbf75b..f405ebcbd1 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -448,6 +448,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" +def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): + # A Vulkan install (marker asset carries 'vulkan') must re-assert + # UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to + # CUDA/ROCm and silently replaces the Vulkan build. + install_dir = tmp_path / "llama.cpp" + binary = _write_install( + install_dir, + "b9493", + repo = "ggml-org/llama.cpp", + asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz", + ) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + def _on_start(cmd): + _write_install( + install_dir, + "b9518", + repo = "ggml-org/llama.cpp", + asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz", + ) + + popen_kwargs: dict = {} + _patch_installer_popen( + monkeypatch, + lines = ["installed\n"], + on_start = _on_start, + captured_kwargs = popen_kwargs, + ) + + assert upd.start_update()["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1" + + def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") @@ -594,9 +636,10 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): - # CPU installs come from ggml-org. Re-running into the same install-dir/repo - # reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU - # detection) is reserved for setup.sh's arm64 rescue and must not appear here. + # Legacy CPU installs recorded a ggml-org marker (new installs use the fork). + # Re-running into the same install-dir/repo reproduces the same CPU bundle; + # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's + # arm64 rescue and must not appear here. cmd = _capture_install_cmd( monkeypatch, tmp_path, diff --git a/studio/backend/tests/test_llama_cpp_vulkan_probe.py b/studio/backend/tests/test_llama_cpp_vulkan_probe.py new file mode 100644 index 0000000000..92aaab4873 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_vulkan_probe.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Vulkan free-VRAM reader regression tests on a synthetic probe output. + +Covers the post-probe handling in +``LlamaCppBackend._get_gpu_free_memory_vulkan``: + + * integrated GPUs (probe reports is_igpu=1) leave a flat per-device host + margin matching llama.cpp's --fit-target, so context auto-sizing can't + over-commit shared RAM, and report total 0 (shared RAM is not a budget), + * discrete GPUs (is_igpu=0) keep their free untouched and pass their real + total through so the fit can reserve absolute headroom, + * an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged + (ggml applies it), not stripped or filtered in Python -- the probe reports + ggml's compact ordinal, which load_model pins with ``--device Vulkan``. + +The ggml Vulkan library is never loaded: subprocess.run is mocked to emit +the tab-separated lines the real ``_vulkan_probe.py`` would print. +""" + +from __future__ import annotations + +import subprocess +import sys +import types as _types +from pathlib import Path +from unittest import mock + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import importlib as _importlib # noqa: E402 + + +def _maybe_stub(name: str, builder): + try: + _importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = _types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", lambda: _types.ModuleType("structlog")) + +from core.inference import llama_cpp as _llama_mod # noqa: E402 +from core.inference.llama_cpp import ( # noqa: E402 + LlamaCppBackend, + _llama_lib_dir, + _vulkan_lib_filename, +) + +MIB = 1024 * 1024 +GIB = 1024 * MIB + + +def _make_vulkan_install(tmp_path: Path) -> str: + """A binary whose sibling dir holds the Vulkan ggml lib, so the + reader's ``is_vulkan_backend`` sibling-file check passes.""" + bindir = tmp_path / "build" / "bin" + bindir.mkdir(parents = True) + binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server") + binary.write_bytes(b"stub") + (bindir / _vulkan_lib_filename()).write_bytes(b"stub") + return str(binary) + + +def _mock_probe(rows: list[str], captured_env: dict | None = None): + """Patch subprocess.run so the _vulkan_probe.py call returns ``rows`` + (already tab-formatted), recording the env it was launched with.""" + real_run = subprocess.run + + def fake_run(cmd, *args, **kwargs): + if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd): + if captured_env is not None: + captured_env.clear() + captured_env.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess( + args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = "" + ) + return real_run(cmd, *args, **kwargs) + + return mock.patch("subprocess.run", side_effect = fake_run) + + +def _row( + idx: int, + free_bytes: int, + is_igpu: int, + total_bytes: int = 0, +) -> str: + return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}" + + +def test_integrated_gpu_leaves_host_margin(tmp_path): + binary = _make_vulkan_install(tmp_path) + # iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target). + # total stays 0: shared system RAM is not a VRAM budget for the fit. + rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus + + +def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path): + binary = _make_vulkan_install(tmp_path) + # 6 GiB free on a partially occupied 24 GiB card: free is untouched and the + # real total flows through so the fit reserves absolute headroom (CUDA/ROCm + # parity) instead of the looser free*frac budget. + rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus + + +def test_large_discrete_gpu_is_untouched(tmp_path): + binary = _make_vulkan_install(tmp_path) + # A 48 GiB discrete card stays untouched regardless of size; only the + # iGPU flag triggers the host margin, never a VRAM/RAM ratio. + rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus + + +def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch): + # The mask is NOT stripped or filtered in Python: ggml parses it in raw + # physical-device space while this probe reports the compact post-filter + # ordinal, so mixing spaces would be wrong. It is passed through unchanged + # so ggml applies it to the same device list the launch will enumerate. + binary = _make_vulkan_install(tmp_path) + monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1") + captured: dict = {} + rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows, captured_env = captured): + LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured + + +def test_vulkan_pin_args_uses_device_names_not_env_mask(): + # Pin by compact device name via --device (the space the probe reports and + # the registry names), never by writing a compact ordinal into the raw + # GGML_VK_VISIBLE_DEVICES index space. + assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"] + assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"] + assert LlamaCppBackend._vulkan_pin_args(None) == [] + assert LlamaCppBackend._vulkan_pin_args([]) == [] + + +def test_vulkan_only_build_is_detected(tmp_path): + binary = _make_vulkan_install(tmp_path) + assert LlamaCppBackend._is_vulkan_backend(binary) is True + + +def test_multi_backend_build_is_not_vulkan_only(tmp_path): + # A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be + # treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan + # device; defer to the CUDA/HIP path instead. + binary = _make_vulkan_install(tmp_path) + cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so" + (_llama_lib_dir(binary) / cuda).write_bytes(b"stub") + assert LlamaCppBackend._is_vulkan_backend(binary) is False + + +@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX") +def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path): + # create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root + # when it cannot symlink; _find_llama_server_binary returns that root entrypoint, + # so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else + # _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently + # never engage on a valid Vulkan install. + import os + + binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib + bindir = Path(binary).parent + wrapper = tmp_path / "llama-server" + wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n') + os.chmod(wrapper, 0o755) + assert _llama_lib_dir(str(wrapper)) == bindir + assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index 24866bd03e..b4666e3d18 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel(): asyncio.run(_run(False)) asyncio.run(_run(True)) + + +def test_stream_stall_timeout_callable_re_resolved_each_read(): + # The OpenAI passthrough passes a callable so the stall bound can switch to + # the short post-terminal grace mid-stream; it must be re-resolved per read, + # not captured once at generator start. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + values = iter([100.0, 2.0]) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 3: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 1, + post_first_item_read_timeout_s = lambda: next(values, 5.0), + ): + seen.append(response.request.extensions["timeout"].get("read")) + + assert len(seen) == 3 + # The callable is resolved right after the first item (arming the + # post-first window) and again before each later read, consuming + # successive values. + assert seen[0] == 100.0 + assert 1.0 <= seen[1] <= 2.0 + assert 4.0 <= seen[2] <= 5.0 + + asyncio.run(_run()) + + +def test_stream_stall_timeout_disabled_clears_read_timeout(): + # UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so + # the callable returns None. Once a chunk has arrived the leftover + # first-token read timeout must be cleared, else a long post-first-chunk gap + # trips a stale deadline the operator asked to turn off. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 2: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 5, + post_first_item_read_timeout_s = lambda: None, + ): + seen.append(response.request.extensions["timeout"].get("read")) + + # The first-token path armed a finite read timeout; after the first chunk + # with the guard disabled, it is cleared to None on every subsequent read. + assert seen == [None, None], seen + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index ac4088fb25..55a3198a6b 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -4,6 +4,8 @@ import sys import types from types import SimpleNamespace +import pytest + class _DummyMetal: @staticmethod @@ -185,6 +187,129 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri assert isinstance(backend._tokenizer, _DummyTokenizer) +def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch): + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + from core.inference.mlx_inference import MLXInferenceBackend + + group = SimpleNamespace(size = lambda: 2, rank = lambda: 0) + config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False) + for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")): + calls.clear() + assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group) + _, kwargs = calls.pop() + assert kwargs["text_only"] is False and kwargs[group_key] is group + + calls.clear() + singleton = SimpleNamespace(size = lambda: 1, rank = lambda: 0) + assert MLXInferenceBackend().load_model( + config, parallel_mode = "tensor", distributed_group = singleton + ) + assert not {"tensor_group", "pipeline_group"} & set(calls.pop()[1]) + + config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True) + with pytest.raises(ValueError, match = "LoRA adapter repos"): + MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group) + + +@pytest.mark.parametrize("accepts_backend", (True, False)) +def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _init_mlx_distributed + + group = SimpleNamespace(rank = lambda: 1, size = lambda: 2) + calls = [] + + def _init(**kwargs): + calls.append(kwargs) + if kwargs and not accepts_backend: + raise TypeError("backend keyword unsupported") + return group + + sys.modules["mlx.core"].distributed = SimpleNamespace(init = _init) + monkeypatch.setenv("MLX_JACCL_COORDINATOR", "127.0.0.1:12345") + monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json") + + assert _init_mlx_distributed() == (group, 1, 2) + assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}]) + + +def test_worker_share_object_receives_distributed_payload(monkeypatch): + from core.inference import worker + + shared_obj = {"type": "turn", "text": "hi"} + payload = worker._encode_share_object(shared_obj) + + def _array(value): + val = value.item() if hasattr(value, "item") else value + return SimpleNamespace( + item = lambda: val, + tolist = lambda: list(val) if hasattr(val, "__iter__") else [val], + ) + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.uint8 = "uint8" + mlx_core.array = _array + mlx_core.zeros = lambda *_a, **_k: _array([]) + + def _all_sum(value, group = None): + value = value.item() if hasattr(value, "item") else value + return _array(len(payload)) if value == 0 else _array(payload) + + mlx_core.distributed = SimpleNamespace(all_sum = _all_sum) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 1, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": None}, + SimpleNamespace(put = responses.append), + ) + + response = responses[0] + assert response["object"] == shared_obj + + +def test_worker_share_object_oversize_notifies_peers(monkeypatch): + from core.inference import worker + + calls = [] + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.array = lambda value, **_kwargs: SimpleNamespace(item = lambda: value) + mlx_core.eval = lambda value: value + mlx_core.distributed = SimpleNamespace( + all_sum = lambda value, group = None: calls.append(value.item()) or value + ) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setattr(worker, "_SHARE_OBJECT_MAX_BYTES", 8) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 0, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": {"text": "too long"}}, + SimpleNamespace(put = responses.append), + ) + + assert calls == [worker._SHARE_OBJECT_ERROR_SIZE] + assert responses[0]["type"] == "share_error" + + # Regression: generate_chat_response must accept the four template kwargs # (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route # layer can forward UI toggles. The old signature raised diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index d5e8d13652..86e528ae67 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path: from hub.utils.download_manifest import ExpectedFile from hub.utils.gguf import is_mtp_drafter_path -from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files +from hub.utils.gguf_plan import ( + build_gguf_variant_plans, + plan_from_expected_files, + preferred_mtp_sibling, +) from utils.models.model_config import ( _is_mtp_drafter, detect_gguf_model, @@ -37,6 +41,8 @@ from utils.models.model_config import ( DRAFTER_CASES = [ ("mtp-gemma-4-12b-it.gguf", True), ("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True), + # New-scheme MTP/ copies carry the mtp- basename prefix too. + ("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True), ("foo/MTP/bar.gguf", True), ("gemma-4-12b-it-Q8_0.gguf", False), # Baked-in Qwen MTP repos: the head is inside the main GGUF, the file @@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path): assert detect_gguf_model(str(copy)) is None # Selecting the MTP dir itself must not surface the copies as models. assert detect_gguf_model(str(sub)) is None + + +# ── Root drafter wins over new-scheme MTP/ copies ──────────────────── +# The MTP/ copies were renamed to share the mtp- basename prefix (e.g. +# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the +# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes +# lowercase, so the subdir path would otherwise win). + +NEW_SCHEME_SIBLINGS = [ + _sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"), + _sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"), + _sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"), + _sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"), + _sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"), + _sib("mmproj-F16.gguf", 500, "mmproj"), +] + + +def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies(): + picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS) + assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf" + + +def test_variant_plans_new_scheme_uses_root_drafter(): + plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS) + assert set(plans) == {"q4_k_m", "q8_0"} + for plan in plans.values(): + assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames + assert not any("MTP/" in name for name in plan.target_filenames) + assert "drafter" in plan.companion_hashes + # Download size = main + mmproj + root drafter (not the 200-byte BF16 copy). + assert plans["q4_k_m"].download_size_bytes == 4_600 + + +def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch): + # _pick_mtp is nested; capture it via the companion-download seam. + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) # online: skip reuse probe + captured = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + captured["pick"] = pick + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + + repo_files = [ + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "MTP/mtp-gemma-4-E4B-it-Q4_0.gguf", + "MTP/mtp-gemma-4-E4B-it-Q8_0.gguf", + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mmproj-F16.gguf", + "mtp-gemma-4-E4B-it.gguf", + ] + assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf" + + +# ── Reuse an on-disk drafter offline; fetch fresh online ───────────── + + +def _seed_snapshot(tmp_path, names): + snap = tmp_path / "snap" + for rel in names: + f = snap / rel + f.parent.mkdir(parents = True, exist_ok = True) + f.write_bytes(b"x") + return snap + + +def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch): + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mtp-gemma-4-E4B-it.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "mmproj-F16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch): + # Pre-fix build may have fetched only the MTP/ copy; reuse it offline. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf" + + +def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch): + # A newer partial snapshot holds only the MTP/ copy; an older one has the + # root. Must still return the small root, not the large subdir copy. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"]) + snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch): + # Two snapshots both hold a root drafter; newest-first order must win so a + # fresh main GGUF is not paired with a stale drafter revision. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"]) + oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).parent.parent.name == "newest" + + +def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): + # Online, do not reuse a cached copy: go to the download path so a changed + # drafter is refetched (hf_hub_download checks the current revision). + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + reached = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + reached["hit"] = True + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None + assert reached.get("hit") is True diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 4499881c4d..e24e2ca451 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -79,9 +79,11 @@ from huggingface_hub import constants as hf_constants from core.inference.llama_cpp import ( LlamaCppBackend, + _cached_colocated_split_main, _gguf_files_for_variant, _hf_offline_if_dns_dead, _probe_dns_dead, + _resolve_repo_id_casing, ) from utils.models.model_config import ( _detect_gguf_from_hf_cache, @@ -217,7 +219,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch( "huggingface_hub.list_repo_files", @@ -239,6 +241,214 @@ class TestGgufVariantFileResolution: assert downloaded == ["tinyllamas/stories260K.gguf"] assert out == "/fake/ggml-org/models/tinyllamas/stories260K.gguf" + def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial( + self, monkeypatch, hf_cache + ): + # Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download + # resumes the partial current-ref download and revalidates the revision instead + # of serving an older snapshot's same-name blob. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache( + hf_cache, + repo, + {"model-UD-Q4_K_XL.gguf": 4}, + snapshot_sha = "a" * 40, + ) + _build_cache( + hf_cache, + repo, + {"mtp-model.gguf": 1}, + snapshot_sha = "b" * 40, + ) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf", "mtp-model.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it( + self, monkeypatch, hf_cache + ): + # Case-variant cross-dir reuse is offline-only; online the canonical repo id + # resolves up front and hf_hub_download fetches the current revision. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + gguf_file = "gemma-4-E2B-it-UD-Q4_K_XL.gguf" + snap = _build_cache( + hf_cache, + canonical_repo, + {gguf_file: 4}, + snapshot_sha = "a" * 40, + ) + lower_snap = _build_cache( + hf_cache, + requested_repo, + {"mtp-gemma-4-E2B-it.gguf": 1}, + snapshot_sha = "b" * 40, + ) + os.utime(lower_snap, (2000, 2000)) + os.utime(snap, (1000, 1000)) + seen_repos: list[str] = [] + + def fake_list_repo_files(repo_id, token = None): + seen_repos.append(repo_id) + return [gguf_file] + + def fake_get_paths_info( + repo_id, + paths, + token = None, + ): + seen_repos.append(repo_id) + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fake_cache(repo_id, filename, *args, **kwargs): + seen_repos.append(repo_id) + return str(snap / filename) if repo_id == canonical_repo else None + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", fake_cache), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = requested_repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(snap / gguf_file) + assert seen_repos + + def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache): + # Online, an older same-name snapshot must not be served (it may be a stale + # revision); hf_hub_download is called so the current revision is fetched and + # its etag revalidated. + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + downloaded.append(filename) + return f"/fresh/{filename}" + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert downloaded == ["model-UD-Q4_K_XL.gguf"] + assert out == "/fresh/model-UD-Q4_K_XL.gguf" + + def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): + # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline + # cache reuse must trigger for those too, otherwise the earlier Hub calls run + # offline while this branch still attempts hf_hub_download and the cached GGUF + # cannot load. + monkeypatch.setenv("HF_HUB_OFFLINE", "true") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_companion_resolves_from_case_variant_snapshot_offline( + self, monkeypatch, hf_cache + ): + # Offline, resolve_cached_repo_id_case can keep a partial lower-case spelling, + # so the companion (mmproj) must resolve from whichever case-variant snapshot + # actually holds it rather than being dropped by an hf_hub_download on the + # wrong casing. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + snap = _build_cache(hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40) + # A partial lower-case dir exists so casing resolution keeps the requested spelling. + _build_cache(hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40) + + _offline_exc = type("OfflineModeIsEnabled", (Exception,), {}) + + def fake_list_repo_files(repo_id, token = None): + raise _offline_exc("offline") + + def fail_download(*_args, **_kwargs): + raise AssertionError("should resolve the companion from cache, not download") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_mmproj(hf_repo = requested_repo) + + assert out == str(snap / "mmproj-F16.gguf") + def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path): backend = LlamaCppBackend() downloaded: list[str] = [] @@ -264,7 +474,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), @@ -279,6 +489,48 @@ class TestGgufVariantFileResolution: assert downloaded == files assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF" + def test_download_refetches_split_gguf_when_shards_span_snapshots(self, monkeypatch, hf_cache): + # The cached main shard lives in an older snapshot; its sibling shard is only + # in a newer, separate snapshot. Reusing the main shard alone would leave + # llama.cpp unable to resolve the sibling, so the whole set must be re-fetched + # together (co-located) rather than served split across snapshot dirs. + backend = LlamaCppBackend() + repo = "org/split" + files = [ + "model-Q4_K_M-00001-of-00002.gguf", + "model-Q4_K_M-00002-of-00002.gguf", + ] + _build_cache(hf_cache, repo, {files[0]: 4}, snapshot_sha = "a" * 40) + _build_cache(hf_cache, repo, {files[1]: 4}, snapshot_sha = "b" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "Q4_K_M") + + assert downloaded == files + assert out == f"/fake/{repo}/{files[0]}" + def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" @@ -315,6 +567,21 @@ class TestIterHfCacheSnapshots: out = list(_iter_hf_cache_snapshots("unsloth/multi")) assert [p.name for p in out] == ["b" * 40, "a" * 40] + def test_skips_snapshot_when_mtime_is_unavailable(self, hf_cache, monkeypatch): + stale = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) + good = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + original_stat = Path.stat + + def flaky_stat(self, *args, **kwargs): + if self == stale: + raise FileNotFoundError(str(self)) + return original_stat(self, *args, **kwargs) + + monkeypatch.setattr(Path, "stat", flaky_stat) + + out = list(_iter_hf_cache_snapshots("unsloth/multi")) + assert out == [good] + def test_repo_id_match_is_case_insensitive(self, hf_cache): _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1}) # Lookup with different org/name casing still resolves @@ -347,6 +614,87 @@ class TestListGgufVariantsFromCache: assert _list_gguf_variants_from_hf_cache("unsloth/absent") is None +class TestCachedColocatedSplitMain: + def test_prefers_older_complete_snapshot_over_newer_partial(self, hf_cache): + # Newer snapshot has only shard 1; older snapshot has the complete set. The + # complete older snapshot must win so the split GGUF can load co-located. + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + old = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100, shard2: 100}, snapshot_sha = "a" * 40 + ) + new = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + main = _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) + assert main is not None + assert main.startswith(str(old)) + + def test_returns_none_when_shards_span_snapshots(self, hf_cache): + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + a = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40) + b = _build_cache(hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40) + os.utime(a, (1000, 1000)) + os.utime(b, (2000, 2000)) + + assert _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) is None + + +class TestResolveRepoIdCasing: + def test_maps_to_canonical_casing(self, monkeypatch): + monkeypatch.setattr( + "utils.paths.resolve_cached_repo_id_case", + lambda repo: "unsloth/Gemma-4-GGUF" if repo.lower() == "unsloth/gemma-4-gguf" else repo, + ) + # A companion download passed the resolved id reads the same cache entry + # as the main GGUF instead of missing it under the requested casing. + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/Gemma-4-GGUF" + + def test_passthrough_on_resolver_error(self, monkeypatch): + def boom(_repo): + raise RuntimeError("resolver unavailable") + + monkeypatch.setattr("utils.paths.resolve_cached_repo_id_case", boom) + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/gemma-4-gguf" + + def test_companion_only_newer_snapshot_does_not_shadow_real_variants(self, hf_cache): + # A newer snapshot holds only a vision projector fetched on demand, + # while the quant files live in an older snapshot. The newer snapshot + # must not shadow the real variants; the vision flag carries over. + old = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"vision-Q4_K_M.gguf": 100}, + snapshot_sha = "a" * 40, + ) + new = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"mmproj-vision-F16.gguf": 10}, + snapshot_sha = "b" * 40, + ) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert [v.quant for v in variants] == ["Q4_K_M"] + assert has_vision is True + + def test_companion_only_cache_returns_empty_variants_with_vision(self, hf_cache): + # Only a vision projector is cached anywhere: report the vision flag + # with an empty variant list rather than None. + _build_cache(hf_cache, "unsloth/vision-GGUF", {"mmproj-vision-F16.gguf": 10}) + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert variants == [] + assert has_vision is True + + class TestListGgufVariantsOffline: def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 7d2e2213b3..d02a2a4f7e 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -3037,3 +3037,60 @@ def test_acquire_swap_gate_is_cancellation_safe(): inference_route._auto_switch_process_lock.release() asyncio.run(asyncio.wait_for(main(), timeout = 5)) + + +def test_no_model_loaded_detail_appends_hint_only_when_off(monkeypatch): + # The "no model loaded" errors point at the opt-in auto-switch toggle so a + # request naming a listed-but-unloaded model is self-explanatory -- but only + # when it's off. With it on the name simply didn't resolve, so no hint. + base = "No GGUF model loaded. Load a GGUF model first." + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + off = inference_route._no_model_loaded_detail(base) + assert off.startswith(base) + assert "Model auto-switch" in off and "Settings > API" in off + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert inference_route._no_model_loaded_detail(base) == base + + +def _run_responses_stream_no_model(monkeypatch, *, enabled, active_model_name): + # Drive _responses_stream's GGUF-not-loaded guard: llama backend unloaded, + # inference backend maybe holding a non-GGUF model. Returns the 400 detail. + from fastapi import HTTPException + from models.inference import ResponsesRequest, ChatMessage + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: _FakeBackend(loaded_id = None) + ) + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: type("_B", (), {"active_model_name": active_model_name})(), + ) + payload = ResponsesRequest(model = "unsloth/Qwen3.5-4B-GGUF", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route._responses_stream(payload, messages, None)) + assert exc.value.status_code == 400 + return exc.value.detail + + +def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeypatch): + # Streaming /v1/responses shares the GGUF-only 400 with the other "no model + # loaded" sites, so the auto-switch hint attaches whenever the toggle is + # off -- including while a non-GGUF model is active, since auto-switch + # evicts it to load a resolved GGUF (_maybe_auto_switch_model's resolver + # branch has no active-model guard, unlike its reload-stash branch). Only + # the toggle being on suppresses it. + hinted = _run_responses_stream_no_model(monkeypatch, enabled = False, active_model_name = None) + assert "Model auto-switch" in hinted + + on = _run_responses_stream_no_model(monkeypatch, enabled = True, active_model_name = None) + assert "Model auto-switch" not in on + + non_gguf_loaded = _run_responses_stream_no_model( + monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct" + ) + assert "Model auto-switch" in non_gguf_loaded diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 05e017ba7f..910818d7d8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -8,6 +8,7 @@ import sys import asyncio import json import threading +import time from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -29,7 +30,17 @@ from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from core.inference.api_monitor import ApiMonitor +from core.inference.llama_admission import ( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + get_llama_admission_queue, + reset_llama_admission_queues, +) from routes.inference import ( + _aclose_stream_resources, _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, @@ -38,24 +49,69 @@ from routes.inference import ( _coalesce_consecutive_user_turns, _drop_empty_assistant_sentinels, _effective_max_tokens, + _effective_openai_max_tokens, + _effective_openai_max_tokens_from_values, _extract_content_parts, _friendly_error, _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, + _normalize_openai_passthrough_sse_line, + _openai_compat_stream_stall_timeout, + _openai_llama_admission_capacity, _openai_messages_for_gguf_chat, + _openai_passthrough_sse_line_terminal_state, + _openai_passthrough_upstream_headers, _openai_passthrough_non_streaming, _openai_passthrough_stream, + _responses_stream, + _openai_stream_error_sse, _openai_stream_usage_chunk, + _openai_admission_wait_stream_chunks, + _wait_for_openai_admission_non_streaming, _proxy_to_external_provider, _SameTaskStreamingResponse, + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, _set_or_prepend_system_message, openai_completions, openai_embeddings, openai_chat_completions, ) -from state.tool_policy import reset_tool_policy +from state.tool_policy import reset_tool_policy, set_tool_policy + + +@pytest.fixture(autouse = True) +def _reset_admission_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_aclose_stream_resources_attempts_remaining_closes_after_cancel(): + class Closeable: + def __init__(self, *, cancel = False): + self.cancel = cancel + self.closed = False + + async def aclose(self): + self.closed = True + if self.cancel: + raise asyncio.CancelledError() + + async def _run(): + iterator = Closeable(cancel = True) + resp = Closeable() + client = Closeable() + + with pytest.raises(asyncio.CancelledError): + await _aclose_stream_resources(iterator = iterator, resp = resp, client = client) + + assert iterator.closed + assert resp.closed + assert client.closed + + asyncio.run(_run()) class TestFriendlyUpstreamError: @@ -548,6 +604,263 @@ class TestChatCompletionRequestToolFields: assert "n > 1 is not supported" in entry["error"] assert monitor.active_count() == 0 + def test_client_tools_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must not fall through to the standard GGUF path") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + self._assert_unsupported_param(resp, "tools") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + + def test_client_tools_use_passthrough_capability_when_tool_loop_is_disabled(self, monkeypatch): + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop must stay disabled") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch): + # DiffusionGemma forces supports_tools off while passthrough stays + # available (#6851): enable_tools=True must not steal client tools + # from the passthrough into a Studio tool loop that cannot run. + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop cannot run on a non-tool backend") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "enable_tools": True, + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_tool_choice_none_allows_tool_catalog_without_tool_template(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + "tool_choice": "none", + }, + ) + + assert resp.status_code == 200 + assert resp.json()["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + + def test_tool_call_history_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError( + "tool-call history must not fall through to the standard GGUF path" + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [ + {"role": "user", "content": "use a tool"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "{}"}, + ], + }, + ) + + self._assert_unsupported_param(resp, "messages") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: is_loaded = False @@ -748,11 +1061,32 @@ class TestBuildPassthroughPayloadToolChoice: ) assert body.get("stream_options") == {"include_usage": False} + def test_response_format_without_tools_omits_tool_fields(self): + args = self._args() + args["openai_tools"] = None + + body = _build_passthrough_payload( + **args, + response_format = {"type": "json_object"}, + ) + + assert body["response_format"] == {"type": "json_object"} + assert "tools" not in body + assert "tool_choice" not in body + def test_repetition_penalty_renamed(self): body = _build_passthrough_payload(**self._args(), repetition_penalty = 1.1) assert body.get("repeat_penalty") == 1.1 assert "repetition_penalty" not in body + def test_omitted_passthrough_max_tokens_uses_backend_context(self): + args = self._args() + args["max_tokens"] = None + + body = _build_passthrough_payload(**args, backend_ctx = 4096) + + assert body["max_tokens"] == 4096 + def test_passthrough_body_merges_system_and_developer_messages(self): payload = ChatCompletionRequest( model = "default", @@ -772,6 +1106,74 @@ class TestBuildPassthroughPayloadToolChoice: ] +class TestOpenAIPassthroughSSETerminalState: + def test_done_sentinel(self): + assert _openai_passthrough_sse_line_terminal_state("data: [DONE]") == "done" + + def test_finish_reason_with_space(self): + line = 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_finish_reason_without_space(self): + line = 'data:{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_usage_chunk(self): + line = 'data: {"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "usage" + + def test_error_chunk(self): + line = 'data: {"error":{"message":"boom"}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "error" + + def test_cap_parallel_tool_calls_accepts_no_space_after_data_colon(self): + line = ( + 'data:{"choices":[{"delta":{"tool_calls":[' + '{"index":0,"function":{"name":"a"}},' + '{"index":1,"function":{"name":"b"}}]}}]}' + ) + + capped = _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) + + data = json.loads(capped[len("data:") :].lstrip()) + assert data["choices"][0]["delta"]["tool_calls"] == [ + {"index": 0, "function": {"name": "a"}} + ] + + def test_plain_content_line_is_returned_identically(self): + # The relay dispatches terminal classification on `out_line is raw_line`, + # so the no-mutation path must return the identical string object. + line = 'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}' + assert _normalize_openai_passthrough_sse_line(line) is line + assert _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) is line + + def test_reasoning_key_inside_content_text_keeps_line_identical(self): + # Fast-path substring gate fires, but the parse finds nothing to change: + # the original object must come back so the relay stays byte-identical. + line = ( + 'data: {"choices":[{"index":0,"delta":{"content":' + '"mentions \\"reasoning_content\\" in text"},"finish_reason":null}]}' + ) + assert _normalize_openai_passthrough_sse_line(line) is line + + def test_reasoning_only_delta_gets_empty_content(self): + line = ( + 'data: {"choices":[{"index":0,' + '"delta":{"reasoning_content":"thinking"},' + '"finish_reason":null}]}' + ) + + normalized = _normalize_openai_passthrough_sse_line(line) + + data = json.loads(normalized[len("data:") :].lstrip()) + delta = data["choices"][0]["delta"] + assert delta["reasoning_content"] == "thinking" + assert delta["content"] == "" + + def test_reasoning_normalization_preserves_done_sentinel(self): + assert _normalize_openai_passthrough_sse_line("data: [DONE]") == "data: [DONE]" + + # ===================================================================== # Passthrough reasoning kwargs — enable_thinking / reasoning_effort / # preserve_thinking must reach llama-server via chat_template_kwargs, @@ -892,6 +1294,172 @@ class TestOpenAICompatibilityHelpers: payload = SimpleNamespace(max_tokens = 128, max_completion_tokens = 64) assert _effective_max_tokens(payload) == 64 + def test_openai_compat_max_tokens_returns_none_when_omitted(self): + payload = SimpleNamespace(max_tokens = None, max_completion_tokens = None) + assert _effective_openai_max_tokens(payload) is None + + @pytest.mark.parametrize( + ("payload", "expected"), + [ + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = None), 8192), + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = 256), 256), + ], + ) + def test_openai_compat_explicit_values_pass_through(self, payload, expected): + assert _effective_openai_max_tokens(payload) == expected + + @pytest.mark.parametrize( + ("payload", "param"), + [ + (SimpleNamespace(max_tokens = "128", max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = True, max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = 12.5, max_completion_tokens = None), "max_tokens"), + ( + SimpleNamespace(max_tokens = None, max_completion_tokens = "128"), + "max_completion_tokens", + ), + ], + ) + def test_openai_compat_max_tokens_rejects_non_integer_explicit_values(self, payload, param): + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens(payload) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == param + assert exc.value.detail["error"]["code"] == "invalid_type" + + def test_openai_compat_max_tokens_zero_is_valid_and_negative_rejected(self): + # Legacy completions spec: max_tokens has minimum 0, so 0 must pass + # through; only negatives are invalid_value. + assert _effective_openai_max_tokens_from_values(0) == 0 + + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens_from_values(-1) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["code"] == "invalid_value" + assert exc.value.detail["error"]["param"] == "max_tokens" + + def test_chat_reasoning_chunk_carries_empty_content(self): + from routes.inference import _chat_reasoning_chunk + + line = _chat_reasoning_chunk("chatcmpl-test", 123, "gguf", "thinking...") + chunk = json.loads(line[len("data: ") :]) + delta = chunk["choices"][0]["delta"] + + assert delta["reasoning_content"] == "thinking..." + assert delta["content"] == "" + + def test_passthrough_upstream_headers_include_backend_auth(self): + headers = _openai_passthrough_upstream_headers( + llama_backend = SimpleNamespace(_auth_headers = {"Authorization": "Bearer secret"}), + ) + + assert headers["Authorization"] == "Bearer secret" + assert headers["Connection"] == "close" + + def test_openai_admission_capacity_prefers_backend_effective_slots(self): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + ) + backend = SimpleNamespace(effective_parallel_slots = 3) + + assert _openai_llama_admission_capacity(request, backend) == 3 + + @pytest.mark.parametrize("backend_value", [None, 0, -1, "not-an-int"]) + def test_openai_admission_capacity_falls_back_to_app_state(self, backend_value): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 2)) + ) + backend = SimpleNamespace(effective_parallel_slots = backend_value) + + assert _openai_llama_admission_capacity(request, backend) == 2 + + def test_openai_admission_capacity_falls_back_to_one_without_request(self): + assert _openai_llama_admission_capacity(None, SimpleNamespace()) == 1 + + def test_openai_admission_non_streaming_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for( + _wait_for_openai_admission_non_streaming( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ), + timeout = 0.1, + ) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_admission_stream_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.stream.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + chunks = _openai_admission_wait_stream_chunks( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ) + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for(chunks.__anext__(), timeout = 0.1) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_compat_stream_stall_timeout_uses_default(self, monkeypatch): + monkeypatch.delenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raising = False) + assert _openai_compat_stream_stall_timeout() == 120.0 + + def test_openai_compat_stream_stall_timeout_uses_env_override(self, monkeypatch): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, "4.5") + assert _openai_compat_stream_stall_timeout() == 4.5 + + @pytest.mark.parametrize("raw_value", ["", "not-a-float"]) + def test_openai_compat_stream_stall_timeout_invalid_env_uses_default( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() == 120.0 + + @pytest.mark.parametrize("raw_value", ["0", "-1"]) + def test_openai_compat_stream_stall_timeout_non_positive_env_disables( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() is None + + def test_openai_stream_error_sse_closes_with_done(self): + error = {"error": {"message": "boom"}} + assert _openai_stream_error_sse(error) == ( + 'data: {"error": {"message": "boom"}}\n\n' "data: [DONE]\n\n" + ) + @pytest.mark.parametrize( "finish_reason", ["stop", "length", "tool_calls", "content_filter", "function_call"], @@ -1515,9 +2083,659 @@ class TestGgufVisionToolRouting: assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" assert "".join(d.get("content", "") for d in deltas) == "visible" assert all("" not in d.get("content", "") for d in deltas) + assert all("content" in d for d in deltas if "reasoning_content" in d) [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_stream_queued_request_sends_keepalive_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_close_after_first_chunk_cleans_tracker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-stream-close-cleanup" + + def _generate(**_kwargs): + yield "visible" + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + cancel_id = cancel_id, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert cancel_id in inf_mod._CANCEL_REGISTRY + await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + aclose = getattr(iterator, "aclose", None) + assert aclose is not None + await aclose() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_task_cancel_after_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_queued_request_sends_keepalive_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_task_cancel_after_first_chunk_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_global_enable_tools_does_not_preempt_response_format_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not steal response_format") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + assert "tools" not in captured["body"] + assert "tool_choice" not in captured["body"] + finally: + reset_tool_policy() + + def test_global_enable_tools_does_not_replace_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not replace client tools") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + finally: + reset_tool_policy() + + def test_global_enable_tools_honors_client_tool_choice_none(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + def _tools(**_kwargs): + raise AssertionError("tool_choice='none' must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "do not use tools"}], + tools = client_tools, + tool_choice = "none", + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + finally: + reset_tool_policy() + + def test_enabled_tools_without_enable_tools_keeps_response_format_passthrough( + self, monkeypatch + ): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + enabled_tools = ["web_search"], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + + def test_enabled_tools_without_enable_tools_keeps_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + enabled_tools = ["web_search"], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch): def _generate(**_kwargs): yield "planvisible" @@ -1639,6 +2857,313 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_cancel_id_stops_queued_request_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after cancel_id") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "standard-nonstream-admission-cancel" + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + try: + for _ in range(50): + if cancel_id in inf_mod._CANCEL_REGISTRY: + break + await asyncio.sleep(0.01) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + with pytest.raises(HTTPException) as exc: + await asyncio.wait_for(task, timeout = 0.5) + assert exc.value.status_code == 499 + finally: + if not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_admission_task_cancel_cleans_tracker_and_slot( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-nonstream-task-cancel" + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after task cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + with pytest.raises(asyncio.CancelledError): + await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_cancel_drains_worker_before_releasing_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): import routes.inference as inf_mod @@ -1691,6 +3216,58 @@ class TestGgufVisionToolRouting: assert entry["completion_tokens"] == 3 assert monitor.active_count() == 0 + def test_non_streaming_gguf_cancel_drains_worker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod @@ -1801,7 +3378,12 @@ class TestApiMonitorProviderAndCompletionStreams: async def is_disconnected(self): return False - async def _run_passthrough_stream(self, monkeypatch, lines): + async def _run_passthrough_stream( + self, + monkeypatch, + lines, + stream_options = None, + ): import routes.inference as inf_mod class Request: @@ -1829,6 +3411,7 @@ class TestApiMonitorProviderAndCompletionStreams: model = "default", messages = [ChatMessage(role = "user", content = "hi")], stream = True, + stream_options = stream_options, tools = [ { "type": "function", @@ -1913,6 +3496,134 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured_headers = {} + + async def fake_send(_client, req, *_args, **_kwargs): + captured_headers.update(dict(req.headers)) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + + assert "data: [DONE]\n\n" in "".join(chunks) + assert captured_headers["authorization"] == "Bearer secret" + assert captured_headers["connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_stream_keepalive_while_upstream_headers_are_pending(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, + "_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S", + 0.01, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + + first = await asyncio.wait_for(response.body_iterator.__anext__(), timeout = 0.2) + assert first == ": keep-alive\n\n" + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data: [DONE]\n\n" in body + + asyncio.run(_run()) + def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2055,6 +3766,7 @@ class TestApiMonitorProviderAndCompletionStreams: body = "".join(chunks) assert "data:" in body assert '"error"' in body + assert "data: [DONE]" in body [entry] = monitor.snapshot() assert entry["status"] == "error" assert "bad" in entry["error"] @@ -2117,7 +3829,13 @@ class TestApiMonitorProviderAndCompletionStreams: async for chunk in response.body_iterator ] body = "".join(chunks) - payload = json.loads(body.removeprefix("data: ").strip()) + events = [ + line.removeprefix("data: ") + for line in body.splitlines() + if line.startswith("data: ") + ] + assert events[-1] == "[DONE]" + payload = json.loads(events[0]) assert payload["error"]["code"] == "context_length_exceeded" assert payload["error"]["param"] == "messages" assert isinstance(payload["error"], dict) @@ -2207,6 +3925,105 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_preheader_immediate_context_retry_adopts_delayed_response( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + ok_lines = [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{"content":"OK"},' + '"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + return httpx.Response(400, content = err_body) + await gate.wait() + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in ok_lines: + yield line + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + + assert "OK" in body + assert "context_length_exceeded" not in body + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2647,6 +4464,150 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_completions_omitted_max_tokens_falls_back_to_context(self, monkeypatch): + # With no env knobs set, an omitted max_tokens must forward the + # backend's context length, exactly as on main. + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "ok"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 4096 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_forwards_spec_valid_zero_max_tokens(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": 0} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "", "finish_reason": "length"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 0, + "total_tokens": 1, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 0 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_rejects_non_integer_max_tokens_before_forwarding(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": "128"} + + class UnusedClient: + async def post(self, *_args, **_kwargs): + raise AssertionError("invalid max_tokens must not reach llama-server") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: UnusedClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + with pytest.raises(HTTPException) as exc: + await openai_completions(Request(), current_subject = "test") + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == "max_tokens" + assert exc.value.detail["error"]["code"] == "invalid_type" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_monitor_openai_chunk_records_all_choice_replies(self, monkeypatch): import routes.inference as inf_mod @@ -2793,6 +4754,8 @@ class TestApiMonitorProviderAndCompletionStreams: yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' await asyncio.sleep(3600) + cancel_id = "passthrough-stream-delete-cancel" + monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) @@ -2807,6 +4770,7 @@ class TestApiMonitorProviderAndCompletionStreams: model = "default", messages = [ChatMessage(role = "user", content = "hi")], stream = True, + cancel_id = cancel_id, tools = [ { "type": "function", @@ -2824,6 +4788,7 @@ class TestApiMonitorProviderAndCompletionStreams: SimpleNamespace( base_url = "http://llama.test", context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, _request_reasoning_kwargs = lambda *_args, **_kwargs: None, ), payload, @@ -2835,6 +4800,7 @@ class TestApiMonitorProviderAndCompletionStreams: iterator = response.body_iterator first = await anext(iterator) assert "hello" in first + assert cancel_id in inf_mod._CANCEL_REGISTRY pending = asyncio.create_task(anext(iterator)) await asyncio.sleep(0) @@ -2846,6 +4812,271 @@ class TestApiMonitorProviderAndCompletionStreams: assert entry["status"] == "cancelled" assert entry["reply"] == "hello" assert monitor.active_count() == 0 + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_immediate_task_cancel_releases_admission_and_tracker( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + async def fake_cancel_check(*_args, **_kwargs): + raise asyncio.CancelledError() + + cancel_id = "passthrough-stream-immediate-task-cancel" + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_raise_if_openai_admission_cancelled", + fake_cancel_check, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + backend = SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_stream( + self._Request(), + threading.Event(), + backend, + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_before_inner_first_chunk_runs_cleanup( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + body_holder = {} + cleanup_called = threading.Event() + + async def fake_admitted(*_args, admission_lease, tracker, **_kwargs): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + cleanup_called.set() + + class BlockingBody: + def __init__(self): + self.started = threading.Event() + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + self.started.set() + await asyncio.sleep(3600) + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + await cleanup() + + body = BlockingBody() + body_holder["body"] = body + return _SameTaskStreamingResponse( + body, + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-unstarted-cleanup" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + assert cancel_id in inf_mod._CANCEL_REGISTRY + + blocker.release() + pending = asyncio.create_task(iterator.__anext__()) + for _ in range(100): + if "body" in body_holder: + break + await asyncio.sleep(0.01) + body = body_holder["body"] + assert await asyncio.to_thread(body.started.wait, 1.0) + + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert body_holder["body"].closed + assert cleanup_called.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_after_inner_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_admitted( + *_args, + monitor_id = None, + admission_lease, + tracker, + **_kwargs, + ): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + + async def body(): + try: + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + await asyncio.sleep(3600) + except asyncio.CancelledError: + inf_mod.api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await cleanup() + + return _SameTaskStreamingResponse( + body(), + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-cancel-monitor" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + + blocker.release() + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "hello" in first + + pending = asyncio.create_task(iterator.__anext__()) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 asyncio.run(_run()) @@ -2931,6 +5162,313 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_usage_done_are_separate_sse_events(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}', + ], + stream_options = {"include_usage": True}, + ) + + assert ( + '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2' in result.body + ) + assert "data: [DONE]" in result.body + assert "}\n\ndata: [DONE]\n\n" in result.body + assert "}\ndata: [DONE]\n\n" not in result.body + + asyncio.run(_run()) + + def test_passthrough_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_admitted(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_openai_passthrough_stream_admitted", fail_admitted) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_timeout_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_queue_full_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start when admission queue is full") + + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve( + capacity = 1, + config = LlamaAdmissionConfig(max_queue = 1), + ).lease_nowait() + queued = queue.reserve(capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)) + assert blocker is not None + assert queued.lease_nowait() is None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 429 + finally: + queued.cancel() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_immediate_cancel_stops_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after client cancellation") + + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + cancel_event = threading.Event() + cancel_event.set() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = cancel_event, + ) + + assert exc.value.status_code == 499 + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after admission task cancel") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = threading.Event(), + ) + + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2990,6 +5528,407 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + async def is_disconnected(self): + return False + + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = Request(), + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + cancel_event.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_route_registers_cancel_id(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + async def is_disconnected(self): + return False + + cancel_id = "passthrough-nonstream-cancel-id" + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client) + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + ), + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + cancel_id = cancel_id, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_disconnect_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + def __init__(self): + self.disconnected = False + + async def is_disconnected(self): + return self.disconnected + + client = HangingCancelableClient() + request = Request() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + request.disconnected = True + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_event.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def post(self, *_args, **kwargs): + captured["headers"] = kwargs.get("headers") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "OK" + assert captured["headers"]["Authorization"] == "Bearer secret" + assert captured["headers"]["Connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forces_upstream_stream_false(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **kwargs): + captured["json"] = kwargs.get("json") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + stream_options = {"include_usage": True}, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert captured["json"]["stream"] is False + assert "stream_options" not in captured["json"] + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch): async def _run(): result = await self._run_passthrough_stream( @@ -3009,6 +5948,216 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_finish_without_done_closes_stream_early(self, monkeypatch): + # Some llama-server builds emit the finish chunk and then hold the HTTP + # stream open without sending [DONE]; the terminal classifier must end + # the client stream promptly instead of hanging on the open socket. + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + await asyncio.Event().wait() # upstream never closes + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + + async def _consume(): + return [chunk async for chunk in response.body_iterator] + + chunks = await asyncio.wait_for(_consume(), timeout = 2) + body = "".join(chunks) + + assert '"finish_reason":"stop"' in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stall_after_finish_closes_cleanly(self, monkeypatch): + # include_usage keeps the stream open past the finish chunk waiting for + # the usage chunk; if that never arrives, the post-terminal grace path + # must close with a clean [DONE], not an in-band error. + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + raise httpx.ReadTimeout("usage chunk never arrived") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + stream_options = {"include_usage": True}, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert '"type":"api_error"' not in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_stall_after_data_emits_error(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + raise httpx.ReadTimeout("upstream went silent") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' in body + assert '"finish_reason"' not in body.replace(" ", "") + assert '"type":"api_error"' in body.replace(" ", "") + assert "still processing the prompt" in body + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "still processing the prompt" in entry["error"] + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + class TestApiMonitorSafetensorsUsage: class _Request: @@ -3512,6 +6661,15 @@ class TestApiMonitorAudioInput: class TestResponsesChatTemplateKwargs: _messages = [ChatMessage(role = "user", content = "What is 100 - 67?")] + class _Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/responses") + method = "POST" + + async def is_disconnected(self): + return False + def test_enable_thinking_lifted_from_extra_body(self): payload = ResponsesRequest( model = "qwen-local", @@ -3544,6 +6702,113 @@ class TestResponsesChatTemplateKwargs: chat_req = _build_chat_request(payload, self._messages, stream = False) assert chat_req.enable_thinking is None + def test_responses_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + + queue = get_llama_admission_queue("http://llama.responses.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_responses_stream_cancel_after_created_finalizes_monitor_and_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start after created cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "event: response.created" in first + + with pytest.raises(asyncio.CancelledError): + await iterator.athrow(asyncio.CancelledError()) + + assert get_llama_admission_queue("http://llama.responses.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + # ===================================================================== # GGUF chat-template role alternation: coalesce orphaned user turns left diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py index 83bcc5864a..9b22ab8b05 100644 --- a/studio/backend/tests/test_passthrough_healing.py +++ b/studio/backend/tests/test_passthrough_healing.py @@ -515,6 +515,7 @@ class ScriptedClient: _url, json = None, timeout = None, + headers = None, ): self.posts.append(json) return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)]) diff --git a/studio/backend/tests/test_response_template_markers.py b/studio/backend/tests/test_response_template_markers.py new file mode 100644 index 0000000000..8c813e62f2 --- /dev/null +++ b/studio/backend/tests/test_response_template_markers.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""TEMPLATE_TO_RESPONSES_MAPPER markers must match what the templates render. + +The manual instruction/response markers are the fallback for +train_on_completions when auto-detection is unavailable, so a marker that +never matches the rendered chat template masks every assistant token and the +run dies on the all-labels-masked safety net. Six template families shipped +such markers: + + mistral - "[INST] " / " [/INST]": the surrounding spaces fold into + the neighbouring tokens ("[INST]" is a single special + token in Mistral v0.3), so the padded strings never match. + llama - same space folding, plus llama-2 tokenizes [INST] after + as bare "[" on transformers 5.x while the standalone + encoding gives "▁[", so the marker must anchor on . + starling - trailing space after "GPT4 Correct Assistant:" folds + into the next content token ("▁Hello"). + glm - "[gMASK]" renders once at text start, never before + later user turns; "" is generation scaffolding + that non-final turns render as a lone "". + qwen3-thinking - "" is stripped from non-final assistant turns + (Qwen3-Thinking-2507) or never rendered (QwQ). + zephyr - role tags are plain text, and SentencePiece tokenizes + "<|assistant|>" differently at text start than after + "\\n" mid-conversation; the markers need the leading + newline anchor to tokenize like a real turn boundary. + +Literal assertions run everywhere; the token-level masking checks need the +representative tokenizers plus unsloth_zoo and skip when either is +unavailable (offline CI). +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# model_mappings is dependency-free: load it directly so these tests run +# without the studio venv / package import side effects. +_MM_PATH = Path(_BACKEND_DIR) / "utils" / "datasets" / "model_mappings.py" +_mm_spec = importlib.util.spec_from_file_location("_marker_test_mm", _MM_PATH) +model_mappings = importlib.util.module_from_spec(_mm_spec) +_mm_spec.loader.exec_module(model_mappings) + +T2R = model_mappings.TEMPLATE_TO_RESPONSES_MAPPER + + +# ── Fixed entries: markers derived from what each representative tokenizer +# actually renders (see PR for the token-level derivation). ── +EXPECTED_FIXED = { + "mistral": {"instruction": "[INST]", "response": "[/INST]"}, + "llama": {"instruction": "[INST]", "response": "[/INST]"}, + "starling": {"instruction": "GPT4 Correct User:", "response": "GPT4 Correct Assistant:"}, + "glm": {"instruction": "<|user|>", "response": "<|assistant|>"}, + "qwen3-thinking": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "zephyr": {"instruction": "\n<|user|>\n", "response": "\n<|assistant|>\n"}, +} + +# Spot-pin some known-good entries so a refactor cannot silently change them. +EXPECTED_UNCHANGED = { + "qwen3": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "llama-3.1": { + "instruction": "<|start_header_id|>user<|end_header_id|>\n\n", + "response": "<|start_header_id|>assistant<|end_header_id|>\n\n", + }, + "phi-4": { + "instruction": "<|im_start|>user<|im_sep|>", + "response": "<|im_start|>assistant<|im_sep|>", + }, + "gemma-3": {"instruction": "user\n", "response": "model\n"}, + "gpt-oss": { + "instruction": "<|start|>user<|message|>", + "response": "<|start|>assistant<|channel|>final<|message|>", + }, +} + + +@pytest.mark.parametrize("template", sorted(EXPECTED_FIXED)) +def test_fixed_marker_literals(template): + assert T2R[template] == EXPECTED_FIXED[template] + + +@pytest.mark.parametrize("template", sorted(EXPECTED_UNCHANGED)) +def test_unchanged_marker_literals(template): + assert T2R[template] == EXPECTED_UNCHANGED[template] + + +def test_no_marker_is_empty_or_whitespace(): + for template, parts in T2R.items(): + assert parts["instruction"].strip(), template + assert parts["response"].strip(), template + + +# ── Token-level checks: markers must select exactly the assistant turns on a +# rendered two-turn fixture, and the final EOS label must never be -100. ── + +REPRESENTATIVES = { + "mistral": ["unsloth/mistral-7b-instruct-v0.3"], + "llama": ["unsloth/llama-2-7b-chat"], + "starling": ["unsloth/Starling-LM-7B-beta"], + "glm": ["unsloth/GLM-4.7-Flash"], + "qwen3-thinking": ["unsloth/Qwen3-4B-Thinking-2507", "Qwen/QwQ-32B"], + "zephyr": ["unsloth/zephyr-sft"], +} + +FIXTURE = [ + {"role": "user", "content": "zebra alpha question one?"}, + {"role": "assistant", "content": "grape reply number one."}, + {"role": "user", "content": "zebra beta question two?"}, + {"role": "assistant", "content": "grape reply number two."}, +] + + +def _load_tokenizer(repo): + try: + from transformers import AutoTokenizer + except Exception as e: # pragma: no cover + pytest.skip(f"transformers unavailable: {e}") + try: + return AutoTokenizer.from_pretrained(repo) + except OSError as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + except Exception: + # Tokenizer class newer than this transformers (e.g. GLM-4.7's + # TokenizersBackend): build directly from tokenizer.json. + try: + import json as _json + from huggingface_hub import hf_hub_download + from transformers import PreTrainedTokenizerFast + + with open(hf_hub_download(repo, "tokenizer_config.json"), encoding = "utf-8") as f: + cfg = _json.load(f) + tok_file = hf_hub_download(repo, "tokenizer.json") + + def _tokval(v): + return v["content"] if isinstance(v, dict) else v + + return PreTrainedTokenizerFast( + tokenizer_file = tok_file, + chat_template = cfg.get("chat_template"), + **{ + k: _tokval(cfg[k]) + for k in ("bos_token", "eos_token", "pad_token", "unk_token") + if cfg.get(k) is not None + }, + ) + except Exception as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + + +def _train_on_responses_only(): + try: + from unsloth_zoo.dataset_utils import train_on_responses_only + except Exception as e: + pytest.skip(f"unsloth_zoo unavailable: {e}") + return train_on_responses_only + + +@pytest.mark.parametrize( + "template,repo", + [(t, r) for t, repos in sorted(REPRESENTATIVES.items()) for r in repos], +) +def test_fixed_markers_token_level(template, repo): + tor = _train_on_responses_only() + tok = _load_tokenizer(repo) + parts = T2R[template] + + msgs = [{"role": "system", "content": "You are a terse assistant."}] + FIXTURE + try: + ids = tok.apply_chat_template(msgs, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] # transformers 5.x returns a BatchEncoding + except Exception: + ids = tok.apply_chat_template(FIXTURE, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] + + fn = tor( + None, + instruction_part = parts["instruction"], + response_part = parts["response"], + tokenizer = tok, + return_function = True, + ) + labels = fn({"input_ids": [list(ids)]})["labels"][0] + + n = len(ids) + trained = tok.decode([ids[i] for i in range(n) if labels[i] != -100]) + masked = tok.decode([ids[i] for i in range(n) if labels[i] == -100]) + + # User and system content fully masked + assert "question one" not in trained and "question one" in masked + assert "question two" not in trained and "question two" in masked + assert "terse assistant" not in trained + # EVERY assistant turn trained, not just the last + assert "reply number one" in trained + assert "reply number two" in trained + # The final EOS (last non-whitespace token) must never be -100, or the + # fine-tuned model never learns to stop generating. + i = n - 1 + while i > 0 and tok.decode([ids[i]]).strip() == "": + i -= 1 + assert labels[i] != -100, f"final token {tok.convert_ids_to_tokens(int(ids[i]))!r} is masked" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py new file mode 100644 index 0000000000..300ff92776 --- /dev/null +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unit tests for detect_think_prefill. + +Reasoning templates (Qwen3.6-style) end the generation prompt with an open +``\\n`` so the model starts reasoning immediately. skip_prompt +streaming drops that opening tag, so the safetensors/MLX paths must re-emit +it for the frontend's parser to render a thinking block. +""" + +import os +import sys + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.chat_template_helpers import detect_think_prefill + + +QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n" + + +def test_open_think_prefill_reemitted(): + """Qwen3.6-style enable_thinking=True prompt tail: \\n.""" + assert detect_think_prefill(QWEN_PROMPT + "\n") == "\n" + + +def test_bare_open_think_prefill_reemitted(): + """Prefill without trailing newline still detected.""" + assert detect_think_prefill(QWEN_PROMPT + "") == "" + + +def test_closed_think_prefill_not_reemitted(): + """enable_thinking=False prefills a closed, empty think block.""" + assert detect_think_prefill(QWEN_PROMPT + "\n\n\n\n") == "" + + +def test_prompt_without_think_untouched(): + """Non-reasoning templates produce no prefix.""" + assert detect_think_prefill(QWEN_PROMPT) == "" + + +def test_historical_think_blocks_ignored(): + """A closed think block in a prior assistant turn (preserve_thinking) + must not trigger re-emission when the generation tail is plain.""" + prompt = ( + "<|im_start|>user\nHi!<|im_end|>\n" + "<|im_start|>assistant\n\nprior reasoning\n\n\nHello!<|im_end|>\n" + "<|im_start|>user\nAgain?<|im_end|>\n<|im_start|>assistant\n" + ) + assert detect_think_prefill(prompt) == "" + + +def test_historical_blocks_plus_open_prefill(): + """Prior closed blocks plus a fresh open prefill: only the tail matters.""" + prompt = ( + "<|im_start|>assistant\n\nprior\n\n\nHello!<|im_end|>\n" + "<|im_start|>assistant\n\n" + ) + assert detect_think_prefill(prompt) == "\n" + + +def test_content_after_open_tag_not_reemitted(): + """If non-whitespace follows the tag it is not a plain prefill.""" + assert detect_think_prefill(QWEN_PROMPT + "\npartial reasoning") == "" + + +def test_empty_and_none_prompts(): + assert detect_think_prefill("") == "" + assert detect_think_prefill(None) == "" + + +def test_guard_suppresses_when_close_tag_is_special(): + """If is a special token, skip_special_tokens strips the model's + close tag, so re-emitting the open would leave an unclosed block. Guard off.""" + specials = ["<|im_end|>", "", ""] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "" + + +def test_guard_emits_when_think_not_special(): + specials = ["<|im_end|>", "<|endoftext|>"] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "\n" + + +def test_guard_default_and_empty_keep_emitting(): + assert detect_think_prefill(QWEN_PROMPT + "\n", None) == "\n" + assert detect_think_prefill(QWEN_PROMPT + "\n", []) == "\n" diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 2d3dc5fbff..e4775a10a6 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -32,16 +32,23 @@ def _load_module(monkeypatch): @pytest.mark.parametrize( "torch_version, expected", [ - # torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0, - # independent of the local +cuXXX/+rocm/+cpu suffix or patch level. - ("2.10.0+cu130", "torchao==0.16.0"), + # torch 2.10 on CUDA <= 12 -> 0.16.0 (its cpp is built for torch 2.10.0 and + # loads against the CUDA-12 PyPI wheel). Independent of patch level. + ("2.10.0+cu128", "torchao==0.16.0"), + ("2.10.0+cu126", "torchao==0.16.0"), ("2.10.0+rocm6.4", "torchao==0.16.0"), ("2.10.0+cpu", "torchao==0.16.0"), ("2.10.1", "torchao==0.16.0"), ("2.10.0", "torchao==0.16.0"), - # Pre-release / dev / rc builds: the minor is cleaned of non-digits. + # torch 2.10 on CUDA >= 13 (Blackwell / cu130): 0.16.0's CUDA-12 cpp can't + # load against a CUDA-13 torch (libcudart.so.12 error), so use 0.17.0. + ("2.10.0+cu130", "torchao==0.17.0"), + ("2.10.0+cu140", "torchao==0.17.0"), + # Pre-release / dev / rc builds: the minor is cleaned of non-digits; the + # CUDA tag still decides 0.16.0 vs 0.17.0. ("2.10.0rc1", "torchao==0.16.0"), - ("2.10.0.dev20250804+cu130", "torchao==0.16.0"), + ("2.10.0.dev20250804+cu130", "torchao==0.17.0"), + ("2.10.0.dev20250804+cu128", "torchao==0.16.0"), ("2.10rc1", "torchao==0.16.0"), # torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0. ("2.11.0+cu130", "torchao==0.17.0"), diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 54048a65dd..47c6669f8f 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -6,9 +6,15 @@ empty-chat-template crash) before train(). The real methods are bound onto a lig fake self so the production logic runs against controlled batches.""" import importlib +import json +import os +import queue +import subprocess import sys +import threading import types import unittest +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -184,5 +190,231 @@ class TestChatTemplateRendersEmpty(unittest.TestCase): self.assertFalse(s._chat_template_renders_empty()) +def _clear_trainer_module(package: str): + sys.modules.pop(f"{package}.trainer", None) + pkg = sys.modules.get(package) + if pkg is not None and hasattr(pkg, "trainer"): + delattr(pkg, "trainer") + + +def _set_training_platform(monkeypatch, package: str, backend: str): + training_mod = importlib.import_module(f"{package}.training") + from utils.hardware import hardware as hw + + monkeypatch.setattr(hw, "DEVICE", None) + monkeypatch.setattr( + training_mod.platform, + "system", + lambda: "Darwin" if backend == "mlx" else "Linux", + ) + monkeypatch.setattr( + training_mod.platform, + "machine", + lambda: "arm64" if backend == "mlx" else "x86_64", + ) + + +def _load_trainer_module( + monkeypatch, + backend: str, + package: str = "core.training", +): + _set_training_platform(monkeypatch, package, backend) + _clear_trainer_module(package) + if package in sys.modules: + importlib.reload(sys.modules[package]) + trainer_mod = importlib.import_module(f"{package}.trainer") + training_mod = importlib.import_module(f"{package}.training") + monkeypatch.setattr( + training_mod._MLXTrainerAdapter, + "_activate_transformers_for_model", + lambda self, model_name, hf_token: None, + ) + return trainer_mod + + +class _ExitedProc: + def join(self, timeout = None): + return None + + def is_alive(self): + return False + + +class _TerminableProc: + def __init__(self): + self.terminated = False + self._done = threading.Event() + + def join(self, timeout = None): + self._done.wait(timeout = timeout or 5) + + def is_alive(self): + return not self.terminated + + def terminate(self): + self.terminated = True + self._done.set() + + +def test_unsloth_trainer_dispatches_for_mlx_and_torch(monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + + mlx_trainer = trainer_mod.UnslothTrainer() + + assert type(mlx_trainer).__module__ == "core.training.training" + assert mlx_trainer.get_training_progress().status_message == "Ready to train" + + trainer_mod = _load_trainer_module(monkeypatch, "torch") + + assert trainer_mod.UnslothTrainer().__class__ is trainer_mod.UnslothTrainer + + +def test_cli_mlx_trainer_activates_before_importing_trainer(): + repo_root = Path(__file__).resolve().parents[3] + script = """ +import json +import sys +import unsloth_cli.commands.train as train_cmd +from studio.backend.core.training import training as training_mod +from utils.hardware import hardware as hw + +training_mod.platform.system = lambda: "Darwin" +training_mod.platform.machine = lambda: "arm64" +hw.DEVICE = None +events = [] + +def fake_activate(model_name, hf_token): + events.append({ + "model_name": model_name, + "trainer_loaded": "studio.backend.core.training.trainer" in sys.modules, + }) + +train_cmd._activate_mlx_transformers = fake_activate +trainer = train_cmd._create_cli_trainer("mlx-community/Qwen3-0.6B-4bit", None) +print(json.dumps({ + "trainer_module": type(trainer).__module__, + "events": events, +})) +""" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [str(repo_root), str(repo_root / "studio" / "backend"), env.get("PYTHONPATH", "")] + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd = repo_root, + env = env, + text = True, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + check = True, + ) + payload = json.loads(result.stdout) + + assert payload["trainer_module"] == "studio.backend.core.training.training" + assert payload["events"] == [ + {"model_name": "mlx-community/Qwen3-0.6B-4bit", "trainer_loaded": False} + ] + + +def test_mlx_adapter_builds_config_and_reports_completion(tmp_path, monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + captured = {} + + def fake_run_worker(config, event_queue, stop_queue): + captured["config"] = config + event_queue.put({"type": "progress", "step": 1, "total_steps": 1, "loss": 0.25}) + event_queue.put( + {"type": "complete", "status_message": "done", "output_dir": config["output_dir"]} + ) + + trainer = trainer_mod.UnslothTrainer() + monkeypatch.setattr(trainer, "_run_mlx_worker", fake_run_worker) + + assert trainer.load_model("mlx-community/Qwen3-0.6B-4bit", max_seq_length = 1024) + assert trainer.prepare_model_for_training(use_lora = False) + dataset, eval_dataset = trainer.load_and_format_dataset("org/dataset") + output_dir = tmp_path / "mlx-out" + + assert trainer.start_training( + dataset = dataset, + eval_dataset = eval_dataset, + output_dir = output_dir, + project_name = "Sales Assistant", + max_steps = 1, + learning_rate = 3e-4, + ) + trainer.training_thread.join(timeout = 5) + + progress = trainer.get_training_progress() + config = captured["config"] + assert progress.is_completed + assert progress.output_dir == str(output_dir.resolve()) + progress.status_message = "mutated" + assert trainer.get_training_progress().status_message == "done" + assert config["model_name"] == "mlx-community/Qwen3-0.6B-4bit" + assert config["project_name"] == "Sales Assistant" + assert config["hf_dataset"] == "org/dataset" + assert config["training_type"] == "Full Finetuning" + assert config["load_in_4bit"] is False + assert config["max_seq_length"] == 1024 + assert config["learning_rate"] == 3e-4 + assert config["output_dir"] == str(output_dir.resolve()) + assert config["allow_external_output_dir"] is True + + +def test_mlx_worker_helpers_cover_cli_paths(tmp_path, monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training.worker import ( + _resolve_mlx_local_dataset_files, + _resolve_mlx_output_dir, + ) + + dataset = tmp_path / "train.jsonl" + dataset.write_text('{"text":"hello"}\n', encoding = "utf-8") + monkeypatch.chdir(tmp_path) + + assert _resolve_mlx_local_dataset_files(["train.jsonl"]) == [str(dataset)] + assert _resolve_mlx_output_dir( + {"output_dir": "cli-out", "allow_external_output_dir": True}, + "mlx-community/Qwen3-0.6B-4bit", + ) == str((tmp_path / "cli-out").resolve()) + + +def test_run_mlx_training_process_applies_side_effects_before_hardware_detection(monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training import worker + from utils.hardware import hardware as hw + + order = [] + + def fake_activate(model_name, hf_token): + order.append(("activate", model_name, hf_token)) + + def fake_detect_hardware(): + order.append("detect") + hw.DEVICE = hw.DeviceType.CPU + return hw.DEVICE + + monkeypatch.delenv("HF_HUB_DISABLE_XET", raising = False) + monkeypatch.delenv("HF_HUB_ENABLE_HF_TRANSFER", raising = False) + monkeypatch.setattr(worker, "_activate_transformers_version_or_warn", fake_activate) + monkeypatch.setattr(hw, "detect_hardware", fake_detect_hardware) + + event_queue = queue.Queue() + worker.run_mlx_training_process( + event_queue = event_queue, + stop_queue = queue.Queue(), + config = {"model_name": "mlx-community/Gemma-4-12B", "disable_xet": True}, + ) + + event = event_queue.get_nowait() + assert order == [("activate", "mlx-community/Gemma-4-12B", None), "detect"] + assert os.environ["HF_HUB_DISABLE_XET"] == "1" + assert os.environ["HF_HUB_ENABLE_HF_TRANSFER"] == "0" + assert "MLX training requires Apple Silicon" in event["error"] + + if __name__ == "__main__": unittest.main() diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd094..7e7fc1af48 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -59,7 +59,6 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -88,7 +87,6 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -141,27 +139,6 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): worker._sp.run.assert_not_called() -def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): - statuses: list[str] = [] - install_mock = mock.Mock() - - monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) - monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) - monkeypatch.setattr( - worker, - "_send_status", - lambda queue, message: statuses.append(message), - ) - - worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536) - - install_mock.assert_not_called() - assert len(statuses) == 1 - assert "Blackwell" in statuses[0] - - def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): install_mock = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) diff --git a/studio/backend/tests/test_training_worker_import_discipline.py b/studio/backend/tests/test_training_worker_import_discipline.py new file mode 100644 index 0000000000..a047c91704 --- /dev/null +++ b/studio/backend/tests/test_training_worker_import_discipline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: the training worker must not import ``transformers`` before it activates the +transformers sidecar. + +``core/training/worker.py:run_training_process`` runs a preflight (Xet decision, logging, hardware +detection) and only THEN calls ``_activate_transformers_version`` -> ``activate_transformers_for_subprocess``, +which prepends the correct ``.venv_t5_*`` (5.x) sidecar to ``sys.path``. Because activation only edits +``sys.path``, it is a no-op for any module already cached in ``sys.modules``. So if the preflight imports +``transformers`` (directly or transitively via ``unsloth_zoo``), the default 4.57.x gets pinned before +the sidecar is on the path -- and 5.x models (Qwen3.5, GLM-4.7, gemma-4) then fail to load their +tokenizer/config ("Tokenizer class TokenizersBackend does not exist"). + +This regression shipped once when ``utils/hf_xet_fallback.py`` eagerly imported ``unsloth_zoo`` (which +imports ``transformers``) at module load; the worker imports that shim during preflight to decide the +Xet env flip (see issue #6951). This test locks the invariant in a fresh interpreter. It is CPU-only, +needs no network/GPU/weights/sidecars, so it runs in the standard ``studio-backend-ci`` matrix. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend + +# Mirrors run_training_process's imports that run BEFORE _activate_transformers_version (worker.py); +# keep in sync. torch-dependent imports are optional (a no-torch CI shard skips them) but must still +# not drag in transformers. +_PREFLIGHT_SNIPPET = r""" +import sys + +# worker.py: from utils.hf_xet_fallback import child_should_disable_xet (+ call it) +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) + +# worker.py: from loggers.config import LogConfig +from loggers.config import LogConfig # noqa: F401 + +# worker.py: from utils.hardware import hardware (imports torch, not transformers) +try: + from utils.hardware import hardware as _hw # noqa: F401 +except Exception: + pass # torch may be absent in a no-torch shard; the invariant below still applies + +# worker.py: from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend +# (the MLX-dispatch preflight; must also stay clear of transformers). Guarded because it may pull +# unsloth/trl, absent in a minimal shard -- but a partial import that leaked transformers would still +# be caught by the assertion below. +try: + from core.training.training import ( # noqa: F401 + is_apple_silicon_training_platform as _is_apple, + should_use_mlx_training_backend as _use_mlx, + ) +except Exception: + pass + +leaked_tf = sorted(m for m in sys.modules if m == "transformers" or m.startswith("transformers.")) +leaked_zoo = sorted(m for m in sys.modules if m == "unsloth_zoo" or m.startswith("unsloth_zoo.")) +assert not leaked_tf, f"transformers imported during worker preflight (before sidecar activation): {leaked_tf}" +assert not leaked_zoo, f"unsloth_zoo imported during worker preflight (before sidecar activation): {leaked_zoo}" +print("PREFLIGHT_CLEAN") +""" + + +def test_worker_preflight_does_not_import_transformers(): + """A fresh interpreter running the worker's pre-activation imports must leave ``transformers`` + (and ``unsloth_zoo``) unimported, so the 5.x sidecar prepend is not defeated by a stale module.""" + result = subprocess.run( + [sys.executable, "-c", _PREFLIGHT_SNIPPET], + cwd = str(_BACKEND_DIR), + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight imported transformers before sidecar activation.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + assert "PREFLIGHT_CLEAN" in result.stdout, result.stdout diff --git a/studio/backend/tests/test_worker_activates_correct_transformers.py b/studio/backend/tests/test_worker_activates_correct_transformers.py new file mode 100644 index 0000000000..fe7b8dd25a --- /dev/null +++ b/studio/backend/tests/test_worker_activates_correct_transformers.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: after the training worker runs its preflight and then activates the transformers +sidecar, the in-process ``transformers`` must be the sidecar version the model requires -- not the +default 4.57.x that the base environment ships. + +The CPU-only "does it choose the correct transformers version" guard, stronger than the pure +import-order check in ``test_training_worker_import_discipline.py``: it runs the REAL tier detection +(``get_transformers_tier``) and REAL activation (``activate_transformers_for_subprocess``) for a +transformers-5.x model (Qwen3.5, tier 530) and asserts the version actually switched. It catches the +whole failure family at once: + + * a stale pre-activation ``transformers`` import (the #6951 / ``TokenizersBackend`` regression: an + already-cached 4.57.x defeats the sidecar's ``sys.path`` prepend), + * a wrong tier selected for a 5.x model, and + * activation not actually swapping the resident module. + +Why the CUDA spoof matters (verified): ``unsloth_zoo``'s eager ``import transformers`` only happens on +its full, GPU-present init path. On a GPU-less runner it silently degrades and never preloads +transformers -- which would MASK the stale-import bug (the check would falsely pass). Spoofing +``torch.cuda`` so ``unsloth_zoo`` believes a GPU is present forces the real init path, exposing the +regression on CPU CI. The spoof mirrors ``tests/_zoo_aggressive_cuda_spoof.py`` but is inlined so the +test is self-contained in the ``studio-backend-ci`` matrix (whose conftest does not apply the shared +spoof). No GPU/network/weights/real sidecar needed: a one-line stub sidecar stands in for the 5.x venv, +so we only assert activation lands on it. + +Proven: passes on the fixed tree (active == 5.3.0) and fails on the buggy tree (active == 4.57.x) on +a simulated GPU-less runner. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend +# Canonical CUDA spoof at the repo root (studio/backend -> studio -> repo root). Loaded by the +# subprocess when present (matches the consolidated CI); absent in a standalone studio checkout, where +# the subprocess falls back to a minimal inline spoof. +_SPOOF_PATH = _BACKEND_DIR.parent.parent / "tests" / "_zoo_aggressive_cuda_spoof.py" + +# Runs in a fresh interpreter with cwd == studio/backend so ``utils.*`` resolves like the worker. +# STUB_HOME (a pytest tmp dir) holds a throwaway ``.venv_t5_530`` sidecar exporting transformers 5.3.0. +_SNIPPET = r""" +import os, sys +sys.path.insert(0, os.getcwd()) + +# CUDA spoof so unsloth_zoo takes its full, transformers-importing init path on a GPU-less runner. +# Without it unsloth_zoo degrades and never preloads transformers, which would MASK the stale-import +# regression under test (verified). Prefer the repo's canonical spoof (single source of truth, and the +# one the consolidated CI already relies on); fall back to a minimal inline spoof so this also works in +# a standalone studio checkout. If torch is absent the fixed tree still passes below; the bug just +# would not be exposable in that shard. +try: + import torch # noqa: F401 + _sp = os.environ.get("SPOOF_PATH") + if _sp and os.path.exists(_sp): + import importlib.util + _spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", _sp) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + _mod.apply() + else: + torch.cuda.is_available = lambda: True + torch.cuda.device_count = lambda: 1 + torch.cuda.current_device = lambda: 0 + torch.cuda.get_device_capability = lambda *a, **k: (8, 0) + torch.cuda.get_device_name = lambda *a, **k: "NVIDIA A100-SPOOFED" + torch.cuda.is_bf16_supported = lambda *a, **k: True + class _Props: + name = "NVIDIA A100-SPOOFED" + major = 8 + minor = 0 + total_memory = 80 * 1024**3 + multi_processor_count = 108 + torch.cuda.get_device_properties = lambda *a, **k: _Props() + torch.cuda.mem_get_info = lambda *a, **k: (0, 80 * 1024**3) +except Exception: + pass +os.environ["UNSLOTH_IS_PRESENT"] = "1" + +# Stub 5.x sidecar: activation only edits sys.path, so a package that merely exports __version__ is +# enough to prove the resident transformers switched to it. +home = os.environ["STUB_HOME"] +pkg = os.path.join(home, ".venv_t5_530", "transformers") +os.makedirs(pkg, exist_ok = True) +with open(os.path.join(pkg, "__init__.py"), "w") as f: + f.write('__version__ = "5.3.0"\n') +os.environ["UNSLOTH_STUDIO_HOME"] = home + +# Faithful worker preflight (worker.py: from utils.hf_xet_fallback import child_should_disable_xet). +# This is the exact stale-import trigger: on the buggy tree it pulls unsloth_zoo -> transformers 4.57.x +# into sys.modules BEFORE activation. +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) +_tf = sys.modules.get("transformers") +preload = _tf.__version__ if _tf is not None else None + +# Real tier detection + real activation, with the 530 sidecar pointed at the stub above. +import utils.transformers_version as tv +tv._VENV_T5_530_DIR = os.path.join(home, ".venv_t5_530") +tv._ensure_venv_t5_530_exists = lambda: True +tier = tv.get_transformers_tier("Qwen/Qwen3.5-9B", None) +tv.activate_transformers_for_subprocess("Qwen/Qwen3.5-9B", None) + +import transformers +print(f"RESULT tier={tier} preload={preload} active={transformers.__version__}") +""" + + +def _parse(stdout: str) -> dict[str, str]: + for line in stdout.splitlines(): + if line.startswith("RESULT "): + return dict(kv.split("=", 1) for kv in line.split()[1:]) + return {} + + +def test_worker_activates_correct_transformers_version(tmp_path): + """The worker's real preflight + activation for a transformers-5.x model (Qwen3.5, tier 530) must + leave the in-process ``transformers`` on the 5.x sidecar. A stale pre-activation import leaves the + default 4.57.x pinned and fails this assertion -- exactly the #6951 ``TokenizersBackend`` regression.""" + result = subprocess.run( + [sys.executable, "-c", _SNIPPET], + cwd = str(_BACKEND_DIR), + env = { + **__import__("os").environ, + "STUB_HOME": str(tmp_path), + **({"SPOOF_PATH": str(_SPOOF_PATH)} if _SPOOF_PATH.exists() else {}), + }, + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight + activation harness crashed.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + parsed = _parse(result.stdout) + assert parsed, f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + + # Correct tier chosen for a transformers-5.x model (pure, deterministic; no network/GPU). + assert parsed["tier"] == "530", ( + f"Wrong transformers tier for Qwen3.5 (expected 530, got {parsed['tier']}). " + "Tier detection regressed." + ) + + # Activation must actually swap the resident transformers to the sidecar version. If a preflight + # import cached 4.57.x first, the sidecar prepend is a no-op and this stays 4.57.x -- the bug. + assert parsed["active"] == "5.3.0", ( + "Sidecar activation did NOT switch the in-process transformers to the model's 5.x version " + f"(active={parsed['active']}, preloaded-before-activation={parsed['preload']}). A pre-activation " + "transformers import (directly or via unsloth_zoo) defeated the sidecar; 5.x models (Qwen3.5, " + "GLM-4.7, gemma-4) then fail with 'Tokenizer class TokenizersBackend does not exist'. See #6951." + ) diff --git a/studio/backend/utils/coding_agents.py b/studio/backend/utils/coding_agents.py new file mode 100644 index 0000000000..f7dd2f8357 --- /dev/null +++ b/studio/backend/utils/coding_agents.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Detect which `unsloth start ` coding-agent CLIs are on PATH. + +The web UI only ever shows the user the "claude" flavor of the `unsloth start` +command (see agent-command.ts), leaving anyone using Codex, OpenCode, and the +other supported agents to manually edit the copied command. This module gives +the frontend a way to ask which of those CLIs are actually installed so it can +default to one the user can run immediately. +""" + +import shutil + +# Keep in sync with the `unsloth start ` subcommands defined in +# unsloth_cli/commands/start.py. Each entry is the exact executable name that +# subcommand launches, so a hit here means `unsloth start ` can find the +# binary on PATH without the user installing anything first. +CODING_AGENTS: tuple[str, ...] = ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def _is_on_path(agent: str) -> bool: + # shutil.which is documented to return None on a miss, but PATH lookups can + # still raise (e.g. a permission error while probing a directory entry); + # this is an advisory check, so a lookup failure should read as "not + # installed" instead of breaking the settings endpoint. + try: + return shutil.which(agent) is not None + except OSError: + return False + + +def detect_installed_coding_agents() -> list[str]: + """Return the subset of CODING_AGENTS whose CLI binary is on PATH. + + Order follows CODING_AGENTS, not discovery order, so callers can treat the + first entry as the preferred default among the installed agents. + """ + return [agent for agent in CODING_AGENTS if _is_on_path(agent)] diff --git a/studio/backend/utils/datasets/completion_masking.py b/studio/backend/utils/datasets/completion_masking.py new file mode 100644 index 0000000000..c7c4a474e3 --- /dev/null +++ b/studio/backend/utils/datasets/completion_masking.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy shared by the CUDA and MLX training paths. + +Decides how train_on_responses_only is applied for a model: chat template +auto-detection first, manual TEMPLATE_TO_RESPONSES_MAPPER markers as the +fallback. gpt-oss included: its quantized checkpoints ship a different +chat template, so only detection from the actual template is reliable. +""" + +from .model_mappings import ( + MODEL_TO_TEMPLATE_MAPPER, + TEMPLATE_TO_RESPONSES_MAPPER, + is_gpt_oss_model_name, +) + + +def lookup_manual_markers(model_name): + """Return (template_name, instruction_part, response_part) from the + manual template table, with None parts when the model or template is + not mapped.""" + template = MODEL_TO_TEMPLATE_MAPPER.get((model_name or "").lower()) + markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template) if template else None + if markers: + return template, markers["instruction"], markers["response"] + return template, None, None + + +def apply_completion_masking( + trainer, + model_name, + train_fn, + num_proc = None, + notify = None, + detect_fn = None, +): + """Apply completion-only masking with auto-detection first and the manual + template table as fallback. + + Args: + trainer: The platform trainer (SFTTrainer or MLXTrainer). + model_name: Model repo id used for table lookup and the gpt-oss + renamed-checkpoint fallback. + train_fn: The platform train_on_responses_only callable. + num_proc: Forwarded to train_fn when not None (CUDA path only). + notify: Optional callback notify(level, message) with level "info" or + "warning" for user-visible progress and warnings. + detect_fn: Marker detector (tokenizer/processor) -> (instruction_part, + response_part). Defaults to unsloth_zoo's get_chat_template_parts, + which raises loudly when the template cannot be parsed. Test seam. + + Returns: + (trainer, applied): the possibly wrapped trainer and whether masking + was applied. When applied is False the trainer is unchanged and + training runs on full sequences. + + Only marker DETECTION failures trigger the table fallback. Exceptions + raised while applying the masking (dataset map, tokenization) propagate + to the caller in both the auto and manual paths, so a real failure stops + the run instead of silently changing the training objective. + """ + if notify is None: + notify = lambda level, message: None + kwargs = {} + if num_proc is not None: + kwargs["num_proc"] = num_proc + + template, instruction_part, response_part = lookup_manual_markers(model_name) + + # gpt-oss goes auto-first: quantized/BF16 checkpoints ship a channel-less + # template, so the manual markers match nothing (zero tokens trained). Auto + # derives markers from whichever template ships, and per the harmony format + # only the final terminator carries stop supervision. Renamed checkpoints + # miss the exact-name table, so give the fallback the gpt-oss markers. + if is_gpt_oss_model_name(model_name) and not (instruction_part and response_part): + markers = TEMPLATE_TO_RESPONSES_MAPPER.get("gpt-oss") + if markers: + template = "gpt-oss" + instruction_part = markers["instruction"] + response_part = markers["response"] + processor = getattr(trainer, "processing_class", None) or getattr(trainer, "tokenizer", None) + # mlx-lm TokenizerWrapper hides underscore attrs, so preset _unsloth_* + # markers are invisible through it. Unwrap to the real tokenizer (as + # zoo's MLX resolver does) before the preset check and detection. + if type(processor).__name__ == "TokenizerWrapper": + wrapped = getattr(processor, "_tokenizer", None) + if wrapped is not None: + processor = wrapped + inner = getattr(processor, "tokenizer", processor) + if hasattr(inner, "_unsloth_input_part") and hasattr(inner, "_unsloth_output_part"): + # Markers preset on the tokenizer; zoo reuses them on a bare call. + trainer = train_fn(trainer, **kwargs) + notify( + "info", + "Train on responses only configured via tokenizer preset markers", + ) + return trainer, True + auto_instruction = auto_response = None + try: + if detect_fn is None: + # Torch-backed import is fine: the MLX train_fn itself requires + # unsloth_zoo.dataset_utils, so a torch-free host cannot mask either way. + from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn + auto_instruction, auto_response = detect_fn(processor) + except Exception as e: + notify( + "warning", + f"Auto-detection of instruction/response markers failed ({e}); " + f"falling back to the template table", + ) + if auto_instruction and auto_response: + trainer = train_fn( + trainer, + instruction_part = auto_instruction, + response_part = auto_response, + **kwargs, + ) + notify( + "info", + "Train on responses only configured via chat template auto-detection", + ) + return trainer, True + + if instruction_part and response_part: + trainer = train_fn( + trainer, + instruction_part = instruction_part, + response_part = response_part, + **kwargs, + ) + notify( + "info", + f"Train on responses only configured with template table markers ({template})", + ) + return trainer, True + + notify( + "warning", + f"'Train on completions' could not be applied for {model_name}: no " + f"auto-detected or mapped instruction/response markers. Training " + f"will run on full sequences (prompts included).", + ) + return trainer, False diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 9d2c983aed..65ba4b4688 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -485,9 +485,11 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No "" suffix: Qwen3-Thinking-2507 strips it from non-final turns + # and QwQ renders none, so a marker holding it masks those responses. "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", @@ -525,29 +527,39 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user<|im_sep|>", "response": "<|im_start|>assistant<|im_sep|>", }, + # No surrounding spaces: in Mistral v0.3 they fold into neighbouring text + # tokens ("[INST]"/"[/INST]" are single special tokens), so padded strings + # never match and everything masks. Same for Llama-2's SentencePiece. "mistral": { - "instruction": "[INST] ", - "response": " [/INST]", + "instruction": "[INST]", + "response": "[/INST]", }, "llama": { - "instruction": "[INST] ", - "response": " [/INST]", + # -anchored: llama-2 tokenizes [INST] after as bare "[" on + # transformers 5.x (standalone gives space-prefixed "▁["), so an + # unanchored marker misses every turn boundary there. + "instruction": "[INST]", + "response": "[/INST]", }, "chatml": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # Leading newline required: Zephyr's role tags are plain text, and + # SentencePiece tokenizes "<|assistant|>" differently at text start than + # after "\n". Without the "\n" anchor the markers never match real + # turns, so every assistant token masks. "zephyr": { - "instruction": "<|user|>\n", - "response": "<|assistant|>\n", + "instruction": "\n<|user|>\n", + "response": "\n<|assistant|>\n", }, "unsloth": { - "instruction": ">>> User: ", - "response": ">>> Assistant: ", + "instruction": ">>> User:", + "response": ">>> Assistant:", }, "vicuna": { - "instruction": "USER: ", - "response": "ASSISTANT: ", + "instruction": "USER:", + "response": "ASSISTANT:", }, "alpaca": { "instruction": "### Instruction:\n", @@ -573,16 +585,21 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No trailing space: SentencePiece folds it into the next content token + # ("▁Hello"), so the padded marker never matches and masks everything. "starling": { - "instruction": "GPT4 Correct User: ", - "response": "GPT4 Correct Assistant: ", + "instruction": "GPT4 Correct User:", + "response": "GPT4 Correct Assistant:", }, "yi-chat": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # "[gMASK]" appears once at text start, so a marker holding it matches + # no later user turn; "" is scaffolding GLM-4.x renders as a lone + # "" on non-final turns, so "<|assistant|>" never matches. "glm": { - "instruction": "[gMASK]<|user|>", - "response": "<|assistant|>", + "instruction": "<|user|>", + "response": "<|assistant|>", }, } diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 2dd2247396..9bc4a60fad 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -6,6 +6,16 @@ Re-exports the shared API and injects Studio's marker-aware cache purge (``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` marker semantics on the HTTP retry. + +Import discipline: ``unsloth_zoo``'s ``__init__`` eagerly imports ``transformers``. The workers +import this shim at startup (to decide the per-worker Xet env flip) *before* activating the model's +``transformers`` sidecar. Activation only prepends the sidecar to ``sys.path``, so a ``transformers`` +already cached in ``sys.modules`` (via an eager ``unsloth_zoo`` import here) wins -- pinning the +default 4.57.x and regressing Qwen3.5 / GLM-4.7 / gemma-4 training with +``Tokenizer class TokenizersBackend does not exist``. So the shared backend is loaded **lazily** +(``_load_shared``), only on first use of a heavy download helper, i.e. after the sidecar is active. +``child_should_disable_xet`` and the ``DEFAULT_*`` constants are defined locally so importing them +never triggers the heavy load. """ from __future__ import annotations @@ -13,161 +23,230 @@ from __future__ import annotations import threading from typing import Any, Callable, Optional -_shared_import_error = None -try: - import unsloth_zoo.hf_xet_fallback as _shared - _shared_available = True -except Exception as _exc: # noqa: BLE001 - any import failure must degrade, not crash - # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less Studio - # host. The download helper needs none of it, so retry via the light UNSLOTH_ZOO_DISABLE_GPU_INIT - # path before giving up. - _shared_import_error = _exc - import os as _os +# Defaults mirror unsloth_zoo.hf_xet_fallback; plain literals so they resolve (including as +# default args below) without importing unsloth_zoo/transformers. +DEFAULT_GRACE_PERIOD = 10.0 +DEFAULT_HEARTBEAT_INTERVAL = 30.0 +DEFAULT_STALL_TIMEOUT = 180.0 - _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") - _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" - try: - import unsloth_zoo.hf_xet_fallback as _shared - _shared_available = True - _shared_import_error = None - except Exception as _exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF downloads - _shared_import_error = _exc2 - _shared_available = False - finally: - if _prev_gpu_init is None: - _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) - else: - _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init +# --- lazy shared-backend loader ---------------------------------------------------------------- +_shared: Any = None +_shared_available: Optional[bool] = None # None = not yet attempted +_shared_import_error: Optional[BaseException] = None +_load_lock = threading.Lock() -if _shared_available: - # Bind by assignment so each public name shares one module-level binding with the degraded branch. - DEFAULT_GRACE_PERIOD = _shared.DEFAULT_GRACE_PERIOD - DEFAULT_HEARTBEAT_INTERVAL = _shared.DEFAULT_HEARTBEAT_INTERVAL - DEFAULT_STALL_TIMEOUT = _shared.DEFAULT_STALL_TIMEOUT - DownloadStallError = _shared.DownloadStallError - child_should_disable_xet = _shared.child_should_disable_xet - get_hf_download_state = _shared.get_hf_download_state - start_watchdog = _shared.start_watchdog - _shared_hf_hub_download_with_xet_fallback = _shared.hf_hub_download_with_xet_fallback - _shared_snapshot_download_with_xet_fallback = _shared.snapshot_download_with_xet_fallback -else: - # Degrade instead of crashing Studio: plain HF downloads, stall watchdog disabled. Thin stubs, - # not a second copy of the orchestration; recovery returns once unsloth_zoo is upgraded. - import logging as _logging - _logging.getLogger(__name__).warning( - "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " - "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " - "re-enable automatic Xet -> HTTP download recovery.", - _shared_import_error, - ) +def _load_shared() -> bool: + """Import ``unsloth_zoo.hf_xet_fallback`` on demand; return True if available. Deferred so + importing this module at worker startup does not pull transformers in before the sidecar is + activated. Degrades (returns False) rather than crashing when unsloth_zoo is unavailable.""" + global _shared, _shared_available, _shared_import_error + if _shared_available is not None: + return _shared_available + with _load_lock: + if _shared_available is not None: + return _shared_available + try: + import unsloth_zoo.hf_xet_fallback as shared - DEFAULT_HEARTBEAT_INTERVAL = 30.0 - DEFAULT_STALL_TIMEOUT = 180.0 - DEFAULT_GRACE_PERIOD = 10.0 + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less + # host. The download helper needs none of it, so retry via UNSLOTH_ZOO_DISABLE_GPU_INIT. + _shared_import_error = exc + import os as _os - class DownloadStallError(RuntimeError): - """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" + try: + import unsloth_zoo.hf_xet_fallback as shared - def child_should_disable_xet(config: dict) -> bool: - return bool(config.get("disable_xet")) + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF + _shared_import_error = exc2 + _shared_available = False + import logging as _logging - def get_hf_download_state(*args: Any, **kwargs: Any) -> None: - return None # unmeasurable -> the (absent) watchdog never fires + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) + return False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init - def start_watchdog( - *, - on_heartbeat: "Optional[Callable[[str], None]]" = None, - interval: float = DEFAULT_HEARTBEAT_INTERVAL, - xet_disabled: bool = False, - **kwargs: Any, - ) -> "threading.Event": - # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline - # is not tripped during a long download. - stop = threading.Event() - if on_heartbeat is None: - return stop - transport = "https" if xet_disabled else "xet" - def _beat() -> None: - while not stop.wait(interval): - try: - on_heartbeat(f"Downloading ({transport} transport)...") - except Exception: - pass +def child_should_disable_xet(config: dict) -> bool: + """Single source of truth for the per-worker Xet env flip (mirrors + ``unsloth_zoo.hf_xet_fallback.child_should_disable_xet``). Deliberately lightweight: importing or + calling it must NOT pull in unsloth_zoo/transformers, so the worker can decide before activating + the transformers sidecar (see the module docstring).""" + return bool(config.get("disable_xet")) - threading.Thread( - target = _beat, - daemon = True, - name = "hf-xet-degraded-heartbeat", - ).start() + +# --- degraded stubs (used only when unsloth_zoo is unavailable) ------------------------------- +class _DegradedDownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + + +def _degraded_get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires + + +def _degraded_start_watchdog( + *, + on_heartbeat: "Optional[Callable[[str], None]]" = None, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + xet_disabled: bool = False, + **kwargs: Any, +) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. + stop = threading.Event() + if on_heartbeat is None: return stop + transport = "https" if xet_disabled else "xet" - def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: - return cancel_event is not None and cancel_event.is_set() + def _beat() -> None: + while not stop.wait(interval): + try: + on_heartbeat(f"Downloading ({transport} transport)...") + except Exception: + pass - def _shared_hf_hub_download_with_xet_fallback( - repo_id: str, - filename: str, - token: Optional[str], - *, - repo_type: str = "model", - revision: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - cancel_event: "Optional[threading.Event]" = None, - **_ignored: Any, - ) -> str: - # Keep the cancellation contract: do not start or return a download once cancelled. - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() + return stop - from huggingface_hub import hf_hub_download - path = hf_hub_download( - repo_id = repo_id, - filename = filename, - token = token, - repo_type = repo_type, - revision = revision, - cache_dir = cache_dir, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path +def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() - def _shared_snapshot_download_with_xet_fallback( - repo_id: str, - *, - revision: Optional[str] = None, - token: Optional[str] = None, - repo_type: str = "model", - cache_dir: Optional[str] = None, - allow_patterns: Optional[Any] = None, - ignore_patterns: Optional[Any] = None, - force_download: bool = False, - cancel_event: "Optional[threading.Event]" = None, - **_ignored: Any, - ) -> str: - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - from huggingface_hub import snapshot_download +def _degraded_hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - path = snapshot_download( - repo_id = repo_id, - repo_type = repo_type, - revision = revision, - token = token, - cache_dir = cache_dir, - allow_patterns = allow_patterns, - ignore_patterns = ignore_patterns, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path + from huggingface_hub import hf_hub_download + + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + token = token, + repo_type = repo_type, + revision = revision, + cache_dir = cache_dir, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +def _degraded_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + + from huggingface_hub import snapshot_download + + path = snapshot_download( + repo_id = repo_id, + repo_type = repo_type, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +# --- lazy attribute access for the heavy shared API ------------------------------------------- +# ``DownloadStallError`` (class identity matters for ``except``), ``start_watchdog`` and +# ``get_hf_download_state`` come from the shared backend when available, else the degraded stubs. +# Resolved via PEP 562 ``__getattr__`` so ``from utils.hf_xet_fallback import X`` triggers the load +# only for these heavy names, not for ``child_should_disable_xet`` / ``DEFAULT_*``. +_DEGRADED_ATTRS = { + "DownloadStallError": _DegradedDownloadStallError, + "start_watchdog": _degraded_start_watchdog, + "get_hf_download_state": _degraded_get_hf_download_state, +} + +# Annotation-only declarations for the three names above: they bind NO value, so lookup still misses +# and PEP 562 ``__getattr__`` resolves them lazily -- but ruff/pyflakes see them as defined, so listing +# them in ``__all__`` does not trip F822 (while F822 still catches a real typo elsewhere in the list). +DownloadStallError: type +start_watchdog: Any +get_hf_download_state: Any + + +def __getattr__(name: str) -> Any: + if name in _DEGRADED_ATTRS: + if _load_shared(): + return getattr(_shared, name) + return _DEGRADED_ATTRS[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# Indirection seam the public wrappers call (and tests monkeypatch): lazy-load the shared backend, +# then dispatch to it or the degraded stub. The ``_shared_*`` names preserve the pre-refactor contract. +def _shared_hf_hub_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.hf_hub_download_with_xet_fallback + if _load_shared() + else _degraded_hf_hub_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +def _shared_snapshot_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.snapshot_download_with_xet_fallback + if _load_shared() + else _degraded_snapshot_download_with_xet_fallback + ) + return impl(*args, **kwargs) __all__ = [ diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index c16ae91467..1bcbfbf95a 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path logger.info("llama update: installing", cmd = " ".join(cmd)) # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") + # Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm + # box would otherwise re-route and silently replace the Vulkan build. + # Re-assert it via the same env flag setup uses (mirrors + # _rocm_install_args). + if asset and "vulkan" in asset.lower(): + env["UNSLOTH_FORCE_VULKAN"] = "1" proc = subprocess.Popen( cmd, stdout = subprocess.PIPE, diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5d8458e5f0..281ca24281 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1617,36 +1617,60 @@ def _iter_hf_cache_snapshots(repo_id: str): cache_dir = Path(hf_constants.HF_HUB_CACHE) target = f"models--{repo_id.replace('/', '--')}".lower() - repo_dir: Optional[Path] = None + repo_dirs: list[Path] = [] try: if not cache_dir.is_dir(): return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: - repo_dir = entry - break + repo_dirs.append(entry) except OSError: return - if repo_dir is None: + if not repo_dirs: return - snapshots = repo_dir / "snapshots" - try: - if not snapshots.is_dir(): - return - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] - except OSError: + snap_dirs: list[Path] = [] + for repo_dir in repo_dirs: + snapshots = repo_dir / "snapshots" + try: + if snapshots.is_dir(): + for snap_dir in snapshots.iterdir(): + try: + if snap_dir.is_dir(): + snap_dirs.append(snap_dir) + except OSError: + continue + except OSError: + continue + if not snap_dirs: return - snap_dirs.sort(key = lambda s: s.stat().st_mtime, reverse = True) - yield from snap_dirs + snap_dirs_with_mtime = [] + for snap_dir in snap_dirs: + try: + snap_dirs_with_mtime.append((snap_dir.stat().st_mtime, snap_dir)) + except OSError: + continue + snap_dirs_with_mtime.sort(key = lambda item: item[0], reverse = True) + yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime) def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: - """Variants from the local HF cache snapshot, or None if not cached.""" + """Variants from the local HF cache snapshot, or None if not cached. + + A newer snapshot can hold only a companion file (for example a vision + projector fetched on demand) while the quant files live in an older + snapshot. Returning the first snapshot that merely reports a vision flag + would shadow those real variants, so keep scanning older snapshots for + actual variants and carry the vision flag across snapshots. + """ + any_vision = False for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) - if variants or has_vision: - return variants, has_vision + any_vision = any_vision or has_vision + if variants: + return variants, any_vision + if any_vision: + return [], True return None diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 2a63caa5b2..a69673f081 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -28,7 +28,9 @@ Strategy: sys.path swap using the same directories pre-installed by setup.sh. """ +import ast import importlib +import importlib.util import json import structlog from loggers import get_logger @@ -173,6 +175,7 @@ _TRANSFORMERS_530_ARCHITECTURES: set[str] = { "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", "Glm4MoeLiteForCausalLM", + "Lfm2MoeForCausalLM", "Lfm2VlForConditionalGeneration", } _TRANSFORMERS_530_MODEL_TYPES: set[str] = { @@ -183,6 +186,7 @@ _TRANSFORMERS_530_MODEL_TYPES: set[str] = { "qwen3_moe", "qwen3_next", "glm4_moe_lite", + "lfm2_moe", "lfm2_vl", } @@ -870,6 +874,116 @@ def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None: return _config_json_cache.get(_token_cache_key(model_name, hf_token)) +# --- Static tier from CONFIG_MAPPING_NAMES (AST only: no import/network/exec) --- +# A model_type absent from an overlay's mapping can't load there. Parse each sidecar's +# config map from source and pick the lowest tier that ships it, so a new arch routes +# correctly with no per-model table edit. Only ever upgrades default, never lowers. +_config_mapping_cache: dict[str, frozenset[str]] = {} + + +def _overlay_transformers_dir(tier: str) -> str | None: + """transformers source dir for a tier, located without importing it.""" + if tier != "default": + root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier) + src = os.path.join(root, "transformers") if root else None + return src if src and _safe_is_dir(Path(src)) else None + # default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one + # is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir. + sidecars = tuple( + os.path.abspath(d) + os.sep + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR) + ) + candidates = [] + try: + spec = importlib.util.find_spec("transformers") + if spec and spec.origin: + candidates.append(os.path.dirname(spec.origin)) + except Exception: + pass + candidates += [os.path.join(e, "transformers") for e in sys.path if e] + for c in candidates: + if _safe_is_dir(Path(c)) and not os.path.abspath(c).startswith(sidecars): + return c + return None + + +def _mapping_first_keys(value: ast.AST) -> set[str]: + """First keys of a dict literal, or of an OrderedDict(...)/dict(...)/.update(...) + built from 2-tuple lists and **{...} unpacking.""" + + def keys_of(node): + if isinstance(node, ast.Dict): + return list(node.keys) + if isinstance(node, (ast.List, ast.Tuple)): + return [ + el.elts[0] for el in node.elts if isinstance(el, (ast.Tuple, ast.List)) and el.elts + ] + return [] + + nodes = keys_of(value) + if isinstance(value, ast.Call): + for a in value.args: + nodes += keys_of(a) + for kw in value.keywords: # **{...} unpacking has kw.arg is None + if kw.arg is None: + nodes += keys_of(kw.value) + return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)} + + +def _config_model_types(tier: str) -> frozenset[str]: + """model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py).""" + cached = _config_mapping_cache.get(tier) + if cached is not None: + return cached + tdir = _overlay_transformers_dir(tier) + if tdir is None: + return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads + keys: set[str] = set() + for rel in ("models/auto/configuration_auto.py", "models/auto/auto_mappings.py"): + path = Path(tdir) / rel + if not _safe_is_file(path): + continue + try: + tree = ast.parse(path.read_text(encoding = "utf-8")) + for node in ast.walk(tree): + # direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets + ): + keys |= _mapping_first_keys(node.value) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + fn = node.value.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "update" + and isinstance(fn.value, ast.Name) + and fn.value.id == "CONFIG_MAPPING_NAMES" + ): + keys |= _mapping_first_keys(node.value) + except Exception: + continue + result = frozenset(keys) + _config_mapping_cache[tier] = result + return result + + +def _tier_from_config_mapping(cfg: dict) -> str | None: + """Lowest tier whose transformers ships cfg's model_type, or None if unknown.""" + model_type = cfg.get("model_type") + if not isinstance(model_type, str): + for key in _NESTED_CONFIG_KEYS: + sub = cfg.get(key) + if isinstance(sub, dict) and isinstance(sub.get("model_type"), str): + model_type = sub["model_type"] + break + if not isinstance(model_type, str): + return None + for tier in sorted(_TIER_RANK, key = _TIER_RANK.get): + if model_type in _config_model_types(tier): + return tier + return None + + # --- AutoConfig probe: general tier resolution for ambiguous models ---------- # When the cheap signals only say "needs some 5.x", parse config.json with the built-in # parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond @@ -1210,6 +1324,14 @@ def get_transformers_tier( match, ) return tier + static = _tier_from_config_mapping(cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static local_tc = Path(model_name) / "tokenizer_config.json" if _safe_is_file(local_tc) and _check_tokenizer_config_needs_v5(model_name, hf_token): if not probe: @@ -1269,6 +1391,18 @@ def get_transformers_tier( return override logger.info("Transformers tier 530 selected for %s (config.json check)", model_name) return "530" + # _load_config_json (not the cache-only reader) so a config served from the hub + # cache during a transient outage still feeds the mapping resolver. + remote_cfg = _load_config_json(model_name, hf_token) + if remote_cfg is not None: + static = _tier_from_config_mapping(remote_cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static if _check_tokenizer_config_needs_v5(model_name, hf_token): if not probe: return "530" diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83c..1b5926fd49 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -26,11 +26,14 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele def has_blackwell_gpu() -> bool: """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell). - Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels - fail to load, so callers use this to skip the flash-attn install path. Cached - for the process lifetime; tests mocking nvidia-smi must call + Cached for the process lifetime; tests mocking nvidia-smi must call ``has_blackwell_gpu.cache_clear()`` first. """ + # Detection disabled for now: Dao-AILab ships Blackwell (sm_100+) flash-attn + # wheels and url_exists() already gates resolution, so we no longer skip + # flash-attn on Blackwell. The nvidia-smi probe below is kept for possible + # future arch-based gating; drop this early return to re-enable it. + return False exe = shutil.which("nvidia-smi") if not exe: return False @@ -117,6 +120,19 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non return env +# torch 2.11 has no native prebuilt wheels for flash-attn / causal-conv1d / mamba +# yet, but their torch 2.10 CUDA wheels load and pass the projects' own test suites +# on torch 2.11 (verified on B200: FA2 fwd/bwd, causal-conv1d, and mamba selective +# scan all match reference). Reuse the torch 2.10 wheels on torch 2.11 so a 2.11 +# install still gets these prebuilt accelerators instead of building from source. +_PREBUILT_WHEEL_TORCH_MM = {"2.11": "2.10"} + + +def prebuilt_wheel_torch_mm(torch_mm: str) -> str: + """Map a torch major.minor to the one whose prebuilt accelerator wheels to use.""" + return _PREBUILT_WHEEL_TORCH_MM.get(torch_mm, torch_mm) + + def direct_wheel_url( *, filename_prefix: str, @@ -130,7 +146,7 @@ def direct_wheel_url( filename = ( f"{filename_prefix}-{package_version}" - f"+cu{env['cuda_major']}torch{env['torch_mm']}" + f"+cu{env['cuda_major']}torch{prebuilt_wheel_torch_mm(env['torch_mm'])}" f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}" f"-{env['platform_tag']}.whl" ) @@ -152,7 +168,7 @@ def flash_attn_package_version(torch_mm: str) -> str | None: def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None: if env is None: return None - package_version = flash_attn_package_version(env["torch_mm"]) + package_version = flash_attn_package_version(prebuilt_wheel_torch_mm(env["torch_mm"])) if package_version is None: return None return direct_wheel_url( diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index f59a952b3a..1a74b38524 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -81,7 +81,6 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, TooltipContent, @@ -97,6 +96,7 @@ import { createChatProject, deleteChatProject, deleteChatItem, + listStoredChatThreads, moveChatItemToProject, renameChatItem, renameChatProject, @@ -582,7 +582,14 @@ export function AppSidebar() { useEffect(() => { if (!pendingRename) return; const match = allChatItems.find((i) => i.id === pendingRename.id); - if (match && match.title === pendingRename.title) setPendingRename(null); + if (!match || match.title !== pendingRename.title) return; + queueMicrotask(() => { + setPendingRename((current) => + current?.id === pendingRename.id && current.title === pendingRename.title + ? null + : current, + ); + }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); const [projectNameDraft, setProjectNameDraft] = useState(""); @@ -680,12 +687,6 @@ export function AppSidebar() { useState(null); const [deleteProjectFiles, setDeleteProjectFiles] = useState(false); - useEffect(() => { - if (confirmingDelete?.kind !== "project") { - setDeleteProjectFiles(false); - } - }, [confirmingDelete]); - async function commitDelete() { const target = confirmingDelete; if (!target) return; @@ -1572,9 +1573,6 @@ export function AppSidebar() { > {t("shell.navigation.api")} - - {t("common.new")} - } diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx index 823696693a..331a06a4c6 100644 --- a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -18,7 +18,7 @@ import { useExternalProvidersStore, } from "@/features/chat"; import { cn } from "@/lib/utils"; -import { FileDatabaseIcon } from "@hugeicons/core-free-icons"; +import { HelpCircleIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useMessage, useMessageTiming } from "@assistant-ui/react"; import type { FC, ReactNode } from "react"; @@ -341,7 +341,7 @@ export const MessageResponseDetailsSheet: FC<{ diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 9890c5f574..ff0351883d 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -393,6 +393,7 @@ function ModelRow({ vramEst, gpuGb, tooltipText, + hubUrl, optionProps, onArrowDownIntoChildren, capabilities, @@ -409,6 +410,10 @@ function ModelRow({ vramEst?: number; gpuGb?: number; tooltipText?: ReactNode; + /** Hugging Face address (e.g. "huggingface.co/owner/name") for online/Hub + * rows; surfaced on hover so their repo id / URL is discoverable the same + * way local rows show an on-disk path. Omit to show no address line. */ + hubUrl?: string; optionProps?: ModelRowOptionProps; onArrowDownIntoChildren?: () => boolean; /** Capability override (HF rows have tags); falls back to name detection. */ @@ -546,30 +551,41 @@ function ModelRow({ ); - if (vramTooltipText) { - return ( - - {content} - - {label} - {vramTooltipText} - - - ); - } + // Optional Hugging Face address line for online/Hub rows, rendered under + // whichever tooltip shows so the repo id / URL is always visible on hover. + const hubUrlLine = hubUrl ? ( + + {hubUrl} + + ) : null; - if (tooltipText) { + const tooltipBody = vramTooltipText ? ( + <> + {label} + {vramTooltipText} + {hubUrlLine} + + ) : tooltipText ? ( + <> + {tooltipText} + {hubUrlLine} + + ) : hubUrl ? ( + <> + {label} + {hubUrlLine} + + ) : null; + + if (tooltipBody) { return ( - + {content} - {tooltipText} + {tooltipBody} ); @@ -1193,6 +1209,13 @@ function localPathTooltip(name: string, path: string): ReactNode { ); } +/** Hugging Face address for an online/Hub row, or undefined when the repo id is + * missing so the row shows no (empty) address line on hover. */ +function hubRepoUrl(id: string | null | undefined): string | undefined { + const trimmed = id?.trim(); + return trimmed ? `huggingface.co/${trimmed}` : undefined; +} + /** Whether a local model is an MLX build (name hint). MLX runs on Mac only, so * callers gate visibility on the host being a Mac. */ function localModelIsMlx(m: LocalModelInfo): boolean { @@ -2462,6 +2485,7 @@ export function HubModelPicker({
{ strokeWidth={1.75} className="size-icon" /> - Export as Markdown + Export as markdown { className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground" > diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx index bce4bf2831..0a51875de9 100644 --- a/studio/frontend/src/components/floating-monitor.tsx +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -3,27 +3,35 @@ import { Button } from "@/components/ui/button"; import { Progress } from "@/components/ui/progress"; -import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useMonitorOverlayStore } from "@/features/settings"; import { useSystemInfo } from "@/hooks/use-system"; import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react"; -import { motion } from "motion/react"; -import { useRef } from "react"; +import { AnimatePresence, motion, useDragControls } from "motion/react"; +import { type PointerEvent, useMemo, useState } from "react"; function clampPercent(value: number): number { return Math.max(0, Math.min(100, value)); } function usageIndicatorClass(percent: number): string { - if (percent >= 90) return "bg-destructive"; - if (percent >= 70) return "bg-amber-500"; + if (percent >= 90) { + return "bg-destructive"; + } + if (percent >= 70) { + return "bg-amber-500"; + } return "bg-primary"; } function usageTextClass(percent: number): string { - if (percent >= 90) return "text-destructive"; - if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + if (percent >= 90) { + return "text-destructive"; + } + if (percent >= 70) { + return "text-amber-600 dark:text-amber-400"; + } return "text-primary"; } @@ -39,9 +47,18 @@ export function FloatingMonitor() { const { isOpen, setIsOpen } = useMonitorOverlayStore(); const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); - const constraintsRef = useRef(null); + const [constraintsElement, setConstraintsElement] = + useState(null); + const constraintsRef = useMemo( + () => ({ current: constraintsElement }), + [constraintsElement], + ); + const dragControls = useDragControls(); - if (!isOpen) return null; + function startDrag(event: PointerEvent) { + event.preventDefault(); + dragControls.start(event); + } const ramTotal = systemInfo.memory?.total_gb ?? 0; const ramAvailable = systemInfo.memory?.available_gb ?? 0; @@ -64,99 +81,109 @@ export function FloatingMonitor() { const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; return ( -
- -
-
- - - {t("settings.resources.liveMonitor.title")} - -
-
-
- -
- - -
-
- - + {isOpen && ( +
-
-
- {t("settings.resources.liveMonitor.ram")} - - {Math.round(ramPercent)}% - -
-
- {formatGiB(ramUsed)} / {formatGiB(ramTotal)} -
- -
- - {hasGpu && ( -
-
- - {t("settings.resources.liveMonitor.vram")}{" "} - {devices.length > 1 - ? `(${devices.length} GPUs)` - : `(${devices[0].name ?? "GPU"})`} + +
+
+ + + {t("settings.resources.liveMonitor.title")} - +
+
- {Math.round(vramPercent)}% - + +
+ +
-
- {formatGiB(vramUsed)} / {formatGiB(vramTotal)} -
-
- )} - - -
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGiB(ramUsed)} / {formatGiB(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGiB(vramUsed)} / {formatGiB(vramTotal)} +
+ +
+ )} +
+
+
+ )} + ); } diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 2ca6ceddf7..ca062bf93c 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2,6 +2,10 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { getAuthToken } from "@/features/auth"; +import { + loadRememberedLoadSettings, + rememberedLoadSettingsKey, +} from "@/components/assistant-ui/model-selector/remembered-load-settings"; import { projectHasSources } from "@/features/rag/api/rag-api"; import { apiUrl } from "@/lib/api-base"; import { parseParamCountB } from "@/lib/model-size"; @@ -63,11 +67,17 @@ import { listStoredChatThreads, updateStoredChatThread, } from "../utils/chat-history-storage"; +import { + readLastLocalModelLoad, + recordLastLocalModelLoad, + type LastLocalModelKind, +} from "../utils/last-local-model-load"; import { getImageInputUnavailableReason } from "../utils/image-input-support"; import { hasClosedThinkTag, parseAssistantContent, } from "../utils/parse-assistant-content"; +import { resolveLoadMaxSeqLength } from "../presets/preset-policy"; import { generateAudio, listCachedGguf, @@ -1309,6 +1319,30 @@ const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi; const GGUF_KNOWN_QUANT_RE = /(UD-)?(MXFP[0-9]+(?:_[A-Z0-9]+)*|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?|TQ[0-9]+_[0-9]+|Q[0-9]+_K_[A-Z]+|Q[0-9]+_[0-9]+|Q[0-9]+_K|BF16|F16|F32)/i; +type AutoLoadCandidate = { + id: string; + kind: LastLocalModelKind; + ggufVariant: string | null; + maxSeqLength: number; + successLabel: string; +}; + +function autoLoadCandidateKey( + kind: LastLocalModelKind, + id: string, + ggufVariant?: string | null, +): string { + return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`; +} + +function findCachedRepo( + repos: T[], + id: string, +): T | undefined { + const normalized = id.toLowerCase(); + return repos.find((repo) => repo.repo_id.toLowerCase() === normalized); +} + function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean { const normalized = filename.replace(/\\/g, "/").toLowerCase(); const separatorIndex = normalized.lastIndexOf("/"); @@ -1357,14 +1391,18 @@ async function autoLoadSmallestModel(): Promise<{ const hfToken = store.hfToken || null; const trustRemoteCode = store.params.trustRemoteCode ?? false; const specSettings = resolveSpeculativeSettingsForLoad(); + const lastLoaded = readLastLocalModelLoad(); const toastId = toast("Loading a model…", { - description: "Auto-selecting the smallest downloaded model.", + description: lastLoaded + ? "Loading last used model." + : "Auto-selecting the smallest downloaded model.", duration: 5000, closeButton: true, }); let blockedByTrustRemoteCode = false; let hadNonTrustFailure = false; let loadAttempts = 0; + const skippedAutoLoadCandidates = new Set(); async function canAutoLoad(payload: { model_path: string; @@ -1389,12 +1427,224 @@ async function autoLoadSmallestModel(): Promise<{ } return true; } + + async function loadAutoLoadCandidate( + candidate: AutoLoadCandidate, + ): Promise { + if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) { + return false; + } + const currentStore = useChatRuntimeStore.getState(); + const remembered = loadRememberedLoadSettings( + rememberedLoadSettingsKey({ + id: candidate.id, + ggufVariant: candidate.ggufVariant, + }), + ); + const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ + modelId: candidate.id, + ggufVariant: candidate.ggufVariant, + isGguf: candidate.kind === "gguf", + customContextLength: remembered?.contextLength ?? null, + ggufContextLength: null, + currentCheckpoint: currentStore.params.checkpoint, + activeGgufVariant: currentStore.activeGgufVariant, + maxSeqLength: candidate.maxSeqLength, + presetSource: currentStore.activePresetSource, + }); + const effectiveSpeculativeType = + remembered?.speculativeType ?? specSettings.speculativeType; + const effectiveSpecDraftNMax = + remembered?.specDraftNMax ?? specSettings.specDraftNMax; + if ( + !(await canAutoLoad({ + model_path: candidate.id, + max_seq_length: effectiveMaxSeqLength, + is_lora: false, + gguf_variant: candidate.ggufVariant, + })) + ) { + skippedAutoLoadCandidates.add( + autoLoadCandidateKey(candidate.kind, candidate.id, candidate.ggufVariant), + ); + return false; + } + loadAttempts += 1; + const loadResp = await loadModel({ + model_path: candidate.id, + hf_token: hfToken, + max_seq_length: effectiveMaxSeqLength, + load_in_4bit: true, + is_lora: false, + gguf_variant: candidate.ggufVariant, + trust_remote_code: trustRemoteCode, + cache_type_kv: remembered?.kvCacheDtype ?? null, + speculative_type: effectiveSpeculativeType, + spec_draft_n_max: effectiveSpecDraftNMax, + tensor_parallel: remembered?.tensorParallel ?? false, + }); + saveSpeculativeType(effectiveSpeculativeType); + useChatRuntimeStore + .getState() + .setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined); + const store = useChatRuntimeStore.getState(); + store.setModelRequiresTrustRemoteCode( + loadResp.requires_trust_remote_code ?? false, + ); + store.setParams({ + ...store.params, + maxTokens: + candidate.kind === "gguf" + ? loadResp.context_length ?? 131072 + : effectiveMaxSeqLength, + }); + const autoModel: ChatModelSummary = { + id: candidate.id, + name: loadResp.display_name ?? candidate.id, + isVision: loadResp.is_vision ?? false, + isLora: loadResp.is_lora ?? false, + isGguf: loadResp.is_gguf ?? candidate.kind === "gguf", + isAudio: loadResp.is_audio ?? false, + audioType: loadResp.audio_type ?? null, + hasAudioInput: loadResp.has_audio_input ?? false, + }; + if (!store.models.some((m) => m.id === candidate.id)) { + store.setModels([...store.models, autoModel]); + } + if (candidate.kind === "gguf") { + useChatRuntimeStore.setState({ + ggufContextLength: loadResp.context_length ?? 131072, + ggufMaxContextLength: + loadResp.max_context_length ?? loadResp.context_length ?? 131072, + ggufNativeContextLength: loadResp.native_context_length ?? null, + supportsReasoning: loadResp.supports_reasoning ?? false, + reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, + reasoningEnabled: loadResp.supports_reasoning ?? false, + ...reasoningCapsFromLoad(loadResp), + supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false, + supportsTools: loadResp.supports_tools ?? false, + ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + tensorParallel: loadResp.tensor_parallel ?? false, + loadedTensorParallel: loadResp.tensor_parallel ?? false, + defaultChatTemplate: loadResp.chat_template ?? null, + chatTemplateOverride: null, + loadedChatTemplateOverride: null, + loadedIsMultimodal: isMultimodalResponse(loadResp), + loadedIsDiffusion: loadResp.is_diffusion ?? false, + ...resolveLoadedSpeculativeSettings(loadResp), + }); + } else { + useChatRuntimeStore.setState({ + supportsReasoning: loadResp.supports_reasoning ?? false, + reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, + reasoningEnabled: loadResp.supports_reasoning ?? false, + ...reasoningCapsFromLoad(loadResp), + supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false, + supportsTools: loadResp.supports_tools ?? false, + ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + tensorParallel: loadResp.tensor_parallel ?? false, + loadedTensorParallel: loadResp.tensor_parallel ?? false, + defaultChatTemplate: loadResp.chat_template ?? null, + chatTemplateOverride: null, + loadedChatTemplateOverride: null, + ...resolveLoadedSpeculativeSettings(loadResp), + loadedIsMultimodal: isMultimodalResponse(loadResp), + loadedIsDiffusion: loadResp.is_diffusion ?? false, + }); + } + if (!(loadResp.is_lora ?? false)) { + recordLastLocalModelLoad({ + id: candidate.id, + kind: candidate.kind, + ggufVariant: candidate.ggufVariant, + }); + } + toast.success(candidate.successLabel, { id: toastId }); + return true; + } try { const [ggufRepos, modelRepos] = await Promise.all([ listCachedGguf().catch(() => []), listCachedModels().catch(() => []), ]); + if (lastLoaded) { + if (lastLoaded.kind === "gguf") { + const repo = findCachedRepo(ggufRepos, lastLoaded.id); + if (repo && lastLoaded.ggufVariant) { + try { + const variants = await listGgufVariants(repo.repo_id); + const variant = variants.variants.find( + (entry) => + entry.downloaded && + entry.quant?.toLowerCase() === + lastLoaded.ggufVariant?.toLowerCase() && + isAutoLoadableGgufVariant(entry), + ); + if (variant) { + toast("Loading last used model…", { + id: toastId, + description: `${repo.repo_id} (${variant.quant})`, + duration: 5000, + }); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "gguf", + ggufVariant: variant.quant, + maxSeqLength: 0, + successLabel: `Loaded ${repo.repo_id} (${variant.quant})`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + autoLoadCandidateKey("gguf", repo.repo_id, lastLoaded.ggufVariant), + ); + } + } + } else { + const repo = findCachedRepo(modelRepos, lastLoaded.id); + if (repo) { + try { + toast("Loading last used model…", { + id: toastId, + description: repo.repo_id, + duration: 5000, + }); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "model", + ggufVariant: null, + maxSeqLength: store.params.maxSeqLength, + successLabel: `Loaded ${repo.repo_id}`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + autoLoadCandidateKey("model", repo.repo_id), + ); + } + } + } + toast("Loading a model…", { + id: toastId, + description: "Auto-selecting the smallest downloaded model.", + duration: 5000, + }); + } + // GGUF first: smallest-total-size repo, then its smallest variant. if (ggufRepos.length > 0) { const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes); @@ -1408,82 +1658,23 @@ async function autoLoadSmallestModel(): Promise<{ if (downloaded.length > 0) { const variant = downloaded[0]; if ( - !(await canAutoLoad({ - model_path: repo.repo_id, - max_seq_length: 0, - is_lora: false, - gguf_variant: variant.quant, - })) + skippedAutoLoadCandidates.has( + autoLoadCandidateKey("gguf", repo.repo_id, variant.quant), + ) ) { continue; } - loadAttempts += 1; - const loadResp = await loadModel({ - model_path: repo.repo_id, - hf_token: hfToken, - max_seq_length: 0, - load_in_4bit: true, - is_lora: false, - gguf_variant: variant.quant, - trust_remote_code: trustRemoteCode, - speculative_type: specSettings.speculativeType, - spec_draft_n_max: specSettings.specDraftNMax, - }); - saveSpeculativeType(specSettings.speculativeType); - useChatRuntimeStore - .getState() - .setCheckpoint(repo.repo_id, variant.quant); - const store = useChatRuntimeStore.getState(); - store.setModelRequiresTrustRemoteCode( - loadResp.requires_trust_remote_code ?? false, - ); - store.setParams({ - ...store.params, - maxTokens: loadResp.context_length ?? 131072, - }); - // Add to store so the selector shows the name. - const autoModel: ChatModelSummary = { - id: repo.repo_id, - name: loadResp.display_name ?? repo.repo_id, - isVision: loadResp.is_vision ?? false, - isLora: loadResp.is_lora ?? false, - isGguf: loadResp.is_gguf ?? false, - isAudio: loadResp.is_audio ?? false, - audioType: loadResp.audio_type ?? null, - hasAudioInput: loadResp.has_audio_input ?? false, - }; - const existingModels = store.models; - if (!existingModels.some((m) => m.id === repo.repo_id)) { - store.setModels([...existingModels, autoModel]); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "gguf", + ggufVariant: variant.quant, + maxSeqLength: 0, + successLabel: `Loaded ${repo.repo_id} (${variant.quant})`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; } - useChatRuntimeStore.setState({ - ggufContextLength: loadResp.context_length ?? 131072, - ggufMaxContextLength: - loadResp.max_context_length ?? - loadResp.context_length ?? - 131072, - supportsReasoning: loadResp.supports_reasoning ?? false, - reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, - reasoningEnabled: loadResp.supports_reasoning ?? false, - ...reasoningCapsFromLoad(loadResp), - supportsPreserveThinking: - loadResp.supports_preserve_thinking ?? false, - supportsTools: loadResp.supports_tools ?? false, - ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), - kvCacheDtype: loadResp.cache_type_kv ?? null, - loadedKvCacheDtype: loadResp.cache_type_kv ?? null, - tensorParallel: loadResp.tensor_parallel ?? false, - loadedTensorParallel: loadResp.tensor_parallel ?? false, - defaultChatTemplate: loadResp.chat_template ?? null, - chatTemplateOverride: null, - loadedChatTemplateOverride: null, - loadedIsMultimodal: isMultimodalResponse(loadResp), - ...resolveLoadedSpeculativeSettings(loadResp), - }); - toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { - id: toastId, - }); - return { loaded: true, blockedByTrustRemoteCode: false }; } } catch { hadNonTrustFailure = true; @@ -1501,64 +1692,23 @@ async function autoLoadSmallestModel(): Promise<{ if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; try { if ( - !(await canAutoLoad({ - model_path: repo.repo_id, - max_seq_length: 4096, - is_lora: false, - gguf_variant: null, - })) + skippedAutoLoadCandidates.has( + autoLoadCandidateKey("model", repo.repo_id), + ) ) { continue; } - loadAttempts += 1; - const sfLoadResp = await loadModel({ - model_path: repo.repo_id, - hf_token: hfToken, - max_seq_length: 4096, - load_in_4bit: true, - is_lora: false, - gguf_variant: null, - trust_remote_code: trustRemoteCode, - speculative_type: specSettings.speculativeType, - spec_draft_n_max: specSettings.specDraftNMax, - }); - saveSpeculativeType(specSettings.speculativeType); - useChatRuntimeStore.getState().setCheckpoint(repo.repo_id); - const store = useChatRuntimeStore.getState(); - store.setModelRequiresTrustRemoteCode( - sfLoadResp.requires_trust_remote_code ?? false, - ); - store.setParams({ ...store.params, maxTokens: 4096 }); - useChatRuntimeStore.setState({ - supportsReasoning: sfLoadResp.supports_reasoning ?? false, - reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false, - reasoningEnabled: sfLoadResp.supports_reasoning ?? false, - ...reasoningCapsFromLoad(sfLoadResp), - supportsPreserveThinking: - sfLoadResp.supports_preserve_thinking ?? false, - supportsTools: sfLoadResp.supports_tools ?? false, - // Parity with the GGUF branch above. - ...resolveToolsEnabledOnLoad(sfLoadResp.supports_tools ?? false), - defaultChatTemplate: sfLoadResp.chat_template ?? null, - chatTemplateOverride: null, - loadedChatTemplateOverride: null, - ...resolveLoadedSpeculativeSettings(sfLoadResp), - }); - const sfModel: ChatModelSummary = { - id: repo.repo_id, - name: sfLoadResp.display_name ?? repo.repo_id, - isVision: sfLoadResp.is_vision ?? false, - isLora: sfLoadResp.is_lora ?? false, - isGguf: sfLoadResp.is_gguf ?? false, - }; - if (!store.models.some((m) => m.id === repo.repo_id)) { - store.setModels([...store.models, sfModel]); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "model", + ggufVariant: null, + maxSeqLength: 4096, + successLabel: `Loaded ${repo.repo_id}`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; } - useChatRuntimeStore.setState({ - loadedIsMultimodal: isMultimodalResponse(sfLoadResp), - }); - toast.success(`Loaded ${repo.repo_id}`, { id: toastId }); - return { loaded: true, blockedByTrustRemoteCode: false }; } catch { hadNonTrustFailure = true; continue; @@ -1650,6 +1800,11 @@ async function autoLoadSmallestModel(): Promise<{ loadedIsMultimodal: isMultimodalResponse(loadResp), ...resolveLoadedSpeculativeSettings(loadResp), }); + recordLastLocalModelLoad({ + id: "unsloth/Qwen3.5-4B-MTP-GGUF", + kind: "gguf", + ggufVariant: "UD-Q4_K_XL", + }); toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId }); return { loaded: true, blockedByTrustRemoteCode: false }; } catch { diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 798c1658f9..0a4342f208 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -44,6 +44,7 @@ import { mergeBackendRecommendedInference, resolveLoadMaxSeqLength, } from "../presets/preset-policy"; +import { recordLastLocalModelLoad } from "../utils/last-local-model-load"; import { isMultimodalResponse, } from "../types/api"; @@ -818,6 +819,23 @@ export function useChatModelRuntime() { } } await refresh({ signal: abortCtrl.signal }); + if ( + !isLora && + !(loadResponse.is_lora ?? false) && + !nativePathToken && + !isLocalModelPath(modelId) && + !isExternalModelId(modelId) + ) { + if (loadResponse.is_gguf || isGguf || ggufVariant) { + recordLastLocalModelLoad({ + id: modelId, + kind: "gguf", + ggufVariant: ggufVariant ?? null, + }); + } else { + recordLastLocalModelLoad({ id: modelId, kind: "model" }); + } + } // A successful load owns the shared (pick-unscoped) settings fields, // so any surviving stage is stale: the just-loaded pick itself, or a // pick queued for a different model mid-load whose knobs this load diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 7cd9611c71..d070ed15de 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -36,6 +36,7 @@ export { ChatSearchDialog } from "./components/chat-search-dialog"; export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; export type { ProjectRecord } from "./types"; export { clearAllChats, countAllChats } from "./utils/clear-all-chats"; +export { listStoredChatThreads } from "./utils/chat-history-storage"; export { ArtifactCard } from "./artifacts/artifact-card"; export { useChatArtifactsStore, diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 9386650fee..60788c23bd 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -3,16 +3,19 @@ import { getInferenceStatus } from "../api/chat-api"; import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { clampReasoningEffortToLevels } from "../provider-capabilities"; import { CHAT_REASONING_ENABLED_KEY, - loadOptionalBool, type ReasoningEffort, type ReasoningStyle, + loadOptionalBool, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; -import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api"; -import { clampReasoningEffortToLevels } from "../provider-capabilities"; +import { + type InferenceStatusResponse, + isMultimodalResponse, +} from "../types/api"; import type { ChatModelSummary } from "../types/runtime"; type LocalReasoningEffort = Extract; @@ -31,7 +34,10 @@ export function normalizeSpeculativeType( return "ngram"; } if (s === "mtp+ngram") return "mtp+ngram"; - const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const parts = s + .split(",") + .map((p) => p.trim()) + .filter(Boolean); const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); const hasNgram = parts.some( (p) => p === "ngram" || p === "ngram-mod" || p === "ngram-simple", @@ -197,6 +203,12 @@ export function applyActiveModelStatusToStore( ggufContextLength: currentGgufContextLength, ggufMaxContextLength, ggufNativeContextLength, + // A non-GGUF status must also drop a stale native-path token: without this the + // isGguf OR (activeGgufVariant || activeNativePathToken || ggufContextLength) + // stays true after switching from a native GGUF to a transformers model, so a + // Codex-only detection would auto-select for a model its preflight rejects. A real + // GGUF load reports is_gguf: true, so its token is preserved (the load path owns it). + ...(status.is_gguf ? {} : { activeNativePathToken: null }), modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false, defaultChatTemplate: nextDefaultChatTemplate, loadedIsMultimodal: isMultimodalResponse(status), @@ -245,7 +257,7 @@ export function applyActiveModelStatusToStore( const mid = checkpointId.toLowerCase(); if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); - if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { + if (sizeMatch && Number.parseFloat(sizeMatch[1]) < 9) { reasoningDefault = false; } } @@ -281,8 +293,7 @@ export async function tryAdoptServerActiveModel(): Promise { } // Re-check after the await: keep a checkpoint the user picked meanwhile. - const previousCheckpoint = - useChatRuntimeStore.getState().params.checkpoint; + const previousCheckpoint = useChatRuntimeStore.getState().params.checkpoint; if (previousCheckpoint) { return true; } diff --git a/studio/frontend/src/features/chat/utils/last-local-model-load.ts b/studio/frontend/src/features/chat/utils/last-local-model-load.ts new file mode 100644 index 0000000000..099386fbc7 --- /dev/null +++ b/studio/frontend/src/features/chat/utils/last-local-model-load.ts @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export type LastLocalModelKind = "gguf" | "model"; + +export type LastLocalModelLoad = { + id: string; + kind: LastLocalModelKind; + ggufVariant: string | null; + loadedAt: number; +}; + +const STORAGE_KEY = "unsloth.last-local-model-load.v1"; + +function storage(): Storage | null { + try { + return typeof localStorage === "undefined" ? null : localStorage; + } catch { + return null; + } +} + +function isLastLocalModelKind(value: unknown): value is LastLocalModelKind { + return value === "gguf" || value === "model"; +} + +export function readLastLocalModelLoad(): LastLocalModelLoad | null { + try { + const raw = storage()?.getItem(STORAGE_KEY); + if (!raw) { + return null; + } + const parsed = JSON.parse(raw) as Partial; + if ( + typeof parsed.id !== "string" || + !parsed.id.trim() || + !isLastLocalModelKind(parsed.kind) || + typeof parsed.loadedAt !== "number" + ) { + return null; + } + if ( + parsed.kind === "gguf" && + (typeof parsed.ggufVariant !== "string" || !parsed.ggufVariant.trim()) + ) { + return null; + } + return { + id: parsed.id, + kind: parsed.kind, + ggufVariant: + typeof parsed.ggufVariant === "string" ? parsed.ggufVariant : null, + loadedAt: parsed.loadedAt, + }; + } catch { + return null; + } +} + +export function recordLastLocalModelLoad(input: { + id: string; + kind: LastLocalModelKind; + ggufVariant?: string | null; +}): void { + const id = input.id.trim(); + if (!id) { + return; + } + const ggufVariant = input.ggufVariant?.trim() || null; + if (input.kind === "gguf" && !ggufVariant) { + return; + } + try { + storage()?.setItem( + STORAGE_KEY, + JSON.stringify({ + id, + kind: input.kind, + ggufVariant: input.kind === "gguf" ? ggufVariant : null, + loadedAt: Date.now(), + } satisfies LastLocalModelLoad), + ); + } catch { + // Ignore disabled storage / quota errors; auto-load falls back to size order. + } +} diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 4fe24b1f6d..f4e8167cb3 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -494,3 +494,13 @@ export async function removeUnstructuredFile( throw new Error("Failed to remove file"); } } + +export async function removeUnstructuredBlock(blockId: string): Promise { + const res = await authFetch( + `${DATA_DESIGNER_API_BASE}/seed/unstructured-block/${encodeURIComponent(blockId)}`, + { method: "DELETE" }, + ); + if (!res.ok && res.status !== 404) { + throw new Error("Failed to remove uploaded files"); + } +} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 0ed7eeed75..53f8566090 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -53,6 +53,11 @@ import { inspectSeedDataset, inspectSeedUpload, } from "../../api"; +import { useRecipeStudioStore } from "../../stores/recipe-studio"; +import { + makeUnstructuredUploadUid, + resolveUnstructuredUploadBlockId, +} from "../../utils/config-factories"; import { resolveImagePreview } from "../../utils/image-preview"; import type { GithubItemType, @@ -597,6 +602,41 @@ export function SeedDialog({ const mode = config.seed_source_type ?? "hf"; const previewEmpty = getPreviewEmptyStateCopy(mode); + const queueUploadCleanup = useRecipeStudioStore( + (state) => state.queueUploadCleanup, + ); + + // config.id collides across recipes (ids reset to n1 on import); use a + // stable per-block uid instead. Generate one synchronously so the first + // rendered drop zone cannot upload under a legacy node id. + const uploadUid = config.unstructured_upload_uid?.trim() ?? ""; + const unstructuredFileCount = config.unstructured_file_ids?.length ?? 0; + const generatedUploadUidRef = useRef(null); + if ( + mode === "unstructured" && + !uploadUid && + unstructuredFileCount === 0 && + generatedUploadUidRef.current === null + ) { + generatedUploadUidRef.current = makeUnstructuredUploadUid(); + } + const uploadBlockId = resolveUnstructuredUploadBlockId({ + configId: config.id, + uploadUid, + generatedUploadUid: generatedUploadUidRef.current, + unstructuredFileCount, + }); + + useEffect(() => { + if (mode !== "unstructured") return; + if (uploadUid) return; + if (unstructuredFileCount > 0) return; + const nextUid = + generatedUploadUidRef.current ?? makeUnstructuredUploadUid(); + generatedUploadUidRef.current = nextUid; + onUpdate({ unstructured_upload_uid: nextUid }); + }, [mode, uploadUid, unstructuredFileCount, onUpdate]); + const prevModeRef = useRef(mode); useEffect(() => { const prevMode = prevModeRef.current; @@ -720,6 +760,11 @@ export function SeedDialog({ subset: config.hf_subset?.trim() || undefined, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -730,6 +775,7 @@ export function SeedDialog({ hf_split: response.split ?? "", hf_subset: response.subset ?? "", local_file_name: "", + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -754,6 +800,11 @@ export function SeedDialog({ content_base64: payload, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -765,6 +816,7 @@ export function SeedDialog({ hf_subset: "", hf_split: "", local_file_name: localFile.name, + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -789,7 +841,7 @@ export function SeedDialog({ const { chunkSize, chunkOverlap } = resolveChunking(config); const response = await inspectSeedUpload({ - block_id: config.id, + block_id: uploadBlockId, file_ids: fileIds, file_names: fileNames, preview_size: 10, @@ -827,7 +879,18 @@ export function SeedDialog({ setIsInspecting(false); } }, - [config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles], + [ + config, + getCurrentLoadKey, + localFile, + mode, + onUpdate, + queueUploadCleanup, + unstructuredFiles, + unstructuredFileCount, + uploadBlockId, + uploadUid, + ], ); useEffect(() => { @@ -997,7 +1060,7 @@ export function SeedDialog({ {mode === "unstructured" && ( (null); const filesRef = useRef(files); + const blockIdRef = useRef(blockId); + const mountedRef = useRef(true); const [isDragOver, setIsDragOver] = useState(false); useEffect(() => { filesRef.current = files; - }, [files]); + blockIdRef.current = blockId; + }, [files, blockId]); + useEffect(() => () => { + mountedRef.current = false; + }, []); const totalSize = files.reduce((sum, f) => sum + f.size, 0); @@ -134,15 +140,32 @@ export function UnstructuredDropZone({ if (entry.status === "uploading" && entry.abortController) { entry.abortController.abort(); } - if ( + const needsServerRemove = entry.id && entry.status === "ok" && - !deletedIdsRef.current.has(entry.id) - ) { - deletedIdsRef.current.add(entry.id); - void removeUnstructuredFile(blockId, entry.id).catch(() => {}); - } + !deletedIdsRef.current.has(entry.id); onFilesChange((prev) => prev.filter((_, i) => i !== index)); + if (!needsServerRemove) return; + deletedIdsRef.current.add(entry.id); + removeUnstructuredFile(blockId, entry.id).catch(() => { + // Skip if the drop zone unmounted or its block changed: the id no + // longer belongs here and restoring would leak it into another block. + if (!mountedRef.current || blockIdRef.current !== blockId) return; + // Still exists server-side (counts toward quota); restore it at its + // original position. + deletedIdsRef.current.delete(entry.id); + onFilesChange((prev) => { + const next = [...prev]; + next.splice(Math.min(index, next.length), 0, { + id: entry.id, + name: entry.name, + size: entry.size, + status: "ok", + error: "Remove failed — try again", + }); + return next; + }); + }); }, [blockId, onFilesChange], ); diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts index 0417a6dd22..2d91272469 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts @@ -4,11 +4,13 @@ import { useCallback, useEffect, useMemo, useState } from "react"; import { toastError, toastSuccess } from "@/shared/toast"; import { normalizeNonEmptyName } from "@/utils"; +import { removeUnstructuredBlock } from "../api"; import { buildSignature, copyTextToClipboard, formatSavedLabel, } from "../executions/execution-helpers"; +import { useRecipeStudioStore } from "../stores/recipe-studio"; import { importRecipePayload, type RecipeSnapshot } from "../utils/import"; import type { RecipePayloadResult } from "../utils/payload/types"; @@ -72,7 +74,10 @@ function stripApiKeys(value: unknown): unknown { !Array.isArray(output.env) ) { output.env = Object.fromEntries( - Object.keys(output.env as Record).map((envKey) => [envKey, ""]), + Object.keys(output.env as Record).map((envKey) => [ + envKey, + "", + ]), ); } return output; @@ -82,10 +87,7 @@ function inferHfRepoIdFromPath(pathValue: unknown): string { if (typeof pathValue !== "string") { return ""; } - const parts = pathValue - .trim() - .split("/") - .filter(Boolean); + const parts = pathValue.trim().split("/").filter(Boolean); if (parts.length >= 3 && parts[0] === "datasets") { return `${parts[1]}/${parts[2]}`; } @@ -126,8 +128,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { typeof ui?.seed_source_type === "string" ? ui.seed_source_type : null; const sourceType = typeof source?.seed_type === "string" ? source.seed_type : null; - const shouldResetHfState = - sourceType === "hf" || uiSourceType === "hf"; + const shouldResetHfState = sourceType === "hf" || uiSourceType === "hf"; const shouldResetLocalState = sourceType === "local" || sourceType === "unstructured" || @@ -144,6 +145,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -165,6 +167,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -174,6 +177,43 @@ function sanitizeSeedForShare(payload: unknown): unknown { return root; } +// Delete queued upload directories once a save stops referencing them, so a +// reload before autosave can never leave the saved recipe pointing at +// already-deleted files. Skips any uid the just-saved payload still uses. +function drainQueuedUploadCleanups( + savedPayload: RecipePayloadResult["payload"], +): void { + const pending = useRecipeStudioStore.getState().pendingUploadCleanups; + if (pending.length === 0) { + return; + } + const ui = + savedPayload && typeof savedPayload === "object" + ? (savedPayload as { ui?: Record }).ui + : undefined; + const savedUid = + ui && typeof ui.unstructured_upload_uid === "string" + ? ui.unstructured_upload_uid + : ""; + const ready = pending.filter((uid) => uid !== savedUid); + if (ready.length === 0) { + return; + } + for (const uid of ready) { + void removeUnstructuredBlock(uid) + .then(() => { + useRecipeStudioStore.setState((state) => ({ + pendingUploadCleanups: state.pendingUploadCleanups.filter( + (pendingUid) => pendingUid !== uid, + ), + })); + }) + .catch((error) => { + console.warn("Failed to clean up uploaded documents:", error); + }); + } +} + export function useRecipePersistence({ recipeId, initialRecipeName, @@ -202,8 +242,10 @@ export function useRecipePersistence({ () => buildSignature(normalizedWorkflowName, currentPayload), [currentPayload, normalizedWorkflowName], ); - const isDirty = savedSignature.length > 0 && currentSignature !== savedSignature; - const saveTone: SaveTone = !isDirty && Boolean(lastSavedAt) ? "success" : "error"; + const isDirty = + savedSignature.length > 0 && currentSignature !== savedSignature; + const saveTone: SaveTone = + !isDirty && Boolean(lastSavedAt) ? "success" : "error"; const savedAtLabel = formatSavedLabel(lastSavedAt); useEffect(() => { @@ -214,7 +256,9 @@ export function useRecipePersistence({ setLastSavedAt(initialSavedAt); setCopied(false); - const parsed = importRecipePayload(JSON.stringify(initialPayload)); + const parsed = importRecipePayload(JSON.stringify(initialPayload), { + preserveUnstructuredUploads: true, + }); if (parsed.snapshot) { loadRecipe(parsed.snapshot); } else { @@ -252,6 +296,7 @@ export function useRecipePersistence({ }); setLastSavedAt(result.updatedAt); setSavedSignature(buildSignature(nextName, currentPayload)); + drainQueuedUploadCleanups(currentPayload); } catch (error) { console.error("Save recipe failed:", error); toastError("Save failed", "Could not save recipe."); @@ -270,11 +315,28 @@ export function useRecipePersistence({ return () => window.clearTimeout(timeoutId); }, [isDirty, persistRecipe, saveLoading]); + // Drain queued cleanups even when autosave is skipped: a net-zero edit (add + // then remove an unstructured seed before the 800ms debounce) keeps isDirty + // false, so the autosave effect never drains and the queued uid leaks its + // upload dir. Not-dirty means currentPayload equals the saved recipe, and + // drain skips the uid it still references, so only dirs no saved recipe + // points at are deleted (keeps the save-first invariant). + useEffect(() => { + if (!initialRecipeReady || isDirty || saveLoading) { + return; + } + drainQueuedUploadCleanups(currentPayload); + }, [currentPayload, initialRecipeReady, isDirty, saveLoading]); + const copyRecipe = useCallback(async (): Promise => { setCopied(false); try { - const safePayload = sanitizeSeedForShare(stripApiKeys(payloadResult.payload)); - const ok = await copyTextToClipboard(JSON.stringify(safePayload, null, 2)); + const safePayload = sanitizeSeedForShare( + stripApiKeys(payloadResult.payload), + ); + const ok = await copyTextToClipboard( + JSON.stringify(safePayload, null, 2), + ); if (!ok) { throw new Error("Clipboard not available."); } diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index 8659cad4fb..a1ff72ee9b 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -36,6 +36,7 @@ import { } from "../utils/handles"; import type { RecipeSnapshot } from "../utils/import"; import { getLayoutedElements } from "../utils/layout"; +import { makeUnstructuredUploadUid } from "../utils/config-factories"; import { centerModelInfraNodes, optimizeModelInfraEdgeHandles, @@ -76,6 +77,12 @@ type RecipeStudioState = { nextId: number; nextY: number; fitViewTick: number; + // Upload-uid directories whose owning block dropped them; server-side + // deletion is deferred until a save no longer references them, so a + // reload before autosave cannot leave a saved recipe pointing at + // deleted files. + pendingUploadCleanups: string[]; + queueUploadCleanup: (uid: string) => void; setSheetOpen: (open: boolean) => void; setSheetView: (view: SheetView) => void; setProcessors: (processors: RecipeProcessorConfig[]) => void; @@ -137,6 +144,7 @@ const INITIAL_STATE = { nextId: 3, nextY: 280, fitViewTick: 0, + pendingUploadCleanups: [], } satisfies Pick< RecipeStudioState, | "nodes" @@ -154,6 +162,7 @@ const INITIAL_STATE = { | "nextId" | "nextY" | "fitViewTick" + | "pendingUploadCleanups" >; function buildAddedNodeState( @@ -269,6 +278,20 @@ function isModelSemanticEdge( ); } +// Upload uid of a seed block whose server-side directory becomes orphaned +// when the block drops it. Only uid directories qualify (single owner); +// legacy node-id directories can be shared by other recipes. +function seedUploadCleanupUid(config: NodeConfig | undefined): string | null { + if (!config || config.kind !== "seed") { + return null; + } + const uid = config.unstructured_upload_uid?.trim(); + if (!uid || !config.unstructured_file_ids?.length) { + return null; + } + return uid; +} + export const useRecipeStudioStore = create((set, get) => ({ ...INITIAL_STATE, setSheetOpen: (open) => set({ sheetOpen: open }), @@ -278,6 +301,12 @@ export const useRecipeStudioStore = create((set, get) => ({ setDialogOpen: (open) => set({ dialogOpen: open }), setExecutionLocked: (locked) => set({ executionLocked: locked }), resetRecipe: () => set(INITIAL_STATE), + queueUploadCleanup: (uid) => + set((state) => + state.pendingUploadCleanups.includes(uid) + ? state + : { pendingUploadCleanups: [...state.pendingUploadCleanups, uid] }, + ), selectConfig: (id) => set({ activeConfigId: id, dialogOpen: false }), openConfig: (id) => set({ activeConfigId: id, dialogOpen: true }), setLayoutDirection: (direction) => @@ -383,7 +412,18 @@ export const useRecipeStudioStore = create((set, get) => ({ } return buildAddedNodeState(state, "sampler", type, position, openDialog); }), - addSeedNode: (type, position, openDialog = true) => + addSeedNode: (type, position, openDialog = true) => { + const current = get(); + if (!current.executionLocked) { + // The reset below clears the block's upload uid and file list; queue + // its server-side directory for deletion after the next save. + const uid = seedUploadCleanupUid( + Object.values(current.configs).find((config) => config.kind === "seed"), + ); + if (uid) { + current.queueUploadCleanup(uid); + } + } set((state) => { if (state.executionLocked) { return state; @@ -413,6 +453,8 @@ export const useRecipeStudioStore = create((set, get) => ({ hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + unstructured_upload_uid: + nextSourceType === "unstructured" ? makeUnstructuredUploadUid() : "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -446,7 +488,8 @@ export const useRecipeStudioStore = create((set, get) => ({ activeConfigId: existing.id, dialogOpen: openDialog, }; - }), + }); + }, addLlmNode: (type, position, openDialog = true) => set((state) => { if (state.executionLocked) { @@ -699,6 +742,9 @@ export const useRecipeStudioStore = create((set, get) => ({ dialogOpen: false, sheetView: "root", fitViewTick: state.fitViewTick + 1, + // Queued cleanups belong to the previous recipe; draining them after + // a save of this one could delete files its saved payload still uses. + pendingUploadCleanups: [], })), setAuxNodePosition: (id, position) => set((state) => { @@ -786,6 +832,17 @@ export const useRecipeStudioStore = create((set, get) => ({ set(applyUpdate); }, onNodesChange: (changes) => { + const current = get(); + if (!current.executionLocked) { + for (const change of changes) { + if (change.type === "remove") { + const uid = seedUploadCleanupUid(current.configs[change.id]); + if (uid) { + current.queueUploadCleanup(uid); + } + } + } + } const applyNodesChange = (state: RecipeStudioState) => { if (state.executionLocked) { return state; diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index b8ed13f70b..9231c5f7a3 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -340,6 +340,8 @@ export type SeedConfig = { hf_token?: string; hf_endpoint?: string; local_file_name?: string; + // ui-only: stable per-block id for uploads, since node ids collide across imports + unstructured_upload_uid?: string; unstructured_file_ids?: string[]; unstructured_file_names?: string[]; unstructured_file_sizes?: number[]; diff --git a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts index 76fc2c38ee..d47bac8858 100644 --- a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts +++ b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts @@ -20,6 +20,46 @@ import type { } from "../types"; import { nextName } from "./naming"; +export function makeUnstructuredUploadUid(): string { + if (typeof globalThis.crypto?.randomUUID === "function") { + return globalThis.crypto.randomUUID().replace(/-/g, "").toLowerCase(); + } + if (typeof globalThis.crypto?.getRandomValues === "function") { + const bytes = new Uint8Array(16); + globalThis.crypto.getRandomValues(bytes); + return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join( + "", + ); + } + let uid = ""; + while (uid.length < 32) { + uid += Math.floor(Math.random() * 0x100000000) + .toString(16) + .padStart(8, "0"); + } + return uid.slice(0, 32); +} + +export function resolveUnstructuredUploadBlockId({ + configId, + uploadUid, + generatedUploadUid, + unstructuredFileCount, +}: { + configId: string; + uploadUid: string; + generatedUploadUid: string | null; + unstructuredFileCount: number; +}): string { + if (uploadUid) { + return uploadUid; + } + if (generatedUploadUid) { + return generatedUploadUid; + } + return unstructuredFileCount > 0 ? configId : ""; +} + export function makeSamplerConfig( id: string, samplerType: SamplerType, @@ -368,6 +408,9 @@ export function makeSeedConfig( hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + ...(seedSourceType === "unstructured" + ? { unstructured_upload_uid: makeUnstructuredUploadUid() } + : {}), unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts index ac881d0373..54df2ffd50 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts @@ -16,11 +16,7 @@ import type { } from "../../types"; import { buildEdges } from "./edges"; import { isRecord, parseJson, readString } from "./helpers"; -import { - parseColumn, - parseModelConfig, - parseModelProvider, -} from "./parsers"; +import { parseColumn, parseModelConfig, parseModelProvider } from "./parsers"; import { parseSeedConfig } from "./parsers/seed-config-parser"; import { buildNodes, parseUi } from "./ui"; import type { ImportResult } from "./types"; @@ -43,6 +39,7 @@ type UiInput = { seed_drop_columns?: unknown; seed_preview_rows?: unknown; local_file_name?: unknown; + unstructured_upload_uid?: unknown; unstructured_file_ids?: unknown; unstructured_file_names?: unknown; unstructured_file_sizes?: unknown; @@ -51,6 +48,10 @@ type UiInput = { advanced_open_by_node?: unknown; }; +type ImportRecipePayloadOptions = { + preserveUnstructuredUploads?: boolean; +}; + type UiMarkdownNoteNode = { name: string; markdown: string; @@ -90,7 +91,7 @@ function parseProcessors(input: unknown): RecipeProcessorConfig[] { ? templateRaw : isRecord(templateRaw) ? JSON.stringify(templateRaw, null, 2) - : "{\n \"text\": \"{{ column_name }}\"\n}"; + : '{\n "text": "{{ column_name }}"\n}'; processors.push({ id: `p${index + 1}`, // biome-ignore lint/style/useNamingConvention: api schema @@ -135,9 +136,7 @@ function parseSeedDropColumns(input: unknown): string[] { return Array.from(values); } -function parseMcpProviders( - input: unknown, -): Map { +function parseMcpProviders(input: unknown): Map { const providers = new Map(); if (!Array.isArray(input)) { return providers; @@ -156,13 +155,12 @@ function parseMcpProviders( const args = Array.isArray(item.args) ? item.args.map((value) => String(value)) : []; - const envPairs = - isRecord(item.env) - ? Object.entries(item.env).map(([key, value]) => ({ - key: String(key), - value: String(value), - })) - : []; + const envPairs = isRecord(item.env) + ? Object.entries(item.env).map(([key, value]) => ({ + key: String(key), + value: String(value), + })) + : []; providers.set(name, { id: `mcp-${index + 1}`, name, @@ -209,7 +207,8 @@ function parseToolConfigs(input: unknown): Map { allow_tools: allowTools, // biome-ignore lint/style/useNamingConvention: api schema max_tool_call_turns: - item.max_tool_call_turns === null || item.max_tool_call_turns === undefined + item.max_tool_call_turns === null || + item.max_tool_call_turns === undefined ? "5" : String(item.max_tool_call_turns), // biome-ignore lint/style/useNamingConvention: api schema @@ -257,7 +256,9 @@ function parseUiMarkdownNoteNodes(input: unknown): UiMarkdownNoteNode[] { return noteNodes; } -function parseUiToolProfileNodes(input: unknown): Map> { +function parseUiToolProfileNodes( + input: unknown, +): Map> { const toolProfiles = new Map>(); if (!Array.isArray(input)) { return toolProfiles; @@ -312,9 +313,15 @@ function parseAdvancedOpenByNode(input: unknown): Record { return out; } -type AdvancedOpenConfig = LlmConfig | SamplerConfig | SeedConfig | ValidatorConfig; +type AdvancedOpenConfig = + | LlmConfig + | SamplerConfig + | SeedConfig + | ValidatorConfig; -function isAdvancedOpenConfig(config: NodeConfig): config is AdvancedOpenConfig { +function isAdvancedOpenConfig( + config: NodeConfig, +): config is AdvancedOpenConfig { return ( config.kind === "llm" || config.kind === "sampler" || @@ -350,7 +357,8 @@ function buildToolProfileConfig( .map((providerName) => mcpProvidersByName.get(providerName)) .flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])), // biome-ignore lint/style/useNamingConvention: ui schema - fetched_tools_by_provider: fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, + fetched_tools_by_provider: + fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, // biome-ignore lint/style/useNamingConvention: api schema allow_tools: [...(canonical.allow_tools ?? [])], // biome-ignore lint/style/useNamingConvention: api schema @@ -360,7 +368,10 @@ function buildToolProfileConfig( }; } -export function importRecipePayload(input: string): ImportResult { +export function importRecipePayload( + input: string, + options: ImportRecipePayloadOptions = {}, +): ImportResult { const parsed = parseJson(input); if (!parsed.data || !isRecord(parsed.data)) { return { @@ -369,9 +380,9 @@ export function importRecipePayload(input: string): ImportResult { }; } - const recipe = (isRecord(parsed.data.recipe) - ? parsed.data.recipe - : parsed.data) as RecipeInput; + const recipe = ( + isRecord(parsed.data.recipe) ? parsed.data.recipe : parsed.data + ) as RecipeInput; const ui = isRecord(parsed.data.ui) ? (parsed.data.ui as UiInput) : null; if (!Array.isArray(recipe.columns)) { @@ -410,21 +421,36 @@ export function importRecipePayload(input: string): ImportResult { .map((row) => ({ ...row })) : undefined; const uiLocalFileName = readString(ui?.local_file_name) ?? undefined; - // Preserve file IDs/names from saved recipes (cleared at share time by sanitizeSeedForShare) - const uiUnstructuredFileIds: string[] = Array.isArray(ui?.unstructured_file_ids) - ? (ui.unstructured_file_ids as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileNames: string[] = Array.isArray(ui?.unstructured_file_names) - ? (ui.unstructured_file_names as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileSizes: number[] = Array.isArray(ui?.unstructured_file_sizes) - ? (ui.unstructured_file_sizes as number[]).filter((v): v is number => typeof v === "number") - : []; + const preserveUnstructuredUploads = + options.preserveUnstructuredUploads === true; + const uiUnstructuredUploadUid = preserveUnstructuredUploads + ? (readString(ui?.unstructured_upload_uid) ?? undefined) + : undefined; + const uiUnstructuredFileIds: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_ids) + ? (ui.unstructured_file_ids as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileNames: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_names) + ? (ui.unstructured_file_names as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileSizes: number[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_sizes) + ? (ui.unstructured_file_sizes as number[]).filter( + (v): v is number => typeof v === "number", + ) + : []; const uiUnstructuredChunkSize = readStringNumber(ui?.unstructured_chunk_size); const uiUnstructuredChunkOverlap = readStringNumber( ui?.unstructured_chunk_overlap, ); - const uiAdvancedOpenByNode = parseAdvancedOpenByNode(ui?.advanced_open_by_node); + const uiAdvancedOpenByNode = parseAdvancedOpenByNode( + ui?.advanced_open_by_node, + ); const uiMarkdownNotes = parseUiMarkdownNoteNodes(ui?.nodes); const uiToolProfilesByName = parseUiToolProfileNodes(ui?.nodes); @@ -459,11 +485,13 @@ export function importRecipePayload(input: string): ImportResult { : payloadSeedDropColumns, seed_preview_rows: uiSeedPreviewRows, local_file_name: uiLocalFileName, + unstructuredUploadUid: uiUnstructuredUploadUid, unstructuredFileIds: uiUnstructuredFileIds, unstructuredFileNames: uiUnstructuredFileNames, unstructuredFileSizes: uiUnstructuredFileSizes, unstructured_chunk_size: uiUnstructuredChunkSize, unstructured_chunk_overlap: uiUnstructuredChunkOverlap, + preserveUnstructuredUploads, }); if (seedConfig) { applyAdvancedOpen(seedConfig, uiAdvancedOpenByNode); @@ -567,12 +595,7 @@ export function importRecipePayload(input: string): ImportResult { const { layouts, auxNodes, edges: uiEdges, layoutDirection } = parseUi(ui); const resolvedLayoutDirection = layoutDirection ?? "LR"; const nodes = buildNodes(configs, layouts); - const edges = buildEdges( - configs, - nameToId, - uiEdges, - resolvedLayoutDirection, - ); + const edges = buildEdges(configs, nameToId, uiEdges, resolvedLayoutDirection); const auxNodePositions = Object.fromEntries( auxNodes.flatMap((item) => { const llmId = nameToId.get(item.llm); @@ -583,10 +606,7 @@ export function importRecipePayload(input: string): ImportResult { }), ); - const maxY = nodes.reduce( - (acc, node) => Math.max(acc, node.position.y), - 0, - ); + const maxY = nodes.reduce((acc, node) => Math.max(acc, node.position.y), 0); return { errors: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts index 21eadb3195..939205fe6d 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts @@ -197,17 +197,26 @@ export function parseSeedConfig( seed_drop_columns?: string[]; seed_preview_rows?: Record[]; local_file_name?: string; + unstructuredUploadUid?: string; unstructuredFileIds?: string[]; unstructuredFileNames?: string[]; unstructuredFileSizes?: number[]; unstructured_chunk_size?: string; unstructured_chunk_overlap?: string; + preserveUnstructuredUploads?: boolean; }, ): SeedConfig | null { if (!seedConfigRaw) { return null; } - const parsed = parseSeedSettings(seedConfigRaw); + const parsed = { ...parseSeedSettings(seedConfigRaw) }; + if ( + parsed.seed_source_type === "unstructured" && + options?.preserveUnstructuredUploads !== true + ) { + parsed.hf_path = ""; + parsed.resolved_paths = []; + } let sourceType: SeedSourceType = "hf"; if (parsed.seed_source_type === "hf") { sourceType = "hf"; @@ -230,6 +239,9 @@ export function parseSeedConfig( ...(options?.local_file_name !== undefined ? { local_file_name: options.local_file_name } : {}), + ...(options?.unstructuredUploadUid + ? { unstructured_upload_uid: options.unstructuredUploadUid } + : {}), ...(options?.unstructuredFileIds !== undefined ? { unstructured_file_ids: options.unstructuredFileIds } : {}), diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts index 34c3c34274..9b2ad5b2b5 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts @@ -440,6 +440,9 @@ export function buildRecipePayload( unstructured_file_names: firstSeed.unstructured_file_names, unstructured_file_sizes: firstSeed.unstructured_file_sizes, }), + ...(firstSeed?.unstructured_upload_uid?.trim() && { + unstructured_upload_uid: firstSeed.unstructured_upload_uid, + }), ...(firstSeed && firstSeed.unstructured_chunk_size !== undefined && { unstructured_chunk_size: firstSeed.unstructured_chunk_size, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts index 902ea796b4..763e68c1fb 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts @@ -71,6 +71,8 @@ export type RecipePayload = { seed_preview_rows?: Record[]; local_file_name?: string; // biome-ignore lint/style/useNamingConvention: api schema + unstructured_upload_uid?: string; + // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_ids?: string[]; // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_names?: string[]; diff --git a/studio/frontend/src/features/settings/api/coding-agents.ts b/studio/frontend/src/features/settings/api/coding-agents.ts new file mode 100644 index 0000000000..ae371b2d3a --- /dev/null +++ b/studio/frontend/src/features/settings/api/coding-agents.ts @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type CodingAgentsInfo = { + // Every agent `unsloth start` supports, in the CLI's declared order. + agents: string[]; + // Subset of `agents` whose CLI binary was found on PATH by the backend. + detected: string[]; +}; + +type ApiCodingAgentsInfo = { + agents: string[]; + detected: string[]; +}; + +// Which CLIs are on PATH is environment state, not a persisted setting -- it +// can change any time the user installs something new, so this only +// de-duplicates concurrent in-flight calls (e.g. React strict-mode's double +// mount) rather than caching the result across the module's lifetime. Every +// fresh call (each time a settings panel mounts) re-checks PATH for real. +let inFlightInfo: Promise | null = null; + +function fromApi(info: ApiCodingAgentsInfo): CodingAgentsInfo { + return { agents: info.agents, detected: info.detected }; +} + +async function fetchCodingAgents(): Promise { + const res = await authFetch("/api/settings/coding-agents"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load installed coding agents"), + ); + } + return fromApi(await res.json()); +} + +export async function loadCodingAgents(): Promise { + inFlightInfo ??= fetchCodingAgents().finally(() => { + inFlightInfo = null; + }); + return inFlightInfo; +} diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts index 9e87922970..38b2d73c3b 100644 --- a/studio/frontend/src/features/settings/components/agent-command.ts +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -12,7 +12,7 @@ const DEFAULT_AGENT = "claude"; // URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is // "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below. -function normalizeHost(host: string): string { +export function normalizeHost(host: string): string { const lower = host.toLowerCase(); return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; } @@ -26,7 +26,7 @@ function isDefaultLocalHost(host: string): boolean { } // Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. -function isLoopbackHost(host: string): boolean { +export function isLoopbackHost(host: string): boolean { if (host === "localhost" || host === "::1") return true; const octets = host.split("."); return ( diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index c43dd0f219..b3396c8d9b 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -16,6 +16,7 @@ import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { useChatRuntimeStore } from "@/features/chat"; import { useT } from "@/i18n"; import type { TranslationKey } from "@/i18n"; +import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { Tick02Icon } from "@/lib/tick-icon"; import { cn } from "@/lib/utils"; @@ -25,14 +26,15 @@ import { InformationCircleIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useEffect, useMemo, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { Streamdown } from "streamdown"; +import { loadCodingAgents } from "../api/coding-agents"; import { type OpenAIAutoSwitchSettings, loadOpenAIAutoSwitchSettings, updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; -import { buildAgentCommand } from "./agent-command"; +import { buildAgentCommand, isLoopbackHost, normalizeHost } from "./agent-command"; type ExampleType = | "curl" @@ -114,6 +116,30 @@ const DOC_LINKS = [ { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; +// Falls back to this list until the backend's installed-CLI check resolves; +// kept in sync with the `unsloth start ` subcommands and with +// CODING_AGENTS in studio/backend/utils/coding_agents.py. +const DEFAULT_AGENTS = [ + "claude", + "codex", + "openclaw", + "opencode", + "hermes", + "pi", +]; +// The agent selection resets to this whenever an auto-pick is no longer +// trustworthy (leaving loopback, or the only compatible detected agent +// stops being compatible) rather than lingering on a stale choice. +const DEFAULT_AGENT = "claude"; +const AGENT_LABELS: Record = { + claude: "Claude Code", + codex: "Codex", + openclaw: "OpenClaw", + opencode: "OpenCode", + hermes: "Hermes", + pi: "Pi", +}; + const j = (s: string): string => JSON.stringify(s); const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); const psSingle = (s: string): string => s.replace(/'/g, "''"); @@ -399,6 +425,17 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } +// Backend PATH detection is only safe in the desktop app, where the UI owns +// the local backend. A browser loopback URL may be an SSH/local port forward. +function canUseLocalAgentDetection(base: string): boolean { + if (!isTauri) return false; + try { + return isLoopbackHost(normalizeHost(new URL(base).hostname)); + } catch { + return false; + } +} + const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -443,7 +480,18 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); const [copiedAgent, setCopiedAgent] = useState(false); + const [agent, setAgent] = useState(DEFAULT_AGENT); + const [availableAgents, setAvailableAgents] = + useState(DEFAULT_AGENTS); + const [detectedAgents, setDetectedAgents] = useState([]); + // True once the user has picked an agent themselves; guards the detection + // effect below from clobbering that choice if it resolves afterward. + const agentPickedByUserRef = useRef(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + const origin = typeof window !== "undefined" ? window.location.origin : ""; + const base = + useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const localAgentDetection = canUseLocalAgentDetection(base); // null while loading; the same setting the General tab exposes (shared cache). const [autoSwitch, setAutoSwitch] = useState( null, @@ -454,6 +502,78 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { void fetchDeviceType({ force: true }); }, []); + // Fetching is the only job of this effect: populate availableAgents/ + // detectedAgents (or clear them). Which agent gets auto-picked from that + // list is derived separately below, so it can react to the loaded model + // changing too, not just a fresh fetch. + useEffect(() => { + // Browser loopback URLs can be SSH/local forwards, so only the desktop app + // may use backend PATH checks to mark or auto-pick local agents. + if (!localAgentDetection) { + setDetectedAgents([]); + // A previously auto-picked agent was only ever verified against the + // Studio backend's PATH, which is meaningless now that this panel no + // longer targets a loopback base -- don't leave it selected, but + // never touch a choice the user made by hand. + if (!agentPickedByUserRef.current) { + setAgent(DEFAULT_AGENT); + } + return; + } + + let cancelled = false; + void loadCodingAgents() + .then((info) => { + if (cancelled) return; + setAvailableAgents(info.agents); + setDetectedAgents(info.detected); + }) + .catch(() => { + // Best-effort: keep the default agent list and let the user pick manually. + }); + return () => { + cancelled = true; + }; + }, [localAgentDetection]); + + // Single source of truth for the auto-picked agent, re-derived whenever + // the detected list or the loaded model's GGUF-ness changes -- in either + // direction. `codex` needs a GGUF model (unsloth_cli's + // _require_gguf_for_codex exits otherwise), so it's only preferred once + // the loaded model actually qualifies; loading a GGUF model *after* a + // non-GGUF-gated fallback picked something else re-steers back to codex + // just as loading a non-GGUF model steers away from it. Never overrides a + // choice the user made by hand. + // activeGgufVariant alone only covers an HF-repo GGUF pick (a specific + // quant variant string) -- a direct local .gguf file (custom folder / + // LM Studio / drag-drop) is just as much a GGUF the codex preflight would + // accept, but never has a "variant" to report, and would otherwise read as + // non-GGUF here. activeNativePathToken covers the drag-drop/picked-file + // case; ggufContextLength is only ever populated when the backend's + // /api/inference/status last reported is_gguf: true for the active model + // (see applyActiveModelStatusToStore), so together these three cover every + // path a model can be GGUF through, matching the same is_gguf-or-equivalent + // check hasGgufSource applies to a staged pick. + const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + const activeNativePathToken = useChatRuntimeStore((s) => s.activeNativePathToken); + const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + useEffect(() => { + if (agentPickedByUserRef.current) return; + if (detectedAgents.length === 0) return; + const isGguf = + activeGgufVariant != null || activeNativePathToken != null || ggufContextLength != null; + const preferred = detectedAgents.find((a) => a !== "codex" || isGguf); + if (preferred) { + setAgent(preferred); + } else if (agent === "codex" && !isGguf) { + // codex was auto-picked while a GGUF model was active and it's the + // only detected agent; now that the model isn't GGUF anymore, nothing + // detected is actually runnable, so fall back to the default instead + // of leaving a codex command unsloth_cli will reject. + setAgent(DEFAULT_AGENT); + } + }, [agent, detectedAgents, activeGgufVariant, activeNativePathToken, ggufContextLength]); + useEffect(() => { let cancelled = false; void loadOpenAIAutoSwitchSettings() @@ -470,9 +590,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const model = useLoadedModelName(); const key = apiKey || KEY_PLACEHOLDER; - const origin = typeof window !== "undefined" ? window.location.origin : ""; - const base = - useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( @@ -481,8 +598,8 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); // Agent command must target the server the panel shows, not the :8888 default. const agentCommand = useMemo( - () => buildAgentCommand(base, key, os), - [base, key, os], + () => buildAgentCommand(base, key, os, agent), + [base, key, os, agent], ); const osAware = OS_AWARE[lang]; @@ -710,6 +827,42 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {t("settings.apiKeys.codingAgentsHint")} +
+ {availableAgents.map((id) => { + const installed = detectedAgents.includes(id); + const active = agent === id; + return ( + + ); + })} +
{agentCommand} @@ -727,7 +880,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
- {t("settings.apiKeys.codingAgentsSwap")} + {detectedAgents.length > 0 + ? t("settings.apiKeys.codingAgentsDetectedHint", { + agents: detectedAgents + .map((id) => AGENT_LABELS[id] ?? id) + .join(", "), + }) + : t("settings.apiKeys.codingAgentsSwap")}
diff --git a/studio/frontend/src/features/settings/index.ts b/studio/frontend/src/features/settings/index.ts index 364ca1611f..3fefd8c63a 100644 --- a/studio/frontend/src/features/settings/index.ts +++ b/studio/frontend/src/features/settings/index.ts @@ -7,6 +7,7 @@ export { savePersonalization, } from "./api/personalization"; export { setTheme, useTheme } from "./stores/theme-store"; +export { useMonitorOverlayStore } from "./stores/monitor-overlay-store"; export type { Personalization, PersonalizationAppearance, diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 0d201edc0d..d8a0d45d3f 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -56,6 +56,7 @@ const TABS: TabDef[] = [ id: "resources", labelKey: "settings.tabs.resources", icon: CpuIcon, + badgeKey: "common.new", }, { id: "chat", @@ -72,7 +73,6 @@ const TABS: TabDef[] = [ id: "connections", labelKey: "settings.tabs.connections", icon: CloudIcon, - badgeKey: "common.new", }, { id: "about", labelKey: "settings.tabs.about", icon: HelpCircleIcon }, ]; diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index b67fd5ca1d..e0cb8030ae 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -409,7 +409,6 @@ export const en = { description: "Access Unsloth via the OpenAI-compatible API.", readDocs: "Read the API docs", noAccess: "No API access yet.", - newBadge: "New", accessTokens: "Access tokens", loadError: "Couldn't load API access.", createError: "Couldn't create access token.", @@ -445,6 +444,8 @@ export const en = { codingAgentsHint: "Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.", codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.", + codingAgentDetected: "Installed on this machine", + codingAgentsDetectedHint: "Detected on this machine: {agents}.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts index f2020f3cfb..08b8d4f4d3 100644 --- a/studio/frontend/src/i18n/locales/ja.ts +++ b/studio/frontend/src/i18n/locales/ja.ts @@ -296,7 +296,6 @@ export const ja = { description: "OpenAI互換 API を介して Unsloth にアクセスします。", readDocs: "API ドキュメントを読む", noAccess: "まだ API アクセス権がありません。", - newBadge: "新規", accessTokens: "アクセストークン", loadError: "API アクセス権を読み込めませんでした。", createError: "アクセストークンを作成できませんでした。", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts index c261e4ed0c..494a98ec32 100644 --- a/studio/frontend/src/i18n/locales/pt-br.ts +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -363,7 +363,6 @@ export const ptBR = { "Acesse o Unsloth por meio da API compatível com OpenAI.", readDocs: "Leia a documentação da API", noAccess: "Nenhum acesso à API ainda.", - newBadge: "Novo", accessTokens: "Tokens de acesso", loadError: "Não foi possível carregar o acesso à API.", createError: "Não foi possível criar o token de acesso.", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index f6dc265fc7..dda8e017ab 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -267,7 +267,6 @@ export const zhCN = { description: "通过兼容 OpenAI 的 API 以编程方式访问 Unsloth。", readDocs: "阅读 API 文档", noAccess: "还没有 API 访问权限。", - newBadge: "新", accessTokens: "访问 token", loadError: "无法加载 API 访问权限。", createError: "无法创建访问 token。", diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 86a9634048..edf9875602 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -173,7 +173,11 @@ function looksLikeMathBody(body: string): boolean { * (`**$X$**`, `__$X$__`) are always math: LLMs use that for "bold math" * and the heuristic would otherwise reject prose-shaped bodies like "90 - x". */ -function hasInlineMathCloser(content: string, offset: number): boolean { +function hasInlineMathCloser( + content: string, + offset: number, + mathRegions: Array<[number, number]>, +): boolean { const MAX_SPAN = 200; const limit = Math.min(content.length, offset + 1 + MAX_SPAN); for (let i = offset + 1; i < limit; i++) { @@ -181,6 +185,9 @@ function hasInlineMathCloser(content: string, offset: number): boolean { if (c === "\n") return false; if (c !== "$") continue; if (content[i - 1] === "\\") continue; + // A `$` opening a generated span (from `\(...\)`) is not a currency closer; + // pairing with it would swallow the price into math (`$5 + x \(y\)`). + if (isInRegion(i, mathRegions)) return false; if (content[i + 1] === "$") { i++; continue; @@ -294,7 +301,22 @@ function convertLatexDelimiters(content: string): { continue; } append(content.slice(last, match.index)); - const wrapped = isDisplay ? `\n$$\n${body}\n$$\n` : `$${body}$`; + let wrapped: string; + if (isDisplay) { + // Keep the opener's leading indentation so a `$$` block inside a list item + // stays in the container instead of breaking out at column 0. Only when the + // opener is whitespace-prefixed, so inline `text \[x\]` keeps column 0. + const lineStart = + match.index > 0 ? content.lastIndexOf("\n", match.index - 1) + 1 : 0; + const prefix = content.slice(lineStart, match.index); + const indent = /^\s*$/.test(prefix) ? prefix : ""; + // Indent every body line, not just the first, so multi-line display math + // (`\[a\nb\]`) stays wholly inside the container. + const inner = indent ? body.replace(/\n/g, `\n${indent}`) : body; + wrapped = `\n${indent}$$\n${indent}${inner}\n${indent}$$\n`; + } else { + wrapped = `$${body}$`; + } const start = append(wrapped); mathRegions.push([start, offset]); last = matchEnd; @@ -334,7 +356,7 @@ export function preprocessLaTeX(content: string): string { if (isInRegion(offset, mathRegions)) { return match; } - if (hasInlineMathCloser(text, offset)) { + if (hasInlineMathCloser(text, offset, mathRegions)) { return match; } return "\\" + match; diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index e40cb3083e..ca1fe79efa 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -10,6 +10,7 @@ import argparse import atexit import errno import fnmatch +import glob import hashlib import json import os @@ -165,9 +166,9 @@ def env_int( # errors. Only use "master" temporarily when the latest release is missing # support for a new model architecture. DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", "latest") -# Default published repo for prebuilt release resolution. Linux uses -# Unsloth prebuilts; setup.sh/setup.ps1 pass --published-repo explicitly -# for macOS/Windows to override with ggml-org/llama.cpp when needed. +# Default published repo for prebuilt release resolution. Every host plans +# its prebuilt against the Unsloth fork; setup.sh/setup.ps1 pass it via +# --published-repo. ggml-org is reachable only via an explicit override. DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" DEFAULT_PUBLISHED_TAG = os.environ.get("UNSLOTH_LLAMA_RELEASE_TAG") DEFAULT_PUBLISHED_MANIFEST_ASSET = os.environ.get( @@ -265,6 +266,7 @@ class HostInfo: has_physical_nvidia: bool has_usable_nvidia: bool has_rocm: bool = False + has_intel_gpu: bool = False rocm_gfx_target: str | None = None # (major, minor) from platform.mac_ver(); None off macOS or if unparseable. # Skips a macos prebuilt whose minimum-OS exceeds this host. @@ -1284,162 +1286,6 @@ def synthetic_checksums_for_release( ) -def parse_direct_linux_release_bundle( - repo: str, release: dict[str, Any] -) -> PublishedReleaseBundle | None: - release_tag = release.get("tag_name") - if not isinstance(release_tag, str) or not release_tag: - return None - - assets = release_asset_map(release) - artifacts: list[PublishedLlamaArtifact] = [] - inferred_labels: list[str] = [] - - linux_asset_re = re.compile( - r"^app-(?P