Merge remote-tracking branch 'origin/main' into cuda-torch-index-override
# Conflicts: # install.ps1 # install.sh # studio/install_python_stack.py
This commit is contained in:
Daniel Han
commit
afe4da63a0
247 changed files with 48195 additions and 5050 deletions
164
.github/scripts/agent-guides-drive.sh
vendored
164
.github/scripts/agent-guides-drive.sh
vendored
|
|
@ -154,7 +154,12 @@ raw_env() { # $1 = var name -> value (one shlex-quote layer stripped)
|
|||
# writers as a side effect (it writes each agent's relocated session config).
|
||||
parse_connect() {
|
||||
local raw="$LOGS_DIR/connect-${AGENT}.txt"
|
||||
if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
|
||||
# CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their
|
||||
# config (which now prompts by default), so the file-edit test opts into auto-approval
|
||||
# here, the same intent as claude/codex's per-call bypass flags.
|
||||
local yolo=()
|
||||
[ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo)
|
||||
if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
|
||||
cat_redacted "$raw"
|
||||
guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero"
|
||||
fi
|
||||
|
|
@ -371,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
|
||||
|
|
@ -394,7 +399,10 @@ case "$MODE" in
|
|||
T2='Run hello.py with python and show me the exact output.'
|
||||
|
||||
# The start.py recipe writers + crosscheck must see the repo; run them
|
||||
# from the repo root BEFORE cd-ing into the scratch work dir.
|
||||
# from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw
|
||||
# gate tool approval through their config (prompting by default), so file-edit
|
||||
# opts them into auto-approval to run edits/commands headlessly.
|
||||
case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac
|
||||
parse_connect
|
||||
crosscheck_contract
|
||||
# File-edit needs real tools, so we cannot zero them as in connection.
|
||||
|
|
@ -441,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
|
||||
|
|
@ -519,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
|
||||
# <agent> ...`, 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
|
||||
|
|
|
|||
3
.github/workflows/consolidated-tests-ci.yml
vendored
3
.github/workflows/consolidated-tests-ci.yml
vendored
|
|
@ -364,6 +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
|
||||
|
|
|
|||
170
.github/workflows/local-agent-guides-ci.yml
vendored
170
.github/workflows/local-agent-guides-ci.yml
vendored
|
|
@ -471,6 +471,176 @@ jobs:
|
|||
redacted-configs/
|
||||
retention-days: 7
|
||||
|
||||
# ═════════════════════════════════════════════════════════════════════
|
||||
# Job: resume
|
||||
# Does a conversation started with `unsloth start <agent>` 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}#<REDACTED>#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
|
||||
|
|
|
|||
4
.github/workflows/lockfile-audit.yml
vendored
4
.github/workflows/lockfile-audit.yml
vendored
|
|
@ -60,11 +60,11 @@ jobs:
|
|||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
|
|
|
|||
78
.github/workflows/ossf.yml
vendored
Normal file
78
.github/workflows/ossf.yml
vendored
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
# This workflow uses actions that are not certified by GitHub. They are provided
|
||||
# by a third-party and are governed by separate terms of service, privacy
|
||||
# policy, and support documentation.
|
||||
|
||||
name: Scorecard supply-chain security
|
||||
on:
|
||||
# For Branch-Protection check. Only the default branch is supported. See
|
||||
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection
|
||||
branch_protection_rule:
|
||||
# To guarantee Maintained check is occasionally updated. See
|
||||
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained
|
||||
schedule:
|
||||
- cron: '21 20 * * 0'
|
||||
push:
|
||||
branches: [ "main" ]
|
||||
|
||||
# Declare default permissions as read only.
|
||||
permissions: read-all
|
||||
|
||||
jobs:
|
||||
analysis:
|
||||
name: Scorecard analysis
|
||||
runs-on: ubuntu-latest
|
||||
# `publish_results: true` only works when run from the default branch. conditional can be removed if disabled.
|
||||
if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request'
|
||||
permissions:
|
||||
# Needed to upload the results to code-scanning dashboard.
|
||||
security-events: write
|
||||
# Needed to publish results and get a badge (see publish_results below).
|
||||
id-token: write
|
||||
# Uncomment the permissions below if installing in a private repository.
|
||||
# contents: read
|
||||
# actions: read
|
||||
|
||||
steps:
|
||||
- name: "Checkout code"
|
||||
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: "Run analysis"
|
||||
uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1
|
||||
with:
|
||||
results_file: results.sarif
|
||||
results_format: sarif
|
||||
# (Optional) "write" PAT token. Uncomment the `repo_token` line below if:
|
||||
# - you want to enable the Branch-Protection check on a *public* repository, or
|
||||
# - you are installing Scorecard on a *private* repository
|
||||
# To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional.
|
||||
# repo_token: ${{ secrets.SCORECARD_TOKEN }}
|
||||
|
||||
# Public repositories:
|
||||
# - Publish results to OpenSSF REST API for easy access by consumers
|
||||
# - Allows the repository to include the Scorecard badge.
|
||||
# - See https://github.com/ossf/scorecard-action#publishing-results.
|
||||
# For private repositories:
|
||||
# - `publish_results` will always be set to `false`, regardless
|
||||
# of the value entered here.
|
||||
publish_results: true
|
||||
|
||||
# (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore
|
||||
# file_mode: git
|
||||
|
||||
# Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
|
||||
# format to the repository Actions tab.
|
||||
- name: "Upload artifact"
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: SARIF file
|
||||
path: results.sarif
|
||||
retention-days: 5
|
||||
|
||||
# Upload the results to GitHub's code scanning dashboard (optional).
|
||||
# Commenting out will disable upload of results to your repo's Code Scanning dashboard
|
||||
- name: "Upload to code-scanning"
|
||||
uses: github/codeql-action/upload-sarif@v3
|
||||
with:
|
||||
sarif_file: results.sarif
|
||||
6
.github/workflows/security-audit.yml
vendored
6
.github/workflows/security-audit.yml
vendored
|
|
@ -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]
|
||||
|
|
|
|||
44
.github/workflows/studio-inference-smoke.yml
vendored
44
.github/workflows/studio-inference-smoke.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
44
.github/workflows/studio-mac-inference-smoke.yml
vendored
44
.github/workflows/studio-mac-inference-smoke.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
212
.github/workflows/studio-windows-inference-smoke.yml
vendored
212
.github/workflows/studio-windows-inference-smoke.yml
vendored
|
|
@ -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<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
|
||||
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
|
||||
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
|
||||
|
||||
Write-Host "ProgramFiles simulation root: $pf"
|
||||
Write-Host "ProgramFiles(x86) simulation root: $pfx86"
|
||||
if ($blocked.Count -gt 0) {
|
||||
Write-Host "Removed build-tool PATH dirs:"
|
||||
$blocked | Sort-Object | ForEach-Object { Write-Host " $_" }
|
||||
} else {
|
||||
Write-Host "No cmake or cl.exe PATH dirs found to remove."
|
||||
}
|
||||
("HIDDEN_CMAKE=" + ($hidden -join '|')) | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
|
||||
|
||||
- name: Assert Visual Studio + CMake are genuinely undetectable
|
||||
shell: pwsh
|
||||
run: |
|
||||
$ErrorActionPreference = 'Stop'
|
||||
# Set in-script: the runner does not apply step-level env keys with
|
||||
# parentheses (`ProgramFiles(x86)`), so vswhere still found VS.
|
||||
if (-not $env:NO_BUILD_TOOLS_PROGRAMFILES) { Write-Error "NO_BUILD_TOOLS_* env missing (Prepare step did not run?)"; exit 1 }
|
||||
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
|
||||
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
|
||||
$env:Path = $env:NO_BUILD_TOOLS_PATH
|
||||
. (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1')
|
||||
$setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1'
|
||||
foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Find-VsBuildTools')) {
|
||||
|
|
@ -1394,6 +1463,10 @@ jobs:
|
|||
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
|
||||
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
|
||||
run: |
|
||||
# Set in-script (see the assert step); child processes inherit these.
|
||||
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
|
||||
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
|
||||
$env:Path = $env:NO_BUILD_TOOLS_PATH
|
||||
New-Item -ItemType Directory -Force -Path logs | Out-Null
|
||||
$ProgressPreference = 'SilentlyContinue'
|
||||
& ./install.ps1 --local --no-torch *>&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<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
|
||||
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
|
||||
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
|
||||
|
||||
- name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS)
|
||||
env:
|
||||
|
|
@ -1577,25 +1663,34 @@ jobs:
|
|||
echo "Windows CUDA and ROCm prebuilts are available -- GPU users get them without compiling."
|
||||
|
||||
- name: The prebuilt resolver runs without Visual Studio
|
||||
shell: pwsh
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
$ErrorActionPreference = 'Stop'
|
||||
# pwsh: bash cannot export `ProgramFiles(x86)`; set in-script so the
|
||||
# python child inherits the overrides.
|
||||
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
|
||||
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
|
||||
$env:Path = $env:NO_BUILD_TOOLS_PATH
|
||||
# Resolver-only (no GPU on hosted runners, so the host resolves to the
|
||||
# CPU bundle). The point is that resolution needs no compiler/VS.
|
||||
python -m pip install --upgrade huggingface_hub
|
||||
python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > /tmp/resolve.json || {
|
||||
echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; }
|
||||
cat /tmp/resolve.json
|
||||
echo "Prebuilt resolver ran with no Visual Studio present."
|
||||
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:
|
||||
|
|
@ -1610,6 +1705,13 @@ jobs:
|
|||
- name: Install Pester v5
|
||||
shell: pwsh
|
||||
run: |
|
||||
# PSGallery is intermittently absent from the repository list on GitHub's Windows
|
||||
# runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the
|
||||
# name 'PSGallery' was found." Re-register the default gallery first so the policy
|
||||
# change and module install below always have a repository to target.
|
||||
if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) {
|
||||
Register-PSRepository -Default -ErrorAction SilentlyContinue
|
||||
}
|
||||
Set-PSRepository PSGallery -InstallationPolicy Trusted
|
||||
Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser
|
||||
Import-Module Pester -MinimumVersion 5.5.0
|
||||
|
|
|
|||
|
|
@ -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-<tag>-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
|
||||
|
|
|
|||
86
.github/workflows/version-compat-ci.yml
vendored
86
.github/workflows/version-compat-ci.yml
vendored
|
|
@ -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:
|
||||
|
|
|
|||
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -11,6 +11,8 @@ outputs/
|
|||
exports/
|
||||
/datasets/
|
||||
studio/backend/assets/datasets/
|
||||
# Generated async worker / reviewer transcripts (never part of the product).
|
||||
studio/backend/async_task_outputs/
|
||||
unsloth_training_checkpoints/
|
||||
*.gguf
|
||||
*.safetensors
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
|
|
|||
34
install.ps1
34
install.ps1
|
|
@ -473,6 +473,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 {
|
||||
|
|
@ -492,6 +503,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] } } }
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -2254,7 +2266,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.6.9" "unsloth-zoo>=2026.6.7" }
|
||||
$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
|
||||
|
|
@ -2268,7 +2280,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.6.9" "unsloth-zoo>=2026.6.7" }
|
||||
$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
|
||||
|
|
@ -2299,7 +2311,7 @@ exit 0
|
|||
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
|
||||
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
|
||||
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -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.
|
||||
|
|
@ -2312,7 +2324,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 $CpuFallbackIndexUrl }
|
||||
$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 $CpuFallbackIndexUrl }
|
||||
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)
|
||||
|
|
@ -2329,7 +2341,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)
|
||||
|
|
@ -2341,7 +2353,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.6.9" "unsloth-zoo>=2026.6.7" }
|
||||
$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 }
|
||||
|
|
@ -2353,7 +2365,7 @@ exit 0
|
|||
}
|
||||
}
|
||||
} elseif ($StudioLocalInstall) {
|
||||
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
|
||||
$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" }
|
||||
}
|
||||
|
|
@ -2381,7 +2393,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.6.7" "unsloth>=2026.6.9" --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)
|
||||
|
|
@ -2412,7 +2424,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
|
||||
|
|
@ -2428,7 +2440,7 @@ exit 0
|
|||
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
|
||||
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -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)
|
||||
|
|
@ -2440,7 +2452,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)
|
||||
|
|
|
|||
210
install.sh
210
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=$?
|
||||
|
|
@ -1446,8 +1452,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then
|
|||
SKIP_TORCH=true
|
||||
fi
|
||||
|
||||
# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for
|
||||
# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file
|
||||
# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover.
|
||||
_MLX_LM_EXCLUDE_ARG=""
|
||||
|
||||
# Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file).
|
||||
if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
|
||||
_MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3"
|
||||
_OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt"
|
||||
if [ -f "$_OVERRIDES_FILE" ]; then
|
||||
# uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace
|
||||
|
|
@ -1481,6 +1493,81 @@ elif [ "$OS" = "macos" ]; then
|
|||
fi
|
||||
tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none"
|
||||
|
||||
# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in
|
||||
# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded
|
||||
# to 10s. Defined here so the reroute below can use it before _run_bounded exists.
|
||||
_WSL_AMD_GPU_NAME_CACHE=""
|
||||
_wsl_amd_gpu_name() {
|
||||
if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then
|
||||
[ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1
|
||||
printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0
|
||||
fi
|
||||
command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; }
|
||||
_wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name"
|
||||
if command -v timeout >/dev/null 2>&1; then
|
||||
_wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
|
||||
else
|
||||
_wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
|
||||
fi
|
||||
if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi
|
||||
_WSL_AMD_GPU_NAME_CACHE="-"; return 1
|
||||
}
|
||||
|
||||
# ── Bounded command runner ──
|
||||
# Runs a command under a 10s timeout when the `timeout` binary is available,
|
||||
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
|
||||
# driver init or after a reset) from hanging the installer: a timed-out probe
|
||||
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
|
||||
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
|
||||
_run_bounded() {
|
||||
if command -v timeout >/dev/null 2>&1; then
|
||||
timeout 10 "$@"
|
||||
else
|
||||
"$@"
|
||||
fi
|
||||
}
|
||||
|
||||
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
|
||||
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
|
||||
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
|
||||
# var, so the probes below cannot see the distinction on their own.
|
||||
_cvd_hides_nvidia() {
|
||||
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
|
||||
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
|
||||
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
|
||||
}
|
||||
|
||||
# ── NVIDIA usable-GPU helper ──
|
||||
# Returns 0 (true) if an NVIDIA GPU is present and usable.
|
||||
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
|
||||
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
|
||||
# -- handles PATH gaps, subprocess timeouts, and driver init races that
|
||||
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
|
||||
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
|
||||
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
|
||||
_has_usable_nvidia_gpu() {
|
||||
if _cvd_hides_nvidia; then
|
||||
return 1
|
||||
fi
|
||||
_nvsmi=""
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
_nvsmi="nvidia-smi"
|
||||
elif [ -x "/usr/bin/nvidia-smi" ]; then
|
||||
_nvsmi="/usr/bin/nvidia-smi"
|
||||
fi
|
||||
if [ -n "$_nvsmi" ]; then
|
||||
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
|
||||
return 0
|
||||
fi
|
||||
fi
|
||||
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
|
||||
if [ -d /proc/driver/nvidia/gpus ] && \
|
||||
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04)
|
||||
# with a 24.04 distro present, re-run the install there and stop; else fall through
|
||||
# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before
|
||||
|
|
@ -1491,7 +1578,15 @@ _maybe_reroute_strixhalo_to_2404() {
|
|||
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
|
||||
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
|
||||
[ -e /dev/dxg ] || return 0
|
||||
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
|
||||
# A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on
|
||||
# this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors
|
||||
# CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps.
|
||||
if _has_usable_nvidia_gpu; then return 0; fi
|
||||
# Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes.
|
||||
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
|
||||
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
|
||||
return 0
|
||||
fi
|
||||
# Already ROCm-on-WSL? leave a working GPU alone, whatever the version.
|
||||
if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then
|
||||
return 0
|
||||
|
|
@ -1965,61 +2060,6 @@ _has_amd_rocm_gpu() {
|
|||
return 1
|
||||
}
|
||||
|
||||
# ── Bounded command runner ──
|
||||
# Runs a command under a 10s timeout when the `timeout` binary is available,
|
||||
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
|
||||
# driver init or after a reset) from hanging the installer: a timed-out probe
|
||||
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
|
||||
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
|
||||
_run_bounded() {
|
||||
if command -v timeout >/dev/null 2>&1; then
|
||||
timeout 10 "$@"
|
||||
else
|
||||
"$@"
|
||||
fi
|
||||
}
|
||||
|
||||
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
|
||||
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
|
||||
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
|
||||
# var, so the probes below cannot see the distinction on their own.
|
||||
_cvd_hides_nvidia() {
|
||||
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
|
||||
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
|
||||
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
|
||||
}
|
||||
|
||||
# ── NVIDIA usable-GPU helper ──
|
||||
# Returns 0 (true) if an NVIDIA GPU is present and usable.
|
||||
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
|
||||
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
|
||||
# -- handles PATH gaps, subprocess timeouts, and driver init races that
|
||||
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
|
||||
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
|
||||
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
|
||||
_has_usable_nvidia_gpu() {
|
||||
if _cvd_hides_nvidia; then
|
||||
return 1
|
||||
fi
|
||||
_nvsmi=""
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
_nvsmi="nvidia-smi"
|
||||
elif [ -x "/usr/bin/nvidia-smi" ]; then
|
||||
_nvsmi="/usr/bin/nvidia-smi"
|
||||
fi
|
||||
if [ -n "$_nvsmi" ]; then
|
||||
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
|
||||
return 0
|
||||
fi
|
||||
fi
|
||||
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
|
||||
if [ -d /proc/driver/nvidia/gpus ] && \
|
||||
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
# ── Detect GPU and choose PyTorch index URL ──
|
||||
# Mirrors Get-TorchIndexUrl in install.ps1.
|
||||
# On CPU-only machines this returns the cpu index, avoiding the solver
|
||||
|
|
@ -2199,9 +2239,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.
|
||||
# Lowercase the leaf first so a cased leaf (e.g. gfx120X-all) is recognised.
|
||||
|
|
@ -2432,19 +2472,19 @@ _persist_rocm_wsl_dropin() {
|
|||
fi
|
||||
}
|
||||
|
||||
# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it.
|
||||
_maybe_bootstrap_rocm_wsl() {
|
||||
[ "${OS:-}" = "wsl" ] || return 0
|
||||
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
|
||||
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
|
||||
# Leave any already-usable GPU completely alone (NVIDIA, or working ROCm).
|
||||
if _has_usable_nvidia_gpu; then return 0; fi
|
||||
# "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the
|
||||
# generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and
|
||||
# would skip this bootstrap while the real GPU is still unusable. awk consumes
|
||||
# all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail.
|
||||
# Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000,
|
||||
# the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so
|
||||
# rocminfo isn't SIGPIPE'd like `grep -q` under pipefail.
|
||||
_ensure_rocm_probe_env
|
||||
if command -v rocminfo >/dev/null 2>&1 && \
|
||||
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then
|
||||
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then
|
||||
# rocminfo may work only via the transient env _ensure_rocm_probe_env
|
||||
# just set, which dies with the installer. Persist the drop-in so login
|
||||
# shells (Studio, llama.cpp) inherit it -- else a reinstall over an
|
||||
|
|
@ -2454,9 +2494,12 @@ _maybe_bootstrap_rocm_wsl() {
|
|||
fi
|
||||
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
|
||||
[ -e /dev/dxg ] || return 0
|
||||
# Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match
|
||||
# the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S").
|
||||
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
|
||||
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
|
||||
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
|
||||
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
|
||||
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
|
||||
return 0
|
||||
fi
|
||||
command -v bash >/dev/null 2>&1 || return 0
|
||||
|
||||
# Fast path: already configured (librocdxg present) but launched from a
|
||||
|
|
@ -2474,7 +2517,8 @@ _maybe_bootstrap_rocm_wsl() {
|
|||
fi
|
||||
|
||||
echo ""
|
||||
substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN"
|
||||
_rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU"
|
||||
substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN"
|
||||
substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU."
|
||||
substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)"
|
||||
|
||||
|
|
@ -2852,7 +2896,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.6.9" "unsloth-zoo>=2026.6.7"
|
||||
"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.
|
||||
|
|
@ -2863,9 +2907,11 @@ if [ "$_MIGRATED" = true ]; then
|
|||
run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
|
||||
fi
|
||||
else
|
||||
# Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no
|
||||
# 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.6.9" "unsloth-zoo>=2026.6.7"
|
||||
"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)..."
|
||||
|
|
@ -2888,7 +2934,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_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--index-url "$TORCH_INDEX_URL" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
# torch was actually reinstalled from $TORCH_INDEX_URL now, so the
|
||||
# marker should record it (the preserved-torch case above must not).
|
||||
|
|
@ -3017,7 +3063,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_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
--default-index "$TORCH_INDEX_URL"
|
||||
else
|
||||
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
|
||||
# Record the ACTUAL wheel source for the torch-index marker: this
|
||||
|
|
@ -3049,18 +3095,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_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--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_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--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_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--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
|
||||
|
|
@ -3081,7 +3127,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.6.9" "unsloth-zoo>=2026.6.7"
|
||||
"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
|
||||
|
|
@ -3099,7 +3145,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.6.9" "unsloth-zoo>=2026.6.7"
|
||||
--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..."
|
||||
|
|
@ -3108,7 +3154,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
|
||||
else
|
||||
run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
|
||||
--upgrade-package unsloth -- "$PACKAGE_NAME"
|
||||
--upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-}
|
||||
fi
|
||||
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
|
||||
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
|
||||
|
|
@ -3120,7 +3166,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_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--index-url "$TORCH_INDEX_URL" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
# The repair reinstalled torch from $TORCH_INDEX_URL (the generic
|
||||
# ROCm index), not the Radeon --find-links repo, so record THAT as
|
||||
|
|
@ -3138,7 +3184,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.6.7" "unsloth>=2026.6.9" --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..."
|
||||
|
|
@ -3162,14 +3208,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_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
|
||||
--index-url "$TORCH_INDEX_URL" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
|
||||
# torch was re-landed from $TORCH_INDEX_URL, so record it even on a
|
||||
# migrated venv whose flavor was genuinely wrong (the gfx-switch case
|
||||
|
|
@ -3184,7 +3230,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_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --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_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -73,7 +73,7 @@ triton = [
|
|||
]
|
||||
|
||||
huggingfacenotorch = [
|
||||
"unsloth_zoo>=2026.6.7",
|
||||
"unsloth_zoo>=2026.7.2",
|
||||
"wheel>=0.42.0",
|
||||
"packaging",
|
||||
"numpy",
|
||||
|
|
@ -94,7 +94,7 @@ huggingfacenotorch = [
|
|||
]
|
||||
huggingface = [
|
||||
"unsloth[huggingfacenotorch]",
|
||||
"unsloth_zoo>=2026.6.7",
|
||||
"unsloth_zoo>=2026.7.2",
|
||||
"torchvision",
|
||||
"unsloth[triton]",
|
||||
]
|
||||
|
|
@ -579,7 +579,7 @@ colab-ampere-torch220 = [
|
|||
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
|
||||
]
|
||||
colab-new = [
|
||||
"unsloth_zoo>=2026.6.7",
|
||||
"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",
|
||||
|
|
|
|||
|
|
@ -3,13 +3,14 @@
|
|||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
#
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151)
|
||||
# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX
|
||||
# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT).
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime
|
||||
# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG
|
||||
# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04
|
||||
# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via
|
||||
# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies).
|
||||
# install.sh routes the detected arch to the right ROCm wheels once a runtime exists;
|
||||
# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg).
|
||||
# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by
|
||||
# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the
|
||||
# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent.
|
||||
#
|
||||
# Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with
|
||||
# production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once
|
||||
|
|
@ -34,10 +35,12 @@ set -euo pipefail
|
|||
|
||||
# ── Tunables (override via env) ──────────────────────────────────────────────
|
||||
ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install
|
||||
GFX="gfx1151"
|
||||
# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200).
|
||||
# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch.
|
||||
GFX="${UNSLOTH_WSL_GFX:-}"
|
||||
LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build
|
||||
# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test.
|
||||
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/"
|
||||
# AMD's wheel index for the (optional) smoke test; resolved after arch detection.
|
||||
TORCH_INDEX=""
|
||||
# Optional torch smoke test (throwaway venv). OFF by default: install.sh installs
|
||||
# torch itself into the real venv right after, so a duplicate download is wasteful.
|
||||
SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}"
|
||||
|
|
@ -220,12 +223,12 @@ $SUDO ldconfig
|
|||
say "Persisting ROCm-on-WSL environment"
|
||||
_envfile="/etc/profile.d/unsloth-rocm-wsl.sh"
|
||||
$SUDO tee "$_envfile" >/dev/null <<EOF
|
||||
# >>> Unsloth ROCm-on-WSL (gfx1151) >>>
|
||||
# >>> Unsloth ROCm-on-WSL >>>
|
||||
export HSA_ENABLE_DXG_DETECTION=1
|
||||
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
|
||||
export PATH="${ROCM_DIR}/bin:\${PATH}"
|
||||
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}"
|
||||
# <<< Unsloth ROCm-on-WSL (gfx1151) <<<
|
||||
# <<< Unsloth ROCm-on-WSL <<<
|
||||
EOF
|
||||
# also drop into ~/.bashrc for interactive shells
|
||||
if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then
|
||||
|
|
@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}"
|
|||
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
# ── Step 5: verify the runtime enumerates the GPU ────────────────────────────
|
||||
say "Verifying rocminfo sees ${GFX}"
|
||||
say "Verifying rocminfo enumerates the GPU over DXG"
|
||||
# Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs
|
||||
# rocminfo on first match, which under `set -o pipefail` turns a successful match
|
||||
# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a
|
||||
# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass.
|
||||
# into a pipeline failure.
|
||||
_rocminfo_out="$(rocminfo 2>/dev/null || true)"
|
||||
if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then
|
||||
# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU
|
||||
# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch.
|
||||
_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)"
|
||||
if [ -z "$_detected_gfx" ]; then
|
||||
printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true
|
||||
die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
|
||||
die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
|
||||
fi
|
||||
# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under
|
||||
# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt.
|
||||
if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then
|
||||
die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'."
|
||||
fi
|
||||
GFX="${GFX:-$_detected_gfx}"
|
||||
# Display-only summary: best-effort (|| true) so head's early pipe-close under
|
||||
# `set -o pipefail` can't fail the bootstrap after verification already passed.
|
||||
printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true
|
||||
note "ROCm-on-WSL runtime is live for ${GFX}."
|
||||
|
||||
# ── Step 6 (optional): torch smoke test from the gfx1151 index ───────────────
|
||||
# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ───────
|
||||
if [ "$SMOKE_TEST" = "1" ]; then
|
||||
say "Smoke-testing PyTorch on ${GFX} (throwaway venv)"
|
||||
# Map the detected arch to AMD's repo.amd.com wheel family index.
|
||||
case "$GFX" in
|
||||
gfx1200|gfx1201) _fam="gfx120X-all" ;;
|
||||
gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;;
|
||||
*) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index
|
||||
esac
|
||||
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/"
|
||||
_venv="${HOME}/.unsloth/rocm-smoketest"
|
||||
rm -rf "$_venv"; python3 -m venv "$_venv"
|
||||
"$_venv/bin/pip" install --quiet --upgrade pip
|
||||
# gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py
|
||||
# AMD arch index is primary (torch + triton); PyPI only an extra for pure-py
|
||||
# deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch.
|
||||
"$_venv/bin/pip" install --index-url "$TORCH_INDEX" \
|
||||
--extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \
|
||||
die "torch install from ${TORCH_INDEX} failed."
|
||||
# WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib.
|
||||
_tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)"
|
||||
[ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true
|
||||
"$_venv/bin/python" - <<'PY'
|
||||
import torch
|
||||
ok = torch.cuda.is_available()
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -564,6 +564,12 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
|
|||
for n, tids in b["module_import_targets"].items():
|
||||
if tids & after_used:
|
||||
continue # resolved -> fine
|
||||
# `from __future__ import ...` is a compiler directive, not a runtime
|
||||
# binding: the name (`annotations`, ...) is never loaded, so it can never
|
||||
# "resolve" to a use. Skip it so a legitimately-added future import
|
||||
# (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged.
|
||||
if all(t.startswith("from:__future__:") for t in tids):
|
||||
continue
|
||||
newly_added = bool(tids - before_module_targets)
|
||||
was_used_before = bool(tids & before_used)
|
||||
if newly_added or was_used_before:
|
||||
|
|
@ -588,9 +594,23 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
|
|||
# package object and only *add* submodule attributes (e.g. adding
|
||||
# `import urllib.error` next to `import urllib.request`). Nothing the name
|
||||
# resolved to before is lost, so no reference is re-pointed -- skip it.
|
||||
#
|
||||
# A deliberate *relocation* is also benign and must not block: when a name
|
||||
# keeps its spelling but its import source is moved A -> B in THIS diff (the
|
||||
# old `from A import x` is removed at module level and a new `from B import x`
|
||||
# is added), the swap is intentional, not a silent re-point to a pre-existing
|
||||
# different object. This mirrors the relocation tolerance already applied to
|
||||
# TARGET-MISSING. The dangerous case -- the name now resolving to a target
|
||||
# that already existed before (shadow/clash) -- is NOT exempted.
|
||||
removed_module_targets = before_module_targets - after_module_targets
|
||||
for key, tafter in b["target_by_use"].items():
|
||||
tbefore = a["target_by_use"].get(key)
|
||||
if tbefore and tbefore != tafter and (tbefore - tafter):
|
||||
lost = tbefore - tafter
|
||||
gained = tafter - tbefore
|
||||
relocated = lost <= removed_module_targets and gained <= added_module_targets
|
||||
if relocated:
|
||||
continue
|
||||
findings.append(
|
||||
(
|
||||
"BLOCKER",
|
||||
|
|
|
|||
|
|
@ -235,6 +235,13 @@
|
|||
"min_p": 0.1,
|
||||
"repetition_penalty": 1.0
|
||||
},
|
||||
"deepseek-v4": {
|
||||
"temperature": 1.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"min_p": 0.0,
|
||||
"repetition_penalty": 1.0
|
||||
},
|
||||
"deepseek-r1": {
|
||||
"temperature": 0.6,
|
||||
"top_p": 0.95,
|
||||
|
|
@ -394,7 +401,7 @@
|
|||
"phi-4", "phi-3",
|
||||
"mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral",
|
||||
"devstral", "pixtral",
|
||||
"deepseek-r1", "deepseek-v3", "deepseek-ocr",
|
||||
"deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr",
|
||||
"glm-5", "glm-4",
|
||||
"nemotron",
|
||||
"minimax-m2.7", "minimax-m2.5", "minimax",
|
||||
|
|
|
|||
|
|
@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation.
|
|||
The default get_inference_backend() returns an InferenceOrchestrator that
|
||||
delegates to a subprocess. The original InferenceBackend runs inside the
|
||||
subprocess and can be imported directly from .inference when needed.
|
||||
|
||||
Public names are resolved lazily (PEP 562): importing this package -- or a
|
||||
dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull
|
||||
the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML
|
||||
backend and its Studio dependencies). Those load only when a public name is
|
||||
actually accessed, so standalone helpers stay unit-testable without the full
|
||||
inference stack.
|
||||
"""
|
||||
|
||||
from .orchestrator import InferenceOrchestrator, get_inference_backend
|
||||
from .llama_cpp import LlamaCppBackend
|
||||
|
||||
# Expose InferenceOrchestrator as InferenceBackend for backward compat.
|
||||
InferenceBackend = InferenceOrchestrator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
__all__ = [
|
||||
"InferenceBackend",
|
||||
|
|
@ -21,3 +24,33 @@ __all__ = [
|
|||
"get_inference_backend",
|
||||
"LlamaCppBackend",
|
||||
]
|
||||
|
||||
# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator.
|
||||
_LAZY_ATTRS = {
|
||||
"InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"),
|
||||
"InferenceBackend": ("orchestrator", "InferenceOrchestrator"),
|
||||
"get_inference_backend": ("orchestrator", "get_inference_backend"),
|
||||
"LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
try:
|
||||
submodule, attr = _LAZY_ATTRS[name]
|
||||
except KeyError:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
|
||||
from importlib import import_module
|
||||
|
||||
value = getattr(import_module(f"{__name__}.{submodule}"), attr)
|
||||
globals()[name] = value # cache so later access skips __getattr__
|
||||
return value
|
||||
|
||||
|
||||
def __dir__():
|
||||
return sorted(set(globals()) | set(__all__))
|
||||
|
||||
|
||||
if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names
|
||||
from .llama_cpp import LlamaCppBackend
|
||||
from .orchestrator import InferenceOrchestrator, get_inference_backend
|
||||
InferenceBackend = InferenceOrchestrator
|
||||
|
|
|
|||
110
studio/backend/core/inference/_vulkan_probe.py
Normal file
110
studio/backend/core/inference/_vulkan_probe.py
Normal file
|
|
@ -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 <bindir>``) so the
|
||||
Vulkan instance never lives in the long-running backend process. Loads the
|
||||
bundled ggml Vulkan backend from ``<bindir>`` and prints one
|
||||
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>`` 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())
|
||||
|
|
@ -258,6 +258,10 @@ class AnthropicStreamEmitter:
|
|||
self._open_tool_use_id: Optional[str] = None
|
||||
self._open_tool_args_sent: bool = False
|
||||
self._prev_text: str = ""
|
||||
# Net <think> minus </think> in the text emitted to the client. Tracked
|
||||
# from emitted deltas (not _prev_text, which a final bare shrink clobbers)
|
||||
# so an unclosed reasoning-only block can be balanced before close.
|
||||
self._open_think_tags: int = 0
|
||||
self._usage: dict = {}
|
||||
|
||||
def start(
|
||||
|
|
@ -317,6 +321,7 @@ class AnthropicStreamEmitter:
|
|||
"""Close any open block and emit message_delta + message_stop."""
|
||||
events = []
|
||||
if self._text_block_open or self._open_tool_call_id is not None:
|
||||
events.extend(self._close_open_think())
|
||||
events.append(self._close_block())
|
||||
self._open_tool_call_id = None
|
||||
self._open_tool_use_id = None
|
||||
|
|
@ -344,12 +349,33 @@ class AnthropicStreamEmitter:
|
|||
)
|
||||
return events
|
||||
|
||||
def _close_open_think(self) -> list[str]:
|
||||
"""Emit a ``</think>`` delta when the streamed text left a ``<think>``
|
||||
open. This emitter diffs cumulative snapshots and drops the generator's
|
||||
final bare shrink, so a reasoning-only reply would otherwise end on an
|
||||
unclosed tag. Mirrors the chat route's reasoning extractor, which closes
|
||||
the block on finish; balances the block before it is closed."""
|
||||
if not self._text_block_open or self._open_think_tags <= 0:
|
||||
return []
|
||||
self._open_think_tags = 0
|
||||
return [
|
||||
build_anthropic_sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": self.block_index,
|
||||
"delta": {"type": "text_delta", "text": "</think>"},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
def _handle_content(self, event: dict) -> list[str]:
|
||||
cumulative = event.get("text", "")
|
||||
new_text = cumulative[len(self._prev_text) :]
|
||||
self._prev_text = cumulative
|
||||
if not new_text:
|
||||
return []
|
||||
self._open_think_tags += new_text.count("<think>") - new_text.count("</think>")
|
||||
if not self._text_block_open:
|
||||
events = self._open_text_block()
|
||||
else:
|
||||
|
|
@ -374,6 +400,7 @@ class AnthropicStreamEmitter:
|
|||
|
||||
events = []
|
||||
if self._text_block_open:
|
||||
events.extend(self._close_open_think())
|
||||
events.append(self._close_block())
|
||||
# Defensive: close a stale open tool_use block before starting another.
|
||||
elif self._open_tool_call_id is not None:
|
||||
|
|
@ -452,6 +479,7 @@ class AnthropicStreamEmitter:
|
|||
events.extend(self._open_text_block())
|
||||
# Reset text tracking for the next synthesis turn
|
||||
self._prev_text = ""
|
||||
self._open_think_tags = 0
|
||||
return events
|
||||
|
||||
def _open_text_block(self) -> list[str]:
|
||||
|
|
|
|||
109
studio/backend/core/inference/chat_eos.py
Normal file
109
studio/backend/core/inference/chat_eos.py
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Resolve a chat model's assistant-turn-end stop tokens.
|
||||
|
||||
Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships
|
||||
config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its
|
||||
small chat variants ship no generation_config), so generation runs past the turn
|
||||
and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns.
|
||||
|
||||
Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it
|
||||
actually uses to end a turn), not raw vocab membership: a base/coder model can
|
||||
carry ChatML control tokens in a shared vocab without using them, and a loader
|
||||
may have synced ``eos_token`` to the document terminator. Dependency-light (no
|
||||
torch / unsloth) so it is unit-testable without the full inference stack.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
# Canonical assistant-turn-end markers per chat family.
|
||||
_CHAT_TURN_END_TOKENS = (
|
||||
"<|im_end|>", # ChatML: Qwen, Yi
|
||||
"<|eot_id|>", # Llama 3.x
|
||||
"<|eom_id|>", # Llama 3.x tool turns
|
||||
"<end_of_turn>", # Gemma
|
||||
"<turn|>", # Gemma-4
|
||||
"<|end|>", # Phi
|
||||
"<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's)
|
||||
)
|
||||
# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has
|
||||
# its own streamer, so its eos is left untouched.
|
||||
_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>")
|
||||
|
||||
|
||||
def _eos_id_set(eos_token_id) -> set:
|
||||
if isinstance(eos_token_id, (list, tuple)):
|
||||
return {int(t) for t in eos_token_id if t is not None}
|
||||
if eos_token_id is not None:
|
||||
return {int(eos_token_id)}
|
||||
return set()
|
||||
|
||||
|
||||
def _collect_template_text(chat_template) -> str:
|
||||
"""Flatten a tokenizer ``chat_template`` into one scannable string.
|
||||
|
||||
Usually the template is a single jinja string, but multi-variant models
|
||||
(e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a
|
||||
``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list
|
||||
of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case
|
||||
would skip turn-end detection for those valid models, so gather every string
|
||||
leaf (variant names are harmless: they never contain the markers).
|
||||
"""
|
||||
if isinstance(chat_template, str):
|
||||
return chat_template
|
||||
if isinstance(chat_template, dict):
|
||||
values = chat_template.values()
|
||||
elif isinstance(chat_template, (list, tuple)):
|
||||
values = chat_template
|
||||
else:
|
||||
return ""
|
||||
parts = [_collect_template_text(v) for v in values]
|
||||
return "\n".join(p for p in parts if p)
|
||||
|
||||
|
||||
def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list:
|
||||
"""eos of ``id_tokenizer`` plus any canonical turn-end marker the
|
||||
``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` --
|
||||
the tokenizer generation actually uses.
|
||||
|
||||
Pass the same tokenizer for both at load time. After a mapped ``get_chat_template``
|
||||
pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective
|
||||
template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped
|
||||
template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab
|
||||
folds the turn-end token onto the doc-eos id, and generate_stream re-reads the
|
||||
original tokenizer, so resolving ids on the mapped tokenizer would store the wrong
|
||||
(doc-eos) id and let generation run past the real turn marker."""
|
||||
ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None))
|
||||
template = _collect_template_text(getattr(template_tokenizer, "chat_template", None))
|
||||
if not template or any(h in template for h in _HARMONY_MARKERS):
|
||||
return sorted(ids)
|
||||
unk = getattr(id_tokenizer, "unk_token_id", None)
|
||||
for marker in _CHAT_TURN_END_TOKENS:
|
||||
if marker in template:
|
||||
try:
|
||||
tid = id_tokenizer.convert_tokens_to_ids(marker)
|
||||
except Exception:
|
||||
tid = None
|
||||
if tid is not None and tid != unk and int(tid) >= 0:
|
||||
ids.add(int(tid))
|
||||
return sorted(ids)
|
||||
|
||||
|
||||
def resolve_chat_turn_end_eos_ids(tokenizer) -> list:
|
||||
"""tokenizer.eos plus any canonical turn-end marker the model's chat_template
|
||||
actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended
|
||||
to be resolved once at load. Returns eos unchanged for harmony templates."""
|
||||
return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer)
|
||||
|
||||
|
||||
def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]:
|
||||
"""Merged eos_token_id list, or None if ``current_eos`` already covers every
|
||||
resolved turn-end id. Used to repair a model's generation_config at load so
|
||||
every ``.generate()`` path (vision, tool loops) stops at the turn boundary."""
|
||||
if not turn_end_ids:
|
||||
return None
|
||||
current_set = _eos_id_set(current_eos)
|
||||
if set(turn_end_ids) <= current_set:
|
||||
return None
|
||||
return sorted(current_set | set(turn_end_ids))
|
||||
|
|
@ -3,11 +3,96 @@
|
|||
|
||||
"""
|
||||
Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg
|
||||
fallback for templates that reject reasoning/tools args.
|
||||
fallback for templates that reject reasoning/tools args, plus the shared
|
||||
native-chat-template fallback used by the transformers and MLX backends.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
_THINK_OPEN = "<think>"
|
||||
_THINK_CLOSE = "</think>"
|
||||
|
||||
|
||||
def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str:
|
||||
"""Return the trailing open ``<think>`` prefill of a rendered prompt.
|
||||
|
||||
Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation
|
||||
prompt with ``<think>\\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 ``<think>``/``</think>`` 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. ``"<think>\\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 ``<think>\\n\\n</think>``.
|
||||
|
||||
``special_tokens`` is the tokenizer's special-token list. If ``</think>``
|
||||
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__)
|
||||
|
||||
|
||||
def _normalize_tool_call_arguments(messages: list) -> list:
|
||||
"""Coerce each assistant ``tool_calls[].function.arguments`` from a JSON
|
||||
string to a dict.
|
||||
|
||||
The OpenAI wire format carries ``arguments`` as a JSON string, but some chat
|
||||
templates (e.g. the stricter Qwen tool templates shipped with mlx-community
|
||||
checkpoints) iterate ``arguments.items()`` and raise
|
||||
``TypeError: Can only get item pairs from a mapping.`` on the string form
|
||||
when a prior tool call is re-rendered on the next turn. A dict works on both
|
||||
strict and lenient templates, so parse the string; leave non-JSON or non-dict
|
||||
values untouched. Returns the original list unchanged when nothing needed
|
||||
coercing (no copy)."""
|
||||
mutated = False
|
||||
out: list = []
|
||||
for msg in messages:
|
||||
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
|
||||
if not tool_calls:
|
||||
out.append(msg)
|
||||
continue
|
||||
new_calls = []
|
||||
msg_changed = False
|
||||
for call in tool_calls:
|
||||
fn = call.get("function") if isinstance(call, dict) else None
|
||||
args = fn.get("arguments") if isinstance(fn, dict) else None
|
||||
if isinstance(args, str):
|
||||
try:
|
||||
parsed = json.loads(args)
|
||||
except (ValueError, TypeError):
|
||||
parsed = None
|
||||
if isinstance(parsed, dict):
|
||||
call = {**call, "function": {**fn, "arguments": parsed}}
|
||||
msg_changed = True
|
||||
new_calls.append(call)
|
||||
if msg_changed:
|
||||
out.append({**msg, "tool_calls": new_calls})
|
||||
mutated = True
|
||||
else:
|
||||
out.append(msg)
|
||||
return out if mutated else messages
|
||||
|
||||
|
||||
def apply_chat_template_for_generation(
|
||||
tokenizer,
|
||||
|
|
@ -38,21 +123,209 @@ def apply_chat_template_for_generation(
|
|||
attempts.append(dict(reasoning_kwargs))
|
||||
attempts.append({})
|
||||
|
||||
last_exc: Optional[Exception] = None
|
||||
for kwargs in attempts:
|
||||
def _render(msgs: list) -> str:
|
||||
last_exc: Optional[Exception] = None
|
||||
for kwargs in attempts:
|
||||
try:
|
||||
return tokenizer.apply_chat_template(
|
||||
msgs,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kwargs,
|
||||
)
|
||||
except TypeError as e:
|
||||
last_exc = e
|
||||
continue
|
||||
except Exception as e:
|
||||
last_exc = e
|
||||
break
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
|
||||
|
||||
try:
|
||||
return _render(messages)
|
||||
except Exception:
|
||||
# Strict tool templates reject the JSON-string ``arguments`` form via
|
||||
# TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
|
||||
# Original messages render first, so working templates stay byte-identical.
|
||||
normalized = _normalize_tool_call_arguments(messages)
|
||||
if normalized is messages:
|
||||
raise
|
||||
return _render(normalized)
|
||||
|
||||
|
||||
def render_native_template(
|
||||
*,
|
||||
model_info: dict,
|
||||
active_model_name: Optional[str],
|
||||
messages: list,
|
||||
tools: list,
|
||||
enable_thinking: Optional[bool] = None,
|
||||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
apply_fn = None,
|
||||
hf_token: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""Render ``messages`` + ``tools`` with the model's NATIVE chat template.
|
||||
|
||||
Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit
|
||||
the ``tools`` schema, so a tool-calling turn silently stops advertising tools.
|
||||
The native template ships in the model repo and carries the family's
|
||||
tool-calling syntax. It is loaded straight from the repo (bypassing any
|
||||
override on the live tokenizer) and cached on ``model_info``. Returns the
|
||||
rendered prompt only if the native template actually emits the tools (render
|
||||
differs with vs without tools); otherwise ``None``.
|
||||
|
||||
``hf_token`` is the token the model was loaded with -- passed to the repo load
|
||||
so a gated/private model's native template can still be fetched (otherwise the
|
||||
fallback fails silently and keeps the override prompt that dropped tools).
|
||||
|
||||
``trust_remote_code`` is sourced from ``model_info`` (the value the model was
|
||||
actually loaded with) rather than a call-site argument, so the native-template
|
||||
reload uses exactly the consent already granted at load. A custom-code tokenizer
|
||||
repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is
|
||||
passed, so without this the fallback fails silently and keeps the tool-dropping
|
||||
prompt for a model the user already consented to run remote code for. For a LoRA
|
||||
adapter the reload targets the base model, whose remote code was gated and loaded
|
||||
under the same stored flag, so re-passing it executes no unconsented code.
|
||||
"""
|
||||
# ``apply_fn`` lets a backend inject its own render; defaults to the module helper.
|
||||
if apply_fn is None:
|
||||
apply_fn = apply_chat_template_for_generation
|
||||
native_tpl = model_info.get("native_chat_template")
|
||||
if native_tpl is None:
|
||||
# A LoRA adapter's native template lives on the base model, not the adapter id.
|
||||
template_source = model_info.get("base_model") or active_model_name
|
||||
# Re-use the load-time trust_remote_code so a custom-code tokenizer repo can
|
||||
# instantiate its class (the stored flag already covers template_source).
|
||||
trust_remote_code = bool(model_info.get("trust_remote_code", False))
|
||||
try:
|
||||
return tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kwargs,
|
||||
from transformers import AutoTokenizer
|
||||
nt = AutoTokenizer.from_pretrained(
|
||||
template_source,
|
||||
token = hf_token if hf_token and hf_token.strip() else None,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
except TypeError as e:
|
||||
last_exc = e
|
||||
continue
|
||||
except Exception as e:
|
||||
last_exc = e
|
||||
break
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
|
||||
native_tpl = nt.chat_template or False
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Could not load native chat template for '%s': %s",
|
||||
template_source,
|
||||
exc,
|
||||
)
|
||||
# A failed fetch is not "no template": leave the sentinel unset so the next
|
||||
# call retries (caching False would pin the tool-dropping override).
|
||||
return None
|
||||
model_info["native_chat_template"] = native_tpl
|
||||
if not native_tpl:
|
||||
return None
|
||||
|
||||
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
|
||||
if tokenizer is None:
|
||||
return None
|
||||
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
|
||||
# Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the
|
||||
# generation lock) races concurrent requests.
|
||||
try:
|
||||
render_tokenizer = copy.copy(tokenizer)
|
||||
render_tokenizer.chat_template = native_tpl
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Could not clone tokenizer for native-template render of '%s': %s",
|
||||
active_model_name,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
try:
|
||||
with_tools = apply_fn(
|
||||
render_tokenizer,
|
||||
messages,
|
||||
tools = tools,
|
||||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
)
|
||||
no_tools = apply_fn(
|
||||
render_tokenizer,
|
||||
messages,
|
||||
tools = None,
|
||||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Native-template tool render failed for '%s': %s",
|
||||
active_model_name,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
return with_tools if with_tools != no_tools else None
|
||||
|
||||
|
||||
def render_with_native_template_fallback(
|
||||
*,
|
||||
formatted_prompt: str,
|
||||
tokenizer,
|
||||
model_info: dict,
|
||||
active_model_name: Optional[str],
|
||||
messages: list,
|
||||
tools: Optional[list],
|
||||
enable_thinking: Optional[bool] = None,
|
||||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
apply_fn = None,
|
||||
hf_token: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Return ``formatted_prompt``, swapping in a native-template render when an
|
||||
override template dropped the ``tools`` schema.
|
||||
|
||||
If ``tools`` were requested but the live render is identical with and without
|
||||
them (detected by comparison, robust against tool names in the system prompt),
|
||||
re-render with the model's native template. Shared by the transformers and MLX
|
||||
backends so both advertise tools consistently. ``hf_token`` is forwarded so a
|
||||
gated/private model's native template can still be fetched."""
|
||||
if not tools:
|
||||
return formatted_prompt
|
||||
if apply_fn is None:
|
||||
apply_fn = apply_chat_template_for_generation
|
||||
# Probe whether the live template dropped the schema. A tools-requiring template
|
||||
# can raise here; on any error keep the valid tools prompt rather than lose it.
|
||||
try:
|
||||
probe_no_tools = apply_fn(
|
||||
tokenizer,
|
||||
messages,
|
||||
tools = None,
|
||||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"No-tools probe failed for '%s'; keeping the existing tools prompt: %s",
|
||||
active_model_name,
|
||||
exc,
|
||||
)
|
||||
return formatted_prompt
|
||||
if formatted_prompt != probe_no_tools:
|
||||
return formatted_prompt # template already emits the tools schema
|
||||
native_prompt = render_native_template(
|
||||
model_info = model_info,
|
||||
active_model_name = active_model_name,
|
||||
messages = messages,
|
||||
tools = tools,
|
||||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
apply_fn = apply_fn,
|
||||
hf_token = hf_token,
|
||||
)
|
||||
if native_prompt:
|
||||
logger.info(
|
||||
"Override template for '%s' dropped tool schemas; using the model's "
|
||||
"native template for this tool-calling turn.",
|
||||
active_model_name,
|
||||
)
|
||||
return native_prompt
|
||||
return formatted_prompt
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import utils.hardware.hardware as hw
|
|||
DEFAULT_MODELS_GGUF = [
|
||||
"unsloth/Qwen3.6-27B-MTP-GGUF",
|
||||
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
|
||||
"unsloth/DeepSeek-V4-Flash-GGUF",
|
||||
"unsloth/gemma-4-E2B-it-GGUF",
|
||||
"unsloth/gemma-4-E4B-it-GGUF",
|
||||
"unsloth/gemma-4-31B-it-GGUF",
|
||||
|
|
@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [
|
|||
DEFAULT_MODELS_STANDARD = [
|
||||
"unsloth/Qwen3.6-27B-MTP-GGUF",
|
||||
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
|
||||
"unsloth/DeepSeek-V4-Flash-GGUF",
|
||||
"unsloth/gemma-4-E2B-it-GGUF",
|
||||
"unsloth/gemma-4-E4B-it-GGUF",
|
||||
"unsloth/gemma-4-31B-it-GGUF",
|
||||
|
|
|
|||
|
|
@ -27,6 +27,11 @@ from utils.hardware import (
|
|||
from core.inference.audio_codecs import AudioCodecManager
|
||||
from core.inference.runtime_context import runtime_context_length
|
||||
from core.inference.message_content import content_to_text
|
||||
from core.inference.chat_eos import (
|
||||
chat_eos_repair,
|
||||
resolve_chat_turn_end_eos_ids_using,
|
||||
)
|
||||
from core.inference.presence_penalty import _make_presence_penalty_processor
|
||||
from io import StringIO
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
|
|
@ -210,6 +215,50 @@ class InferenceBackend:
|
|||
# API uses -1 to disable top-k; transformers uses 0.
|
||||
return 0 if top_k < 0 else top_k
|
||||
|
||||
def _resolve_chat_eos(self, model_name: str) -> None:
|
||||
"""Resolve this chat model's assistant-turn-end stop tokens once at load,
|
||||
cache them in model_info, and repair generation_config so every
|
||||
``.generate()`` path stops at the turn boundary.
|
||||
|
||||
Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with
|
||||
``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no
|
||||
``generation_config.json``, so paths that read ``generation_config`` (the
|
||||
vision path, tool loops) run past the turn and loop. Turn-end markers are
|
||||
derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids),
|
||||
so base/coder models and harmony templates are left untouched.
|
||||
"""
|
||||
info = self.models.get(model_name) or {}
|
||||
model = info.get("model")
|
||||
container = info.get("tokenizer")
|
||||
tokenizer = getattr(container, "tokenizer", container) # unwrap processors
|
||||
if model is None or tokenizer is None:
|
||||
return
|
||||
# Vision models carry the chat_template on the processor, not the inner
|
||||
# tokenizer. Read markers from whichever has one, but resolve ids on the
|
||||
# generation tokenizer, else the vision path misses the turn-end token.
|
||||
template_source = container if getattr(container, "chat_template", None) else tokenizer
|
||||
try:
|
||||
turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
|
||||
except Exception as e: # never block a load on eos resolution
|
||||
logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
|
||||
return
|
||||
info["chat_turn_end_eos_ids"] = turn_end_ids
|
||||
|
||||
gen = getattr(model, "generation_config", None)
|
||||
if gen is None:
|
||||
return
|
||||
repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids)
|
||||
if repaired is None:
|
||||
return
|
||||
previous = gen.eos_token_id
|
||||
gen.eos_token_id = repaired
|
||||
logger.info(
|
||||
"Repaired generation_config.eos_token_id for %s: %s -> %s",
|
||||
model_name,
|
||||
previous,
|
||||
repaired,
|
||||
)
|
||||
|
||||
def load_model(
|
||||
self,
|
||||
config: ModelConfig,
|
||||
|
|
@ -221,6 +270,9 @@ class InferenceBackend:
|
|||
gpu_ids: Optional[list[int]] = None,
|
||||
) -> bool:
|
||||
"""Load any model: base, LoRA adapter, text, or vision."""
|
||||
# Keep the token so the native-template fallback can fetch a
|
||||
# gated model's repo template later during generation.
|
||||
self._hf_token = hf_token
|
||||
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
|
||||
if max_seq_length <= 0:
|
||||
max_seq_length = 2048
|
||||
|
|
@ -231,6 +283,8 @@ class InferenceBackend:
|
|||
# Already loaded?
|
||||
if model_name in self.models and self.models[model_name].get("model"):
|
||||
logger.info(f"Model {model_name} already loaded")
|
||||
if hf_token:
|
||||
self.models[model_name]["hf_token"] = hf_token
|
||||
self.active_model_name = model_name
|
||||
return True
|
||||
|
||||
|
|
@ -246,6 +300,14 @@ class InferenceBackend:
|
|||
)
|
||||
|
||||
self.models[model_name] = {
|
||||
# Per-model token: the native-template fallback must use the
|
||||
# token this model was loaded with, not whichever loaded last.
|
||||
"hf_token": hf_token,
|
||||
# Per-model consent: the native-template reload must re-use the
|
||||
# exact trust_remote_code this model (and a LoRA's base) was loaded
|
||||
# with, so a custom-code tokenizer repo can be re-fetched without
|
||||
# executing any code the user did not already consent to.
|
||||
"trust_remote_code": trust_remote_code,
|
||||
"is_vision": config.is_vision,
|
||||
"is_lora": config.is_lora,
|
||||
"is_audio": config.is_audio,
|
||||
|
|
@ -496,6 +558,7 @@ class InferenceBackend:
|
|||
max_seq_length,
|
||||
)
|
||||
|
||||
self._resolve_chat_eos(model_name)
|
||||
self._load_chat_template_info(model_name)
|
||||
|
||||
self.active_model_name = model_name
|
||||
|
|
@ -766,9 +829,11 @@ class InferenceBackend:
|
|||
preserve_thinking: Optional[bool] = None,
|
||||
max_tool_iterations: int = 25,
|
||||
auto_heal_tool_calls: bool = True,
|
||||
nudge_tool_calls: Optional[bool] = None,
|
||||
tool_call_timeout: int = 300,
|
||||
session_id: Optional[str] = None,
|
||||
rag_scope: Optional[dict] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
):
|
||||
"""Run an agentic tool loop on top of ``generate_chat_response``.
|
||||
|
||||
|
|
@ -802,6 +867,7 @@ class InferenceBackend:
|
|||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
|
||||
initial = list(messages)
|
||||
|
|
@ -815,6 +881,7 @@ class InferenceBackend:
|
|||
execute_tool = execute_tool,
|
||||
cancel_event = cancel_event,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
nudge_tool_calls = nudge_tool_calls,
|
||||
max_tool_iterations = max_tool_iterations,
|
||||
tool_call_timeout = tool_call_timeout,
|
||||
session_id = session_id,
|
||||
|
|
@ -837,12 +904,14 @@ class InferenceBackend:
|
|||
enable_thinking: Optional[bool] = None,
|
||||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Generate response for text or vision models (lock held by background thread).
|
||||
|
||||
``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking``
|
||||
are forwarded into ``apply_chat_template`` so templates that understand them
|
||||
(Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls.
|
||||
``presence_penalty`` matches the GGUF sampling path (0 disables it).
|
||||
"""
|
||||
yield from self._generate_chat_response_inner(
|
||||
messages = messages,
|
||||
|
|
@ -859,6 +928,7 @@ class InferenceBackend:
|
|||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
|
||||
def _generate_chat_response_inner(
|
||||
|
|
@ -878,6 +948,7 @@ class InferenceBackend:
|
|||
enable_thinking: Optional[bool] = None,
|
||||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Inner generation logic, called by generate_chat_response and
|
||||
generate_with_adapter_control.
|
||||
|
|
@ -917,6 +988,7 @@ class InferenceBackend:
|
|||
max_new_tokens,
|
||||
repetition_penalty,
|
||||
cancel_event = cancel_event,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
return
|
||||
else:
|
||||
|
|
@ -946,6 +1018,22 @@ class InferenceBackend:
|
|||
tokenizer,
|
||||
chat_template = template_name,
|
||||
)
|
||||
# The mapper installs the effective template only now, at generate
|
||||
# time, so re-resolve and UNION into the load-time cache (never
|
||||
# overwrite). get_chat_template can return a remapped tokenizer
|
||||
# (turn-end folded onto doc-eos) while generate_stream reads the
|
||||
# original, so take marker strings from the mapped template but
|
||||
# resolve their ids on the original.
|
||||
try:
|
||||
_gen_tok = model_info.get("tokenizer") or tokenizer
|
||||
refreshed = resolve_chat_turn_end_eos_ids_using(
|
||||
getattr(tokenizer, "tokenizer", tokenizer),
|
||||
getattr(_gen_tok, "tokenizer", _gen_tok),
|
||||
)
|
||||
existing = model_info.get("chat_turn_end_eos_ids") or []
|
||||
model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
|
||||
else:
|
||||
logger.info(
|
||||
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
|
||||
|
|
@ -975,6 +1063,27 @@ class InferenceBackend:
|
|||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
)
|
||||
|
||||
# If tools were requested but the (possibly overridden) template ignored
|
||||
# them, fall back to the model's native template (shared with MLX).
|
||||
from core.inference.chat_template_helpers import (
|
||||
render_with_native_template_fallback,
|
||||
)
|
||||
|
||||
formatted_prompt = render_with_native_template_fallback(
|
||||
formatted_prompt = formatted_prompt,
|
||||
tokenizer = tokenizer,
|
||||
model_info = model_info,
|
||||
active_model_name = self.active_model_name,
|
||||
messages = template_messages,
|
||||
tools = tools,
|
||||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
apply_fn = self._apply_chat_template_for_generation,
|
||||
hf_token = model_info.get("hf_token"),
|
||||
)
|
||||
|
||||
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
|
||||
except Exception as e:
|
||||
logger.error(f"Error applying chat template: {e}")
|
||||
|
|
@ -992,6 +1101,7 @@ class InferenceBackend:
|
|||
repetition_penalty,
|
||||
cancel_event = cancel_event,
|
||||
_adapter_state = _adapter_state,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
|
||||
def _generate_vision_response(
|
||||
|
|
@ -1006,6 +1116,7 @@ class InferenceBackend:
|
|||
max_new_tokens,
|
||||
repetition_penalty,
|
||||
cancel_event = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Handle vision model generation with true token-by-token streaming."""
|
||||
model_info = self.models[self.active_model_name]
|
||||
|
|
@ -1067,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 <think> 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
|
||||
|
||||
|
|
@ -1095,6 +1215,14 @@ class InferenceBackend:
|
|||
top_k = top_k,
|
||||
min_p = min_p,
|
||||
)
|
||||
# Presence penalty (GGUF parity) for VLM chat.
|
||||
_vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None
|
||||
if _vision_input_ids is not None:
|
||||
_pp = _make_presence_penalty_processor(
|
||||
presence_penalty, int(_vision_input_ids.shape[1])
|
||||
)
|
||||
if _pp is not None:
|
||||
generation_kwargs["logits_processor"] = _pp
|
||||
|
||||
err: dict[str, str] = {}
|
||||
|
||||
|
|
@ -1114,7 +1242,11 @@ class InferenceBackend:
|
|||
thread = threading.Thread(target = generate_fn)
|
||||
thread.start()
|
||||
|
||||
output = ""
|
||||
output = think_prefix
|
||||
# Emit the prefilled <think> 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
|
||||
|
|
@ -1323,11 +1455,13 @@ class InferenceBackend:
|
|||
repetition_penalty: float = 1.0,
|
||||
cancel_event = None,
|
||||
_adapter_state = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Generate a streaming text response (text models only).
|
||||
|
||||
_adapter_state: if not None, the background thread toggles adapters
|
||||
before model.generate(), under _generation_lock.
|
||||
``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it).
|
||||
"""
|
||||
if not self.active_model_name:
|
||||
yield "Error: No active model"
|
||||
|
|
@ -1346,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 <think> 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 <think> tags
|
||||
|
|
@ -1382,11 +1526,18 @@ class InferenceBackend:
|
|||
min_p = min_p,
|
||||
repetition_penalty = repetition_penalty,
|
||||
do_sample = temperature > 0,
|
||||
eos_token_id = tokenizer.eos_token_id,
|
||||
# Resolved once at load (chat_template-derived turn-end tokens).
|
||||
eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
|
||||
pad_token_id = tokenizer.eos_token_id
|
||||
if tokenizer.pad_token_id is None
|
||||
else tokenizer.pad_token_id,
|
||||
)
|
||||
# Presence penalty (GGUF parity); prompt_len excludes prompt tokens.
|
||||
_pp = _make_presence_penalty_processor(
|
||||
presence_penalty, int(inputs["input_ids"].shape[1])
|
||||
)
|
||||
if _pp is not None:
|
||||
generation_kwargs["logits_processor"] = _pp
|
||||
if cancel_event is not None:
|
||||
from transformers.generation.stopping_criteria import (
|
||||
StoppingCriteria,
|
||||
|
|
@ -1422,7 +1573,11 @@ class InferenceBackend:
|
|||
thread = threading.Thread(target = generate_fn)
|
||||
thread.start()
|
||||
|
||||
output = ""
|
||||
output = think_prefix
|
||||
# Emit the prefilled <think> 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
|
||||
|
|
|
|||
368
studio/backend/core/inference/llama_admission.py
Normal file
368
studio/backend/core/inference/llama_admission.py
Normal file
|
|
@ -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()
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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,87 @@ 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.
|
||||
|
||||
generate_step calls processors as ``fn(tokens, logits)`` with ``tokens`` the
|
||||
full running sequence; the first call is prompt-only, so latch that length
|
||||
and penalize only after it.
|
||||
"""
|
||||
state = {"prompt_len": None}
|
||||
|
||||
def _processor(tokens, logits):
|
||||
if state["prompt_len"] is None:
|
||||
# First call is prompt-only; latch its length.
|
||||
state["prompt_len"] = int(tokens.shape[0])
|
||||
return logits
|
||||
generated = tokens[state["prompt_len"] :]
|
||||
if generated.size == 0:
|
||||
return logits
|
||||
import mlx.core as mx
|
||||
|
||||
vocab = logits.shape[-1]
|
||||
# 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 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]
|
||||
return logits
|
||||
|
||||
return _processor
|
||||
|
||||
|
||||
class MLXInferenceBackend:
|
||||
def __init__(self):
|
||||
self.models = {}
|
||||
|
|
@ -49,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
|
||||
|
|
@ -57,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 = {}
|
||||
|
|
@ -101,16 +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 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}': "
|
||||
|
|
@ -129,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
|
||||
|
|
@ -143,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:
|
||||
|
|
@ -168,18 +287,25 @@ class MLXInferenceBackend:
|
|||
|
||||
self.active_model_name = model_name
|
||||
self.models[model_name] = {
|
||||
# Per-model token for the native-template fallback (matches transformers).
|
||||
"hf_token": hf_token,
|
||||
# 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,
|
||||
"processor": self._processor,
|
||||
"is_vision": is_vision,
|
||||
"is_lora": getattr(config, "is_lora", False),
|
||||
# For a LoRA adapter the native chat template lives on the base model.
|
||||
"base_model": getattr(config, "base_model", None)
|
||||
if getattr(config, "is_lora", False)
|
||||
else None,
|
||||
"is_audio": False,
|
||||
"audio_type": None,
|
||||
"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)
|
||||
|
|
@ -237,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()
|
||||
|
|
@ -264,12 +393,12 @@ 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,
|
||||
preserve_thinking = None,
|
||||
presence_penalty = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
if self._model is None:
|
||||
raise RuntimeError("No model loaded")
|
||||
|
|
@ -277,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})
|
||||
|
|
@ -294,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)
|
||||
)
|
||||
|
|
@ -317,6 +444,7 @@ class MLXInferenceBackend:
|
|||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
else:
|
||||
yield from self._generate_text(
|
||||
|
|
@ -332,6 +460,7 @@ class MLXInferenceBackend:
|
|||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
|
||||
def _generate_text(
|
||||
|
|
@ -349,12 +478,15 @@ class MLXInferenceBackend:
|
|||
enable_thinking = None,
|
||||
reasoning_effort = None,
|
||||
preserve_thinking = None,
|
||||
presence_penalty = 0.0,
|
||||
):
|
||||
from mlx_lm import stream_generate
|
||||
from mlx_lm.sample_utils import make_sampler, make_logits_processors
|
||||
|
||||
from core.inference.chat_template_helpers import (
|
||||
apply_chat_template_for_generation,
|
||||
detect_think_prefill,
|
||||
render_with_native_template_fallback,
|
||||
)
|
||||
|
||||
prompt = apply_chat_template_for_generation(
|
||||
|
|
@ -368,6 +500,34 @@ class MLXInferenceBackend:
|
|||
if prompt is None:
|
||||
raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
|
||||
|
||||
# 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,
|
||||
tokenizer = self._tokenizer,
|
||||
model_info = model_info,
|
||||
active_model_name = self.active_model_name,
|
||||
messages = messages,
|
||||
tools = tools,
|
||||
enable_thinking = enable_thinking,
|
||||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
hf_token = model_info.get("hf_token"),
|
||||
)
|
||||
|
||||
# An open <think> 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,
|
||||
|
|
@ -375,15 +535,21 @@ class MLXInferenceBackend:
|
|||
min_p = float(min_p or 0.0),
|
||||
min_tokens_to_keep = 1,
|
||||
)
|
||||
# Only build a logits processor for a non-trivial repetition penalty.
|
||||
logits_processors = None
|
||||
# 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,
|
||||
1.0,
|
||||
):
|
||||
logits_processors = make_logits_processors(
|
||||
repetition_penalty = float(repetition_penalty),
|
||||
logits_processors.extend(
|
||||
make_logits_processors(
|
||||
repetition_penalty = float(repetition_penalty),
|
||||
)
|
||||
)
|
||||
if presence_penalty:
|
||||
logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
|
||||
if not logits_processors:
|
||||
logits_processors = None
|
||||
|
||||
token_ids = []
|
||||
logger.info(
|
||||
|
|
@ -410,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
|
||||
|
|
@ -449,6 +614,7 @@ class MLXInferenceBackend:
|
|||
enable_thinking = None,
|
||||
reasoning_effort = None,
|
||||
preserve_thinking = None,
|
||||
presence_penalty = 0.0,
|
||||
):
|
||||
from mlx_vlm import stream_generate as vlm_stream
|
||||
|
||||
|
|
@ -457,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
|
||||
|
|
@ -479,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 <think> 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,
|
||||
|
|
@ -496,10 +666,23 @@ class MLXInferenceBackend:
|
|||
top_k = int(top_k or 0),
|
||||
min_p = float(min_p or 0.0),
|
||||
)
|
||||
if repetition_penalty is not None and float(repetition_penalty) not in (
|
||||
_rep_active = repetition_penalty is not None and float(repetition_penalty) not in (
|
||||
0.0,
|
||||
1.0,
|
||||
):
|
||||
)
|
||||
if presence_penalty:
|
||||
# Presence needs a custom processor: pass the full list (repetition +
|
||||
# presence) instead of the repetition_penalty shortcut so both apply.
|
||||
from mlx_lm.sample_utils import make_logits_processors
|
||||
|
||||
_vlm_processors = []
|
||||
if _rep_active:
|
||||
_vlm_processors.extend(
|
||||
make_logits_processors(repetition_penalty = float(repetition_penalty))
|
||||
)
|
||||
_vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
|
||||
vlm_kwargs["logits_processors"] = _vlm_processors
|
||||
elif _rep_active:
|
||||
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
|
||||
|
||||
with self._generation_lock:
|
||||
|
|
@ -534,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):
|
||||
|
|
|
|||
|
|
@ -45,6 +45,22 @@ _DISPATCH_STOP_TIMEOUT = 5.0
|
|||
_DISPATCH_IDLE_TIMEOUT = 30.0
|
||||
_DISPATCH_DRAIN_TIMEOUT = 5.0
|
||||
|
||||
# Max wait for a cancelled generation to release _gen_lock before unload_model
|
||||
# tears the subprocess down. Only bounds a wedged worker.
|
||||
_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:
|
||||
"""
|
||||
|
|
@ -60,7 +76,13 @@ class InferenceOrchestrator:
|
|||
self._cmd_queue: Any = None
|
||||
self._resp_queue: Any = None
|
||||
self._cancel_event: Any = None # mp.Event — set to cancel generation
|
||||
# Set for the whole unload; the worker never clears it (unlike _cancel_event),
|
||||
# so a generate queued behind the cancelled one is skipped, not run.
|
||||
self._drain_event: Any = None
|
||||
self._gen_lock = threading.Lock() # Serializes generation
|
||||
# Set during a switch so a generation winning the _gen_lock handoff bails
|
||||
# instead of starting on the outgoing model.
|
||||
self._unload_pending = False
|
||||
|
||||
# Dispatcher state for compare mode (adapter-controlled requests):
|
||||
# bypass _gen_lock, send commands directly, read from per-request
|
||||
|
|
@ -69,6 +91,12 @@ class InferenceOrchestrator:
|
|||
self._mailbox_lock = threading.Lock()
|
||||
self._dispatcher_thread: Optional[threading.Thread] = None
|
||||
self._dispatcher_stop = threading.Event()
|
||||
# Serializes dispatcher start/stop. _generate_dispatched (compare mode) bypasses
|
||||
# _gen_lock, so two concurrent compare requests can both reach _start_dispatcher;
|
||||
# without this lock both could observe no live dispatcher and each spawn one,
|
||||
# orphaning the extra thread (self._dispatcher_thread tracks only the last). The
|
||||
# orphan later steals the "unloaded" reply off resp_queue and hangs unload_model.
|
||||
self._dispatcher_lifecycle_lock = threading.Lock()
|
||||
|
||||
# Local state mirrors (updated from subprocess responses)
|
||||
self.active_model_name: Optional[str] = None
|
||||
|
|
@ -92,13 +120,11 @@ class InferenceOrchestrator:
|
|||
|
||||
@property
|
||||
def default_models(self) -> list[str]:
|
||||
# Wait up to 5s for background HF fetch
|
||||
self._top_models_ready.wait(timeout = 5)
|
||||
top_gguf = self._top_gguf_cache or []
|
||||
top_hub = self._top_hub_cache or []
|
||||
# Curated static defaults first, then HF download-ranked to backfill.
|
||||
# Send extras so the frontend keeps 4 per category after removing
|
||||
# downloaded ones.
|
||||
# Never wait for the remote Hugging Face ranking during startup. Chat's
|
||||
# first /api/models/list needs curated defaults immediately; the
|
||||
# background fetch backfills extra choices on later calls.
|
||||
result: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for m in self._static_models + top_gguf + top_hub:
|
||||
|
|
@ -159,6 +185,7 @@ class InferenceOrchestrator:
|
|||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
self._cancel_event = _CTX.Event()
|
||||
self._drain_event = _CTX.Event()
|
||||
|
||||
self._proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
|
|
@ -167,6 +194,7 @@ class InferenceOrchestrator:
|
|||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"cancel_event": self._cancel_event,
|
||||
"drain_event": self._drain_event,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
|
|
@ -228,6 +256,7 @@ class InferenceOrchestrator:
|
|||
self._cmd_queue = None
|
||||
self._resp_queue = None
|
||||
self._cancel_event = None
|
||||
self._drain_event = None
|
||||
logger.info("Inference subprocess shut down")
|
||||
|
||||
def _cleanup(self):
|
||||
|
|
@ -409,6 +438,7 @@ class InferenceOrchestrator:
|
|||
enable_thinking: Optional[bool] = None,
|
||||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> dict:
|
||||
"""Build the 'generate' command shared by the locked and dispatched paths."""
|
||||
cmd = {
|
||||
|
|
@ -423,6 +453,7 @@ class InferenceOrchestrator:
|
|||
"min_p": min_p,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"repetition_penalty": repetition_penalty,
|
||||
"presence_penalty": presence_penalty,
|
||||
}
|
||||
# Only forward template kwargs the caller set, for older worker compat.
|
||||
if use_adapter is not None:
|
||||
|
|
@ -456,12 +487,20 @@ class InferenceOrchestrator:
|
|||
cancel ack from that same source so stale events don't leak into the
|
||||
next request.
|
||||
"""
|
||||
# Latch this stream's subprocess/queue: if a wedged worker is torn down and a
|
||||
# later load spawns a fresh one, bail rather than re-block on the new queue
|
||||
# under _gen_lock (deadlock).
|
||||
initial_proc = self._proc
|
||||
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 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
|
||||
|
||||
|
|
@ -471,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":
|
||||
|
|
@ -486,40 +525,63 @@ 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
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Dispatcher — per-request mailbox routing for compare mode
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _start_dispatcher(self) -> None:
|
||||
def _start_dispatcher(self) -> bool:
|
||||
"""Start the dispatcher thread if not already running.
|
||||
|
||||
The dispatcher reads the shared resp_queue and routes responses to
|
||||
per-request mailbox queues, letting multiple adapter-controlled
|
||||
(compare) requests be in-flight without holding _gen_lock.
|
||||
"""
|
||||
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
|
||||
return
|
||||
|
||||
self._dispatcher_stop.clear()
|
||||
self._dispatcher_thread = threading.Thread(
|
||||
target = self._dispatcher_loop,
|
||||
daemon = True,
|
||||
name = "inference-dispatcher",
|
||||
)
|
||||
self._dispatcher_thread.start()
|
||||
logger.debug("Dispatcher thread started")
|
||||
The whole check-then-spawn runs under _dispatcher_lifecycle_lock so
|
||||
concurrent compare requests (which bypass _gen_lock) can't both observe
|
||||
no live dispatcher and each spawn one. Returns True only for the caller
|
||||
that actually started a new thread; False if one was already alive.
|
||||
"""
|
||||
with self._dispatcher_lifecycle_lock:
|
||||
# Refuse to start while an unload is in progress. unload_model sets
|
||||
# _unload_pending under this same lock before it stops the idle
|
||||
# dispatcher, so a start queued behind that stop observes the unload
|
||||
# here and bails. Without this a fresh dispatcher would be spawned
|
||||
# after the stop, become the resp_queue reader, and consume the
|
||||
# worker's "unloaded" reply (unroutable, so dropped) before
|
||||
# unload_model's _wait_response sees it -- hanging the unload 300s.
|
||||
if self._unload_pending:
|
||||
return False
|
||||
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
|
||||
return False
|
||||
|
||||
self._dispatcher_stop.clear()
|
||||
self._dispatcher_thread = threading.Thread(
|
||||
target = self._dispatcher_loop,
|
||||
daemon = True,
|
||||
name = "inference-dispatcher",
|
||||
)
|
||||
self._dispatcher_thread.start()
|
||||
logger.debug("Dispatcher thread started")
|
||||
return True
|
||||
|
||||
def _stop_dispatcher(self) -> None:
|
||||
"""Signal the dispatcher to stop and wait for it."""
|
||||
if self._dispatcher_thread is None:
|
||||
return
|
||||
self._dispatcher_stop.set()
|
||||
self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT)
|
||||
self._dispatcher_thread = None
|
||||
logger.debug("Dispatcher thread stopped")
|
||||
"""Signal the dispatcher to stop and wait for it.
|
||||
|
||||
Runs under _dispatcher_lifecycle_lock (paired with _start_dispatcher) so
|
||||
a stop can't interleave with a concurrent start. Callers must NOT hold
|
||||
_mailbox_lock here: this joins the dispatcher, and the dispatcher loop
|
||||
takes _mailbox_lock, so holding it would deadlock the join.
|
||||
"""
|
||||
with self._dispatcher_lifecycle_lock:
|
||||
if self._dispatcher_thread is None:
|
||||
return
|
||||
self._dispatcher_stop.set()
|
||||
self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT)
|
||||
self._dispatcher_thread = None
|
||||
logger.debug("Dispatcher thread stopped")
|
||||
|
||||
def _dispatcher_loop(self) -> None:
|
||||
"""Background loop: read resp_queue → route to mailboxes by request_id."""
|
||||
|
|
@ -581,6 +643,7 @@ class InferenceOrchestrator:
|
|||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
stats_holder: Optional[dict] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Dispatched generation — sends command without holding _gen_lock.
|
||||
|
||||
|
|
@ -589,15 +652,32 @@ 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
|
||||
# expected_model check).
|
||||
expected_model = self.active_model_name
|
||||
|
||||
# Switch in flight (unload waiting on _gen_lock). This path bypasses the lock,
|
||||
# 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 GenStreamError("Error: model is being unloaded")
|
||||
return
|
||||
|
||||
# Ensure dispatcher is running
|
||||
self._start_dispatcher()
|
||||
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
|
||||
# _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned
|
||||
# the thread, so at most one dispatcher ever exists even when two compare requests race
|
||||
# here. Derive dispatcher_preexisting from that atomic result (not a separate unlocked
|
||||
# is_alive() read): if THIS call started the dispatcher and then bails on a racing
|
||||
# unload, it must stop it again (see the unloading bail below).
|
||||
started = self._start_dispatcher()
|
||||
dispatcher_preexisting = not started
|
||||
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
|
|
@ -617,6 +697,7 @@ class InferenceOrchestrator:
|
|||
min_p = min_p,
|
||||
max_new_tokens = max_new_tokens,
|
||||
repetition_penalty = repetition_penalty,
|
||||
presence_penalty = presence_penalty,
|
||||
use_adapter = use_adapter,
|
||||
tools = tools,
|
||||
enable_thinking = enable_thinking,
|
||||
|
|
@ -624,17 +705,49 @@ class InferenceOrchestrator:
|
|||
preserve_thinking = preserve_thinking,
|
||||
)
|
||||
|
||||
# Create mailbox BEFORE sending command
|
||||
# Create the mailbox BEFORE sending, rechecking _unload_pending under
|
||||
# _mailbox_lock: an unload sets _unload_pending before _wait_dispatcher_idle
|
||||
# reads _mailboxes under the same lock, so either the idle check sees this
|
||||
# mailbox (and tears the dispatcher down) or we see the unload and bail.
|
||||
# Registering after would orphan the mailbox and hang the compare stream forever.
|
||||
mailbox: queue.Queue = queue.Queue()
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes[request_id] = mailbox
|
||||
# _unload_pending alone is not enough: an unload that ran fully since
|
||||
# _start_dispatcher clears it in its finally and stops the dispatcher, so it
|
||||
# reads False here though the dispatcher is gone and the model swapped. Also
|
||||
# bail when the active model changed or the dispatcher died: a mailbox with no
|
||||
# dispatcher to route gen_done/gen_error hangs the compare stream.
|
||||
dispatcher_alive = (
|
||||
self._dispatcher_thread is not None and self._dispatcher_thread.is_alive()
|
||||
)
|
||||
unloading = (
|
||||
self._unload_pending
|
||||
or self.active_model_name != expected_model
|
||||
or not dispatcher_alive
|
||||
)
|
||||
if not unloading:
|
||||
self._mailboxes[request_id] = mailbox
|
||||
# When bailing without a mailbox, note whether any OTHER compare request still
|
||||
# routes through the dispatcher; if none and this call started it, stop it below.
|
||||
orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes
|
||||
if unloading:
|
||||
# A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was
|
||||
# stopped, then set _unload_pending. The one we just started would otherwise
|
||||
# linger with no mailboxes, race unload_model's _wait_response for the "unloaded"
|
||||
# reply off resp_queue, and drop it as unroutable -- hanging the unload 300s. Stop
|
||||
# it here so the unload stays the sole resp_queue reader. Outside _mailbox_lock:
|
||||
# _stop_dispatcher joins the dispatcher, which itself takes that lock.
|
||||
if orphaned_dispatcher:
|
||||
self._stop_dispatcher()
|
||||
yield GenStreamError("Error: model is being unloaded")
|
||||
return
|
||||
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
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):
|
||||
|
|
@ -676,14 +789,18 @@ class InferenceOrchestrator:
|
|||
return
|
||||
logger.warning("Timed out draining mailbox after cancel")
|
||||
|
||||
def _wait_dispatcher_idle(self) -> None:
|
||||
def _wait_dispatcher_idle(self) -> bool:
|
||||
"""Wait for all dispatched requests to complete, then stop dispatcher.
|
||||
|
||||
Called by _generate_inner before the _gen_lock path so the dispatcher
|
||||
thread isn't competing for resp_queue reads.
|
||||
Returns True if the dispatcher was stopped (all mailboxes drained, or no
|
||||
dispatcher was running), and False if it was left running because compare
|
||||
requests were still active after _DISPATCH_IDLE_TIMEOUT.
|
||||
|
||||
Called before the _gen_lock path so the dispatcher thread isn't competing
|
||||
for resp_queue reads.
|
||||
"""
|
||||
if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive():
|
||||
return
|
||||
return True
|
||||
|
||||
# Wait for all mailboxes to be emptied (dispatched requests complete)
|
||||
deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT
|
||||
|
|
@ -704,8 +821,62 @@ class InferenceOrchestrator:
|
|||
"leaving dispatcher running for compare requests",
|
||||
len(self._mailboxes),
|
||||
)
|
||||
else:
|
||||
self._stop_dispatcher()
|
||||
return False
|
||||
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
|
||||
|
|
@ -722,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.
|
||||
|
||||
|
|
@ -747,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,
|
||||
|
|
@ -772,6 +950,19 @@ class InferenceOrchestrator:
|
|||
)
|
||||
|
||||
for attempt in range(2):
|
||||
# Stop-loading (/unload -> cancel_load) aborts a load by discarding this
|
||||
# model's loading marker. cancel_load only kills a live child; if the cancel
|
||||
# lands before any child exists (GPU placement, or between retries) there is
|
||||
# nothing to kill, and without this check the loop would spawn a worker and
|
||||
# load the model after /unload reported it unloaded. Observe removal and stop.
|
||||
if model_name not in self.loading_models:
|
||||
logger.info(
|
||||
"Load for '%s' was cancelled before spawn; not starting a worker",
|
||||
model_name,
|
||||
)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
return False
|
||||
logger.info(
|
||||
"Spawning fresh inference subprocess for '%s' "
|
||||
"(transformers %s.x, attempt %d/2%s)",
|
||||
|
|
@ -783,6 +974,22 @@ class InferenceOrchestrator:
|
|||
sub_config["disable_xet"] = disable_xet
|
||||
self._spawn_subprocess(sub_config)
|
||||
|
||||
# A cancel can land after the pre-spawn recheck but while _spawn_subprocess
|
||||
# is still creating the queues/process. cancel_load runs off the lifecycle
|
||||
# gate, so its _shutdown_subprocess can see _proc still None and no-op,
|
||||
# orphaning this fresh worker; the load would then wait for "loaded" and
|
||||
# publish a model /unload reported unloaded, over a live subprocess nothing
|
||||
# reaps. Recheck now the child exists and tear it down before publishing.
|
||||
if model_name not in self.loading_models:
|
||||
logger.info(
|
||||
"Load for '%s' was cancelled during spawn; tearing the worker down",
|
||||
model_name,
|
||||
)
|
||||
self._shutdown_subprocess(timeout = 5)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
return False
|
||||
|
||||
try:
|
||||
resp = self._wait_response("loaded")
|
||||
except DownloadStallError:
|
||||
|
|
@ -803,8 +1010,31 @@ class InferenceOrchestrator:
|
|||
)
|
||||
|
||||
if resp.get("success"):
|
||||
# A cancel can land while we were parked in _wait_response above.
|
||||
# cancel_load (off the lifecycle gate) discards this model's loading
|
||||
# marker BEFORE its teardown, so a Stop-loading that fired after the
|
||||
# worker queued "loaded" (which we can still consume during cancel_load's
|
||||
# shutdown window) shows up here only as the marker's removal. Without
|
||||
# this recheck we would publish active_model_name/models for a model
|
||||
# /unload reported cancelled, over a subprocess cancel_load just killed;
|
||||
# its post-teardown re-clear cannot undo a publish that lands after it
|
||||
# returns. Observe the removal and abort; cancel_load owns teardown.
|
||||
if model_name not in self.loading_models:
|
||||
logger.info(
|
||||
"Load for '%s' was cancelled while waiting for 'loaded'; "
|
||||
"not publishing the cancelled model",
|
||||
model_name,
|
||||
)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
return False
|
||||
model_info = resp.get("model_info", {})
|
||||
self.active_model_name = model_info.get("identifier", model_name)
|
||||
# A load always spawns a fresh subprocess holding only this model, so
|
||||
# mirror that. A lingering stale name would pass unload_model's "not in
|
||||
# self.models" guard, and the worker's absent-name fallback would unload
|
||||
# its *active* model, not the already-gone one.
|
||||
self.models = {}
|
||||
self.models[self.active_model_name] = {
|
||||
"is_vision": model_info.get("is_vision", False),
|
||||
"is_lora": model_info.get("is_lora", False),
|
||||
|
|
@ -837,17 +1067,65 @@ class InferenceOrchestrator:
|
|||
self.models.clear()
|
||||
raise
|
||||
|
||||
def unload_model(self, model_name: str) -> bool:
|
||||
"""Unload a model from the subprocess."""
|
||||
if model_name in self.loading_models:
|
||||
logger.info(
|
||||
"Cancelling in-flight load for model '%s' by terminating subprocess",
|
||||
def cancel_load(self, model_name: str) -> bool:
|
||||
"""Abort an in-flight load by terminating its subprocess.
|
||||
|
||||
Returns True if a load for ``model_name`` (matched case-insensitively) was
|
||||
cancelled, False if nothing was loading under that name. This only tears the
|
||||
loading subprocess down -- it sends no command to a worker -- so, unlike the
|
||||
rest of ``unload_model``, it is safe to run WITHOUT the inference lifecycle
|
||||
gate. ``/unload`` calls it off-gate so the "stop loading" button can interrupt
|
||||
a safetensors load that holds the gate for its whole (multi-minute) duration;
|
||||
a gated cancel could never preempt that load.
|
||||
"""
|
||||
target = model_name
|
||||
if target not in self.loading_models:
|
||||
target = next(
|
||||
(m for m in self.loading_models if m.lower() == model_name.lower()),
|
||||
model_name,
|
||||
)
|
||||
self._shutdown_subprocess(timeout = 0.5)
|
||||
self.loading_models.discard(model_name)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
if target not in self.loading_models:
|
||||
return False
|
||||
logger.info(
|
||||
"Cancelling in-flight load for model '%s' by terminating subprocess",
|
||||
target,
|
||||
)
|
||||
# Discard the loading marker (and clear local state) BEFORE the teardown, not
|
||||
# after. cancel_load runs off the lifecycle gate, alongside a load_model that
|
||||
# rechecks this marker before each spawn. But _shutdown_subprocess can block (~1s
|
||||
# tearing a live child down and joining the dispatcher), so clearing only after
|
||||
# leaves a window where load_model reads the marker still set, passes its pre-spawn
|
||||
# recheck, and loads the model after /unload reported it cancelled. Clear first.
|
||||
self.loading_models.discard(target)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
self._shutdown_subprocess(timeout = 0.5)
|
||||
# Clear the local mirrors again AFTER the teardown. A racing off-gate load_model
|
||||
# may still be parked in _wait_response("loaded"): its worker already queued a
|
||||
# "loaded" reply, so during the shutdown window above (the 0.5s settle before the
|
||||
# response queue is drained and nulled) that thread can consume it and repopulate
|
||||
# active_model_name/models, undoing the pre-teardown clear. _shutdown_subprocess
|
||||
# nulls the queue but not the mirrors, so without this second clear /unload reports
|
||||
# success while the backend still advertises a killed model. The nulled queue lets
|
||||
# no further "loaded" through, so re-clearing here wipes any repopulation.
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
return True
|
||||
|
||||
def unload_model(self, model_name: str) -> bool:
|
||||
"""Unload a model from the subprocess."""
|
||||
# active_model_name can differ in case from the client's raw /unload name (the
|
||||
# load path canonicalizes casing). Match case-insensitively and use the canonical
|
||||
# spelling so the guard, unload command, and cleanup below hit the loaded model.
|
||||
if (
|
||||
self.active_model_name is not None
|
||||
and model_name != self.active_model_name
|
||||
and model_name.lower() == self.active_model_name.lower()
|
||||
):
|
||||
model_name = self.active_model_name
|
||||
# In-flight load: tear its subprocess down (shared loading-cancel logic; no
|
||||
# worker command sent).
|
||||
if self.cancel_load(model_name):
|
||||
return True
|
||||
|
||||
if not self._ensure_subprocess_alive():
|
||||
|
|
@ -857,30 +1135,93 @@ class InferenceOrchestrator:
|
|||
self.active_model_name = None
|
||||
return True
|
||||
|
||||
try:
|
||||
self._send_cmd(
|
||||
{
|
||||
"type": "unload",
|
||||
"model_name": model_name,
|
||||
}
|
||||
)
|
||||
resp = self._wait_response("unloaded")
|
||||
|
||||
# Update local state
|
||||
# Nothing loaded under this name: don't unload a stale model. The worker falls
|
||||
# back to unloading its *active* model when the name is absent, so a stale unload
|
||||
# (lost a race to a concurrent load) would hit the wrong one.
|
||||
if model_name != self.active_model_name and model_name not in self.models:
|
||||
self.models.pop(model_name, None)
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
|
||||
logger.info("Model '%s' unloaded from subprocess", model_name)
|
||||
return True
|
||||
|
||||
except Exception as exc:
|
||||
logger.error("Error unloading model '%s': %s", model_name, exc)
|
||||
# Clear local state anyway
|
||||
self.models.pop(model_name, None)
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
return False
|
||||
# The subprocess runs commands sequentially, so a bare unload queues behind a
|
||||
# running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker
|
||||
# polls each token), then take _gen_lock as sole resp_queue reader (like GGUF).
|
||||
#
|
||||
# Set _unload_pending under _dispatcher_lifecycle_lock so it is ordered ahead of
|
||||
# the dispatcher stop that _wait_dispatcher_idle runs under the same lock: a
|
||||
# compare request's _start_dispatcher queued behind that stop then observes the
|
||||
# unload and refuses to spawn a fresh dispatcher that would eat the "unloaded"
|
||||
# reply off resp_queue. This is a standalone acquisition (no _gen_lock held yet),
|
||||
# so it keeps the _gen_lock -> _dispatcher_lifecycle_lock order and can't deadlock.
|
||||
with self._dispatcher_lifecycle_lock:
|
||||
self._unload_pending = True
|
||||
# Cancelling only the running generation isn't enough: the worker clears
|
||||
# cancel_event at each generate start, so a queued one would clear it and run the
|
||||
# outgoing model to completion. drain_event, never cleared, makes any generate
|
||||
# dequeued during the unload skip.
|
||||
if self._drain_event is not None:
|
||||
self._drain_event.set()
|
||||
try:
|
||||
self._cancel_generation()
|
||||
acquired = self._gen_lock.acquire(timeout = _UNLOAD_GEN_LOCK_TIMEOUT)
|
||||
if not acquired:
|
||||
# Wedged worker: tear the subprocess down to free the GPU (next load respawns).
|
||||
logger.warning(
|
||||
"Unload: generation did not yield %.1fs after cancel; "
|
||||
"shutting the inference subprocess down to free the model",
|
||||
_UNLOAD_GEN_LOCK_TIMEOUT,
|
||||
)
|
||||
self._shutdown_subprocess(timeout = 5)
|
||||
self.models.pop(model_name, None)
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
return True
|
||||
|
||||
try:
|
||||
# Stop the compare-mode dispatcher so it can't consume the "unloaded" reply
|
||||
# off resp_queue before we do. A dispatched generation bypasses _gen_lock, so
|
||||
# a wedged one slips past the acquire above; if the dispatcher is still active
|
||||
# it owns resp_queue and the queued unload hangs _wait_response behind the
|
||||
# stuck generate. Mirror the wedged locked path: tear the subprocess down.
|
||||
if not self._wait_dispatcher_idle():
|
||||
logger.warning(
|
||||
"Unload: compare-mode dispatcher still active after idle "
|
||||
"wait; shutting the inference subprocess down to free the model"
|
||||
)
|
||||
self._shutdown_subprocess(timeout = 5)
|
||||
self.models.pop(model_name, None)
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
return True
|
||||
# Drop stale tokens so they can't be read as the unload reply.
|
||||
self._drain_queue()
|
||||
self._send_cmd(
|
||||
{
|
||||
"type": "unload",
|
||||
"model_name": model_name,
|
||||
}
|
||||
)
|
||||
self._wait_response("unloaded")
|
||||
|
||||
self.models.pop(model_name, None)
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
|
||||
logger.info("Model '%s' unloaded from subprocess", model_name)
|
||||
return True
|
||||
|
||||
except Exception as exc:
|
||||
logger.error("Error unloading model '%s': %s", model_name, exc)
|
||||
# Clear local state anyway
|
||||
self.models.pop(model_name, None)
|
||||
if self.active_model_name == model_name:
|
||||
self.active_model_name = None
|
||||
return False
|
||||
finally:
|
||||
self._gen_lock.release()
|
||||
finally:
|
||||
self._unload_pending = False
|
||||
if self._drain_event is not None:
|
||||
self._drain_event.clear()
|
||||
|
||||
def generate_chat_response(
|
||||
self,
|
||||
|
|
@ -899,6 +1240,7 @@ class InferenceOrchestrator:
|
|||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
stats_holder: Optional[dict] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Generate response, streaming tokens from subprocess.
|
||||
|
||||
|
|
@ -908,6 +1250,8 @@ class InferenceOrchestrator:
|
|||
|
||||
``stats_holder``: caller-owned dict; on gen_done its "stats" key gets
|
||||
the worker's usage/timings. Request-scoped to avoid cross-stream reads.
|
||||
|
||||
``presence_penalty`` matches the GGUF sampling path (0 disables it).
|
||||
"""
|
||||
yield from self._generate_inner(
|
||||
messages = messages,
|
||||
|
|
@ -926,6 +1270,7 @@ class InferenceOrchestrator:
|
|||
reasoning_effort = reasoning_effort,
|
||||
preserve_thinking = preserve_thinking,
|
||||
stats_holder = stats_holder,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
|
||||
def generate_chat_completion_with_tools(
|
||||
|
|
@ -945,6 +1290,7 @@ class InferenceOrchestrator:
|
|||
preserve_thinking: Optional[bool] = None,
|
||||
max_tool_iterations: int = 25,
|
||||
auto_heal_tool_calls: bool = True,
|
||||
nudge_tool_calls: Optional[bool] = None,
|
||||
tool_call_timeout: int = 300,
|
||||
session_id: Optional[str] = None,
|
||||
rag_scope: Optional[dict] = None,
|
||||
|
|
@ -952,6 +1298,7 @@ class InferenceOrchestrator:
|
|||
bypass_permissions: bool = False,
|
||||
use_adapter: Optional[Union[bool, str]] = None,
|
||||
stats_holder: Optional[dict] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
**_unused,
|
||||
):
|
||||
"""Run the safetensors agentic tool loop in the parent process,
|
||||
|
|
@ -987,6 +1334,7 @@ class InferenceOrchestrator:
|
|||
preserve_thinking = preserve_thinking,
|
||||
# last turn wins, like the GGUF tool loop
|
||||
stats_holder = stats_holder,
|
||||
presence_penalty = presence_penalty,
|
||||
)
|
||||
if use_adapter is not None:
|
||||
yield from self.generate_with_adapter_control(
|
||||
|
|
@ -1007,6 +1355,7 @@ class InferenceOrchestrator:
|
|||
execute_tool = execute_tool,
|
||||
cancel_event = cancel_event,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
nudge_tool_calls = nudge_tool_calls,
|
||||
max_tool_iterations = max_tool_iterations,
|
||||
tool_call_timeout = tool_call_timeout,
|
||||
session_id = session_id,
|
||||
|
|
@ -1053,6 +1402,7 @@ class InferenceOrchestrator:
|
|||
reasoning_effort: Optional[str] = None,
|
||||
preserve_thinking: Optional[bool] = None,
|
||||
stats_holder: Optional[dict] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Inner generation logic — sends command to subprocess, yields tokens.
|
||||
|
||||
|
|
@ -1060,12 +1410,13 @@ 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
|
||||
|
||||
# Drain any prior compare-mode dispatcher so we can read resp_queue.
|
||||
self._wait_dispatcher_idle()
|
||||
|
|
@ -1074,6 +1425,14 @@ class InferenceOrchestrator:
|
|||
# consume and drop each other's token events. Hold _gen_lock across the
|
||||
# cmd build + send + whole stream so we stay the sole resp_queue reader.
|
||||
with self._gen_lock:
|
||||
# Recheck under the lock: an unload we raced may have cleared/swapped the model.
|
||||
# _unload_pending resets after the lock releases, so it can read False by now;
|
||||
# the active-model check catches that handoff and a reload that swapped models,
|
||||
# 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 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
|
||||
cmd = self._build_generate_cmd(
|
||||
|
|
@ -1087,6 +1446,7 @@ class InferenceOrchestrator:
|
|||
min_p = min_p,
|
||||
max_new_tokens = max_new_tokens,
|
||||
repetition_penalty = repetition_penalty,
|
||||
presence_penalty = presence_penalty,
|
||||
use_adapter = use_adapter,
|
||||
tools = tools,
|
||||
enable_thinking = enable_thinking,
|
||||
|
|
@ -1097,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(
|
||||
|
|
@ -1141,53 +1501,62 @@ class InferenceOrchestrator:
|
|||
raise RuntimeError("Inference subprocess is not running")
|
||||
if not self.active_model_name:
|
||||
raise RuntimeError("No active model")
|
||||
expected_model = self.active_model_name
|
||||
|
||||
request_id = str(uuid.uuid4())
|
||||
# Serialize under _gen_lock (sole resp_queue reader) and refuse to start on the
|
||||
# outgoing model once an unload is pending, like the text and audio-input paths.
|
||||
# Without this a concurrent /audio/generate could run TTS on a model being switched.
|
||||
with self._gen_lock:
|
||||
# Recheck under the lock (see _generate_inner): a raced unload/switch may have
|
||||
# cleared or swapped the model while we waited.
|
||||
if self._unload_pending or self.active_model_name != expected_model:
|
||||
raise RuntimeError("model is being unloaded")
|
||||
|
||||
cmd = {
|
||||
"type": "generate_audio",
|
||||
"request_id": request_id,
|
||||
"text": text,
|
||||
"temperature": temperature,
|
||||
"top_p": top_p,
|
||||
"top_k": top_k,
|
||||
"min_p": min_p,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"repetition_penalty": repetition_penalty,
|
||||
}
|
||||
if use_adapter is not None:
|
||||
cmd["use_adapter"] = use_adapter
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
self._send_cmd(cmd)
|
||||
cmd = {
|
||||
"type": "generate_audio",
|
||||
"request_id": request_id,
|
||||
"text": text,
|
||||
"temperature": temperature,
|
||||
"top_p": top_p,
|
||||
"top_k": top_k,
|
||||
"min_p": min_p,
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"repetition_penalty": repetition_penalty,
|
||||
}
|
||||
if use_adapter is not None:
|
||||
cmd["use_adapter"] = use_adapter
|
||||
|
||||
# Wait for audio_done or audio_error
|
||||
deadline = time.monotonic() + 120.0
|
||||
while time.monotonic() < deadline:
|
||||
remaining = max(0.1, deadline - time.monotonic())
|
||||
resp = self._read_resp(timeout = min(remaining, 1.0))
|
||||
self._send_cmd(cmd)
|
||||
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError(self._subprocess_crash_message("audio generation"))
|
||||
continue
|
||||
deadline = time.monotonic() + 120.0
|
||||
while time.monotonic() < deadline:
|
||||
remaining = max(0.1, deadline - time.monotonic())
|
||||
resp = self._read_resp(timeout = min(remaining, 1.0))
|
||||
|
||||
rtype = resp.get("type", "")
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError(self._subprocess_crash_message("audio generation"))
|
||||
continue
|
||||
|
||||
if rtype == "audio_done":
|
||||
wav_bytes = base64.b64decode(resp["wav_base64"])
|
||||
sample_rate = resp["sample_rate"]
|
||||
return wav_bytes, sample_rate
|
||||
rtype = resp.get("type", "")
|
||||
|
||||
if rtype == "audio_error":
|
||||
raise RuntimeError(resp.get("error", "Audio generation failed"))
|
||||
if rtype == "audio_done":
|
||||
wav_bytes = base64.b64decode(resp["wav_base64"])
|
||||
sample_rate = resp["sample_rate"]
|
||||
return wav_bytes, sample_rate
|
||||
|
||||
if rtype == "error":
|
||||
raise RuntimeError(resp.get("error", "Unknown error"))
|
||||
if rtype == "audio_error":
|
||||
raise RuntimeError(resp.get("error", "Audio generation failed"))
|
||||
|
||||
if rtype == "status":
|
||||
continue
|
||||
if rtype == "error":
|
||||
raise RuntimeError(resp.get("error", "Unknown error"))
|
||||
|
||||
raise RuntimeError("Timeout waiting for audio generation (120s)")
|
||||
if rtype == "status":
|
||||
continue
|
||||
|
||||
raise RuntimeError("Timeout waiting for audio generation (120s)")
|
||||
|
||||
def generate_whisper_response(
|
||||
self,
|
||||
|
|
@ -1247,13 +1616,20 @@ 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
|
||||
|
||||
with self._gen_lock:
|
||||
# Recheck under the lock (see _generate_inner): a raced unload/switch may have
|
||||
# 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 GenStreamError("Error: model is being unloaded")
|
||||
return
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
# numpy array -> list for mp.Queue serialization
|
||||
|
|
@ -1279,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(
|
||||
|
|
|
|||
|
|
@ -29,10 +29,25 @@ import os
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Optional
|
||||
|
||||
from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal
|
||||
from core.inference.tool_loop_controller import coerce_tool_arguments
|
||||
from core.tool_healing import parse_tool_calls_from_text
|
||||
|
||||
# Only the formats this healer's parser can promote -- narrower than the loops'
|
||||
# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare
|
||||
# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a
|
||||
# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in.
|
||||
_HEAL_SIGNALS = (
|
||||
"<tool_call>",
|
||||
"<|tool_call>",
|
||||
"<function=",
|
||||
"[TOOL_CALLS]",
|
||||
)
|
||||
|
||||
|
||||
def _has_heal_signal(text: str) -> bool:
|
||||
return any(s in text for s in _HEAL_SIGNALS)
|
||||
|
||||
|
||||
# Read once at import (same convention as the other UNSLOTH_* switches).
|
||||
_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1"
|
||||
# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process
|
||||
|
|
@ -44,7 +59,7 @@ def nudge_enabled(request_flag: Optional[bool]) -> bool:
|
|||
return _NUDGE_DEFAULT if request_flag is None else bool(request_flag)
|
||||
|
||||
|
||||
_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS)
|
||||
_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS)
|
||||
# A suspected-but-unclosed tool block larger than this is declared a false
|
||||
# alarm and flushed, bounding memory on a model rambling XML-lookalike text.
|
||||
_MAX_HOLD_CHARS = 64 * 1024
|
||||
|
|
@ -198,7 +213,7 @@ def heal_openai_message_events(
|
|||
if not isinstance(msg, dict) or msg.get("tool_calls"):
|
||||
return None
|
||||
content = msg.get("content")
|
||||
if not isinstance(content, str) or not has_tool_signal(content):
|
||||
if not isinstance(content, str) or not _has_heal_signal(content):
|
||||
return None
|
||||
parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True)
|
||||
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
|
||||
|
|
@ -248,7 +263,7 @@ def heal_openai_message(
|
|||
|
||||
def _earliest_signal(buffer: str) -> int:
|
||||
best = -1
|
||||
for signal in TOOL_XML_SIGNALS:
|
||||
for signal in _HEAL_SIGNALS:
|
||||
index = buffer.find(signal)
|
||||
if index >= 0 and (best < 0 or index < best):
|
||||
best = index
|
||||
|
|
@ -275,7 +290,7 @@ def _partial_signal_suffix(buffer: str) -> int:
|
|||
"""Length of the longest buffer suffix that is a proper prefix of a signal."""
|
||||
for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1):
|
||||
tail = buffer[-length:]
|
||||
if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS):
|
||||
if any(signal.startswith(tail) for signal in _HEAL_SIGNALS):
|
||||
return length
|
||||
return 0
|
||||
|
||||
|
|
@ -340,9 +355,10 @@ class StreamToolCallHealer:
|
|||
events.append(("text", emit))
|
||||
self._buffer = self._buffer[len(self._buffer) - keep :]
|
||||
return events
|
||||
# HOLD: handle the FIRST complete block per pass so events keep
|
||||
# document order (a later declared call must not overtake an
|
||||
# earlier undeclared one flushing as text).
|
||||
# HOLD: drain the first contiguous run per pass so events keep document
|
||||
# order (a later declared call must not overtake an earlier undeclared one
|
||||
# flushing as text). A run is one markup call OR a whole Mistral [TOOL_CALLS]
|
||||
# array of contiguous spans, so later calls in it are not stranded as text.
|
||||
parsed, spans = parse_tool_calls_from_text(
|
||||
self._buffer,
|
||||
id_offset = self._id_offset,
|
||||
|
|
@ -363,26 +379,32 @@ class StreamToolCallHealer:
|
|||
self._holding = False
|
||||
continue
|
||||
return events
|
||||
start, end = spans[0]
|
||||
promoted = _promote(
|
||||
[parsed[0]],
|
||||
self._allowed,
|
||||
id_offset = self._id_offset,
|
||||
tool_schemas = self._tool_schemas,
|
||||
)
|
||||
if promoted:
|
||||
if start:
|
||||
events.append(("text", self._buffer[:start]))
|
||||
events.append(("tool_call", promoted[0]))
|
||||
self._id_offset += 1
|
||||
# Drop exactly the promoted markup span; everything else
|
||||
# (leading text, later blocks) stays and is rescanned.
|
||||
self._buffer = self._buffer[end:]
|
||||
else:
|
||||
# Undeclared or unusable name: its markup is DATA, flush it
|
||||
# (and anything before it) verbatim, then rescan the rest.
|
||||
events.append(("text", self._buffer[:end]))
|
||||
self._buffer = self._buffer[end:]
|
||||
pos = 0
|
||||
run_end = spans[0][1]
|
||||
for order, (call, (start, end)) in enumerate(zip(parsed, spans)):
|
||||
# Stop at the first gap or incomplete trailing block: leave it for the
|
||||
# next pass to re-hold and stream incrementally, not flush as text early.
|
||||
if order and start != run_end:
|
||||
break
|
||||
promoted = _promote(
|
||||
[call],
|
||||
self._allowed,
|
||||
id_offset = self._id_offset,
|
||||
tool_schemas = self._tool_schemas,
|
||||
)
|
||||
if promoted:
|
||||
# Flush any leading text, then drop the promoted markup span.
|
||||
if self._buffer[pos:start]:
|
||||
events.append(("text", self._buffer[pos:start]))
|
||||
events.append(("tool_call", promoted[0]))
|
||||
self._id_offset += 1
|
||||
else:
|
||||
# Undeclared/unusable name: markup is DATA, flush it (and prior text) verbatim.
|
||||
events.append(("text", self._buffer[pos:end]))
|
||||
pos = end
|
||||
run_end = end
|
||||
# Everything past the drained run (later blocks) stays and is rescanned.
|
||||
self._buffer = self._buffer[run_end:]
|
||||
self._holding = False
|
||||
|
||||
def finalize(self) -> list:
|
||||
|
|
@ -508,7 +530,7 @@ def nudge_should_retry(
|
|||
if not message or message.get("tool_calls"):
|
||||
return False
|
||||
text = message.get("content")
|
||||
if not isinstance(text, str) or not has_tool_signal(text):
|
||||
if not isinstance(text, str) or not _has_heal_signal(text):
|
||||
return False
|
||||
return not _heal_would_promote(text, allowed_tools, tools)
|
||||
|
||||
|
|
|
|||
49
studio/backend/core/inference/presence_penalty.py
Normal file
49
studio/backend/core/inference/presence_penalty.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Presence-penalty logits helpers for the safetensors/MLX inference paths.
|
||||
|
||||
Kept in a dependency-light leaf module (torch + transformers only, no unsloth /
|
||||
peft) so the pure logic can be imported and unit-tested without pulling in the
|
||||
full inference backend. ``core.inference.inference`` re-exports these for the
|
||||
runtime generate paths.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def apply_presence_penalty(input_ids, scores, penalty: float, prompt_len: int):
|
||||
"""OpenAI/llama.cpp presence penalty: subtract ``penalty`` once per distinct
|
||||
completion token (positions >= prompt_len; prompt excluded, multiplicity
|
||||
ignored, negatives raise). In place; zero is a no-op."""
|
||||
if not penalty:
|
||||
return scores
|
||||
vocab_size = scores.shape[-1]
|
||||
for b in range(input_ids.shape[0]):
|
||||
generated = input_ids[b, prompt_len:]
|
||||
if generated.numel() == 0:
|
||||
continue
|
||||
seen = torch.unique(generated)
|
||||
# Bound generated ids to the valid range [0, vocab_size). Real completion
|
||||
# tokens are always in range, so this is a zero-regression safety net that
|
||||
# drops any stray out-of-range or negative id before indexing (mirrors the
|
||||
# MLX path's bound). Filtering both ends avoids indexing scores with a
|
||||
# negative id (which would silently wrap to the wrong row).
|
||||
seen = seen[(seen >= 0) & (seen < vocab_size)]
|
||||
if seen.numel():
|
||||
scores[b, seen] = scores[b, seen] - penalty
|
||||
return scores
|
||||
|
||||
|
||||
def _make_presence_penalty_processor(penalty: float, prompt_len: int):
|
||||
"""``LogitsProcessorList`` for ``apply_presence_penalty``; ``None`` at zero penalty (generate call stays byte-identical)."""
|
||||
if not penalty:
|
||||
return None
|
||||
from transformers import LogitsProcessor, LogitsProcessorList
|
||||
|
||||
class _PresencePenaltyLogitsProcessor(LogitsProcessor):
|
||||
@torch.no_grad()
|
||||
def __call__(self, input_ids, scores):
|
||||
return apply_presence_penalty(input_ids, scores, penalty, prompt_len)
|
||||
|
||||
return LogitsProcessorList([_PresencePenaltyLogitsProcessor()])
|
||||
|
|
@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via
|
|||
``core.inference.tools``.
|
||||
"""
|
||||
|
||||
import bisect
|
||||
import re
|
||||
import threading
|
||||
from typing import Callable, Generator, Optional
|
||||
|
|
@ -21,14 +22,38 @@ from typing import Callable, Generator, Optional
|
|||
from loggers import get_logger
|
||||
|
||||
from core.inference.tool_call_parser import (
|
||||
_TOOL_ALL_PATS,
|
||||
_GEMMA_BARE_TC_PREFIX_RE,
|
||||
_GEMMA_BARE_TC_RE,
|
||||
_TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS,
|
||||
_TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS,
|
||||
_balanced_brace_end,
|
||||
_strip_function_xml_calls,
|
||||
_strip_gemma_wrapperless_calls,
|
||||
_strip_glm_calls,
|
||||
_strip_mistral_closed_calls,
|
||||
_strip_mistral_reasoning,
|
||||
BUDGET_EXHAUSTED_NUDGE,
|
||||
MAX_ACT_REPROMPTS,
|
||||
RAG_MAX_SEARCHES_PER_TURN,
|
||||
RAG_SEARCH_CAP_NUDGE,
|
||||
TOOL_XML_SIGNALS,
|
||||
is_short_intent_without_action,
|
||||
parse_tool_calls_from_text,
|
||||
reprompt_to_act_message,
|
||||
strip_leading_bare_json_call,
|
||||
strip_llama3_leading_sentinels,
|
||||
strip_tool_markup,
|
||||
)
|
||||
|
||||
# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated
|
||||
# pattern lists, so the safetensors streaming strip stays aligned with the parser.
|
||||
from core.tool_healing import (
|
||||
_REHEARSAL_TAIL_STRIP_RE,
|
||||
_strip_bracket_tag_calls,
|
||||
_think_spans_outside_tool_markup,
|
||||
apply_tool_strip_patterns,
|
||||
strip_outside_think,
|
||||
)
|
||||
from core.inference.tool_loop_controller import (
|
||||
ToolLoopController,
|
||||
coerce_tool_arguments,
|
||||
|
|
@ -50,19 +75,213 @@ logger = get_logger(__name__)
|
|||
# Buffer cap while disambiguating a possible tool-call prefix.
|
||||
_MAX_BUFFER_CHARS = 32
|
||||
|
||||
# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances.
|
||||
_MAX_BARE_JSON_BUFFER = 16384
|
||||
|
||||
|
||||
# No grammar constraint here (unlike llama-server's lazy grammar): collapse
|
||||
# exact-duplicate calls and cap the count so a runaway turn cannot fan out.
|
||||
_MAX_TOOL_CALLS_PER_TURN = 8
|
||||
|
||||
|
||||
def _active_tool_names(active_tools: list[dict]) -> list[str]:
|
||||
names = [
|
||||
(tool.get("function") or {}).get("name")
|
||||
for tool in active_tools
|
||||
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
|
||||
]
|
||||
return [name for name in names if name]
|
||||
|
||||
|
||||
def _active_tool_names(active_tools: list[dict]) -> list[str]:
|
||||
names = [
|
||||
(tool.get("function") or {}).get("name")
|
||||
for tool in active_tools
|
||||
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
|
||||
]
|
||||
return [name for name in names if name]
|
||||
|
||||
|
||||
# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal;
|
||||
# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held.
|
||||
_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?")
|
||||
|
||||
|
||||
def _is_rehearsal_prefix(
|
||||
stripped: str,
|
||||
active_tools: list[dict],
|
||||
*,
|
||||
unrestricted: bool = False,
|
||||
) -> bool:
|
||||
"""True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]``
|
||||
rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space
|
||||
means prose. Unrestricted mode accepts any identifier; else NAME must be active."""
|
||||
if not stripped or any(ch.isspace() for ch in stripped):
|
||||
return False
|
||||
if unrestricted:
|
||||
return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None
|
||||
for name in _active_tool_names(active_tools):
|
||||
if stripped == name or f"{name}[ARGS]".startswith(stripped):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _held_rehearsal_tail_len(
|
||||
text: str,
|
||||
active_tools: list[dict],
|
||||
*,
|
||||
unrestricted: bool = False,
|
||||
) -> int:
|
||||
"""Length of a trailing bare tool-name token that may be a split rehearsal call
|
||||
(``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it
|
||||
instead of leaking the name. Returns 0 for ordinary prose."""
|
||||
i = len(text)
|
||||
while i > 0 and not text[i - 1].isspace():
|
||||
i -= 1
|
||||
tail = text[i:]
|
||||
return (
|
||||
len(tail)
|
||||
if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted)
|
||||
else 0
|
||||
)
|
||||
|
||||
|
||||
def _rehearsal_name_start(
|
||||
candidate: str,
|
||||
signal_pos: int,
|
||||
active_tools: list[dict],
|
||||
*,
|
||||
unrestricted: bool = False,
|
||||
) -> int:
|
||||
"""For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding
|
||||
bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the
|
||||
signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode)."""
|
||||
if not candidate.startswith("[ARGS]", signal_pos):
|
||||
return signal_pos
|
||||
j = signal_pos
|
||||
while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"):
|
||||
j -= 1
|
||||
if j < signal_pos and (
|
||||
unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools)
|
||||
):
|
||||
return j
|
||||
return signal_pos
|
||||
|
||||
|
||||
def _earliest_tool_signal(
|
||||
candidate: str,
|
||||
signals,
|
||||
active_tools: list[dict],
|
||||
*,
|
||||
unrestricted: bool = False,
|
||||
) -> int:
|
||||
"""Index where the turn's first genuine tool-call boundary begins, or -1.
|
||||
|
||||
Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal
|
||||
only when an active tool name (any name in unrestricted mode) precedes it, so a
|
||||
literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a
|
||||
real ``NAME[ARGS]`` the boundary is pulled back to NAME."""
|
||||
best = -1
|
||||
for sig in signals:
|
||||
if sig != "[ARGS]":
|
||||
p = candidate.find(sig)
|
||||
if p >= 0 and (best < 0 or p < best):
|
||||
best = p
|
||||
continue
|
||||
from_idx = 0
|
||||
while True:
|
||||
p = candidate.find("[ARGS]", from_idx)
|
||||
if p < 0:
|
||||
break
|
||||
name_start = _rehearsal_name_start(
|
||||
candidate, p, active_tools, unrestricted = unrestricted
|
||||
)
|
||||
if name_start < p:
|
||||
# Genuine ``NAME[ARGS]``: the boundary is the start of NAME.
|
||||
if best < 0 or name_start < best:
|
||||
best = name_start
|
||||
break
|
||||
# Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found.
|
||||
from_idx = p + len("[ARGS]")
|
||||
return best
|
||||
|
||||
|
||||
def _has_genuine_tool_signal(
|
||||
candidate: str,
|
||||
signals,
|
||||
active_tools: list[dict],
|
||||
*,
|
||||
unrestricted: bool = False,
|
||||
) -> bool:
|
||||
"""True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``.
|
||||
|
||||
Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only
|
||||
when an active tool name (any in unrestricted mode) precedes it. Mirrors the
|
||||
``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not
|
||||
drain inactive-name prose."""
|
||||
for sig in signals:
|
||||
if sig == "[ARGS]":
|
||||
if (
|
||||
_earliest_tool_signal(
|
||||
candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted
|
||||
)
|
||||
>= 0
|
||||
):
|
||||
return True
|
||||
continue
|
||||
if sig in candidate:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def strip_tool_markup_streaming(
|
||||
text: str,
|
||||
*,
|
||||
auto_heal_tool_calls: bool = True,
|
||||
tool_protocol_active: bool = False,
|
||||
enabled_tool_names: Optional[set] = None,
|
||||
) -> str:
|
||||
"""Strip open-ended tool XML from display text without trimming whitespace."""
|
||||
"""Strip open-ended tool XML from display text without trimming whitespace.
|
||||
|
||||
Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so
|
||||
streaming and final display agree: balanced strips first (nested JSON removed whole),
|
||||
then the guarded function-XML / GLM scans that close at each call's REAL terminator so
|
||||
literal markup inside argument values is data and trailing prose survives. Reasoning
|
||||
``<think>`` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must
|
||||
not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names``
|
||||
keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose,
|
||||
not a call), matching the parse / detection active-tool gate."""
|
||||
if not (auto_heal_tool_calls or tool_protocol_active):
|
||||
return text
|
||||
for pat in _TOOL_ALL_PATS:
|
||||
text = pat.sub("", text)
|
||||
return text
|
||||
|
||||
# Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the
|
||||
# ``<think>`` channel) so raw reasoning does not leak into streamed display; an unclosed
|
||||
# leading block is held (dropped to EOF) until its closer streams in.
|
||||
text = _strip_mistral_reasoning(text)
|
||||
|
||||
def _seg(segment: str, is_last: bool) -> str:
|
||||
# Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced
|
||||
# strips first, then the guarded function-XML / GLM scans, then the regex arms
|
||||
# (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last
|
||||
# segment (a bare ``foo[ARGS]`` before <think> is prose). Rehearsal strips are name-gated.
|
||||
seg = _strip_mistral_closed_calls(segment)
|
||||
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
|
||||
if is_last:
|
||||
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
|
||||
seg = _strip_function_xml_calls(seg, final = is_last)
|
||||
seg = _strip_glm_calls(seg, final = is_last)
|
||||
pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS
|
||||
for pat in pats:
|
||||
seg = pat.sub("", seg)
|
||||
if is_last:
|
||||
seg = apply_tool_strip_patterns(
|
||||
seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names
|
||||
)
|
||||
return seg
|
||||
|
||||
# Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then
|
||||
# regrows the cumulative text, corrupting append-by-length consumers.
|
||||
return strip_outside_think(text, _seg)
|
||||
|
||||
|
||||
def _strip_tool_markup_final(
|
||||
|
|
@ -70,10 +289,11 @@ def _strip_tool_markup_final(
|
|||
*,
|
||||
auto_heal_tool_calls: bool,
|
||||
tool_protocol_active: bool = False,
|
||||
enabled_tool_names: Optional[set] = None,
|
||||
) -> str:
|
||||
if not (auto_heal_tool_calls or tool_protocol_active):
|
||||
return text
|
||||
return strip_tool_markup(text, final = True)
|
||||
return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names)
|
||||
|
||||
|
||||
def _status_for_tool(tool_name: str, arguments: dict) -> str:
|
||||
|
|
@ -81,25 +301,76 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str:
|
|||
return status_for_tool(tool_name, arguments)
|
||||
|
||||
|
||||
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool:
|
||||
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
|
||||
probe = strip_llama3_leading_sentinels(text.lstrip())
|
||||
if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)):
|
||||
return False
|
||||
return strip_leading_bare_json_call(probe, enabled_tool_names) != probe
|
||||
|
||||
|
||||
_FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
|
||||
_TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"')
|
||||
# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional
|
||||
# render-html card fires for bracket-tag serializations too.
|
||||
_MISTRAL_RENDER_NAME_RE = re.compile(
|
||||
r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)"
|
||||
)
|
||||
_REHEARSAL_RENDER_NAME_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?=\{)")
|
||||
|
||||
|
||||
def _detect_render_html_tool_start(content: str) -> bool:
|
||||
"""Return True when the first drained tool call is clearly render_html."""
|
||||
function_match = _FUNCTION_SIGNAL_RE.search(content)
|
||||
tool_call_index = content.find("<tool_call>")
|
||||
if not function_match and tool_call_index < 0:
|
||||
"""Return True when the FIRST tool call in ``content`` is clearly render_html.
|
||||
|
||||
Covers every serialization the loop executes (XML ``<function=>`` / ``<tool_call>``,
|
||||
Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a
|
||||
render_html marker inside another call's argument is treated as data. Markers inside
|
||||
a ``<think>`` / ``[THINK]`` block are dropped since the parser skips them."""
|
||||
think_spans = _think_spans_outside_tool_markup(content)
|
||||
_think_starts = [s for s, _e in think_spans]
|
||||
|
||||
def _in_think(pos: int) -> bool:
|
||||
if not think_spans:
|
||||
return False
|
||||
i = bisect.bisect_right(_think_starts, pos) - 1
|
||||
return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1]
|
||||
|
||||
def _first_outside(start: int, finder) -> int:
|
||||
# First occurrence at/after ``start`` that is not inside a think span.
|
||||
pos = finder(start)
|
||||
while pos >= 0 and _in_think(pos):
|
||||
pos = finder(pos + 1)
|
||||
return pos
|
||||
|
||||
candidates: list[tuple[int, str]] = []
|
||||
for fm in _FUNCTION_SIGNAL_RE.finditer(content):
|
||||
if not _in_think(fm.start()):
|
||||
candidates.append((fm.start(), fm.group(1)))
|
||||
break
|
||||
tc = _first_outside(0, lambda i: content.find("<tool_call>", i))
|
||||
if tc >= 0:
|
||||
nm = _TOOL_CALL_NAME_RE.search(content[tc:])
|
||||
candidates.append((tc, nm.group(1) if nm else ""))
|
||||
mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i))
|
||||
if mt >= 0:
|
||||
mm = _MISTRAL_RENDER_NAME_RE.match(content, mt)
|
||||
if mm:
|
||||
candidates.append((mt, mm.group(1)))
|
||||
else:
|
||||
# Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the
|
||||
# first call through the parser (it reads top-level names).
|
||||
arr_calls = parse_tool_calls_from_text(content[mt:])
|
||||
if arr_calls:
|
||||
candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or ""))
|
||||
for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content):
|
||||
if not _in_think(rm.start(1)):
|
||||
candidates.append((rm.start(1), rm.group(1)))
|
||||
break
|
||||
|
||||
if not candidates:
|
||||
return False
|
||||
|
||||
if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index):
|
||||
return function_match.group(1) == "render_html"
|
||||
|
||||
if tool_call_index >= 0:
|
||||
name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:])
|
||||
return bool(name_match and name_match.group(1) == "render_html")
|
||||
|
||||
return False
|
||||
_pos, name = min(candidates, key = lambda c: c[0])
|
||||
return name == "render_html"
|
||||
|
||||
|
||||
def _coerce_arguments_with_provenance(
|
||||
|
|
@ -149,6 +420,7 @@ def run_safetensors_tool_loop(
|
|||
execute_tool: Callable[..., str],
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
auto_heal_tool_calls: bool = True,
|
||||
nudge_tool_calls: Optional[bool] = None,
|
||||
max_tool_iterations: int = 25,
|
||||
tool_call_timeout: int = 300,
|
||||
session_id: Optional[str] = None,
|
||||
|
|
@ -188,8 +460,17 @@ def run_safetensors_tool_loop(
|
|||
for _ev in _auto["events"]:
|
||||
yield _ev
|
||||
conversation.extend(_auto["messages"])
|
||||
# Autoinject ran a KB search outside the controller, so it counts as an
|
||||
# executed tool for the plan-without-action gate.
|
||||
rag_autoinjected = bool(_auto)
|
||||
|
||||
unrestricted_tools = not tools
|
||||
# Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the
|
||||
# ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted.
|
||||
_enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools))
|
||||
# Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent
|
||||
# one-shot), else its repeat is stripped but never drained and the turn ends blank.
|
||||
_detect_tools = [] if unrestricted_tools else list(tools or [])
|
||||
tool_controller = ToolLoopController(
|
||||
tools = None if unrestricted_tools else tools,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
|
|
@ -198,6 +479,14 @@ def run_safetensors_tool_loop(
|
|||
kb_search_count = 0
|
||||
final_attempt_done = False
|
||||
next_call_id = 0
|
||||
reprompt_count = 0
|
||||
# A denied tool confirmation must not be answered with a plan-without-action
|
||||
# re-prompt (which would raise the confirmation gate again).
|
||||
tool_denied = False
|
||||
# Real tool-call turns completed. Only turns that actually executed a tool count
|
||||
# against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a
|
||||
# plan-without-action re-prompt) must not consume budget, matching the GGUF loop.
|
||||
_executed_tool_iters = 0
|
||||
|
||||
def _tool_succeeded(tool_name: str) -> bool:
|
||||
key_prefix = f"{tool_name}:"
|
||||
|
|
@ -215,9 +504,13 @@ def run_safetensors_tool_loop(
|
|||
_state_streaming = 1
|
||||
_state_draining = 2
|
||||
|
||||
for iteration in range(max_tool_iterations + 1):
|
||||
# Reserve re-prompt slots so they don't eat the caller's tool budget.
|
||||
_extra_iters = MAX_ACT_REPROMPTS if max_tool_iterations > 0 else 0
|
||||
for iteration in range(max_tool_iterations + _extra_iters + 1):
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return
|
||||
# Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget.
|
||||
_turn_executed_real_tool = False
|
||||
|
||||
if final_attempt_done:
|
||||
active_tools: list[dict] = []
|
||||
|
|
@ -229,6 +522,8 @@ def run_safetensors_tool_loop(
|
|||
|
||||
tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools))
|
||||
tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else ()
|
||||
# Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
|
||||
_enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools))
|
||||
|
||||
detect_state = _state_buffering
|
||||
content_buffer = ""
|
||||
|
|
@ -304,17 +599,18 @@ def run_safetensors_tool_loop(
|
|||
|
||||
if detect_state == _state_streaming:
|
||||
candidate = cumulative_display + delta
|
||||
signal_pos = -1
|
||||
for sig in tool_xml_signals:
|
||||
p = candidate.find(sig)
|
||||
if p >= 0 and (signal_pos < 0 or p < signal_pos):
|
||||
signal_pos = p
|
||||
# Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is
|
||||
# pulled back to NAME so the name is not flushed.
|
||||
signal_pos = _earliest_tool_signal(
|
||||
candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools
|
||||
)
|
||||
if signal_pos >= 0:
|
||||
before_tool = candidate[:signal_pos]
|
||||
cleaned_before = strip_tool_markup_streaming(
|
||||
before_tool,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = tool_protocol_active,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
if len(cleaned_before) > len(last_emitted):
|
||||
last_emitted = cleaned_before
|
||||
|
|
@ -345,10 +641,20 @@ def run_safetensors_tool_loop(
|
|||
cumulative_display,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = tool_protocol_active,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
if len(cleaned) > len(last_emitted):
|
||||
last_emitted = cleaned
|
||||
yield {"type": "content", "text": cleaned}
|
||||
# Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives;
|
||||
# released by later prose or the end-of-stream flush.
|
||||
if tool_protocol_active:
|
||||
_hold = _held_rehearsal_tail_len(
|
||||
cleaned, _detect_tools, unrestricted = unrestricted_tools
|
||||
)
|
||||
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
|
||||
else:
|
||||
emit = cleaned
|
||||
if len(emit) > len(last_emitted):
|
||||
last_emitted = emit
|
||||
yield {"type": "content", "text": emit}
|
||||
continue
|
||||
|
||||
# BUFFERING: hold until we know it is not a tool call.
|
||||
|
|
@ -366,6 +672,92 @@ def run_safetensors_tool_loop(
|
|||
if sig.startswith(stripped):
|
||||
is_prefix = True
|
||||
break
|
||||
# Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS]
|
||||
# counts only with an active NAME so prose is not drained into a no-op.
|
||||
if sig == "[ARGS]":
|
||||
if (
|
||||
_earliest_tool_signal(
|
||||
stripped,
|
||||
("[ARGS]",),
|
||||
_detect_tools,
|
||||
unrestricted = unrestricted_tools,
|
||||
)
|
||||
>= 0
|
||||
):
|
||||
is_match = True
|
||||
break
|
||||
elif sig.startswith("[") and sig in stripped:
|
||||
is_match = True
|
||||
break
|
||||
|
||||
# Split rehearsal: hold the bare name until its [ARGS] arrives and matches above.
|
||||
is_rehearsal_prefix = False
|
||||
if (
|
||||
not is_match
|
||||
and not is_prefix
|
||||
and tool_protocol_active
|
||||
and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools)
|
||||
):
|
||||
is_prefix = True
|
||||
is_rehearsal_prefix = True
|
||||
|
||||
# Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML
|
||||
# signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses
|
||||
# as a call, else stream as content. Non-call text is always recovered downstream.
|
||||
bare_probe = strip_llama3_leading_sentinels(stripped)
|
||||
if (
|
||||
not is_match
|
||||
and not is_prefix
|
||||
and tool_protocol_active
|
||||
and bare_probe.startswith("{")
|
||||
):
|
||||
if _balanced_brace_end(bare_probe, 0) is None:
|
||||
if len(stripped) < _MAX_BARE_JSON_BUFFER:
|
||||
continue # object still open -- keep buffering
|
||||
elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names):
|
||||
# Oversized still-open ENABLED-tool call: stop holding (memory bound) but
|
||||
# DRAIN instead of leaking the raw prefix; a giant ordinary JSON answer still streams.
|
||||
detect_state = _state_draining
|
||||
continue
|
||||
elif parse_tool_calls_from_text(
|
||||
content_buffer,
|
||||
id_offset = next_call_id,
|
||||
allow_incomplete = auto_heal_tool_calls,
|
||||
enabled_tool_names = _enabled_tool_names,
|
||||
):
|
||||
# Closed object that parses as a bare-JSON call -- drain silently.
|
||||
detect_state = _state_draining
|
||||
continue
|
||||
# Closed non-call object (or oversized non-call) -- stream as text.
|
||||
|
||||
# Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry:
|
||||
# buffer it here or it streams raw until the end-of-turn safety net.
|
||||
# ``(?<!\w)`` keeps "recall:" out; the prefix regex is whitespace-tolerant.
|
||||
if (
|
||||
not is_match
|
||||
and not is_prefix
|
||||
and tool_protocol_active
|
||||
and (
|
||||
"call:".startswith(stripped)
|
||||
or _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None
|
||||
or _GEMMA_BARE_TC_RE.match(stripped) is not None
|
||||
)
|
||||
):
|
||||
if _GEMMA_BARE_TC_RE.match(stripped):
|
||||
detect_state = _state_draining
|
||||
continue
|
||||
# A ``call:`` / ``call:partial_name`` prefix with no ``{`` yet: keep
|
||||
# buffering the variable-length name instead of leaking ``call:longname``.
|
||||
# Names can exceed 32 chars (OpenAI 64, MCP longer), so a fixed cap would
|
||||
# flush real calls raw. The prefix regex self-terminates on ordinary prose
|
||||
# and the ``{`` drains above; bound generously like the bare-JSON path.
|
||||
if _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None:
|
||||
if len(stripped) < _MAX_BARE_JSON_BUFFER:
|
||||
continue
|
||||
detect_state = _state_draining
|
||||
continue
|
||||
if len(stripped) < _MAX_BUFFER_CHARS:
|
||||
continue # bare "call:" prefix still forming
|
||||
|
||||
if is_match:
|
||||
# Tool signal -- flush any visible prefix before DRAINING
|
||||
|
|
@ -375,6 +767,7 @@ def run_safetensors_tool_loop(
|
|||
cumulative_display,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = tool_protocol_active,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
if len(cleaned) > len(last_emitted):
|
||||
last_emitted = cleaned
|
||||
|
|
@ -398,7 +791,8 @@ def run_safetensors_tool_loop(
|
|||
"arguments": {},
|
||||
"provenance": _tool_event_provenance(provisional = True),
|
||||
}
|
||||
elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS:
|
||||
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
|
||||
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
|
||||
continue
|
||||
else:
|
||||
detect_state = _state_streaming
|
||||
|
|
@ -407,57 +801,121 @@ def run_safetensors_tool_loop(
|
|||
cumulative_display,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = tool_protocol_active,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
if len(cleaned) > len(last_emitted):
|
||||
last_emitted = cleaned
|
||||
yield {"type": "content", "text": cleaned}
|
||||
# Same trailing-name hold as STREAMING for this first flush out of BUFFERING.
|
||||
if tool_protocol_active:
|
||||
_hold = _held_rehearsal_tail_len(
|
||||
cleaned, _detect_tools, unrestricted = unrestricted_tools
|
||||
)
|
||||
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
|
||||
else:
|
||||
emit = cleaned
|
||||
if len(emit) > len(last_emitted):
|
||||
last_emitted = emit
|
||||
yield {"type": "content", "text": emit}
|
||||
|
||||
# Stream finished -- resolve what we collected.
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return
|
||||
|
||||
if detect_state == _state_buffering:
|
||||
# Buffer never resolved -- tool XML or plain content?
|
||||
# Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal
|
||||
# ``foo[ARGS]{...}`` is not parsed.
|
||||
stripped = content_buffer.lstrip()
|
||||
_bare_eos = strip_llama3_leading_sentinels(stripped)
|
||||
if (
|
||||
stripped
|
||||
and tool_protocol_active
|
||||
and any(sig in stripped for sig in tool_xml_signals)
|
||||
and _has_genuine_tool_signal(
|
||||
stripped,
|
||||
tool_xml_signals,
|
||||
_detect_tools,
|
||||
unrestricted = unrestricted_tools,
|
||||
)
|
||||
):
|
||||
detect_state = _state_draining
|
||||
elif tool_protocol_active and _looks_like_enabled_bare_json(
|
||||
_bare_eos, _enabled_tool_names
|
||||
):
|
||||
# A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary
|
||||
# JSON answer falls through to the else and streams as content, GGUF parity).
|
||||
detect_state = _state_draining
|
||||
else:
|
||||
# Drain and fall through to STREAMING so the intent re-prompt + safety-net parser
|
||||
# still fire on short emissions like "Let me search." that never exit BUFFERING.
|
||||
if content_buffer:
|
||||
cumulative_display += content_buffer
|
||||
yield {
|
||||
"type": "content",
|
||||
"text": _strip_tool_markup_final(
|
||||
cumulative_display,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = False,
|
||||
),
|
||||
}
|
||||
yield {"type": "status", "text": ""}
|
||||
return
|
||||
cleaned = strip_tool_markup(
|
||||
cumulative_display, final = True, enabled_tool_names = _enabled_tool_names
|
||||
)
|
||||
if len(cleaned) > len(last_emitted):
|
||||
last_emitted = cleaned
|
||||
yield {"type": "content", "text": cleaned}
|
||||
detect_state = _state_streaming
|
||||
|
||||
if detect_state == _state_streaming:
|
||||
# No tool detected mid-stream -- check for late tool XML.
|
||||
safety_tc = None
|
||||
saw_tool_signal = tool_protocol_active and any(
|
||||
sig in content_accum for sig in tool_xml_signals
|
||||
# Run the parser even with no XML signal (the Llama-3.2 bare-JSON form carries none); it's
|
||||
# strict so plain answers stay untouched. Mirrors GGUF.
|
||||
safety_tc = parse_tool_calls_from_text(
|
||||
content_accum,
|
||||
id_offset = next_call_id,
|
||||
allow_incomplete = auto_heal_tool_calls,
|
||||
enabled_tool_names = _enabled_tool_names,
|
||||
)
|
||||
if saw_tool_signal:
|
||||
safety_tc = parse_tool_calls_from_text(
|
||||
content_accum,
|
||||
id_offset = next_call_id,
|
||||
allow_incomplete = auto_heal_tool_calls,
|
||||
)
|
||||
if not safety_tc:
|
||||
# Final answer: if a literal tool marker in prose was stripped
|
||||
# during streaming but did not parse as a real call, restore the
|
||||
# raw cumulative text for core callers. Route-level cleanup can
|
||||
# still apply the Auto-Heal display policy.
|
||||
if saw_tool_signal and content_accum:
|
||||
# Re-prompt once on plan-without-action, before any tool runs
|
||||
# (GGUF loop parity). The retry is gated on nudge_tool_calls so
|
||||
# Studio callers (which send True) always nudge, while API callers
|
||||
# who omit the flag keep today's no-reprompt behavior (opt-in).
|
||||
stripped_answer = content_accum.strip()
|
||||
if (
|
||||
auto_heal_tool_calls
|
||||
and nudge_tool_calls
|
||||
and active_tools
|
||||
and reprompt_count < MAX_ACT_REPROMPTS
|
||||
and not rag_autoinjected
|
||||
and not tool_denied
|
||||
and not any(record.executed for record in tool_controller.history)
|
||||
and is_short_intent_without_action(stripped_answer)
|
||||
):
|
||||
reprompt_count += 1
|
||||
logger.info(
|
||||
"Safetensors re-prompt %d/%d: model responded without "
|
||||
"calling tools (%d chars)",
|
||||
reprompt_count,
|
||||
MAX_ACT_REPROMPTS,
|
||||
len(stripped_answer),
|
||||
)
|
||||
conversation.append({"role": "assistant", "content": stripped_answer})
|
||||
tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool"
|
||||
conversation.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": reprompt_to_act_message(tool_hint),
|
||||
}
|
||||
)
|
||||
# Empty status clears the badge and resets the route's
|
||||
# per-turn text cursor before the re-prompted turn streams.
|
||||
yield {"type": "status", "text": ""}
|
||||
continue
|
||||
|
||||
# Final answer. If a literal tool marker in prose was buffered but
|
||||
# never parsed as a call, restore the raw text so the prose surfaces
|
||||
# in full; route-level cleanup still applies the Auto-Heal policy.
|
||||
if content_accum and any(sig in content_accum for sig in tool_xml_signals):
|
||||
yield {"type": "content", "text": content_accum}
|
||||
else:
|
||||
# Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real
|
||||
# prose, release it.
|
||||
final_clean = strip_tool_markup_streaming(
|
||||
cumulative_display,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = tool_protocol_active,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
if len(final_clean) > len(last_emitted):
|
||||
yield {"type": "content", "text": final_clean}
|
||||
yield {"type": "status", "text": ""}
|
||||
return
|
||||
tool_calls = safety_tc
|
||||
|
|
@ -465,31 +923,41 @@ def run_safetensors_tool_loop(
|
|||
content_accum,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = True,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
logger.info(
|
||||
"Safetensors safety net: parsed %d tool call(s) from streamed content",
|
||||
len(tool_calls),
|
||||
)
|
||||
else:
|
||||
# DRAINING: parse tool calls out of full content.
|
||||
# DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the
|
||||
# ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to
|
||||
# drain here: a spent one-shot (render_html) is off the active list but its re-emitted
|
||||
# ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of
|
||||
# being dropped into a blank continuation.
|
||||
tool_calls = parse_tool_calls_from_text(
|
||||
content_accum,
|
||||
id_offset = next_call_id,
|
||||
allow_incomplete = auto_heal_tool_calls,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
if not tool_calls:
|
||||
# Parser found nothing. Auto-Heal-enabled display cleanup
|
||||
# strips unparseable tool XML; disabled Auto-Heal preserves
|
||||
# the raw text so literal/malformed markup stays visible.
|
||||
if content_accum:
|
||||
yield {
|
||||
"type": "content",
|
||||
"text": _strip_tool_markup_final(
|
||||
content_accum,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = False,
|
||||
),
|
||||
}
|
||||
_drain_text = _strip_tool_markup_final(
|
||||
content_accum,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = False,
|
||||
enabled_tool_names = _enabled_tool_names,
|
||||
)
|
||||
# Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment
|
||||
# (plain JSON answers are left untouched); off keeps it visible per the strict contract.
|
||||
if tool_protocol_active and auto_heal_tool_calls:
|
||||
_drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names)
|
||||
if _drain_text:
|
||||
yield {"type": "content", "text": _drain_text}
|
||||
if provisional_render_html_started and not provisional_resolved:
|
||||
provisional_resolved = True
|
||||
yield {
|
||||
|
|
@ -505,10 +973,14 @@ def run_safetensors_tool_loop(
|
|||
content_accum,
|
||||
auto_heal_tool_calls = auto_heal_tool_calls,
|
||||
tool_protocol_active = True,
|
||||
enabled_tool_names = _enabled_names_gate,
|
||||
)
|
||||
|
||||
if tool_calls:
|
||||
next_call_id += len(tool_calls)
|
||||
# Strip a leading bare-JSON call from the kept content so it isn't replayed as text or
|
||||
# next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers.
|
||||
content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names)
|
||||
|
||||
if final_attempt_done:
|
||||
# Final-answer turn re-called a tool -- stop the loop.
|
||||
|
|
@ -517,6 +989,27 @@ def run_safetensors_tool_loop(
|
|||
yield {"type": "status", "text": ""}
|
||||
return
|
||||
|
||||
# Collapse exact-duplicate calls and cap the count (runaway-turn guard).
|
||||
if tool_calls:
|
||||
seen_keys: set = set()
|
||||
deduped: list = []
|
||||
for _tc in tool_calls:
|
||||
_fn = _tc.get("function", {}) or {}
|
||||
_key = (_fn.get("name", ""), str(_fn.get("arguments", "")))
|
||||
if _key in seen_keys:
|
||||
continue
|
||||
seen_keys.add(_key)
|
||||
deduped.append(_tc)
|
||||
if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN:
|
||||
break
|
||||
if len(deduped) != len(tool_calls):
|
||||
logger.info(
|
||||
"Safetensors: collapsed %d repeated tool call(s) in one turn to %d",
|
||||
len(tool_calls),
|
||||
len(deduped),
|
||||
)
|
||||
tool_calls = deduped
|
||||
|
||||
assistant_msg: dict = {"role": "assistant", "content": content_text}
|
||||
assistant_appended = False
|
||||
|
||||
|
|
@ -593,6 +1086,7 @@ def run_safetensors_tool_loop(
|
|||
"result": TOOL_REJECTED_MESSAGE,
|
||||
"provenance": decision.provenance,
|
||||
}
|
||||
tool_denied = True
|
||||
denied_message = {
|
||||
"role": "tool",
|
||||
"name": decision.tool_name,
|
||||
|
|
@ -634,6 +1128,8 @@ def run_safetensors_tool_loop(
|
|||
completion = tool_controller.record_result(decision, result)
|
||||
if provisional_match:
|
||||
provisional_resolved = True
|
||||
# A tool ran this turn, so it counts against the caller's budget.
|
||||
_turn_executed_real_tool = True
|
||||
yield completion.tool_end_event()
|
||||
conversation.append(completion.tool_message())
|
||||
|
||||
|
|
@ -646,7 +1142,11 @@ def run_safetensors_tool_loop(
|
|||
if not unrestricted_tools and not tool_controller.active_tools():
|
||||
final_attempt_done = True
|
||||
continue
|
||||
if iteration + 1 >= max_tool_iterations and not final_attempt_done:
|
||||
# Count only turns that executed a tool against the cap; a no-op correction turn doesn't
|
||||
# consume budget so the model gets its nudge and another tool-enabled turn (GGUF parity).
|
||||
if _turn_executed_real_tool:
|
||||
_executed_tool_iters += 1
|
||||
if _executed_tool_iters >= max_tool_iterations and not final_attempt_done:
|
||||
# Budget exhausted; nudge a final plain answer.
|
||||
final_attempt_done = True
|
||||
conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE})
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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()
|
||||
|
||||
|
|
@ -406,6 +425,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool:
|
||||
"""Skip a generate queued behind a cancelled one during an unload.
|
||||
|
||||
The parent sets ``drain_event`` for the whole unload. Because the parent's
|
||||
per-token ``cancel_event`` is cleared at the start of every generate, a cancel
|
||||
set while this generate was still queued would otherwise be lost when it is
|
||||
dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so
|
||||
the parent's stream/mailbox drains fast and the switch stays fast, and report
|
||||
the generate was skipped so the caller does not clear the cancel or run it.
|
||||
"""
|
||||
if drain_event is None or not drain_event.is_set():
|
||||
return False
|
||||
request_id = cmd.get("request_id", "")
|
||||
logger.info("Skipping generate for request %s: unload draining", request_id)
|
||||
_send_response(
|
||||
resp_queue,
|
||||
{
|
||||
"type": "gen_done",
|
||||
"request_id": request_id,
|
||||
"cancelled": True,
|
||||
"stats": None,
|
||||
},
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
|
||||
"""Handle a generate command: stream tokens back via resp_queue.
|
||||
|
||||
|
|
@ -431,6 +476,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
|
|||
"min_p": cmd.get("min_p", 0.0),
|
||||
"max_new_tokens": cmd.get("max_new_tokens", 256),
|
||||
"repetition_penalty": cmd.get("repetition_penalty", 1.0),
|
||||
"presence_penalty": cmd.get("presence_penalty", 0.0),
|
||||
"cancel_event": cancel_event,
|
||||
}
|
||||
|
||||
|
|
@ -494,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", "")
|
||||
|
|
@ -632,7 +739,14 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
|
|||
)
|
||||
|
||||
|
||||
def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None:
|
||||
def run_inference_process(
|
||||
*,
|
||||
cmd_queue: Any,
|
||||
resp_queue: Any,
|
||||
cancel_event,
|
||||
config: dict,
|
||||
drain_event = None,
|
||||
) -> None:
|
||||
"""Subprocess entrypoint. Persistent — runs the command loop until shutdown.
|
||||
|
||||
Args:
|
||||
|
|
@ -640,6 +754,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
|
|||
resp_queue: mp.Queue for sending responses to parent.
|
||||
cancel_event: mp.Event the parent sets to cancel generation.
|
||||
config: Initial configuration dict with model info.
|
||||
drain_event: mp.Event the parent sets for the duration of an unload. Unlike
|
||||
cancel_event (cleared at the start of every generate), it is never cleared
|
||||
here, so a generate still queued behind a cancelled one is skipped rather
|
||||
than run — the cancel survives the queue handoff.
|
||||
"""
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
|
||||
|
|
@ -682,9 +800,29 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
|
|||
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..."},
|
||||
|
|
@ -715,8 +853,19 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
|
|||
cmd_type = cmd.get("type", "")
|
||||
try:
|
||||
if cmd_type == "generate":
|
||||
if _drain_skip_generate(cmd, resp_queue, drain_event):
|
||||
continue
|
||||
cancel_event.clear()
|
||||
# Re-check the drain after clearing: the parent sets drain_event
|
||||
# then cancel_event for an unload, so if that pair landed between
|
||||
# the check above and this clear, the clear just erased the unload's
|
||||
# cancel. Skip here so the outgoing model is not run to completion,
|
||||
# which would stall the switch until the dispatcher idle-timeout.
|
||||
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)
|
||||
|
|
@ -918,9 +1067,21 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
|
|||
|
||||
try:
|
||||
if cmd_type == "generate":
|
||||
if _drain_skip_generate(cmd, resp_queue, drain_event):
|
||||
continue
|
||||
cancel_event.clear()
|
||||
# Re-check the drain after clearing: the parent sets drain_event then
|
||||
# cancel_event for an unload, so if that pair landed between the check
|
||||
# above and this clear, the clear just erased the unload's cancel. Skip
|
||||
# here so the outgoing model is not run to completion, which would stall
|
||||
# the switch until the dispatcher idle-timeout tears the subprocess down.
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None:
|
|||
install_torchao_windows_rocm_stub()
|
||||
|
||||
|
||||
class UnsafeEmbeddingModelError(RuntimeError):
|
||||
"""Raised when the embedding model repo is flagged unsafe. A distinct type so the
|
||||
llama-server fallback paths re-raise it instead of masking a security block as a
|
||||
routine ST failure."""
|
||||
|
||||
|
||||
def _ambient_hf_token() -> str | None:
|
||||
"""The HF token the loader itself would use (HF_TOKEN env or the cached login), so
|
||||
the scan can reach a gated/private repo instead of failing open. None if unavailable."""
|
||||
try:
|
||||
from huggingface_hub import get_token
|
||||
return get_token()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
|
||||
"""The module directories a SentenceTransformer load reads weights from, taken from
|
||||
the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``).
|
||||
ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the
|
||||
security scan: a flagged pickle directly under one must block. Returns () on any
|
||||
failure (no modules.json, offline, malformed) so the guard never bricks the embedder.
|
||||
"""
|
||||
try:
|
||||
import json
|
||||
|
||||
from utils.paths import is_local_path
|
||||
|
||||
if is_local_path(name):
|
||||
from pathlib import Path
|
||||
from utils.paths import normalize_path
|
||||
|
||||
path = Path(normalize_path(name)).expanduser() / "modules.json"
|
||||
if not path.is_file():
|
||||
return ()
|
||||
data = json.loads(path.read_text())
|
||||
else:
|
||||
from huggingface_hub import hf_hub_download
|
||||
from huggingface_hub.utils import EntryNotFoundError
|
||||
|
||||
try:
|
||||
local = hf_hub_download(name, "modules.json", token = token or None)
|
||||
except EntryNotFoundError:
|
||||
return ()
|
||||
data = json.loads(open(local).read())
|
||||
subdirs = []
|
||||
for module in data or ():
|
||||
sub = str((module or {}).get("path", "")).strip().strip("/")
|
||||
if sub:
|
||||
subdirs.append(sub)
|
||||
return tuple(dict.fromkeys(subdirs))
|
||||
except Exception:
|
||||
return ()
|
||||
|
||||
|
||||
def _guard_model_security(name: str) -> None:
|
||||
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
|
||||
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
|
||||
/settings gate (a name can also arrive via env/default); local paths and unreachable
|
||||
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
|
||||
"""
|
||||
try:
|
||||
from utils.security import evaluate_file_security, security_load_subdirs
|
||||
|
||||
token = _ambient_hf_token()
|
||||
# Union the audio-model load roots with the ST module dirs so a flagged pickle
|
||||
# directly under a Transformer module dir (0_Transformer/) blocks instead of
|
||||
# passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
|
||||
)
|
||||
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
|
||||
except Exception:
|
||||
return
|
||||
if blocked:
|
||||
raise UnsafeEmbeddingModelError(
|
||||
f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
|
||||
"scan; refusing to load. Set a different RAG embedding model."
|
||||
)
|
||||
|
||||
|
||||
def _get(model_name: str | None = None):
|
||||
"""Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16
|
||||
for a ~1.5x speedup at negligible accuracy loss."""
|
||||
|
|
@ -75,6 +156,7 @@ def _get(model_name: str | None = None):
|
|||
|
||||
device = _device()
|
||||
logger.info("loading embedding model %s on %s", name, device)
|
||||
_guard_model_security(name)
|
||||
_model = SentenceTransformer(
|
||||
name, device = device, model_kwargs = {"torch_dtype": "float16"}
|
||||
)
|
||||
|
|
@ -159,6 +241,8 @@ class _SentenceTransformersBackend:
|
|||
):
|
||||
try:
|
||||
return _st_encode(texts, model_name = model_name, normalize = normalize)
|
||||
except UnsafeEmbeddingModelError:
|
||||
raise # a security block must hard-fail, not fall back to llama-server
|
||||
except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure
|
||||
# ST loaded but this encode blew up; swap the process to the llama-server
|
||||
# embedder (so later encodes stay in one space) and retry.
|
||||
|
|
@ -222,6 +306,8 @@ def _build_st_backend_or_fallback():
|
|||
try:
|
||||
backend.warm(model_name = None)
|
||||
return backend
|
||||
except UnsafeEmbeddingModelError:
|
||||
raise # a security block must hard-fail, not fall back to llama-server
|
||||
except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure
|
||||
fallback = _try_make_llama_backend()
|
||||
if fallback is None:
|
||||
|
|
@ -290,6 +376,37 @@ def _reset_backend() -> None:
|
|||
_backend_key = None
|
||||
|
||||
|
||||
def active_backend_is_llama() -> bool:
|
||||
"""True when this process actually embeds via the llama-server (GGUF) backend.
|
||||
|
||||
Reflects the ACTUAL built backend once one exists: an ``auto`` install that
|
||||
resolves to sentence-transformers but then falls back to llama-server at
|
||||
runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or
|
||||
``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so
|
||||
callers gating on the ST pickle must see llama here. Before any backend is
|
||||
built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw
|
||||
key) exactly as a fresh process would. Never raises: a backend probe must not
|
||||
block saving a model."""
|
||||
try:
|
||||
with _backend_lock:
|
||||
backend = _backend
|
||||
if backend is not None:
|
||||
# A backend exists: report what it ACTUALLY is. A concrete
|
||||
# sentence-transformers backend must return False even if the
|
||||
# resolver would now pick llama, so its pickle stays gated. If the
|
||||
# llama import fails we cannot be llama, so fall to the safe False.
|
||||
try:
|
||||
from .embed_llama_server import LlamaServerBackend
|
||||
except Exception: # noqa: BLE001 - llama plumbing import must never block
|
||||
return False
|
||||
return isinstance(backend, LlamaServerBackend)
|
||||
raw = (config.EMBED_BACKEND or "auto").strip().lower()
|
||||
key = _resolve_auto() if raw in _AUTO_ALIASES else raw
|
||||
return key in _LLAMA_ALIASES
|
||||
except Exception: # noqa: BLE001 - a backend probe must never block saving
|
||||
return False
|
||||
|
||||
|
||||
def warm(model_name: str | None = None) -> None:
|
||||
"""Eagerly load the embedder so the first real request isn't slow."""
|
||||
_get_backend().warm(model_name = model_name)
|
||||
|
|
|
|||
|
|
@ -1,38 +1,137 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
#
|
||||
# Bracket-tag, rehearsal, and thinking-block-strip logic adapted from forge
|
||||
# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026
|
||||
# Antoine Zambelli, used under the MIT License.
|
||||
|
||||
"""Lightweight tool-call XML parsing and stripping helpers.
|
||||
"""Lightweight tool-call parsing and stripping helpers.
|
||||
|
||||
External inference servers import this module without pulling in the inference
|
||||
orchestrator, structlog, httpx, or the rest of the studio backend.
|
||||
orchestrator, structlog, httpx, or the rest of the studio backend. Kept in
|
||||
lockstep with ``core/inference/tool_call_parser.py`` so those servers
|
||||
(llama-server wrappers, llama-swap, custom shims) reuse the same logic. Any
|
||||
change here must also land there.
|
||||
|
||||
Handles these serializations (see ``parse_tool_calls_from_text``):
|
||||
|
||||
* ``<tool_call>{json}</tool_call>``
|
||||
* ``<|tool_call>call:name{...}<tool_call|>`` (Gemma)
|
||||
* ``<function=name><parameter=k>v</parameter></function>``
|
||||
* ``[TOOL_CALLS]name{json}`` (Mistral / Devstral fallback)
|
||||
* ``name[ARGS]{json}`` (reasoning-model rehearsal)
|
||||
"""
|
||||
|
||||
# PEP 604 annotations must stay import-safe on Python 3.9 (requires-python >=3.9).
|
||||
from __future__ import annotations
|
||||
|
||||
import bisect
|
||||
import json
|
||||
import re
|
||||
|
||||
# Pre-compiled patterns for tool XML stripping. The hyphen in the name
|
||||
# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues,
|
||||
# issue-number) parse alongside the built-ins.
|
||||
# One nesting level in the strip regexes; deeper may leak markup (still parsed).
|
||||
_BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"
|
||||
|
||||
# Rehearsal ``name[ARGS]{..}`` strips; group 1 = name for tool-list gating. Closed =
|
||||
# complete body, tail = truncated; ``(?<!\[CALL_ID\])`` keeps the v11 call-id from reading as a name.
|
||||
_REHEARSAL_CLOSED_STRIP_RE = re.compile(
|
||||
r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*" + _BRACKETED_JSON_ONE_LEVEL, re.DOTALL
|
||||
)
|
||||
_REHEARSAL_TAIL_STRIP_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?:\{.*)?$", re.DOTALL)
|
||||
|
||||
# Tool-XML strip patterns; hyphen in the name class covers dashed MCP names.
|
||||
# Closed-pair patterns are named so _PAT_REQUIRED_TOKEN can skip a doomed lazy rescan when
|
||||
# the close token is absent: an unguarded ``<tag>.*?</tag>`` rescans to EOF from every opener
|
||||
# (quadratic on a stream of unclosed openers). Also reused by the quote-aware Gemma pre-pass.
|
||||
_TC_JSON_CLOSED_PAT = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
|
||||
_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL)
|
||||
_TC_FUNC_CLOSED_PAT = re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL)
|
||||
_TOOL_CLOSED_PATS = [
|
||||
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
|
||||
re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL),
|
||||
_TC_JSON_CLOSED_PAT,
|
||||
_TC_GEMMA_CLOSED_PAT,
|
||||
re.compile(r"<tool_call\|>"),
|
||||
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
|
||||
_TC_FUNC_CLOSED_PAT,
|
||||
# Mirror the parser regexes: tolerate whitespace and v11 [CALL_ID]/[ARGS] metadata.
|
||||
re.compile(
|
||||
r"\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*"
|
||||
+ _BRACKETED_JSON_ONE_LEVEL,
|
||||
re.DOTALL,
|
||||
),
|
||||
_REHEARSAL_CLOSED_STRIP_RE,
|
||||
# Drop the bare v11 [/TOOL_CALLS] closer the balanced scan leaves behind.
|
||||
re.compile(r"\[/TOOL_CALLS\]"),
|
||||
]
|
||||
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
|
||||
# Bare open markers strip a partial call mid-stream; the rehearsal tail needs `{` or EOF
|
||||
# so prose ``foo[ARGS]`` survives. The XML open-tail forms reach EOF and are reused by
|
||||
# _tool_call_markup_spans (a think tag in an unclosed call's args stays argument data).
|
||||
_TOOL_OPEN_XML_TAIL_PATS = [
|
||||
re.compile(r"<tool_call>.*$", re.DOTALL),
|
||||
re.compile(r"<\|tool_call>.*$", re.DOTALL),
|
||||
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
|
||||
]
|
||||
_TOOL_ALL_PATS = (
|
||||
_TOOL_CLOSED_PATS
|
||||
+ _TOOL_OPEN_XML_TAIL_PATS
|
||||
+ [
|
||||
re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL),
|
||||
_REHEARSAL_TAIL_STRIP_RE,
|
||||
]
|
||||
)
|
||||
|
||||
# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None.
|
||||
_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE})
|
||||
|
||||
# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument
|
||||
# data cannot make the helper truncate the block and its tail.
|
||||
_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT]
|
||||
# A lazy closed-pair pattern whose close token is absent would rescan to EOF from every
|
||||
# opener; skip that doomed (quadratic) pass. Shared by both strip helpers.
|
||||
_PAT_REQUIRED_TOKEN = {
|
||||
_TC_JSON_CLOSED_PAT: "</tool_call>",
|
||||
_TC_GEMMA_CLOSED_PAT: "<tool_call|>",
|
||||
_TC_FUNC_CLOSED_PAT: "</function>",
|
||||
}
|
||||
|
||||
|
||||
def strip_tool_patterns(text: str, patterns) -> str:
|
||||
"""Apply ``patterns`` in order, skipping closed-pair passes with no close token."""
|
||||
for pat in patterns:
|
||||
token = _PAT_REQUIRED_TOKEN.get(pat)
|
||||
if token is not None and token not in text:
|
||||
continue
|
||||
text = pat.sub("", text)
|
||||
return text
|
||||
|
||||
|
||||
def apply_tool_strip_patterns(
|
||||
text: str,
|
||||
patterns,
|
||||
enabled_tool_names = None,
|
||||
) -> str:
|
||||
"""Apply strip ``patterns`` to ``text``. A bare rehearsal ``name[ARGS]{..}`` pattern
|
||||
strips only when ``name`` is an enabled tool (or when ``enabled_tool_names`` is
|
||||
``None``); every other pattern is removed unconditionally. A closed-pair pattern whose
|
||||
close token is absent is skipped so an unclosed-marker stream stays linear."""
|
||||
for pat in patterns:
|
||||
token = _PAT_REQUIRED_TOKEN.get(pat)
|
||||
if token is not None and token not in text:
|
||||
continue
|
||||
if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS:
|
||||
text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text)
|
||||
else:
|
||||
text = pat.sub("", text)
|
||||
return text
|
||||
|
||||
|
||||
# Pre-compiled patterns for tool-call XML parsing.
|
||||
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
|
||||
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{")
|
||||
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{")
|
||||
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
|
||||
_TC_END_TAG_RE = re.compile(r"</tool_call>")
|
||||
_TC_GEMMA_END_TAG_RE = re.compile(r"<tool_call\|>")
|
||||
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
|
||||
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
|
||||
# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it.
|
||||
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>[^\S\n]*")
|
||||
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
|
||||
_GEMMA_QUOTE = '<|"|>'
|
||||
_PARAM_CLOSE_TAG = "</parameter>"
|
||||
|
|
@ -43,7 +142,62 @@ _FUNC_CLOSE_TAG = "</function>"
|
|||
# must be identifier-shaped (start with a letter or underscore); a comma
|
||||
# followed by digits-then-colon is value text such as a timestamp or ratio
|
||||
# (`meet at 10:00, 11:00 tomorrow`), not a new key.
|
||||
_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:")
|
||||
_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:")
|
||||
|
||||
# A candidate starting inside a think block is a rehearsal (block kept so literal tags in
|
||||
# real args survive); ``$`` accepts an unclosed block mid-stream.
|
||||
_THINK_TAG_RE = re.compile(r"<think>.*?(?:</think>|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL)
|
||||
# Bare open/close markers for prefilled-reasoning turns (template opens <think> in the prompt).
|
||||
_THINK_OPEN_RE = re.compile(r"<think>|\[THINK\]")
|
||||
_THINK_CLOSE_RE = re.compile(r"</think>|\[/THINK\]")
|
||||
|
||||
# Mistral canonical array: [TOOL_CALLS] + JSON list of {"name","arguments"} objects.
|
||||
_MISTRAL_ARRAY_RE = re.compile(r"\[TOOL_CALLS\]\s*(?=\[)")
|
||||
|
||||
# Mistral name form + v11 [ARGS]/[CALL_ID] shapes; [CALL_ID] is metadata, not the name,
|
||||
# and hyphens keep dashed MCP names whole.
|
||||
_MISTRAL_BRACKET_RE = re.compile(
|
||||
r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)"
|
||||
)
|
||||
|
||||
# Rehearsal ``name[ARGS]{json}`` (no [TOOL_CALLS]); the lookbehind keeps the v11 call-id
|
||||
# from being taken as the function name.
|
||||
_REHEARSAL_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?=\{)")
|
||||
|
||||
# Above this size skip the balanced scan; the linear regex catch-all bounds pathological output.
|
||||
_MAX_BRACKET_SCAN_CHARS = 1_000_000
|
||||
|
||||
|
||||
def _balanced_json_span(text: str, start: int) -> int | None:
|
||||
"""Return the end index of a balanced JSON object opening at ``start``,
|
||||
or ``None`` if the braces don't balance. Honors escapes and strings.
|
||||
"""
|
||||
if start >= len(text) or text[start] != "{":
|
||||
return None
|
||||
depth = 0
|
||||
in_string = False
|
||||
escape = False
|
||||
for j in range(start, len(text)):
|
||||
ch = text[j]
|
||||
if escape:
|
||||
escape = False
|
||||
continue
|
||||
if ch == "\\":
|
||||
escape = True
|
||||
continue
|
||||
if in_string:
|
||||
if ch == '"':
|
||||
in_string = False
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = True
|
||||
elif ch == "{":
|
||||
depth += 1
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return j
|
||||
return None
|
||||
|
||||
|
||||
def _balanced_brace_end(
|
||||
|
|
@ -109,6 +263,94 @@ def _balanced_bracket_end(src: str, start: int) -> int:
|
|||
return -1
|
||||
|
||||
|
||||
def _decode_array_items(text: str, body_start: int, body_end: int):
|
||||
"""Return ``(objs, ends)`` for each top-level element of the JSON array between
|
||||
``body_start`` (at or before its ``[``) and ``body_end`` (exclusive): the decoded
|
||||
object and its absolute exclusive end offset.
|
||||
|
||||
Decoding element-by-element with ``raw_decode`` tolerates the comma-less object
|
||||
separators the repo's own Mistral/Ollama multi-call templates emit
|
||||
(``[{...}{...}]``; see ollama_template_mappers.py). A single ``json.loads`` of the
|
||||
whole body rejects that form and would drop every call. The ends also tile the
|
||||
region across the calls' spans so a with_spans consumer strips each exactly once."""
|
||||
decoder = json.JSONDecoder()
|
||||
objs: list = []
|
||||
ends: list[int] = []
|
||||
i = text.find("[", body_start)
|
||||
if i < 0:
|
||||
return objs, ends
|
||||
i += 1
|
||||
while i < body_end:
|
||||
while i < body_end and text[i] in " \t\r\n,":
|
||||
i += 1
|
||||
if i >= body_end or text[i] == "]":
|
||||
break
|
||||
try:
|
||||
obj, rel = decoder.raw_decode(text[i:body_end])
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
break
|
||||
i += rel
|
||||
objs.append(obj)
|
||||
ends.append(i)
|
||||
return objs, ends
|
||||
|
||||
|
||||
def _iter_bracket_spans(
|
||||
text: str,
|
||||
start: int = 0,
|
||||
enabled_tool_names = None,
|
||||
):
|
||||
"""Yield ``(span_start, span_end, kind, match)`` for each balanced bracket-tag
|
||||
call from ``start`` on, in document order; ``span_end`` exclusive. ``kind`` is
|
||||
``"array"`` ([TOOL_CALLS] [..]), ``"name"`` ([TOOL_CALLS]name{..}, incl. v11
|
||||
[CALL_ID]/[ARGS]) or ``"rehearsal"`` (name[ARGS]{..}).
|
||||
|
||||
``enabled_tool_names`` (set, or None = unrestricted) gates only the ambiguous
|
||||
bare rehearsal form: name[ARGS]{..} is a call ONLY when ``name`` is enabled, so a
|
||||
prose ``foo[ARGS]{..}`` (foo disabled) is neither parsed nor stripped. Explicit
|
||||
[TOOL_CALLS] markers stay unconditional, keeping parse/strip/detection symmetric.
|
||||
|
||||
Balance-only (no JSON validation) so strip and parse share one scan. The cursor
|
||||
jumps past each consumed span, so a marker inside consumed JSON is never
|
||||
re-matched and each regex re-searches only once its match falls behind: linear."""
|
||||
n = len(text)
|
||||
specs = (
|
||||
("array", _MISTRAL_ARRAY_RE),
|
||||
("name", _MISTRAL_BRACKET_RE),
|
||||
("rehearsal", _REHEARSAL_RE),
|
||||
)
|
||||
nexts = {kind: rx.search(text, start) for kind, rx in specs}
|
||||
cursor = start
|
||||
while cursor < n:
|
||||
for kind, rx in specs:
|
||||
m = nexts[kind]
|
||||
if m is not None and m.start() < cursor:
|
||||
nexts[kind] = rx.search(text, cursor)
|
||||
live = [(kind, m) for kind, m in nexts.items() if m is not None]
|
||||
if not live:
|
||||
return
|
||||
kind, m = min(live, key = lambda km: km[1].start())
|
||||
if kind == "array":
|
||||
end = _balanced_bracket_end(text, m.end())
|
||||
end = None if end < 0 else end
|
||||
else:
|
||||
end = _balanced_json_span(text, m.end())
|
||||
if end is None:
|
||||
# Truncated body: skip and keep scanning; the caller's catch-all strips the tail.
|
||||
cursor = m.end()
|
||||
continue
|
||||
if (
|
||||
kind == "rehearsal"
|
||||
and enabled_tool_names is not None
|
||||
and m.group(1) not in enabled_tool_names
|
||||
):
|
||||
# Inactive-name rehearsal is prose: advance past its body without yielding.
|
||||
cursor = end + 1
|
||||
continue
|
||||
yield (m.start(), end + 1, kind, m)
|
||||
cursor = end + 1
|
||||
|
||||
|
||||
def _split_top_level_commas(src: str) -> list:
|
||||
"""Split on commas that are not inside a nested ``[]``/``{}`` or a string."""
|
||||
parts: list[str] = []
|
||||
|
|
@ -223,7 +465,7 @@ def _quote_gemma_object_keys(src: str) -> str:
|
|||
while i < len(src) and src[i].isspace():
|
||||
i += 1
|
||||
key_name_start = i
|
||||
while i < len(src) and (src[i].isalnum() or src[i] in "_-"):
|
||||
while i < len(src) and (src[i].isalnum() or src[i] in "_-."):
|
||||
i += 1
|
||||
key_name = src[key_name_start:i]
|
||||
colon_pos = i
|
||||
|
|
@ -267,7 +509,8 @@ def _quote_gemma_object_keys(src: str) -> str:
|
|||
json.loads(raw.strip())
|
||||
parts.append(raw)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
parts.append(json.dumps(raw.strip()) if raw.strip() else raw)
|
||||
# Quote bare value; empty ({k:}) becomes "" so json.loads sees {"k":""} not invalid {"k":}.
|
||||
parts.append(json.dumps(raw.strip()))
|
||||
else:
|
||||
parts.append(src[key_start:i])
|
||||
return "".join(parts)
|
||||
|
|
@ -291,9 +534,99 @@ def _inside_open_parameter(content: str, pos: int) -> bool:
|
|||
last_param_start = match.start()
|
||||
if last_param_start < 0:
|
||||
return False
|
||||
last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos)
|
||||
last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos)
|
||||
return last_param_start > max(last_param_close, last_func_close)
|
||||
# The parameter's OWN close tag decides: if it closes after ``pos`` the position is
|
||||
# argument data (even across literal function closes); an unclosed one falls back to func close.
|
||||
own_close = content.find(_PARAM_CLOSE_TAG, last_param_start)
|
||||
if own_close >= 0:
|
||||
return own_close > pos
|
||||
func_close = content.find(_FUNC_CLOSE_TAG, last_param_start)
|
||||
return func_close < 0 or pos < func_close
|
||||
|
||||
|
||||
def _func_close_index(content: str, body_start: int, body: str) -> int:
|
||||
"""Index in ``body`` of the first ``</function>`` that is not argument
|
||||
data (not inside an open parameter value); -1 when every close is data.
|
||||
Taking the LAST close swallowed prose between the real close and a
|
||||
literal ``</function>`` mentioned later in the answer."""
|
||||
idx = body.find(_FUNC_CLOSE_TAG)
|
||||
while idx >= 0:
|
||||
if not _inside_open_parameter(content, body_start + idx):
|
||||
return idx
|
||||
idx = body.find(_FUNC_CLOSE_TAG, idx + 1)
|
||||
return -1
|
||||
|
||||
|
||||
def _trim_param_value(val: str) -> str:
|
||||
"""Trim the single wrapping newline the chat template adds around an XML
|
||||
parameter value, preserving indentation inside VALUE (``str.strip()`` destroyed
|
||||
code/diff argument indentation)."""
|
||||
if val.startswith("\n"):
|
||||
val = val[1:]
|
||||
if val.endswith("\n"):
|
||||
val = val[:-1]
|
||||
return val
|
||||
|
||||
|
||||
def _marker_coverage(content: str, markers) -> list[tuple[int, int]]:
|
||||
"""Coverage ``[start, end]`` per marker, used to skip markers that are another
|
||||
call's data. Closes pair to markers via a per-format stack so an inner close
|
||||
is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with
|
||||
a paired close cover through it (markers before the close are data); balanced
|
||||
without one cover only the braces, so a later sibling is still recovered."""
|
||||
n = len(content)
|
||||
brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0]
|
||||
events = [] # (position, order) with order 0 = braces-done, 1 = close marker
|
||||
for idx, (_start, brace_end, _kind, _m) in enumerate(markers):
|
||||
if brace_end >= 0:
|
||||
events.append((brace_end, 0, _kind, idx))
|
||||
for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)):
|
||||
for cm in close_re.finditer(content):
|
||||
# A close inside another call's balanced braces is quoted data; it
|
||||
# must not pop an earlier close-less marker and swallow a sibling.
|
||||
if any(s < cm.start() < be for s, be in brace_regions):
|
||||
continue
|
||||
events.append((cm.start(), 1, kind, cm.end()))
|
||||
events.sort(key = lambda e: (e[0], e[1]))
|
||||
waiting = {"json": [], "gemma": []}
|
||||
close_end_for: dict[int, int] = {}
|
||||
for _pos, order, kind, payload in events:
|
||||
if order == 0:
|
||||
waiting[kind].append(payload) # marker index, now awaiting its close
|
||||
elif waiting[kind]:
|
||||
close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here
|
||||
coverage = []
|
||||
for idx, (start, brace_end, _kind, _m) in enumerate(markers):
|
||||
if brace_end < 0:
|
||||
coverage.append((start, n))
|
||||
elif idx in close_end_for:
|
||||
coverage.append((start, close_end_for[idx]))
|
||||
else:
|
||||
coverage.append((start, brace_end))
|
||||
return coverage
|
||||
|
||||
|
||||
def _build_markers(content: str):
|
||||
"""JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document
|
||||
order; ``brace_end < 0`` marks an unbalanced (to-EOF) open."""
|
||||
markers = []
|
||||
for start_re, gemma, kind in (
|
||||
(_TC_JSON_START_RE, False, "json"),
|
||||
(_TC_GEMMA_START_RE, True, "gemma"),
|
||||
):
|
||||
for m in start_re.finditer(content):
|
||||
if _inside_open_parameter(content, m.start()):
|
||||
continue
|
||||
brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma)
|
||||
markers.append((m.start(), brace_end, kind, m))
|
||||
markers.sort(key = lambda c: c[0])
|
||||
return markers
|
||||
|
||||
|
||||
def marker_coverage(content: str) -> list[tuple[int, int]]:
|
||||
"""Coverage spans of JSON/Gemma tool markers so other parsers can treat markup
|
||||
inside a marker's coverage (even a marker that failed to parse) as that call's
|
||||
data rather than a sibling call."""
|
||||
return _marker_coverage(content, _build_markers(content))
|
||||
|
||||
|
||||
def parse_tool_calls_from_text(
|
||||
|
|
@ -301,6 +634,7 @@ def parse_tool_calls_from_text(
|
|||
*,
|
||||
id_offset: int = 0,
|
||||
allow_incomplete: bool = True,
|
||||
enabled_tool_names = None,
|
||||
with_spans: bool = False,
|
||||
):
|
||||
"""Parse OpenAI-format tool calls from model text.
|
||||
|
|
@ -309,55 +643,61 @@ def parse_tool_calls_from_text(
|
|||
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
|
||||
<|tool_call>call:web_search{query:"..."}<tool_call|>
|
||||
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
|
||||
[TOOL_CALLS]web_search{"query":"..."} (Mistral / Devstral fallback)
|
||||
web_search[ARGS]{"query":"..."} (reasoning-model rehearsal)
|
||||
|
||||
A call rehearsed inside a ``<think>`` / ``[THINK]`` block is skipped, not
|
||||
executed; the block is kept so a literal tag in a real argument is preserved.
|
||||
|
||||
With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]``
|
||||
is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup
|
||||
in ``content`` (including its close tag when present), so a caller can
|
||||
remove exactly the parsed markup and keep every other byte intact.
|
||||
"""
|
||||
# Candidates starting inside a think block are rehearsals, skipped; blocks are kept, and a
|
||||
# think marker opening inside a call is argument data (excluded from spans).
|
||||
_think_spans = _think_spans_outside_tool_markup(content)
|
||||
_think_starts = [s for s, _e in _think_spans]
|
||||
|
||||
def _in_think(pos: int) -> bool:
|
||||
# Spans are ordered and non-overlapping; bisect gives O(log M) per candidate.
|
||||
i = bisect.bisect_right(_think_starts, pos) - 1
|
||||
return i >= 0 and _think_spans[i][0] <= pos < _think_spans[i][1]
|
||||
|
||||
tool_calls: list[dict] = []
|
||||
call_spans: list[tuple] = []
|
||||
# Collect every supported call format with spans, then emit in document
|
||||
# order. A marker inside another call's argument string is data, not a
|
||||
# separate executable call.
|
||||
parsed_items = [] # (start, span_end, name, arguments)
|
||||
candidates = [] # (start, brace_end, kind, match)
|
||||
for m in _TC_JSON_START_RE.finditer(content):
|
||||
if _inside_open_parameter(content, m.start()):
|
||||
continue
|
||||
end = _balanced_brace_end(content, m.end() - 1)
|
||||
if end >= 0:
|
||||
candidates.append((m.start(), end, "json", m))
|
||||
for m in _TC_GEMMA_START_RE.finditer(content):
|
||||
if _inside_open_parameter(content, m.start()):
|
||||
continue
|
||||
end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True)
|
||||
if end >= 0:
|
||||
candidates.append((m.start(), end, "gemma", m))
|
||||
candidates.sort(key = lambda c: c[0])
|
||||
|
||||
candidate_spans = [(s, e) for s, e, _kind, _m in candidates]
|
||||
for idx, (start, end, kind, m) in enumerate(candidates):
|
||||
if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx):
|
||||
# Collect JSON/Gemma markers; _marker_coverage decides nesting so a marker inside
|
||||
# another call's coverage (even one that failed to parse) is data, not executed. A
|
||||
# marker opening inside a think block is a rehearsal and is skipped.
|
||||
parsed_items = [] # (start, span_end, name, arguments) in document order
|
||||
markers = [mk for mk in _build_markers(content) if not _in_think(mk[0])]
|
||||
coverage = _marker_coverage(content, markers)
|
||||
for idx, (start, brace_end, kind, m) in enumerate(markers):
|
||||
if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx):
|
||||
continue
|
||||
if brace_end < 0:
|
||||
continue # unclosed: not parseable; the fallback still excludes its XML
|
||||
if not allow_incomplete:
|
||||
tail = content[end + 1 :].lstrip()
|
||||
tail = content[brace_end + 1 :].lstrip()
|
||||
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
|
||||
if close_re.match(tail) is None:
|
||||
continue
|
||||
try:
|
||||
if kind == "json":
|
||||
obj = json.loads(content[m.end() - 1 : end + 1])
|
||||
obj = json.loads(content[m.end() - 1 : brace_end + 1])
|
||||
name = obj.get("name", "")
|
||||
arguments = obj.get("arguments", {})
|
||||
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes).
|
||||
arguments = obj.get("arguments")
|
||||
if arguments is None:
|
||||
arguments = obj.get("parameters", {})
|
||||
if isinstance(arguments, dict):
|
||||
arguments = json.dumps(arguments)
|
||||
else:
|
||||
name = m.group(1)
|
||||
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end]))
|
||||
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
span_end = end + 1
|
||||
span_end = brace_end + 1
|
||||
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
|
||||
ws = len(content[span_end:]) - len(content[span_end:].lstrip())
|
||||
close_m = close_re.match(content, span_end + ws)
|
||||
|
|
@ -369,7 +709,8 @@ def parse_tool_calls_from_text(
|
|||
fm
|
||||
for fm in _TC_FUNC_START_RE.finditer(content)
|
||||
if not _inside_open_parameter(content, fm.start())
|
||||
and not any(s <= fm.start() <= e for s, e in candidate_spans)
|
||||
and not _in_think(fm.start())
|
||||
and not any(s <= fm.start() < e for s, e in coverage)
|
||||
]
|
||||
for idx, fm in enumerate(func_starts):
|
||||
func_name = fm.group(1)
|
||||
|
|
@ -382,7 +723,7 @@ def parse_tool_calls_from_text(
|
|||
body_end = len(content)
|
||||
body_end = min(body_end, next_func)
|
||||
body = content[body_start:body_end]
|
||||
close_idx = body.rfind(_FUNC_CLOSE_TAG)
|
||||
close_idx = _func_close_index(content, body_start, body)
|
||||
if close_idx >= 0:
|
||||
span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG)
|
||||
body = body[:close_idx]
|
||||
|
|
@ -404,7 +745,7 @@ def parse_tool_calls_from_text(
|
|||
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
|
||||
else:
|
||||
val = _TC_PARAM_CLOSE_RE.sub("", val)
|
||||
arguments[pm.group(1)] = val.strip()
|
||||
arguments[pm.group(1)] = _trim_param_value(val)
|
||||
else:
|
||||
valid_params = True
|
||||
for pidx, pm in enumerate(param_starts):
|
||||
|
|
@ -422,7 +763,7 @@ def parse_tool_calls_from_text(
|
|||
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
|
||||
else:
|
||||
val = _TC_PARAM_CLOSE_RE.sub("", val)
|
||||
arguments[param_name] = val.strip()
|
||||
arguments[param_name] = _trim_param_value(val)
|
||||
if not valid_params:
|
||||
continue
|
||||
|
||||
|
|
@ -444,19 +785,293 @@ def parse_tool_calls_from_text(
|
|||
}
|
||||
)
|
||||
call_spans.append((start, span_end))
|
||||
|
||||
# Patterns 3+4: Mistral [TOOL_CALLS] and bare rehearsal via one balanced scan in document
|
||||
# order, so a Mistral call and a rehearsal in one message both parse.
|
||||
if not tool_calls:
|
||||
for start, end, kind, m in _iter_bracket_spans(
|
||||
content, enabled_tool_names = enabled_tool_names
|
||||
):
|
||||
if _in_think(start):
|
||||
continue
|
||||
# Extend the region over an immediately-following v11 closer so with_spans consumers strip it too.
|
||||
closer = re.match(r"\s*\[/TOOL_CALLS\]", content[end:])
|
||||
region_end = end + closer.end() if closer else end
|
||||
if kind == "array":
|
||||
# Decode elements individually (comma-tolerant): one json.loads of the whole
|
||||
# body rejects the comma-less multi-call arrays Mistral/Ollama templates emit.
|
||||
payload, item_ends = _decode_array_items(content, m.end(), end)
|
||||
if not payload:
|
||||
continue
|
||||
# Tile the region so every byte belongs to exactly one span; a with_spans consumer
|
||||
# keeps skipped bytes visible and strips promoted markup exactly once.
|
||||
tile_start = start
|
||||
last_span_idx = -1
|
||||
for item_idx, item in enumerate(payload):
|
||||
if not isinstance(item, dict) or "name" not in item:
|
||||
continue
|
||||
args = item.get("arguments", {})
|
||||
if isinstance(args, str):
|
||||
# ``arguments`` may itself be a JSON string (OpenAI spec).
|
||||
try:
|
||||
args = json.loads(args)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
if not isinstance(args, (dict, str)):
|
||||
# ``"arguments": null`` (or any non-object scalar) becomes {} like the
|
||||
# <tool_call> path, not the string "null" auto-heal would mangle to
|
||||
# a bogus {"query":"null"}.
|
||||
args = {}
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": f"call_{id_offset + len(tool_calls)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": item.get("name", ""),
|
||||
# A bare scalar string stays raw (like the <tool_call> path);
|
||||
# json.dumps would double-encode it so the arg healer wraps
|
||||
# "weather" with its literal quotes.
|
||||
"arguments": args if isinstance(args, str) else json.dumps(args),
|
||||
},
|
||||
}
|
||||
)
|
||||
item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end
|
||||
last_span_idx = len(call_spans)
|
||||
call_spans.append((tile_start, item_end))
|
||||
tile_start = item_end
|
||||
if last_span_idx >= 0:
|
||||
tile_start, _tile_end = call_spans[last_span_idx]
|
||||
call_spans[last_span_idx] = (tile_start, region_end)
|
||||
else:
|
||||
try:
|
||||
payload = json.loads(content[m.end() : end])
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
if not isinstance(payload, dict):
|
||||
continue
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": f"call_{id_offset + len(tool_calls)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": m.group(1),
|
||||
"arguments": json.dumps(payload),
|
||||
},
|
||||
}
|
||||
)
|
||||
call_spans.append((start, region_end))
|
||||
|
||||
if with_spans:
|
||||
return tool_calls, call_spans
|
||||
return tool_calls
|
||||
|
||||
|
||||
def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
|
||||
def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str:
|
||||
"""Strip complete [TOOL_CALLS] arrays / name / bare name[ARGS]{..} calls with one
|
||||
balanced forward scan, so nested JSON args are removed whole (a fixed-depth regex
|
||||
left two-level args behind). Truncated tails go to the caller's catch-all. Linear.
|
||||
``enabled_tool_names`` gates the rehearsal form (inactive-name prose kept; None
|
||||
strips every span)."""
|
||||
if len(text) > _MAX_BRACKET_SCAN_CHARS:
|
||||
return text
|
||||
out: list[str] = []
|
||||
cursor = 0
|
||||
for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names):
|
||||
out.append(text[cursor:start])
|
||||
cursor = end
|
||||
out.append(text[cursor:])
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _tool_call_markup_spans(text: str) -> list[tuple[int, int]]:
|
||||
"""Spans of tool-call markup, so a literal <think>/[THINK] inside a call's args is
|
||||
stripped WITH the call, not kept as a reasoning block. Covers closed XML/bracket
|
||||
calls and an unclosed XML call (run via allow_incomplete); without the open-ended
|
||||
span the unclosed call's markup would leak after execution."""
|
||||
# Skip a lazy closed-pair pattern whose close token is absent: its finditer would rescan
|
||||
# to EOF from every opener (quadratic on a stream of unclosed openers).
|
||||
spans = [
|
||||
m.span()
|
||||
for pat in _TOOL_CLOSED_PATS
|
||||
if (_PAT_REQUIRED_TOKEN.get(pat) is None or _PAT_REQUIRED_TOKEN[pat] in text)
|
||||
for m in pat.finditer(text)
|
||||
]
|
||||
spans.extend((start, end) for start, end, _kind, _m in _iter_bracket_spans(text))
|
||||
# An unclosed opener is a real incomplete call only outside closed/bracket spans.
|
||||
for pat in _TOOL_OPEN_XML_TAIL_PATS:
|
||||
for m in pat.finditer(text):
|
||||
if not any(s <= m.start() < e for s, e in spans):
|
||||
spans.append(m.span())
|
||||
return spans
|
||||
|
||||
|
||||
def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]:
|
||||
"""<think>/[THINK] block spans, minus any whose opening marker sits INSIDE a
|
||||
tool-call span (that tag is argument data, not reasoning). Keeping it would drop a
|
||||
real call after it as rehearsed and leak the call's markup. START tested only, so
|
||||
a greedy unclosed <think> past the call is still that call's argument data."""
|
||||
think_spans = [m.span() for m in _THINK_TAG_RE.finditer(text)]
|
||||
call_spans = _tool_call_markup_spans(text)
|
||||
# Prefilled reasoning: the template opens <think> in the prompt, so add a leading span
|
||||
# (0..close) to skip calls rehearsed there; guarded so a stray close in a normal answer is safe.
|
||||
close = _THINK_CLOSE_RE.search(text)
|
||||
if close is not None:
|
||||
opener = _THINK_OPEN_RE.search(text)
|
||||
if (
|
||||
(opener is None or close.start() < opener.start())
|
||||
and not any(cs <= close.start() < ce for cs, ce in call_spans)
|
||||
and any(cs >= close.end() for cs, ce in call_spans)
|
||||
):
|
||||
think_spans = [(0, close.end())] + think_spans
|
||||
if not think_spans:
|
||||
return think_spans
|
||||
if not call_spans:
|
||||
return think_spans
|
||||
return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)]
|
||||
|
||||
|
||||
def strip_outside_think(text: str, strip_segment) -> str:
|
||||
"""Apply ``strip_segment(segment, is_last)`` to visible text around <think>/[THINK]
|
||||
blocks, preserving the blocks verbatim (tool-looking text inside is rehearsal).
|
||||
``is_last`` is True only after the final block, so trailing-tail patterns apply
|
||||
only there. Shared by every strip path so they stay consistent."""
|
||||
# A think marker opening inside a complete call is argument text; excluding it lets the
|
||||
# stripper see the whole call. START-tested, so an unclosed match stays argument data.
|
||||
think_spans = _think_spans_outside_tool_markup(text)
|
||||
if not think_spans:
|
||||
return strip_segment(text, True)
|
||||
pieces: list[str] = []
|
||||
prev = 0
|
||||
for s, e in think_spans:
|
||||
pieces.append(strip_segment(text[prev:s], False))
|
||||
pieces.append(text[s:e])
|
||||
prev = e
|
||||
pieces.append(strip_segment(text[prev:], True))
|
||||
return "".join(pieces)
|
||||
|
||||
|
||||
def _strip_gemma_native_spans(text: str, *, final: bool) -> str:
|
||||
"""Remove complete Gemma-native spans, brace/quote-balanced so a literal
|
||||
``<tool_call|>`` in a quoted argument cannot truncate the span. An incomplete
|
||||
span is dropped to EOF when ``final``, else kept (still streaming)."""
|
||||
out: list[str] = []
|
||||
cursor = 0
|
||||
for match in _TC_GEMMA_START_RE.finditer(text):
|
||||
start = match.start()
|
||||
if start < cursor:
|
||||
continue
|
||||
brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
|
||||
if brace_end < 0:
|
||||
# Unbalanced: nothing completes from here on. Drop the rest if final,
|
||||
# else keep it; stop either way (rescanning would be quadratic).
|
||||
if final:
|
||||
out.append(text[cursor:start])
|
||||
cursor = len(text)
|
||||
break
|
||||
# Junk between } and <tool_call|> is malformed-call markup: strip through
|
||||
# the close, keep text after it. No close anywhere means stop (linear).
|
||||
close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
|
||||
if close is None:
|
||||
if final:
|
||||
out.append(text[cursor:start])
|
||||
cursor = len(text)
|
||||
break
|
||||
out.append(text[cursor:start])
|
||||
cursor = close.end()
|
||||
out.append(text[cursor:])
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _gemma_span_ranges(text: str) -> list:
|
||||
"""``(start, end)`` of each complete Gemma-native span; same walk as
|
||||
``_strip_gemma_native_spans`` without stripping."""
|
||||
ranges: list[tuple] = []
|
||||
cursor = 0
|
||||
for match in _TC_GEMMA_START_RE.finditer(text):
|
||||
start = match.start()
|
||||
if start < cursor:
|
||||
continue
|
||||
brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
|
||||
if brace_end < 0:
|
||||
break
|
||||
close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
|
||||
if close is None:
|
||||
break
|
||||
ranges.append((start, close.end()))
|
||||
cursor = close.end()
|
||||
return ranges
|
||||
|
||||
|
||||
def _strip_closed_blocks_outside_gemma(text: str) -> str:
|
||||
"""Closed JSON/function pre-pass that skips matches starting inside a complete
|
||||
Gemma span: deleting across the span boundary would mangle the Gemma close and
|
||||
truncate the tail. A skipped match resumes at the covering span's end, so a
|
||||
real function-XML call after the span is still stripped."""
|
||||
ranges = _gemma_span_ranges(text)
|
||||
if not ranges:
|
||||
return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS)
|
||||
for pat in _TOOL_CLOSED_BLOCK_PATS:
|
||||
token = _PAT_REQUIRED_TOKEN.get(pat)
|
||||
if token is not None and token not in text:
|
||||
continue
|
||||
out: list[str] = []
|
||||
pos = 0
|
||||
while True:
|
||||
m = pat.search(text, pos)
|
||||
if m is None:
|
||||
out.append(text[pos:])
|
||||
break
|
||||
covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None)
|
||||
if covering is not None:
|
||||
out.append(text[pos : covering[1]])
|
||||
pos = covering[1]
|
||||
continue
|
||||
out.append(text[pos : m.start()])
|
||||
pos = m.end()
|
||||
new_text = "".join(out)
|
||||
if new_text != text:
|
||||
text = new_text
|
||||
ranges = _gemma_span_ranges(text)
|
||||
return text
|
||||
|
||||
|
||||
def _strip_markup_segment(
|
||||
text: str,
|
||||
*,
|
||||
final: bool,
|
||||
enabled_tool_names = None,
|
||||
) -> str:
|
||||
# Bracket-tag calls (Mistral/rehearsal) first via balanced scan (any nesting depth,
|
||||
# rehearsal name-gated); then the quote-aware Gemma-native passes so a literal
|
||||
# <tool_call|> in an argument cannot truncate a block; finally the regex XML/tail sweeps.
|
||||
text = _strip_bracket_tag_calls(text, enabled_tool_names = enabled_tool_names)
|
||||
text = _strip_closed_blocks_outside_gemma(text)
|
||||
text = _strip_gemma_native_spans(text, final = final)
|
||||
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
|
||||
return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names)
|
||||
|
||||
|
||||
def strip_tool_call_markup(
|
||||
text: str,
|
||||
*,
|
||||
final: bool = False,
|
||||
enabled_tool_names = None,
|
||||
) -> str:
|
||||
"""Strip tool-call XML markup from text.
|
||||
|
||||
When ``final`` is False, only fully closed tool-call blocks are removed.
|
||||
When ``final`` is True, trailing incomplete tool-call blocks are removed
|
||||
too, and the result is stripped of surrounding whitespace.
|
||||
|
||||
``<think>`` / ``[THINK]`` reasoning is preserved verbatim (see
|
||||
``strip_outside_think``); the trailing-tail patterns apply only after the
|
||||
last block. ``enabled_tool_names`` keeps an inactive-name ``foo[ARGS]{..}``
|
||||
example visible (it is prose, not a call) so display cleanup matches detection.
|
||||
"""
|
||||
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
|
||||
for pat in patterns:
|
||||
text = pat.sub("", text)
|
||||
return text.strip() if final else text
|
||||
result = strip_outside_think(
|
||||
text,
|
||||
lambda seg, is_last: _strip_markup_segment(
|
||||
seg, final = final and is_last, enabled_tool_names = enabled_tool_names
|
||||
),
|
||||
)
|
||||
return result.strip() if final else result
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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-<step>/`` partials after a cancel.
|
||||
|
||||
Completed ``checkpoint-<int>/`` 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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -11,8 +11,9 @@ cache lifecycle. Two subdirectories:
|
|||
manifests/ <key>.json per-download expected-files manifest
|
||||
cancelled/ <key>.json per-download cancel marker
|
||||
|
||||
The ``<key>`` mirrors HF's cache dir naming so a state file can be
|
||||
eyeballed next to the on-disk repo it describes:
|
||||
The ``<key>`` 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--<owner>--<name> full snapshot
|
||||
models--<owner>--<name>--variant--<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 ".<target>.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()
|
||||
|
|
|
|||
|
|
@ -177,6 +177,7 @@ class GenerateRequest(BaseModel):
|
|||
temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature")
|
||||
top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling")
|
||||
top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling")
|
||||
min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p sampling")
|
||||
max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate")
|
||||
repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty")
|
||||
presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty")
|
||||
|
|
@ -1146,7 +1147,8 @@ class CompletionMessage(BaseModel):
|
|||
"""The assistant's complete response message."""
|
||||
|
||||
role: Literal["assistant"] = "assistant"
|
||||
content: str
|
||||
# ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise.
|
||||
content: Optional[str] = None
|
||||
refusal: Optional[str] = None
|
||||
reasoning_content: Optional[str] = None
|
||||
tool_calls: Optional[list[dict]] = None
|
||||
|
|
@ -1531,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,
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -1581,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.
|
||||
|
|
|
|||
|
|
@ -10,3 +10,9 @@ transformers>=4.57.6
|
|||
# anyio that also ImportErrors on TaskHandle and 500s the server. An override
|
||||
# wins the fight, so force one consistent <4.14 here too.
|
||||
anyio<4.14.0
|
||||
|
||||
# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights
|
||||
# rejects q_norm/k_norm, so those checkpoints fail to load. mlx-lm #1242.
|
||||
# The override also drops it from transitive resolution; keep the >=0.22.0 floor
|
||||
# (mirrors mlx_repair.py _MLX_MIN_VERSIONS) or the resolver could go below it.
|
||||
mlx-lm>=0.22.0,!=0.31.3
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ class ChatThread(BaseModel):
|
|||
projectId: Optional[str] = None
|
||||
archived: bool = False
|
||||
createdAt: int
|
||||
updatedAt: Optional[int] = None
|
||||
openaiCodeExecContainerId: Optional[str] = None
|
||||
anthropicCodeExecContainerId: Optional[str] = None
|
||||
forkedFromThreadId: Optional[str] = None
|
||||
|
|
@ -70,6 +71,7 @@ class ChatThreadPatch(BaseModel):
|
|||
projectId: Optional[str] = None
|
||||
archived: Optional[bool] = None
|
||||
createdAt: Optional[int] = None
|
||||
updatedAt: Optional[int] = None
|
||||
openaiCodeExecContainerId: Optional[str] = None
|
||||
anthropicCodeExecContainerId: Optional[str] = None
|
||||
|
||||
|
|
@ -177,6 +179,7 @@ class ChatSettingsPayload(BaseModel):
|
|||
collapseHtmlArtifacts: Optional[bool] = None
|
||||
allowArtifactNetworkAccess: Optional[bool] = None
|
||||
autoHealToolCalls: Optional[bool] = None
|
||||
nudgeToolCalls: Optional[bool] = None
|
||||
maxToolCallsPerMessage: Optional[int] = Field(default = None, ge = 1)
|
||||
toolCallTimeout: Optional[int] = Field(default = None, ge = 1)
|
||||
|
||||
|
|
@ -251,7 +254,7 @@ async def patch_thread(
|
|||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
patch = payload.model_dump(exclude_unset = True)
|
||||
for field in ("title", "modelType", "modelId", "archived", "createdAt"):
|
||||
for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"):
|
||||
if field in patch and patch[field] is None:
|
||||
raise HTTPException(status_code = 400, detail = f"{field} cannot be null")
|
||||
if patch.get("projectId") and get_chat_project(patch["projectId"]) is None:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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),
|
||||
|
|
@ -260,17 +274,29 @@ def _embedding_model_response() -> EmbeddingModelResponse:
|
|||
)
|
||||
|
||||
|
||||
def _llama_backend_active() -> bool:
|
||||
"""True when this install embeds via the llama-server (GGUF) backend."""
|
||||
from core.rag import config as rag_config
|
||||
from core.rag import embeddings
|
||||
|
||||
def _ambient_hf_token() -> Optional[str]:
|
||||
"""The HF token the loader would use (HF_TOKEN env or the cached login), so a gated
|
||||
repo is scanned rather than failing open. None if unavailable."""
|
||||
try:
|
||||
raw = (rag_config.EMBED_BACKEND or "auto").strip().lower()
|
||||
key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw
|
||||
from huggingface_hub import get_token
|
||||
return get_token()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _llama_backend_active() -> bool:
|
||||
"""True when this install actually embeds via the llama-server (GGUF) backend.
|
||||
|
||||
Delegates to the embeddings module so a runtime fallback from
|
||||
sentence-transformers to llama-server (after a torch/CUDA load or encode
|
||||
failure) is honored: in that state the process loads only inert GGUF, so the
|
||||
ST pickle gate below must not hard-block a repo whose GGUF companion is clean.
|
||||
Before any backend is built this still reflects the resolver."""
|
||||
from core.rag import embeddings
|
||||
try:
|
||||
return embeddings.active_backend_is_llama()
|
||||
except Exception: # noqa: BLE001 - backend probe must never block saving
|
||||
return False
|
||||
return key in embeddings._LLAMA_ALIASES
|
||||
|
||||
|
||||
def _resolves_as_local_gguf(model: str) -> bool:
|
||||
|
|
@ -357,6 +383,8 @@ def update_embedding_model(
|
|||
"""Set the RAG embedding model. Unless ``force`` is set, the repo is verified
|
||||
to be an embedding model via HF metadata; an unverifiable model (wrong type,
|
||||
typo, gated repo, or no network) returns 409 so the UI can offer "save anyway".
|
||||
A repo flagged unsafe by HF's security scan returns 403 instead: a hard block
|
||||
that ``force`` cannot bypass, so the UI must not offer "save anyway".
|
||||
Documents indexed under the previous model must be re-uploaded."""
|
||||
from utils.models import is_embedding_model
|
||||
|
||||
|
|
@ -370,15 +398,51 @@ def update_embedding_model(
|
|||
event = "settings.update_embedding_model_failed",
|
||||
log = logger,
|
||||
) from exc
|
||||
hf_token = (payload.hf_token or "").strip() or None
|
||||
# The env/default model needs no verification; saving it is a no-op override.
|
||||
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
|
||||
# what the backend loads, and HF metadata cannot verify a local path.
|
||||
if (
|
||||
model != default_embedding_model()
|
||||
and not payload.force
|
||||
and not (_llama_backend_active() and _resolves_as_local_gguf(model))
|
||||
):
|
||||
hf_token = (payload.hf_token or "").strip() or None
|
||||
is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model)
|
||||
# The pickle gate only matters for the sentence-transformers backend, which is what
|
||||
# deserializes pickles. On the llama-server backend the embedder loads GGUF files
|
||||
# (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would
|
||||
# wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability
|
||||
# checks below cover that path instead.
|
||||
scan_st_pickle = (
|
||||
model != default_embedding_model() and not is_local_gguf and not _llama_backend_active()
|
||||
)
|
||||
if scan_st_pickle:
|
||||
# Malware/pickle gate before we persist a repo the embedder later loads with
|
||||
# SentenceTransformer. Runs even under force (force only skips the is-embedding
|
||||
# type check for offline/local repos HF cannot verify); local paths and
|
||||
# unreachable scans fail open inside evaluate_file_security.
|
||||
from utils.security import evaluate_file_security, security_load_subdirs
|
||||
from core.rag.embeddings import _st_module_subdirs
|
||||
|
||||
# Fall back to the loader's own token so a gated/private repo is actually scanned
|
||||
# (a token-less scan fails open for exactly the repo that would still load).
|
||||
scan_token = hf_token or _ambient_hf_token()
|
||||
# Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
|
||||
# one blocks instead of passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys(
|
||||
(
|
||||
*security_load_subdirs(model, scan_token),
|
||||
*_st_module_subdirs(model, scan_token),
|
||||
)
|
||||
)
|
||||
)
|
||||
if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
|
||||
# 403, not 409: the client routes every 409 into the forceable "save anyway"
|
||||
# flow, but this block is a hard, non-forceable security refusal.
|
||||
raise HTTPException(
|
||||
status_code = 403,
|
||||
detail = (
|
||||
f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
|
||||
"cannot be used as the embedding model."
|
||||
),
|
||||
)
|
||||
if model != default_embedding_model() and not payload.force and not is_local_gguf:
|
||||
from core.rag import config as rag_config
|
||||
|
||||
# A GGUF-named repo on the llama-server backend is loaded from its .gguf
|
||||
|
|
|
|||
|
|
@ -12,6 +12,79 @@ import time
|
|||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def _fix_torch_cuda_ld_path():
|
||||
"""Prepend torch's bundled CUDA libs to LD_LIBRARY_PATH.
|
||||
|
||||
PyTorch wheels ship their own CUDA runtime (libcudart, libcublas, ...) in
|
||||
``site-packages/nvidia/*/lib``. On Linux the dynamic linker reads
|
||||
LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a
|
||||
pre-existing LD_LIBRARY_PATH pointing at a different system CUDA (e.g.
|
||||
/usr/local/cuda-13/lib64 from conda or a Docker base image) shadows torch's
|
||||
libs and triggers "undefined symbol" errors when torch is imported. Detect
|
||||
torch's lib dirs (without importing torch) and prepend them. Returns True if
|
||||
LD_LIBRARY_PATH was changed.
|
||||
"""
|
||||
if sys.platform != "linux":
|
||||
return False
|
||||
ld_path = os.environ.get("LD_LIBRARY_PATH", "")
|
||||
if not ld_path:
|
||||
return False
|
||||
try:
|
||||
import importlib.util
|
||||
|
||||
spec = importlib.util.find_spec("torch")
|
||||
if not spec or not spec.origin:
|
||||
return False
|
||||
torch_dir = os.path.dirname(spec.origin)
|
||||
site_pkgs = os.path.dirname(torch_dir)
|
||||
nvidia_dir = os.path.join(site_pkgs, "nvidia")
|
||||
|
||||
lib_dirs = []
|
||||
torch_lib = os.path.join(torch_dir, "lib")
|
||||
if os.path.isdir(torch_lib):
|
||||
lib_dirs.append(torch_lib)
|
||||
if os.path.isdir(nvidia_dir):
|
||||
for sub in sorted(os.listdir(nvidia_dir)):
|
||||
lib = os.path.join(nvidia_dir, sub, "lib")
|
||||
if os.path.isdir(lib):
|
||||
lib_dirs.append(lib)
|
||||
if not lib_dirs:
|
||||
return False
|
||||
|
||||
existing = ld_path.split(":")
|
||||
if existing[: len(lib_dirs)] == lib_dirs:
|
||||
return False # already at the front, nothing to do
|
||||
|
||||
torch_set = set(lib_dirs)
|
||||
cleaned = [p for p in existing if p not in torch_set]
|
||||
os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + cleaned)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
_LD_FIXED_SENTINEL = "_UNSLOTH_STUDIO_LD_FIXED"
|
||||
|
||||
|
||||
def _maybe_reexec_for_cuda_ld_path():
|
||||
"""Re-exec once so the dynamic linker sees the corrected LD_LIBRARY_PATH.
|
||||
|
||||
LD_LIBRARY_PATH is read at process start, so editing os.environ in-process
|
||||
cannot fix the running interpreter; a single re-exec is required. Call only
|
||||
from a true entry point (the ``if __name__ == "__main__"`` block), never at
|
||||
import time, because os.execv replaces the whole process (an embedder such
|
||||
as Colab that does ``from run import run_server`` must not be re-exec'd).
|
||||
"""
|
||||
if _LD_FIXED_SENTINEL in os.environ:
|
||||
return
|
||||
if not _fix_torch_cuda_ld_path():
|
||||
return
|
||||
os.environ[_LD_FIXED_SENTINEL] = "1"
|
||||
argv = getattr(sys, "orig_argv", None) or [sys.executable, *sys.argv]
|
||||
os.execv(sys.executable, argv)
|
||||
|
||||
|
||||
# Suppress C-level dependency warnings globally (e.g. SwigPyPacked).
|
||||
os.environ["PYTHONWARNINGS"] = "ignore"
|
||||
|
||||
|
|
@ -1457,6 +1530,12 @@ def _build_arg_parser():
|
|||
|
||||
# For direct execution (also invoked by CLI via os.execvp / subprocess).
|
||||
if __name__ == "__main__":
|
||||
# Correct a conflicting system CUDA on LD_LIBRARY_PATH before torch is
|
||||
# imported (below, via run_server). Re-execs once on Linux so the dynamic
|
||||
# linker uses torch's bundled CUDA libs; no-op on other platforms, when
|
||||
# LD_LIBRARY_PATH is unset or already correct, or after the single re-exec.
|
||||
_maybe_reexec_for_cuda_ld_path()
|
||||
|
||||
import signal
|
||||
import traceback
|
||||
|
||||
|
|
|
|||
|
|
@ -240,6 +240,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
project_id TEXT,
|
||||
archived INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER,
|
||||
openai_code_exec_container_id TEXT,
|
||||
anthropic_code_exec_container_id TEXT,
|
||||
forked_from_thread_id TEXT,
|
||||
|
|
@ -261,6 +262,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
|
|||
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT")
|
||||
if "forked_from_message_id" not in chat_thread_cols:
|
||||
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT")
|
||||
if "updated_at" not in chat_thread_cols:
|
||||
conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER")
|
||||
# Floor at created_at: forked threads copy older ancestor messages,
|
||||
# so the fork's creation time must win over the branch message times.
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE chat_threads SET updated_at = MAX(
|
||||
COALESCE(
|
||||
(
|
||||
SELECT MAX(m.created_at) FROM chat_messages m
|
||||
WHERE m.thread_id = chat_threads.id
|
||||
),
|
||||
created_at
|
||||
),
|
||||
created_at
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS chat_messages (
|
||||
|
|
@ -992,6 +1011,9 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict:
|
|||
"projectId": data.get("project_id") or None,
|
||||
"archived": bool(data["archived"]),
|
||||
"createdAt": data["created_at"],
|
||||
"updatedAt": data.get("updated_at")
|
||||
if data.get("updated_at") is not None
|
||||
else data["created_at"],
|
||||
"openaiCodeExecContainerId": data.get("openai_code_exec_container_id"),
|
||||
"anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"),
|
||||
"forkedFromThreadId": data.get("forked_from_thread_id"),
|
||||
|
|
@ -1039,8 +1061,8 @@ def upsert_chat_thread(thread: dict) -> dict:
|
|||
conn.execute(
|
||||
"""
|
||||
INSERT INTO chat_threads
|
||||
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
title = excluded.title,
|
||||
model_type = excluded.model_type,
|
||||
|
|
@ -1049,6 +1071,7 @@ def upsert_chat_thread(thread: dict) -> dict:
|
|||
project_id = excluded.project_id,
|
||||
archived = excluded.archived,
|
||||
created_at = excluded.created_at,
|
||||
updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at),
|
||||
openai_code_exec_container_id = excluded.openai_code_exec_container_id,
|
||||
anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id,
|
||||
forked_from_thread_id = excluded.forked_from_thread_id,
|
||||
|
|
@ -1063,6 +1086,7 @@ def upsert_chat_thread(thread: dict) -> dict:
|
|||
thread.get("projectId"),
|
||||
1 if thread.get("archived") else 0,
|
||||
int(thread["createdAt"]),
|
||||
int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None,
|
||||
thread.get("openaiCodeExecContainerId"),
|
||||
thread.get("anthropicCodeExecContainerId"),
|
||||
thread.get("forkedFromThreadId"),
|
||||
|
|
@ -1084,6 +1108,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]:
|
|||
"projectId": ("project_id", patch.get("projectId")),
|
||||
"archived": ("archived", 1 if patch.get("archived") else 0),
|
||||
"createdAt": ("created_at", patch.get("createdAt")),
|
||||
"updatedAt": ("updated_at", patch.get("updatedAt")),
|
||||
"openaiCodeExecContainerId": (
|
||||
"openai_code_exec_container_id",
|
||||
patch.get("openaiCodeExecContainerId"),
|
||||
|
|
@ -1155,7 +1180,8 @@ def list_chat_threads(
|
|||
conn = get_connection()
|
||||
try:
|
||||
rows = conn.execute(
|
||||
f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC",
|
||||
f"SELECT * FROM chat_threads {where} "
|
||||
"ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC",
|
||||
values,
|
||||
).fetchall()
|
||||
return [_chat_thread_from_row(row) for row in rows]
|
||||
|
|
@ -1394,6 +1420,44 @@ def _raise_if_chat_message_thread_conflicts(
|
|||
)
|
||||
|
||||
|
||||
def _bump_chat_thread_updated_at(
|
||||
conn: sqlite3.Connection, thread_id: str, message_created_at: int
|
||||
) -> None:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE chat_threads
|
||||
SET updated_at = MAX(COALESCE(updated_at, created_at), ?)
|
||||
WHERE id = ?
|
||||
""",
|
||||
(message_created_at, thread_id),
|
||||
)
|
||||
|
||||
|
||||
def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None:
|
||||
"""Set updated_at from the remaining messages, floored at created_at.
|
||||
|
||||
Unlike the ratchet-only bump, this can lower updated_at -- needed after
|
||||
pruning, which may delete the thread's newest message.
|
||||
"""
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE chat_threads
|
||||
SET updated_at = MAX(
|
||||
COALESCE(
|
||||
(
|
||||
SELECT MAX(m.created_at) FROM chat_messages m
|
||||
WHERE m.thread_id = chat_threads.id
|
||||
),
|
||||
created_at
|
||||
),
|
||||
created_at
|
||||
)
|
||||
WHERE id = ?
|
||||
""",
|
||||
(thread_id,),
|
||||
)
|
||||
|
||||
|
||||
def upsert_chat_message(message: dict) -> dict:
|
||||
conn = get_connection()
|
||||
try:
|
||||
|
|
@ -1432,6 +1496,7 @@ def upsert_chat_message(message: dict) -> dict:
|
|||
int(message["createdAt"]),
|
||||
),
|
||||
)
|
||||
_bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"]))
|
||||
conn.commit()
|
||||
return message
|
||||
except Exception:
|
||||
|
|
@ -1484,6 +1549,12 @@ def sync_chat_messages(
|
|||
for m in messages
|
||||
],
|
||||
)
|
||||
if prune_missing:
|
||||
_recompute_chat_thread_updated_at(conn, thread_id)
|
||||
elif messages:
|
||||
_bump_chat_thread_updated_at(
|
||||
conn, thread_id, max(int(m["createdAt"]) for m in messages)
|
||||
)
|
||||
conn.commit()
|
||||
return list_chat_messages(thread_id)
|
||||
except ChatMessageConflictError:
|
||||
|
|
|
|||
|
|
@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO
|
|||
from types import SimpleNamespace
|
||||
|
||||
|
||||
def _emitter_client_text(events: list[str]) -> str:
|
||||
"""Concatenate the text_delta payloads an SSE event list carries."""
|
||||
text = ""
|
||||
for line in events:
|
||||
for raw in line.split("\n"):
|
||||
raw = raw.strip()
|
||||
if not raw.startswith("data: "):
|
||||
continue
|
||||
data = json.loads(raw[len("data: ") :])
|
||||
delta = data.get("delta", {})
|
||||
if delta.get("type") == "text_delta":
|
||||
text += delta.get("text", "")
|
||||
return text
|
||||
|
||||
|
||||
def test_anthropic_emitter_closes_reasoning_only_think_block():
|
||||
# A reasoning-only reply streams <think>X live then shrinks to bare X at EOF.
|
||||
# This emitter diffs cumulative snapshots and drops the shrink, so without a
|
||||
# closing pass the client text would end on an unclosed <think>. finish()
|
||||
# must balance it.
|
||||
emitter = AnthropicStreamEmitter()
|
||||
events = emitter.start("msg_1", "m")
|
||||
events += emitter.feed({"type": "content", "text": "<think>The capital"})
|
||||
events += emitter.feed({"type": "content", "text": "<think>The capital of France is Paris."})
|
||||
# The generator's final bare-text shrink (dropped by the cumulative diff).
|
||||
events += emitter.feed({"type": "content", "text": "The capital of France is Paris."})
|
||||
events += emitter.finish()
|
||||
|
||||
assert _emitter_client_text(events) == "<think>The capital of France is Paris.</think>"
|
||||
|
||||
|
||||
def test_anthropic_emitter_does_not_double_close_balanced_think():
|
||||
# A reasoning-then-answer reply already closes its own </think>; the balancer
|
||||
# must not append a second one.
|
||||
emitter = AnthropicStreamEmitter()
|
||||
events = emitter.start("msg_1", "m")
|
||||
events += emitter.feed({"type": "content", "text": "<think>Thinking."})
|
||||
events += emitter.feed({"type": "content", "text": "<think>Thinking.</think>Answer."})
|
||||
events += emitter.finish()
|
||||
|
||||
assert _emitter_client_text(events) == "<think>Thinking.</think>Answer."
|
||||
|
||||
|
||||
def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch):
|
||||
import routes.inference as inf_mod
|
||||
|
||||
|
|
@ -889,6 +932,24 @@ class TestAnthropicToolNonStreaming:
|
|||
assert tool_blocks[0]["name"] == "render_html"
|
||||
assert tool_blocks[0]["input"] == {"code": "<!doctype html><html></html>"}
|
||||
|
||||
def test_display_strip_gates_on_declared_tools(self):
|
||||
# A final answer containing NAME[ARGS]{json} is gated on the declared tools: undeclared
|
||||
# ``foo`` markup is prose and survives, the declared web_search rehearsal strips.
|
||||
def _run_gen():
|
||||
yield {
|
||||
"type": "content",
|
||||
"text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.',
|
||||
}
|
||||
|
||||
tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}]
|
||||
response = asyncio.run(
|
||||
_anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools)
|
||||
)
|
||||
body = json.loads(response.body)
|
||||
text = "".join(b["text"] for b in body["content"] if b["type"] == "text")
|
||||
assert 'foo[ARGS]{"x": 1}' in text # inactive name preserved as prose
|
||||
assert "web_search[ARGS]" not in text # active name stripped from display
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Pass-through emitter tests (client-side tool execution path)
|
||||
|
|
@ -1709,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
|
||||
|
|
|
|||
194
studio/backend/tests/test_chat_eos_template_refresh.py
Normal file
194
studio/backend/tests/test_chat_eos_template_refresh.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Mapper models whose own tokenizer ships no chat_template have their turn-end
|
||||
eos resolved at LOAD from an empty template (document eos only). The effective
|
||||
template is installed later, at generate time, via get_chat_template, so the
|
||||
turn-end-eos cache must be refreshed then; otherwise generate_stream runs past
|
||||
the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes).
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND = Path(__file__).resolve().parent.parent
|
||||
if str(_BACKEND) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND))
|
||||
|
||||
# These tests construct InferenceBackend, pulling the full stack. CI may lack
|
||||
# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup
|
||||
# (RuntimeError); skip at module level so collection is not aborted (exit 2).
|
||||
try:
|
||||
from core.inference import inference as inf_mod # noqa: E402
|
||||
from core.inference.inference import InferenceBackend # noqa: E402
|
||||
except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent
|
||||
pytest.skip(
|
||||
f"full inference backend unavailable ({type(exc).__name__}: {exc})",
|
||||
allow_module_level = True,
|
||||
)
|
||||
|
||||
_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
|
||||
_GEMMA = "{% for m in messages %}<start_of_turn>{{m.role}}\n{{m.content}}<end_of_turn>{% endfor %}"
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
def __init__(
|
||||
self,
|
||||
eos_id,
|
||||
chat_template = "",
|
||||
token_ids = None,
|
||||
):
|
||||
self.eos_token_id = eos_id
|
||||
self.chat_template = chat_template
|
||||
self.pad_token_id = eos_id
|
||||
self.unk_token_id = None
|
||||
self._ids = dict(token_ids or {})
|
||||
|
||||
def convert_tokens_to_ids(self, tok):
|
||||
return self._ids.get(tok)
|
||||
|
||||
|
||||
def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch):
|
||||
import utils.datasets as ds
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "unsloth/qwen2.5-0.5b"
|
||||
|
||||
# No chat_template at load, so the cache stored only the document eos, though
|
||||
# <|im_end|> is atomic in the vocab (unused until the mapper installs a template).
|
||||
bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645})
|
||||
model_info = {
|
||||
"tokenizer": bare_tok,
|
||||
"is_vision": False,
|
||||
"chat_turn_end_eos_ids": [151643],
|
||||
}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
# The mapper installs a ChatML template (turns end with <|im_end|>) at generate time.
|
||||
templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645})
|
||||
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok)
|
||||
monkeypatch.setattr(
|
||||
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False
|
||||
)
|
||||
|
||||
# Stub the tail so the generator runs through the refresh without a real model.
|
||||
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
|
||||
monkeypatch.setattr(
|
||||
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
|
||||
)
|
||||
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
|
||||
|
||||
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
|
||||
|
||||
# After the template is applied the cache must include the ChatML turn-end id.
|
||||
assert model_info["chat_turn_end_eos_ids"] == [151643, 151645]
|
||||
|
||||
|
||||
def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch):
|
||||
# Regression: get_chat_template can return a remapped tokenizer (Gemma: <end_of_turn>
|
||||
# folded onto the eos id) while generate_stream re-reads the original. Resolving on
|
||||
# the swap yields a narrower set, so the refresh must UNION, never overwrite.
|
||||
import utils.datasets as ds
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "unsloth/gemma-2b-it"
|
||||
|
||||
# Original tokenizer (used by generate_stream): <end_of_turn>=107 distinct from
|
||||
# eos=1, so the load-time cache resolved to [1, 107].
|
||||
orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 107})
|
||||
model_info = {
|
||||
"tokenizer": orig_tok,
|
||||
"is_vision": False,
|
||||
"chat_turn_end_eos_ids": [1, 107],
|
||||
}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
# Destructively-swapped tokenizer: <end_of_turn> now maps onto eos id 1, so
|
||||
# resolving on it yields only [1] (drops 107).
|
||||
swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 1})
|
||||
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok)
|
||||
monkeypatch.setattr(
|
||||
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False
|
||||
)
|
||||
|
||||
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
|
||||
monkeypatch.setattr(
|
||||
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
|
||||
)
|
||||
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
|
||||
|
||||
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
|
||||
|
||||
# The load-time <end_of_turn>=107 must survive: overwriting with the swapped
|
||||
# [1] would regress and loop past the turn.
|
||||
assert model_info["chat_turn_end_eos_ids"] == [1, 107]
|
||||
|
||||
|
||||
def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch):
|
||||
# Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but
|
||||
# get_chat_template folds it onto the doc-eos id. generate_stream uses the original,
|
||||
# so read marker strings from the mapped template but ids from the original.
|
||||
import utils.datasets as ds
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "01-ai/yi-6b"
|
||||
|
||||
# Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7.
|
||||
orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
|
||||
model_info = {
|
||||
"tokenizer": orig_tok,
|
||||
"is_vision": False,
|
||||
"chat_turn_end_eos_ids": [2],
|
||||
}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
# Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2.
|
||||
remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
|
||||
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok)
|
||||
monkeypatch.setattr(
|
||||
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False
|
||||
)
|
||||
|
||||
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
|
||||
monkeypatch.setattr(
|
||||
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
|
||||
)
|
||||
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
|
||||
|
||||
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
|
||||
|
||||
# The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2.
|
||||
assert model_info["chat_turn_end_eos_ids"] == [2, 7]
|
||||
|
||||
|
||||
class _FakeProcessor:
|
||||
"""A ProcessorMixin-like container: carries the chat_template itself and
|
||||
wraps the real text tokenizer as ``.tokenizer`` (the vision layout)."""
|
||||
|
||||
def __init__(self, chat_template, tokenizer):
|
||||
self.chat_template = chat_template
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
|
||||
def test_resolve_chat_eos_reads_vision_processor_template():
|
||||
# Vision model: the chat_template lives on the processor while the inner tokenizer
|
||||
# ships none. _resolve_chat_eos must read the marker from the processor but resolve
|
||||
# its id on the inner tokenizer, and repair generation_config.
|
||||
from types import SimpleNamespace
|
||||
|
||||
inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"<end_of_turn>": 107})
|
||||
processor = _FakeProcessor(_GEMMA, inner_tok)
|
||||
model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1))
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "unsloth/gemma-3-4b-it"
|
||||
model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
backend._resolve_chat_eos(backend.active_model_name)
|
||||
|
||||
assert model_info["chat_turn_end_eos_ids"] == [1, 107]
|
||||
# generation_config repaired so the vision .generate() path stops at the turn.
|
||||
assert model.generation_config.eos_token_id == [1, 107]
|
||||
|
|
@ -91,6 +91,17 @@ def test_chat_settings_payload_accepts_fast_mode_presets():
|
|||
assert dumped["customPresets"][0]["params"]["fastMode"] is True
|
||||
|
||||
|
||||
def test_chat_settings_payload_accepts_nudge_tool_calls():
|
||||
# extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the
|
||||
# frontend's persisted nudgeToolCalls needs a payload field (like
|
||||
# autoHealToolCalls).
|
||||
payload = chat_history.ChatSettingsPayload.model_validate(
|
||||
{"autoHealToolCalls": True, "nudgeToolCalls": False}
|
||||
)
|
||||
dumped = payload.model_dump(exclude_unset = True)
|
||||
assert dumped == {"autoHealToolCalls": True, "nudgeToolCalls": False}
|
||||
|
||||
|
||||
def test_chat_inference_settings_covers_frontend_persisted_fields():
|
||||
# Drift guard: every InferenceParams field the UI persists (all but
|
||||
# checkpoint) must exist on ChatInferenceSettings, else extra="forbid"
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import sqlite3
|
||||
import threading
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
|
@ -11,6 +12,7 @@ from pathlib import Path
|
|||
import pytest
|
||||
|
||||
from storage import studio_db
|
||||
from utils.paths import studio_db_path
|
||||
|
||||
|
||||
def _reset_studio_db(
|
||||
|
|
@ -108,6 +110,138 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch):
|
|||
assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}]
|
||||
|
||||
|
||||
def test_chat_thread_updated_at_bumps_on_message_writes(tmp_path, monkeypatch):
|
||||
_reset_studio_db(tmp_path, monkeypatch)
|
||||
thread = studio_db.upsert_chat_thread(_thread())
|
||||
assert thread["updatedAt"] == thread["createdAt"]
|
||||
|
||||
studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi"))
|
||||
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
|
||||
|
||||
studio_db.upsert_chat_message(_message("msg-0", 1_600_000_000_000, "old"))
|
||||
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
|
||||
|
||||
studio_db.sync_chat_messages(
|
||||
"thread-1",
|
||||
[_message("msg-2", 1_700_000_001_000, "newer")],
|
||||
)
|
||||
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000
|
||||
|
||||
|
||||
def test_chat_thread_updated_at_recomputed_when_pruning(tmp_path, monkeypatch):
|
||||
_reset_studio_db(tmp_path, monkeypatch)
|
||||
thread = studio_db.upsert_chat_thread(_thread())
|
||||
studio_db.sync_chat_messages(
|
||||
"thread-1",
|
||||
[
|
||||
_message("msg-1", 1_700_000_000_500, "older"),
|
||||
_message("msg-2", 1_700_000_001_000, "newest"),
|
||||
],
|
||||
prune_missing = True,
|
||||
)
|
||||
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000
|
||||
|
||||
# Pruning the newest message must lower updated_at to the remaining one.
|
||||
studio_db.sync_chat_messages(
|
||||
"thread-1",
|
||||
[_message("msg-1", 1_700_000_000_500, "older")],
|
||||
prune_missing = True,
|
||||
)
|
||||
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
|
||||
|
||||
# Pruning every message falls back to created_at.
|
||||
studio_db.sync_chat_messages("thread-1", [], prune_missing = True)
|
||||
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == thread["createdAt"]
|
||||
|
||||
|
||||
def test_chat_thread_updated_at_survives_thread_resave(tmp_path, monkeypatch):
|
||||
_reset_studio_db(tmp_path, monkeypatch)
|
||||
studio_db.upsert_chat_thread(_thread())
|
||||
studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi"))
|
||||
|
||||
studio_db.upsert_chat_thread(_thread())
|
||||
assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500
|
||||
|
||||
|
||||
def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch):
|
||||
_reset_studio_db(tmp_path, monkeypatch)
|
||||
older = _thread("thread-old")
|
||||
older["createdAt"] = 1_700_000_000_000
|
||||
newer = _thread("thread-new")
|
||||
newer["createdAt"] = 1_700_000_100_000
|
||||
studio_db.upsert_chat_thread(older)
|
||||
studio_db.upsert_chat_thread(newer)
|
||||
assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"]
|
||||
|
||||
studio_db.upsert_chat_message(
|
||||
_message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old")
|
||||
)
|
||||
assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"]
|
||||
|
||||
|
||||
def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch):
|
||||
_reset_studio_db(tmp_path, monkeypatch)
|
||||
db_path = studio_db_path()
|
||||
db_path.parent.mkdir(parents = True, exist_ok = True)
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
try:
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE chat_threads (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
title TEXT NOT NULL,
|
||||
model_type TEXT NOT NULL,
|
||||
model_id TEXT,
|
||||
pair_id TEXT,
|
||||
archived INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE chat_messages (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
thread_id TEXT NOT NULL,
|
||||
parent_id TEXT,
|
||||
role TEXT NOT NULL,
|
||||
content_json TEXT NOT NULL,
|
||||
attachments_json TEXT,
|
||||
metadata_json TEXT,
|
||||
created_at INTEGER NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)",
|
||||
("thread-with-msgs", "Old", "base", 1_700_000_000_000),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)",
|
||||
("thread-empty", "Empty", "base", 1_700_000_050_000),
|
||||
)
|
||||
# Fork-like thread: copied ancestor messages predate the thread itself.
|
||||
conn.execute(
|
||||
"INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)",
|
||||
("thread-fork", "Fork", "base", 1_700_000_100_000),
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT INTO chat_messages (id, thread_id, role, content_json, created_at) VALUES (?, ?, ?, ?, ?)",
|
||||
[
|
||||
("m1", "thread-with-msgs", "user", "[]", 1_700_000_001_000),
|
||||
("m2", "thread-with-msgs", "assistant", "[]", 1_700_000_002_000),
|
||||
("m3", "thread-fork", "user", "[]", 1_700_000_001_000),
|
||||
],
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000
|
||||
assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000
|
||||
assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000
|
||||
|
||||
|
||||
def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch):
|
||||
_reset_studio_db(tmp_path, monkeypatch)
|
||||
project = studio_db.upsert_chat_project(_project())
|
||||
|
|
|
|||
157
studio/backend/tests/test_chat_template_tool_arguments.py
Normal file
157
studio/backend/tests/test_chat_template_tool_arguments.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""apply_chat_template_for_generation must coerce assistant tool_call arguments
|
||||
from the OpenAI JSON-string form to a dict before rendering. Strict tool
|
||||
templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and
|
||||
raise "Can only get item pairs from a mapping." on the string form when a prior
|
||||
tool call is re-rendered on the next turn (MLX + transformers paths).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND = Path(__file__).resolve().parent.parent
|
||||
if str(_BACKEND) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND))
|
||||
|
||||
from core.inference.chat_template_helpers import ( # noqa: E402
|
||||
_normalize_tool_call_arguments,
|
||||
apply_chat_template_for_generation,
|
||||
)
|
||||
|
||||
|
||||
def _conv(arguments):
|
||||
return [
|
||||
{"role": "user", "content": "weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"id": "c1",
|
||||
"function": {"name": "web_search", "arguments": arguments},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "name": "web_search", "content": "21C sunny"},
|
||||
]
|
||||
|
||||
|
||||
class _StrictTemplateTokenizer:
|
||||
"""Mimics a strict Qwen tool template: rejects string tool_call arguments."""
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
*,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kw,
|
||||
):
|
||||
for msg in messages:
|
||||
for call in msg.get("tool_calls", []) or []:
|
||||
args = call.get("function", {}).get("arguments")
|
||||
if isinstance(args, str):
|
||||
raise TypeError("Can only get item pairs from a mapping.")
|
||||
return "RENDERED"
|
||||
|
||||
|
||||
def test_string_arguments_are_parsed_to_dict():
|
||||
out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}'))
|
||||
args = out[1]["tool_calls"][0]["function"]["arguments"]
|
||||
assert args == {"query": "sweden"}
|
||||
|
||||
|
||||
def test_dict_arguments_untouched_and_no_copy():
|
||||
conv = _conv({"query": "sweden"})
|
||||
assert _normalize_tool_call_arguments(conv) is conv
|
||||
|
||||
|
||||
def test_non_json_string_left_as_is():
|
||||
out = _normalize_tool_call_arguments(_conv("not json"))
|
||||
assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json"
|
||||
|
||||
|
||||
def test_render_succeeds_on_strict_template_with_string_arguments():
|
||||
# Regression: strict template + string args used to raise.
|
||||
result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}'))
|
||||
assert result == "RENDERED"
|
||||
|
||||
|
||||
class _RecordingTokenizer:
|
||||
"""Lenient template: renders whatever arguments it is given (string or dict)."""
|
||||
|
||||
def __init__(self):
|
||||
self.seen_arguments = None
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
*,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kw,
|
||||
):
|
||||
for msg in messages:
|
||||
for call in msg.get("tool_calls", []) or []:
|
||||
self.seen_arguments = call.get("function", {}).get("arguments")
|
||||
return "RENDERED"
|
||||
|
||||
|
||||
def test_lenient_template_receives_original_string_untouched():
|
||||
# Lenient template must see the exact original string, not a coerced dict.
|
||||
tok = _RecordingTokenizer()
|
||||
apply_chat_template_for_generation(tok, _conv('{"query": "x"}'))
|
||||
assert tok.seen_arguments == '{"query": "x"}'
|
||||
|
||||
|
||||
def test_messages_without_tool_calls_pass_through_unchanged():
|
||||
conv = [{"role": "user", "content": "hi"}]
|
||||
assert _normalize_tool_call_arguments(conv) is conv
|
||||
|
||||
|
||||
class _RaiseExceptionTemplateTokenizer:
|
||||
"""Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via
|
||||
``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError."""
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
*,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kw,
|
||||
):
|
||||
for msg in messages:
|
||||
for call in msg.get("tool_calls", []) or []:
|
||||
args = call.get("function", {}).get("arguments")
|
||||
if isinstance(args, str):
|
||||
raise ValueError(
|
||||
"chat_template: tool_calls[].function.arguments must be a "
|
||||
"JSON object (mapping), not a string."
|
||||
)
|
||||
return "RENDERED"
|
||||
|
||||
|
||||
def test_render_succeeds_on_raise_exception_template_with_string_arguments():
|
||||
# Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce.
|
||||
result = apply_chat_template_for_generation(
|
||||
_RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}')
|
||||
)
|
||||
assert result == "RENDERED"
|
||||
|
||||
|
||||
def test_unrelated_template_error_still_propagates_with_dict_args():
|
||||
# Failure unrelated to string args (dict args, nothing to coerce) must propagate.
|
||||
class _AlwaysRaises:
|
||||
def apply_chat_template(self, messages, **kw):
|
||||
raise ValueError("template is broken")
|
||||
|
||||
with pytest.raises(ValueError, match = "broken"):
|
||||
apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"}))
|
||||
150
studio/backend/tests/test_chat_turn_end_eos.py
Normal file
150
studio/backend/tests/test_chat_turn_end_eos.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and
|
||||
repair generation_config so a chat model whose eos is a bare document terminator
|
||||
(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn
|
||||
boundary instead of running past it and looping. Dependency-light: imported here
|
||||
without the full inference stack.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND = Path(__file__).resolve().parent.parent
|
||||
if str(_BACKEND) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND))
|
||||
|
||||
from core.inference.chat_eos import ( # noqa: E402
|
||||
chat_eos_repair,
|
||||
resolve_chat_turn_end_eos_ids,
|
||||
resolve_chat_turn_end_eos_ids_using,
|
||||
)
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
def __init__(
|
||||
self,
|
||||
eos_id,
|
||||
chat_template = "",
|
||||
token_ids = None,
|
||||
unk_token_id = None,
|
||||
):
|
||||
self.eos_token_id = eos_id
|
||||
self.chat_template = chat_template
|
||||
self.unk_token_id = unk_token_id
|
||||
self._ids = dict(token_ids or {})
|
||||
|
||||
def convert_tokens_to_ids(self, tok):
|
||||
return self._ids.get(tok, self.unk_token_id)
|
||||
|
||||
|
||||
# ---- resolve_chat_turn_end_eos_ids ---------------------------------------
|
||||
|
||||
_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
|
||||
|
||||
|
||||
def test_qwen35_adds_im_end_from_template():
|
||||
# eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046).
|
||||
tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046]
|
||||
|
||||
|
||||
def test_marker_in_vocab_but_not_in_template_is_ignored():
|
||||
# Base/coder model: <|im_end|> is in the vocab but the template does not use
|
||||
# it, so it must not become a stop token.
|
||||
tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [248044]
|
||||
|
||||
|
||||
def test_harmony_template_is_left_untouched():
|
||||
# gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end.
|
||||
harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"
|
||||
tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [200002]
|
||||
|
||||
|
||||
def test_llama3_eot_id_from_template():
|
||||
tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009]
|
||||
|
||||
|
||||
def test_gemma4_turn_marker_from_template():
|
||||
# Gemma-4 ends turns with <turn|> while keeping a document eos, so <turn|> must
|
||||
# be added as a stop token.
|
||||
tok = _FakeTokenizer(
|
||||
1, chat_template = "...<start_of_turn>...<turn|>...", token_ids = {"<turn|>": 106}
|
||||
)
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [1, 106]
|
||||
|
||||
|
||||
def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer():
|
||||
# map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id,
|
||||
# but the original keeps it atomic. Reading marker STRINGS from the template but
|
||||
# IDS on the original recovers the real turn-end id (7), not the doc-eos id (2).
|
||||
template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
|
||||
id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
|
||||
assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7]
|
||||
# Same tokenizer for both reproduces the plain resolve (load-time behaviour).
|
||||
assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2]
|
||||
|
||||
|
||||
def test_list_eos_preserved():
|
||||
tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [1, 2]
|
||||
|
||||
|
||||
def test_missing_marker_maps_to_unk_and_is_skipped():
|
||||
tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0)
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [7]
|
||||
|
||||
|
||||
def test_starling_barred_end_of_turn_from_template():
|
||||
# OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from
|
||||
# Gemma's <end_of_turn>). eos synced to </s>=2, turn marker at 32000.
|
||||
starling = "GPT4 Correct Assistant: hi<|end_of_turn|>"
|
||||
tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000]
|
||||
|
||||
|
||||
def test_dict_chat_template_scans_all_variants():
|
||||
# Hermes-3 style: chat_template is a {name: template} dict. Detection must scan
|
||||
# every variant, not bail because the container is not a plain str.
|
||||
tmpl = {"default": "{{ messages }}", "tool_use": _CHATML}
|
||||
tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
|
||||
|
||||
|
||||
def test_list_of_dicts_chat_template_scans_all_variants():
|
||||
# tokenizer_config.json stores multi-templates as a list of {name, template}.
|
||||
tmpl = [{"name": "default", "template": _CHATML}]
|
||||
tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
|
||||
|
||||
|
||||
def test_dict_harmony_template_left_untouched():
|
||||
# A multi-variant container whose variant is harmony must still be left alone.
|
||||
tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"}
|
||||
tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007})
|
||||
assert resolve_chat_turn_end_eos_ids(tok) == [200002]
|
||||
|
||||
|
||||
# ---- chat_eos_repair ------------------------------------------------------
|
||||
|
||||
|
||||
def test_repair_adds_missing_turn_end():
|
||||
assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046]
|
||||
|
||||
|
||||
def test_repair_from_missing_generation_config_eos():
|
||||
assert chat_eos_repair(None, [248046]) == [248046]
|
||||
|
||||
|
||||
def test_repair_noop_when_already_covered():
|
||||
assert chat_eos_repair([248046, 248044], [248046]) is None
|
||||
|
||||
|
||||
def test_repair_noop_when_no_turn_end_ids():
|
||||
assert chat_eos_repair(248044, []) is None
|
||||
50
studio/backend/tests/test_coding_agents.py
Normal file
50
studio/backend/tests/test_coding_agents.py
Normal file
|
|
@ -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 <agent>` 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"]
|
||||
314
studio/backend/tests/test_completion_masking.py
Normal file
314
studio/backend/tests/test_completion_masking.py
Normal file
|
|
@ -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 "<INS>", "<RES>"
|
||||
|
||||
|
||||
def _detect_fail(processor):
|
||||
raise ValueError(
|
||||
"Unsloth: Could not reliably auto-detect response_part - "
|
||||
"pass instruction_part and response_part."
|
||||
)
|
||||
|
||||
|
||||
_AUTO = {"instruction_part": "<INS>", "response_part": "<RES>"}
|
||||
|
||||
|
||||
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 = "<I>"
|
||||
_unsloth_output_part = "<O>"
|
||||
|
||||
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 = "<I>"
|
||||
_unsloth_output_part = "<O>"
|
||||
|
||||
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 "<INS>", "<RES>"
|
||||
|
||||
_, applied = apply_completion_masking(
|
||||
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect
|
||||
)
|
||||
assert applied is True
|
||||
assert seen == [inner]
|
||||
|
|
@ -65,12 +65,14 @@ def _backend(
|
|||
vocab = 248320,
|
||||
embd = 5120,
|
||||
mla = None,
|
||||
arch = None,
|
||||
):
|
||||
"""Backend with just the dims the compute-buffer estimate reads."""
|
||||
b = LlamaCppBackend.__new__(LlamaCppBackend)
|
||||
b._vocab_size = vocab
|
||||
b._embedding_length = embd
|
||||
b._key_length_mla = mla # non-None -> MLA (compressed attention)
|
||||
b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4')
|
||||
return b
|
||||
|
||||
|
||||
|
|
@ -290,3 +292,62 @@ class TestContextBufferMLA:
|
|||
b = _backend(embd = 6144, mla = 256)
|
||||
est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB
|
||||
assert est <= 4141 * 1.7
|
||||
|
||||
|
||||
class TestContextBufferDSV4:
|
||||
"""DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention
|
||||
compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache).
|
||||
Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit
|
||||
must see this so it does not commit the full 1M train context and OOM (spilling
|
||||
to CPU at ~4 tok/s)."""
|
||||
|
||||
_MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx
|
||||
GIB = 1024**3
|
||||
|
||||
def test_covers_measured_1m_buffer(self):
|
||||
b = _backend(embd = 4096, arch = "deepseek4")
|
||||
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
|
||||
assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}"
|
||||
|
||||
def test_not_wildly_over_at_1m(self):
|
||||
# Within ~1.3x of measured so the fit still grants a large (~256k) context.
|
||||
b = _backend(embd = 4096, arch = "deepseek4")
|
||||
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
|
||||
assert gib <= self._MEASURED_1M_GIB * 1.3
|
||||
|
||||
def test_fires_for_f16_cache(self):
|
||||
# The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must
|
||||
# reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx.
|
||||
dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes(
|
||||
262144, cache_type_kv = "f16"
|
||||
)
|
||||
other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes(
|
||||
262144, cache_type_kv = "f16"
|
||||
)
|
||||
assert dsv4 > 40 * other
|
||||
|
||||
def test_cache_type_independent(self):
|
||||
# Indexer scratch is present for an f16 and a quantized cache alike.
|
||||
b = _backend(embd = 4096, arch = "deepseek4")
|
||||
assert b._compute_buffer_ctx_bytes(
|
||||
262144, cache_type_kv = "f16"
|
||||
) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
|
||||
|
||||
def test_flat_floor_at_small_ctx(self):
|
||||
# ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB).
|
||||
b = _backend(embd = 4096, arch = "deepseek4")
|
||||
assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0
|
||||
|
||||
def test_scales_with_context_and_ubatch(self):
|
||||
b = _backend(embd = 4096, arch = "deepseek4")
|
||||
assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536)
|
||||
assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes(
|
||||
131072, n_ubatch = 256
|
||||
)
|
||||
|
||||
def test_non_dsv4_unchanged(self):
|
||||
# Regression guard: a non-deepseek4 model keeps the mask-only f16 rate.
|
||||
b = _backend(embd = 4096, arch = "llama")
|
||||
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000
|
||||
expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY
|
||||
assert per_tok == pytest.approx(expected, rel = 1e-6)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
181
studio/backend/tests/test_deepseek_v4_thinking_effort.py
Normal file
181
studio/backend/tests/test_deepseek_v4_thinking_effort.py
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""DeepSeek-V4-Flash reasoning toggle: None / High / Max.
|
||||
|
||||
The GGUF template gates thinking with ``enable_thinking`` and only branches
|
||||
``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking).
|
||||
Detection used to return the single level ``['max']``, so the UI collapsed to
|
||||
None / Max and the plain-thinking tier was unreachable. Detection now surfaces
|
||||
``'high'`` as that plain tier, giving None / High / Max. These tests pin the
|
||||
classifier, the GLM-style parity case, and the full request-kwargs -> rendered
|
||||
prompt path for each state (the model itself is too large to load here).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
_backend_root = Path(__file__).resolve().parent.parent
|
||||
if str(_backend_root) not in sys.path:
|
||||
sys.path.insert(0, str(_backend_root))
|
||||
|
||||
|
||||
# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking
|
||||
# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think
|
||||
# fallback. Any non-'max' effort renders as ordinary thinking.
|
||||
DEEPSEEK_V4_TEMPLATE = """
|
||||
{%- if not thinking is defined -%}
|
||||
{%- if enable_thinking is defined -%}
|
||||
{%- set thinking = enable_thinking -%}
|
||||
{%- else -%}
|
||||
{%- set thinking = false -%}
|
||||
{%- endif -%}
|
||||
{%- endif -%}
|
||||
{%- if not reasoning_effort is defined -%}
|
||||
{%- set reasoning_effort = none -%}
|
||||
{%- endif -%}
|
||||
{{- bos_token -}}
|
||||
{%- if thinking and reasoning_effort == 'max' -%}
|
||||
{{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}}
|
||||
{%- endif -%}
|
||||
{%- for message in messages -%}
|
||||
{{- '<|User|>' + (message['content'] or '') -}}
|
||||
{%- endfor -%}
|
||||
{%- if add_generation_prompt -%}
|
||||
{{- '<|Assistant|>' -}}
|
||||
{%- if thinking -%}{{- '<think>' -}}{%- else -%}{{- '</think>' -}}{%- endif -%}
|
||||
{%- endif -%}
|
||||
"""
|
||||
|
||||
|
||||
# GLM-5.2-style: branches on two effort literals, so 'high' already exists as
|
||||
# the sub-'max' tier and detection must leave the pair untouched.
|
||||
GLM_STYLE_TEMPLATE = """
|
||||
{%- if enable_thinking -%}
|
||||
{%- if reasoning_effort == 'high' -%}{{- 'H' -}}
|
||||
{%- elif reasoning_effort == 'max' -%}{{- 'M' -}}
|
||||
{%- endif -%}
|
||||
{%- endif -%}
|
||||
"""
|
||||
|
||||
|
||||
# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped
|
||||
# to deepseek-v4, so this must stay ['max'] (no phantom 'high').
|
||||
NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE
|
||||
|
||||
|
||||
# A template whose sole effort literal is a sub-'max' level: the guard targets
|
||||
# only the ['max']-alone case, so a lone 'high' stays a singleton.
|
||||
HIGH_ONLY_TEMPLATE = """
|
||||
{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%}
|
||||
"""
|
||||
|
||||
|
||||
def _render(template: str, **kwargs) -> str:
|
||||
jinja2 = pytest.importorskip("jinja2")
|
||||
env = jinja2.Environment()
|
||||
tmpl = env.from_string(template)
|
||||
return tmpl.render(bos_token = "<BOS>", add_generation_prompt = True, **kwargs)
|
||||
|
||||
|
||||
# -- Classifier -------------------------------------------------------
|
||||
|
||||
|
||||
def test_deepseek_v4_surfaces_high_as_plain_tier():
|
||||
"""Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show."""
|
||||
from core.inference.llama_cpp import detect_reasoning_flags
|
||||
|
||||
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash")
|
||||
assert flags["supports_reasoning"] is True
|
||||
assert flags["reasoning_style"] == "enable_thinking_effort"
|
||||
assert flags["reasoning_effort_levels"] == ["high", "max"]
|
||||
|
||||
|
||||
def test_glm_style_two_level_template_unchanged():
|
||||
"""A template that already names a sub-'max' tier is left as-is."""
|
||||
from core.inference.llama_cpp import detect_reasoning_flags
|
||||
|
||||
flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2")
|
||||
assert flags["reasoning_style"] == "enable_thinking_effort"
|
||||
assert flags["reasoning_effort_levels"] == ["high", "max"]
|
||||
|
||||
|
||||
def test_synthetic_high_scoped_to_deepseek_v4():
|
||||
"""The same ['max']-only template under a non-deepseek id keeps ['max']."""
|
||||
from core.inference.llama_cpp import detect_reasoning_flags
|
||||
|
||||
flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF")
|
||||
assert flags["reasoning_effort_levels"] == ["max"]
|
||||
|
||||
|
||||
def test_guard_does_not_fire_for_sub_max_singleton():
|
||||
"""The expansion targets only ['max']; a lone 'high' stays a singleton."""
|
||||
from core.inference.llama_cpp import detect_reasoning_flags
|
||||
|
||||
flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only")
|
||||
assert flags["reasoning_effort_levels"] == ["high"]
|
||||
|
||||
|
||||
# -- Request kwargs -> rendered prompt, for each state ----------------
|
||||
|
||||
|
||||
def _kwargs_for(flags: dict, enable_thinking, reasoning_effort):
|
||||
"""Drive the real backend method with a shim carrying the detected flags."""
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
shim = SimpleNamespace(
|
||||
_supports_reasoning = flags["supports_reasoning"],
|
||||
_reasoning_always_on = flags["reasoning_always_on"],
|
||||
_reasoning_style = flags["reasoning_style"],
|
||||
_reasoning_effort_levels = flags["reasoning_effort_levels"],
|
||||
_supports_preserve_thinking = flags["supports_preserve_thinking"],
|
||||
)
|
||||
build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim)
|
||||
return build(enable_thinking, reasoning_effort, None) or {}
|
||||
|
||||
|
||||
def _flags():
|
||||
from core.inference.llama_cpp import detect_reasoning_flags
|
||||
return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash")
|
||||
|
||||
|
||||
def test_none_state_renders_non_thinking():
|
||||
"""UI 'None' -> enable_thinking=false -> closed </think>, no preamble."""
|
||||
kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None)
|
||||
assert kwargs == {"enable_thinking": False}
|
||||
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
|
||||
assert out.endswith("</think>")
|
||||
assert "Absolute maximum" not in out
|
||||
|
||||
|
||||
def test_high_state_renders_plain_thinking():
|
||||
"""UI 'High' -> et=true, effort=high -> open <think>, no max preamble."""
|
||||
kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high")
|
||||
assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"}
|
||||
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
|
||||
assert out.endswith("<think>")
|
||||
assert "Absolute maximum" not in out
|
||||
|
||||
|
||||
def test_max_state_injects_max_preamble():
|
||||
"""UI 'Max' -> et=true, effort=max -> open <think> plus the max preamble."""
|
||||
kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max")
|
||||
assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"}
|
||||
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
|
||||
assert out.endswith("<think>")
|
||||
assert "Absolute maximum" in out
|
||||
|
||||
|
||||
def test_high_effort_alone_enables_thinking():
|
||||
"""API caller sending only reasoning_effort='high' (no enable_thinking) still
|
||||
gets thinking on, so the newly exposed High mode renders correctly."""
|
||||
kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high")
|
||||
assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"}
|
||||
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
|
||||
assert out.endswith("<think>")
|
||||
assert "Absolute maximum" not in out
|
||||
365
studio/backend/tests/test_embedding_model_security_gate.py
Normal file
365
studio/backend/tests/test_embedding_model_security_gate.py
Normal file
|
|
@ -0,0 +1,365 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""The RAG embedding model must pass the malware/pickle gate before it is persisted or
|
||||
loaded. A flagged repo (or any repo saved with force) previously reached
|
||||
SentenceTransformer unscanned, bypassing the normal model-load protections."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types as _types
|
||||
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.settings as settings
|
||||
|
||||
|
||||
class _Decision:
|
||||
def __init__(self, blocked):
|
||||
self.blocked = blocked
|
||||
|
||||
|
||||
def _security_stub(blocked):
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.evaluate_file_security = lambda *a, **k: _Decision(blocked)
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
return mod
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch):
|
||||
# The settings scan unions in the ST module dirs read from modules.json; keep it
|
||||
# offline and deterministic for the endpoint tests that use this fixture.
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
return TestClient(app, raise_server_exceptions = False), saved
|
||||
|
||||
|
||||
def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch):
|
||||
c, saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
r = c.put(
|
||||
"/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True}
|
||||
)
|
||||
# 403, not the forceable 409, so the client does not offer "save anyway".
|
||||
assert r.status_code == 403
|
||||
assert "model" not in saved # force must not persist a flagged repo
|
||||
|
||||
|
||||
def test_flagged_repo_is_blocked_without_force(client, monkeypatch):
|
||||
c, saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
|
||||
assert r.status_code == 403
|
||||
assert "model" not in saved
|
||||
|
||||
|
||||
def test_hard_block_uses_non_forceable_status(client, monkeypatch):
|
||||
# The forceable verification path uses 409; the hard security block must be distinct
|
||||
# (403) so the frontend never routes it into the "save anyway" force flow.
|
||||
c, _saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
|
||||
assert blocked.status_code == 403
|
||||
|
||||
# A verification failure (not-an-embedding-model) stays forceable at 409.
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False)
|
||||
import utils.models as _models
|
||||
|
||||
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
|
||||
unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"})
|
||||
assert unverified.status_code == 409
|
||||
|
||||
|
||||
def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
|
||||
# On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
|
||||
# pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
# force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped.
|
||||
called = {"scanned": False}
|
||||
mod = _types.ModuleType("utils.security")
|
||||
|
||||
def _fail(*a, **k):
|
||||
called["scanned"] = True
|
||||
return _Decision(True)
|
||||
|
||||
mod.evaluate_file_security = _fail
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.put(
|
||||
"/embedding-model",
|
||||
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert called["scanned"] is False # the ST pickle scan never ran on the llama path
|
||||
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
|
||||
|
||||
|
||||
def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch):
|
||||
# auto resolves to sentence-transformers (GPU present) but the embedder fell back to
|
||||
# llama-server at runtime (torch/CUDA load or encode failure), so the process now loads
|
||||
# only inert GGUF. The real _llama_backend_active() must reflect that cached fallback,
|
||||
# so a flagged ST repo with a clean GGUF companion must not be hard-blocked here.
|
||||
import core.rag.embeddings as embeddings
|
||||
from core.rag.embed_llama_server import LlamaServerBackend
|
||||
|
||||
# Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even
|
||||
# though the auto resolver would still say sentence-transformers.
|
||||
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
|
||||
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
|
||||
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
# Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the
|
||||
# real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored.
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
|
||||
called = {"scanned": False}
|
||||
mod = _types.ModuleType("utils.security")
|
||||
|
||||
def _fail(*a, **k):
|
||||
called["scanned"] = True
|
||||
return _Decision(True)
|
||||
|
||||
mod.evaluate_file_security = _fail
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.put(
|
||||
"/embedding-model",
|
||||
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback
|
||||
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
|
||||
|
||||
|
||||
def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch):
|
||||
# active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers
|
||||
# to the resolver (fresh-process behavior) when none has been built yet.
|
||||
import core.rag.embeddings as embeddings
|
||||
import core.rag.config as rag_config
|
||||
from core.rag.embed_llama_server import LlamaServerBackend
|
||||
|
||||
# A cached llama backend wins even when auto would resolve to sentence-transformers.
|
||||
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto")
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
|
||||
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
|
||||
assert embeddings.active_backend_is_llama() is True
|
||||
|
||||
# A cached ST backend reports False even when the resolver now picks llama, so its
|
||||
# pickle stays gated (the cached backend, not the resolver, is what actually embeds).
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
|
||||
monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend())
|
||||
assert embeddings.active_backend_is_llama() is False
|
||||
|
||||
# No cached backend -> the resolver decides, unchanged from before.
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
|
||||
monkeypatch.setattr(embeddings, "_backend", None)
|
||||
assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers
|
||||
|
||||
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
|
||||
assert embeddings.active_backend_is_llama() is True # auto -> llama-server
|
||||
|
||||
# An explicit (non-auto) key is honored verbatim without a cached backend.
|
||||
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server")
|
||||
assert embeddings.active_backend_is_llama() is True
|
||||
|
||||
|
||||
def test_settings_scan_scopes_module_subdirs(monkeypatch):
|
||||
# The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a
|
||||
# pickle directly under one blocks; assert those subdirs reach evaluate_file_security.
|
||||
saved: dict = {}
|
||||
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
|
||||
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
|
||||
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
|
||||
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
|
||||
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
|
||||
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
|
||||
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
|
||||
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(
|
||||
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
|
||||
)
|
||||
seen = {}
|
||||
|
||||
def _capture(*a, **k):
|
||||
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
|
||||
return _Decision(False)
|
||||
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.security_load_subdirs = lambda *a, **k: ()
|
||||
mod.evaluate_file_security = _capture
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.put(
|
||||
"/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True}
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert "0_Transformer" in seen["subdirs"]
|
||||
|
||||
|
||||
def test_clean_repo_saves_under_force(client, monkeypatch):
|
||||
c, saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True})
|
||||
assert r.status_code == 200
|
||||
assert saved.get("model") == "acme/clean-embed"
|
||||
|
||||
|
||||
def test_load_sink_refuses_flagged_model(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
|
||||
import core.rag.embeddings as embeddings
|
||||
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
|
||||
embeddings._guard_model_security("attacker/malicious-embed")
|
||||
|
||||
|
||||
def test_load_sink_allows_clean_model(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
import core.rag.embeddings as embeddings
|
||||
embeddings._guard_model_security("acme/clean-embed") # no raise
|
||||
|
||||
|
||||
def test_sink_threads_ambient_token_into_scan(monkeypatch):
|
||||
# A gated repo set via env/default has no request token; the guard must feed the
|
||||
# loader's own token to the scan, or it fails open for the repo that still loads.
|
||||
seen = {}
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.security_load_subdirs = (
|
||||
lambda name, token = None: seen.setdefault("subdirs_token", token) or ()
|
||||
)
|
||||
mod.evaluate_file_security = lambda *a, **k: seen.setdefault(
|
||||
"scan_token", k.get("hf_token")
|
||||
) or _Decision(False)
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient")
|
||||
embeddings._guard_model_security("acme/gated-embed")
|
||||
assert seen["scan_token"] == "hf_ambient"
|
||||
assert seen["subdirs_token"] == "hf_ambient"
|
||||
|
||||
|
||||
def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch):
|
||||
# A flagged pickle directly under a Transformer module dir (0_Transformer/) must
|
||||
# reach the scan as a load root; assert the guard unions the module dirs into
|
||||
# load_subdirs so evaluate_file_security treats such a pickle as root-level.
|
||||
seen = {}
|
||||
|
||||
def _capture(*a, **k):
|
||||
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
|
||||
return _Decision(False)
|
||||
|
||||
mod = _types.ModuleType("utils.security")
|
||||
mod.security_load_subdirs = lambda name, token = None: ()
|
||||
mod.evaluate_file_security = _capture
|
||||
monkeypatch.setitem(sys.modules, "utils.security", mod)
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
|
||||
)
|
||||
embeddings._guard_model_security("acme/embed-with-module-dir")
|
||||
assert "0_Transformer" in seen["subdirs"]
|
||||
|
||||
|
||||
def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch):
|
||||
# The helper must parse each module's non-empty "path" from a local repo's
|
||||
# modules.json and drop the root-level ("") Transformer entry.
|
||||
import json
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
(tmp_path / "modules.json").write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."},
|
||||
{"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."},
|
||||
{"idx": 2, "name": "2", "path": "", "type": "..."},
|
||||
]
|
||||
)
|
||||
)
|
||||
subdirs = embeddings._st_module_subdirs(str(tmp_path), None)
|
||||
assert subdirs == ("0_Transformer", "1_Pooling")
|
||||
|
||||
|
||||
def test_st_module_subdirs_swallows_errors(monkeypatch):
|
||||
# Any failure (no modules.json, offline, malformed) returns () so the guard never
|
||||
# bricks the embedder.
|
||||
import huggingface_hub
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise RuntimeError("offline")
|
||||
|
||||
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom)
|
||||
assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == ()
|
||||
|
||||
|
||||
def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch):
|
||||
# The ST encode fallback must re-raise a security block, not swap to llama-server.
|
||||
import core.rag.embeddings as embeddings
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise embeddings.UnsafeEmbeddingModelError("flagged")
|
||||
|
||||
monkeypatch.setattr(embeddings, "_st_encode", _boom)
|
||||
monkeypatch.setattr(
|
||||
embeddings,
|
||||
"_switch_to_llama_fallback",
|
||||
lambda err: pytest.fail("security block must not fall back to llama-server"),
|
||||
)
|
||||
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
|
||||
embeddings._SentenceTransformersBackend().encode(["hi"])
|
||||
|
|
@ -1,15 +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
|
||||
|
||||
"""Edge cases in Gemma-native tool-call parsing.
|
||||
|
||||
Covers two failure modes:
|
||||
1. A bare (unquoted) string argument that contains a comma, e.g.
|
||||
``location:New York, NY`` -- the comma must not be treated as the next
|
||||
key boundary, or the whole call is dropped.
|
||||
2. A tool-call marker that appears INSIDE another call's argument string is
|
||||
data, not a real call, so it must not be promoted to a second tool call.
|
||||
"""
|
||||
"""Gemma-native tool-call parsing edge cases: commas inside bare string values,
|
||||
and markers inside another call's argument data staying data."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
|
@ -21,7 +14,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
from core.inference.tool_call_parser import parse_tool_calls_from_text
|
||||
from core.inference.tool_call_parser import (
|
||||
_gemma_parse_value,
|
||||
parse_tool_calls_from_text,
|
||||
)
|
||||
from core.tool_healing import strip_tool_call_markup
|
||||
|
||||
|
||||
def _args(call: dict) -> dict:
|
||||
|
|
@ -40,14 +37,22 @@ def test_bare_string_argument_with_comma_is_kept():
|
|||
def test_normal_multi_key_arguments_still_split():
|
||||
calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}<tool_call|>')
|
||||
assert len(calls) == 1, calls
|
||||
# Numbers stay numeric, bare strings get quoted, an explicit quoted comma
|
||||
# stays inside its value.
|
||||
assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"}
|
||||
|
||||
|
||||
def test_empty_bare_value_becomes_empty_string_not_dropped():
|
||||
# An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON and dropped the call).
|
||||
calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}<tool_call|>")
|
||||
assert len(calls) == 1, calls
|
||||
assert _args(calls[0]) == {"query": "", "unit": "celsius"}
|
||||
|
||||
only = parse_tool_calls_from_text("<|tool_call>call:get{q:}<tool_call|>")
|
||||
assert len(only) == 1, only
|
||||
assert _args(only[0]) == {"q": ""}
|
||||
|
||||
|
||||
def test_bare_value_with_timestamps_after_comma_is_kept():
|
||||
# A comma followed by digits-then-colon (a timestamp/ratio) is value text,
|
||||
# not a new key, so the whole query must be preserved as one argument.
|
||||
# A comma before digits-then-colon (timestamp/ratio) is value text, not a key.
|
||||
calls = parse_tool_calls_from_text(
|
||||
"<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}<tool_call|>"
|
||||
)
|
||||
|
|
@ -55,9 +60,16 @@ def test_bare_value_with_timestamps_after_comma_is_kept():
|
|||
assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"}
|
||||
|
||||
|
||||
def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept():
|
||||
# The wrapper-less Gemma form (no <|tool_call> markers) goes through the
|
||||
# _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE.
|
||||
calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}")
|
||||
assert len(calls) == 1, calls
|
||||
assert calls[0]["function"]["name"] == "web_search"
|
||||
assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"}
|
||||
|
||||
|
||||
def test_marker_inside_json_argument_is_not_a_second_call():
|
||||
# A python call whose `code` argument contains a Gemma marker string. The
|
||||
# marker is data and must not execute as a second `terminal` call.
|
||||
content = (
|
||||
'<tool_call>{"name":"python","arguments":{"code":'
|
||||
'"x = 1 # <|tool_call>call:terminal{command:ls}<tool_call|>"}}</tool_call>'
|
||||
|
|
@ -75,8 +87,6 @@ def test_two_separate_gemma_calls_both_parse():
|
|||
|
||||
|
||||
def test_mixed_format_calls_preserve_document_order():
|
||||
# A Gemma-native call precedes a JSON-format call in the text; tools execute
|
||||
# in returned order, so `create` must come before `read`.
|
||||
content = (
|
||||
"<|tool_call>call:create{path:a}<tool_call|> then "
|
||||
'<tool_call>{"name":"read","arguments":{"path":"a"}}</tool_call>'
|
||||
|
|
@ -86,8 +96,6 @@ def test_mixed_format_calls_preserve_document_order():
|
|||
|
||||
|
||||
def test_json_marker_inside_gemma_argument_is_not_a_second_call():
|
||||
# The reverse of the JSON-outer case: a JSON-style marker inside a Gemma
|
||||
# call's quoted argument is code text, not a second `terminal` call.
|
||||
content = (
|
||||
'<|tool_call>call:python{code:<|"|>'
|
||||
'print(<tool_call>{"name":"terminal","arguments":{"command":"ls"}}</tool_call>)'
|
||||
|
|
@ -98,18 +106,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call():
|
|||
|
||||
|
||||
def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call():
|
||||
# An UNQUOTED Gemma value containing a literal marker: the outer object fails
|
||||
# to normalize (the inner braces/marker break the JSON), but the inner marker
|
||||
# is nested in the outer candidate span, so it must not be promoted to a
|
||||
# standalone `terminal` call. The safe outcome is no executed tool call.
|
||||
# An UNQUOTED Gemma value containing a literal marker: the marker is nested in the outer
|
||||
# candidate span, so it must not be promoted to a standalone `terminal` call (no tool call).
|
||||
content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}<tool_call|>}<tool_call|>"
|
||||
calls = parse_tool_calls_from_text(content)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_bare_string_array_argument_is_quoted():
|
||||
# Gemma may emit an array of bare strings without per-element quotes; they
|
||||
# must be quoted so the call is not dropped.
|
||||
calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}<tool_call|>")
|
||||
assert len(calls) == 1, calls
|
||||
assert _args(calls[0]) == {"labels": ["bug", "ui"]}
|
||||
|
|
@ -123,8 +127,6 @@ def test_array_keeps_numbers_and_quoted_elements():
|
|||
|
||||
|
||||
def test_array_of_objects_is_normalised():
|
||||
# Arrays of objects are a common tool-schema shape; their (unquoted) keys and
|
||||
# bare values must be normalised too, not left verbatim, or the call drops.
|
||||
calls = parse_tool_calls_from_text(
|
||||
"<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}<tool_call|>"
|
||||
)
|
||||
|
|
@ -138,9 +140,6 @@ def test_nested_array_elements_are_normalised():
|
|||
|
||||
|
||||
def test_gemma_marker_inside_xml_parameter_is_not_a_second_call():
|
||||
# An XML-style <function=...> call whose <parameter=code> value contains a
|
||||
# Gemma marker: the marker is the parameter's data, not a separate terminal
|
||||
# call, so only the python call must be returned.
|
||||
content = (
|
||||
"<tool_call><function=python><parameter=code>"
|
||||
"x = 1 # <|tool_call>call:terminal{command:ls}<tool_call|>"
|
||||
|
|
@ -159,3 +158,254 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call():
|
|||
)
|
||||
calls = parse_tool_calls_from_text(content)
|
||||
assert [c["function"]["name"] for c in calls] == ["python"], calls
|
||||
|
||||
|
||||
def test_unclosed_think_literal_inside_tool_argument_does_not_hide_later_call():
|
||||
# A literal <think> inside a completed call's arguments is argument data; both calls must parse.
|
||||
text = '[TOOL_CALLS]a{"x":"literal <think> marker"} b[ARGS]{"y":2}'
|
||||
calls = parse_tool_calls_from_text(text)
|
||||
assert [c["function"]["name"] for c in calls] == ["a", "b"], calls
|
||||
|
||||
|
||||
def test_real_think_block_with_rehearsal_inside_still_skips_only_the_rehearsal():
|
||||
# A genuine reasoning block still hides its rehearsal while a real call after it parses.
|
||||
text = '<think>web_search[ARGS]{"q":"draft"}</think>real[ARGS]{"q":"go"}'
|
||||
calls = parse_tool_calls_from_text(text)
|
||||
assert [c["function"]["name"] for c in calls] == ["real"], calls
|
||||
|
||||
|
||||
def test_wrapperless_nested_object_argument_is_parsed():
|
||||
# skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare.
|
||||
calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}")
|
||||
assert len(calls) == 1
|
||||
assert _args(calls[0]) == {"loc": {"city": "NYC"}, "n": 3}
|
||||
|
||||
|
||||
def test_wrapperless_array_argument_is_parsed():
|
||||
calls = parse_tool_calls_from_text("call:label{labels:[bug,ui],n:2}")
|
||||
assert len(calls) == 1
|
||||
assert _args(calls[0]) == {"labels": ["bug", "ui"], "n": 2}
|
||||
|
||||
|
||||
def test_wrapperless_deeply_nested_object_and_array_are_preserved():
|
||||
# The single-pass parser must keep multi-level nesting (objects inside
|
||||
# objects, arrays inside arrays) intact, not flatten or drop it.
|
||||
calls = parse_tool_calls_from_text(
|
||||
"call:f{loc:{city:NYC,geo:{lat:1,lng:2}},tags:[a,b,[c,d]],n:3}"
|
||||
)
|
||||
assert len(calls) == 1
|
||||
assert _args(calls[0]) == {
|
||||
"loc": {"city": "NYC", "geo": {"lat": 1, "lng": 2}},
|
||||
"tags": ["a", "b", ["c", "d"]],
|
||||
"n": 3,
|
||||
}
|
||||
|
||||
|
||||
def test_gemma_parse_array_advances_on_stray_brace():
|
||||
# Regression: a stray '}' / ']' / ',' where an array element is expected must
|
||||
# not stall _gemma_parse_value at the same index (it looped forever before).
|
||||
from core.inference.tool_call_parser import _gemma_parse_array
|
||||
|
||||
items, end, closed = _gemma_parse_array("[a,}]", 0)
|
||||
assert end == 5 and closed is True # consumed through the closing ']'
|
||||
assert items[0] == "a"
|
||||
|
||||
|
||||
def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping():
|
||||
# Parse keeps the quoted close marker as data; strip removes the whole span.
|
||||
text = '<|tool_call>call:python{code:<|"|>print("<tool_call|>")<|"|>}<tool_call|>'
|
||||
calls = parse_tool_calls_from_text(text)
|
||||
assert len(calls) == 1, calls
|
||||
assert _args(calls[0]) == {"code": 'print("<tool_call|>")'}
|
||||
assert strip_tool_call_markup("before " + text + " after") == "before after"
|
||||
assert strip_tool_call_markup("before " + text + " after", final = True) == "before after"
|
||||
|
||||
|
||||
def test_nested_xml_in_malformed_gemma_call_does_not_execute():
|
||||
# The failed Gemma candidate's span still covers its nested <function=>.
|
||||
text = (
|
||||
"<|tool_call>call:outer{code:<function=terminal><parameter=command>id"
|
||||
"</parameter></function></tool_call>, broken:{x}}<tool_call|>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_unbalanced_gemma_call_with_xml_does_not_execute():
|
||||
# Unclosed braces cover to EOF, so the trailing <function=> is excluded.
|
||||
text = (
|
||||
"<|tool_call>call:outer{code:<function=terminal>"
|
||||
"<parameter=command>id</parameter></function>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_standalone_function_xml_still_parses():
|
||||
text = "<function=terminal><parameter=command>id</parameter></function>"
|
||||
calls = parse_tool_calls_from_text(text)
|
||||
assert [c["function"]["name"] for c in calls] == ["terminal"], calls
|
||||
|
||||
|
||||
def test_xml_between_braces_and_close_marker_does_not_execute():
|
||||
# Coverage runs to the close marker, so <function=> in the gap is data.
|
||||
text = (
|
||||
"<|tool_call>call:outer{broken:{x}}<function=terminal>"
|
||||
"<parameter=command>id</parameter></function><tool_call|>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_balanced_inner_call_inside_unclosed_outer_does_not_execute():
|
||||
text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}<tool_call|>"
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_strip_preserves_text_after_malformed_gemma_close():
|
||||
# Junk before the close is a malformed span: strip through it, keep the tail.
|
||||
text = "pre <|tool_call>call:t{a:1} note <tool_call|> post"
|
||||
assert strip_tool_call_markup(text) == "pre post"
|
||||
assert strip_tool_call_markup(text, final = True) == "pre post"
|
||||
|
||||
|
||||
def test_malformed_closed_gemma_span_is_stripped():
|
||||
assert (
|
||||
strip_tool_call_markup('before <|tool_call>{"name":"x"}<tool_call|> after')
|
||||
== "before after"
|
||||
)
|
||||
|
||||
|
||||
def test_valid_call_after_missing_close_is_recovered():
|
||||
# A close-less call covers only its braces, so the later call is recovered.
|
||||
text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}<tool_call|>"
|
||||
names_inc = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True)
|
||||
]
|
||||
assert "b" in names_inc, names_inc
|
||||
names_strict = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False)
|
||||
]
|
||||
assert names_strict == ["b"], names_strict
|
||||
|
||||
|
||||
def test_strip_non_final_keeps_incomplete_gemma_block():
|
||||
text = "before <|tool_call>call:t{"
|
||||
assert strip_tool_call_markup(text) == text
|
||||
assert strip_tool_call_markup(text, final = True) == "before"
|
||||
|
||||
|
||||
def test_json_call_between_gemma_braces_and_close_does_not_execute():
|
||||
# A JSON call between the outer's braces and its close is covered data.
|
||||
text = (
|
||||
"<|tool_call>call:outer{broken:{x}}"
|
||||
'<tool_call>{"name":"terminal","arguments":{"command":"id"}}</tool_call>'
|
||||
"<tool_call|>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_gemma_call_between_gemma_braces_and_close_does_not_execute():
|
||||
# Same escape with a Gemma-native inner marker.
|
||||
text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}<tool_call|><tool_call|>"
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert "terminal" not in [c["function"]["name"] for c in calls], calls
|
||||
|
||||
|
||||
def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener():
|
||||
# The to-EOF Gemma sweep must not eat visible text after </function>.
|
||||
text = (
|
||||
'before <function=python><parameter=code>print("<|tool_call>")</parameter></function> after'
|
||||
)
|
||||
assert strip_tool_call_markup(text, final = True) == "before after"
|
||||
assert strip_tool_call_markup(text) == "before after"
|
||||
|
||||
|
||||
def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener():
|
||||
# A call-form Gemma opener quoted in a closed block must not truncate it.
|
||||
xml = "<function=python><parameter=code><|tool_call>call:t{</parameter></function>"
|
||||
json_block = (
|
||||
'<tool_call>{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}</tool_call>'
|
||||
)
|
||||
for block in (xml, json_block):
|
||||
text = "before " + block + " after"
|
||||
assert strip_tool_call_markup(text, final = True) == "before after", block
|
||||
assert strip_tool_call_markup(text) == "before after", block
|
||||
|
||||
|
||||
def test_function_sibling_after_close_less_gemma_marker_is_recovered():
|
||||
# The close-less marker covers only its braces; the XML sibling is recovered.
|
||||
text = (
|
||||
"<|tool_call>call:bad{broken:{x}} "
|
||||
"<function=terminal><parameter=command>id</parameter></function>"
|
||||
)
|
||||
for allow_incomplete in (True, False):
|
||||
calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
|
||||
assert [c["function"]["name"] for c in calls] == ["terminal"], calls
|
||||
|
||||
|
||||
def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered():
|
||||
# A close token quoted in the later call must not extend the earlier
|
||||
# close-less marker's coverage over that call.
|
||||
gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|></tool_call><|"|>}<tool_call|>'
|
||||
names = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False)
|
||||
]
|
||||
assert names == ["b"], names
|
||||
json_text = (
|
||||
'<tool_call>{"name":"a","arguments":{}} '
|
||||
'<tool_call>{"name":"b","arguments":{"x":"</tool_call>"}}</tool_call>'
|
||||
)
|
||||
names_j = [
|
||||
c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False)
|
||||
]
|
||||
assert "b" in names_j, names_j
|
||||
|
||||
|
||||
def test_gemma_parse_value_always_advances_on_stray_delimiter():
|
||||
# A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the
|
||||
# index by at least one, or a caller looping on it spins forever at 100% CPU (DoS).
|
||||
for delim in (",", "}", "]"):
|
||||
text = delim + "rest"
|
||||
value, nxt, _explicit = _gemma_parse_value(text, 0)
|
||||
assert nxt > 0, (delim, value, nxt)
|
||||
|
||||
|
||||
def test_malformed_gemma_array_does_not_hang():
|
||||
# ``[},]`` puts a stray ``}`` at the primitive position inside a list body.
|
||||
# On the buggy parser this hangs the server; guard with a wall-clock timeout
|
||||
# so the regression fails loudly instead of blocking CI forever.
|
||||
import threading
|
||||
|
||||
result: dict = {}
|
||||
|
||||
def _run():
|
||||
result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}<tool_call|>")
|
||||
|
||||
t = threading.Thread(target = _run, daemon = True)
|
||||
t.start()
|
||||
t.join(timeout = 10.0)
|
||||
assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed array input"
|
||||
|
||||
|
||||
def test_malformed_gemma_mapping_value_does_not_hang():
|
||||
# A stray ``}`` where a mapping value is expected must also terminate.
|
||||
import threading
|
||||
|
||||
result: dict = {}
|
||||
|
||||
def _run():
|
||||
result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}<tool_call|>")
|
||||
|
||||
t = threading.Thread(target = _run, daemon = True)
|
||||
t.start()
|
||||
t.join(timeout = 10.0)
|
||||
assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input"
|
||||
|
|
|
|||
|
|
@ -1,18 +1,16 @@
|
|||
# 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 utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP
|
||||
transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on.
|
||||
CPU-only, no network, no real subprocess (the per-attempt download seam is
|
||||
monkeypatched).
|
||||
"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback.
|
||||
|
||||
The transport-policy matrix is tested once in unsloth_zoo; here we assert only the
|
||||
Studio seam: re-exporting the shared API and injecting the marker-aware
|
||||
prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types as _types
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -22,9 +20,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
# Stub heavy/unavailable deps before importing the module under test. Use the
|
||||
# real structlog when present; a bare stub left in sys.modules would break later
|
||||
# modules that log at import time.
|
||||
# Stub heavy/unavailable deps before importing the module under test. Use real structlog when present;
|
||||
# a bare stub would break later modules that log at import time.
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
|
@ -34,171 +31,59 @@ except ImportError:
|
|||
sys.modules["structlog"] = _types.ModuleType("structlog")
|
||||
|
||||
import huggingface_hub
|
||||
from huggingface_hub import constants as hf_constants
|
||||
|
||||
try:
|
||||
import unsloth_zoo.hf_xet_fallback as _shared_mod
|
||||
shared = _shared_mod
|
||||
except Exception: # noqa: BLE001 - still collect degraded-path tests when unsloth_zoo is unavailable
|
||||
shared = None
|
||||
|
||||
import utils.hf_xet_fallback as xf
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total.
|
||||
# --------------------------------------------------------------------------- #
|
||||
REPO = "ztest/xet-watchdog"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hf_cache(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _blobs_dir(root: Path, repo_id: str = REPO) -> Path:
|
||||
d = root / f"models--{repo_id.replace('/', '--')}" / "blobs"
|
||||
d.mkdir(parents = True, exist_ok = True)
|
||||
return d
|
||||
|
||||
|
||||
def _wait(
|
||||
predicate,
|
||||
timeout: float = 2.0,
|
||||
step: float = 0.02,
|
||||
) -> bool:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if predicate():
|
||||
return True
|
||||
time.sleep(step)
|
||||
return predicate()
|
||||
|
||||
|
||||
def test_constant_incomplete_fires_stall(hf_cache):
|
||||
blobs = _blobs_dir(hf_cache)
|
||||
(blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows
|
||||
|
||||
calls: list[str] = []
|
||||
stop = xf.start_watchdog(
|
||||
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
|
||||
)
|
||||
try:
|
||||
assert _wait(
|
||||
lambda: len(calls) >= 1, timeout = 3.0
|
||||
), "watchdog never fired on a constant-size .incomplete"
|
||||
finally:
|
||||
stop.set()
|
||||
assert "stalled" in calls[0].lower()
|
||||
|
||||
|
||||
def test_growing_incomplete_never_stalls(hf_cache):
|
||||
blobs = _blobs_dir(hf_cache)
|
||||
part = blobs / "growing.incomplete"
|
||||
part.write_bytes(b"\0" * 1024)
|
||||
|
||||
grow_stop = threading.Event()
|
||||
|
||||
def _grow():
|
||||
size = 1024
|
||||
while not grow_stop.wait(0.05):
|
||||
size += 4096
|
||||
part.write_bytes(b"\0" * size)
|
||||
|
||||
grower = threading.Thread(target = _grow, daemon = True)
|
||||
grower.start()
|
||||
|
||||
calls: list[str] = []
|
||||
stop = xf.start_watchdog(
|
||||
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
|
||||
)
|
||||
try:
|
||||
time.sleep(1.0) # well past stall_timeout, but bytes keep growing
|
||||
assert calls == [], "watchdog fired despite continuous progress"
|
||||
finally:
|
||||
stop.set()
|
||||
grow_stop.set()
|
||||
|
||||
|
||||
def test_no_incomplete_never_stalls(hf_cache):
|
||||
blobs = _blobs_dir(hf_cache)
|
||||
(blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete
|
||||
|
||||
calls: list[str] = []
|
||||
stop = xf.start_watchdog(
|
||||
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
|
||||
)
|
||||
try:
|
||||
time.sleep(0.8)
|
||||
assert calls == [], "watchdog fired with no active .incomplete"
|
||||
finally:
|
||||
stop.set()
|
||||
|
||||
|
||||
def test_stall_fires_at_most_once(hf_cache):
|
||||
blobs = _blobs_dir(hf_cache)
|
||||
(blobs / "frozen.incomplete").write_bytes(b"\0" * 2048)
|
||||
|
||||
calls: list[str] = []
|
||||
stop = xf.start_watchdog(
|
||||
repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2
|
||||
)
|
||||
try:
|
||||
assert _wait(lambda: len(calls) >= 1, timeout = 3.0)
|
||||
time.sleep(0.6) # keep ticking; must not fire again
|
||||
assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1"
|
||||
finally:
|
||||
stop.set()
|
||||
|
||||
|
||||
def test_get_state_empty_cache(hf_cache):
|
||||
assert xf.get_hf_download_state([REPO]) == (0, False)
|
||||
|
||||
|
||||
def test_get_state_absent_cache_root(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache"))
|
||||
assert xf.get_hf_download_state([REPO]) == (0, False)
|
||||
|
||||
|
||||
def test_get_state_skips_local_paths(hf_cache):
|
||||
# Filesystem paths are not HF repo IDs and must be ignored without error.
|
||||
assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False)
|
||||
|
||||
|
||||
def test_get_state_sparse_aware(hf_cache):
|
||||
blobs = _blobs_dir(hf_cache)
|
||||
sparse = blobs / "sparse.incomplete"
|
||||
with open(sparse, "wb") as f:
|
||||
f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks
|
||||
st = sparse.stat()
|
||||
if getattr(st, "st_blocks", 0) == 0:
|
||||
pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable")
|
||||
total, has_incomplete = xf.get_hf_download_state([REPO])
|
||||
assert has_incomplete is True
|
||||
assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Transport policy: cached short-circuit, cancel, error propagation, and the
|
||||
# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn.
|
||||
# --------------------------------------------------------------------------- #
|
||||
DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf"
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _no_real_cache_hit(monkeypatch):
|
||||
"""Default: the cached probe misses; tests override it to force a hit."""
|
||||
def _requires_shared():
|
||||
if shared is None:
|
||||
pytest.skip("unsloth_zoo.hf_xet_fallback is not installed in this environment")
|
||||
|
||||
|
||||
def test_shim_reexports_shared_api():
|
||||
_requires_shared()
|
||||
assert xf.DownloadStallError is shared.DownloadStallError
|
||||
for name in (
|
||||
"start_watchdog",
|
||||
"get_hf_download_state",
|
||||
"child_should_disable_xet",
|
||||
"hf_hub_download_with_xet_fallback",
|
||||
"snapshot_download_with_xet_fallback",
|
||||
):
|
||||
assert hasattr(xf, name), f"shim missing {name}"
|
||||
|
||||
|
||||
def test_child_should_disable_xet_truth_table():
|
||||
assert xf.child_should_disable_xet({"disable_xet": True}) is True
|
||||
assert xf.child_should_disable_xet({"disable_xet": False}) is False
|
||||
assert xf.child_should_disable_xet({}) is False
|
||||
|
||||
|
||||
def test_shim_injects_studio_prepare_on_http_retry(monkeypatch):
|
||||
"""A Xet stall retries over HTTP and the shim runs Studio's marker-aware
|
||||
``prepare_cache_for_transport(..., 'http')`` before the retry."""
|
||||
_requires_shared()
|
||||
for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"):
|
||||
monkeypatch.delenv(var, raising = False)
|
||||
monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None)
|
||||
|
||||
seen_disable_xet = []
|
||||
|
||||
class _FakeAttempt:
|
||||
"""Records calls to the download seam and returns scripted results."""
|
||||
|
||||
def __init__(self, results):
|
||||
self._results = list(results)
|
||||
self.calls = []
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
def fake_attempt(
|
||||
repo_id,
|
||||
filename,
|
||||
token,
|
||||
*,
|
||||
kind,
|
||||
params,
|
||||
token,
|
||||
repo_type,
|
||||
disable_xet,
|
||||
cancel_event,
|
||||
|
|
@ -208,146 +93,277 @@ class _FakeAttempt:
|
|||
on_status,
|
||||
force_download = False,
|
||||
):
|
||||
self.calls.append(
|
||||
_types.SimpleNamespace(
|
||||
repo_id = repo_id,
|
||||
filename = filename,
|
||||
disable_xet = disable_xet,
|
||||
repo_type = repo_type,
|
||||
)
|
||||
)
|
||||
return self._results[len(self.calls) - 1]
|
||||
seen_disable_xet.append(disable_xet)
|
||||
return ("ok", "/cache/model.gguf") if disable_xet else ("stall", None)
|
||||
|
||||
monkeypatch.setattr(shared, "_run_download_attempt", fake_attempt)
|
||||
|
||||
def _install(monkeypatch, results):
|
||||
fake = _FakeAttempt(results)
|
||||
monkeypatch.setattr(xf, "_run_download_attempt", fake)
|
||||
return fake
|
||||
|
||||
|
||||
def test_cached_file_short_circuits(monkeypatch, tmp_path):
|
||||
cached = tmp_path / "cached.gguf"
|
||||
cached.write_bytes(b"\0" * 8)
|
||||
monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached))
|
||||
fake = _install(monkeypatch, []) # must not be called
|
||||
|
||||
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
|
||||
assert out == str(cached)
|
||||
assert fake.calls == [], "spawned a download for an already-cached file"
|
||||
|
||||
|
||||
def test_cancel_before_start_raises_no_attempt(monkeypatch):
|
||||
fake = _install(monkeypatch, [])
|
||||
ev = threading.Event()
|
||||
ev.set()
|
||||
with pytest.raises(RuntimeError, match = "Cancelled"):
|
||||
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev)
|
||||
assert fake.calls == []
|
||||
|
||||
|
||||
def test_nonstall_error_propagates_without_fallback(monkeypatch):
|
||||
fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")])
|
||||
with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"):
|
||||
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
|
||||
assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback"
|
||||
assert fake.calls[0].disable_xet is False
|
||||
|
||||
|
||||
def test_immediate_success_uses_xet_only(monkeypatch):
|
||||
prepared = []
|
||||
monkeypatch.setattr(
|
||||
"hub.utils.download_registry.prepare_cache_for_transport",
|
||||
lambda *a, **k: prepared.append(a),
|
||||
)
|
||||
fake = _install(monkeypatch, [("ok", "/cache/model.gguf")])
|
||||
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
|
||||
assert out == "/cache/model.gguf"
|
||||
assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False
|
||||
assert prepared == [], "no cache prep should run when Xet succeeds first try"
|
||||
|
||||
|
||||
def test_stall_then_http_fallback_succeeds(monkeypatch):
|
||||
prepared = []
|
||||
monkeypatch.setattr(
|
||||
"hub.utils.download_registry.prepare_cache_for_transport",
|
||||
lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)),
|
||||
)
|
||||
fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")])
|
||||
|
||||
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
|
||||
assert out == "/cache/model.gguf"
|
||||
assert len(fake.calls) == 2
|
||||
assert fake.calls[0].disable_xet is False # Xet first
|
||||
assert fake.calls[1].disable_xet is True # HTTP fallback
|
||||
assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry"
|
||||
assert seen_disable_xet == [False, True] # Xet first, then HTTP
|
||||
assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep"
|
||||
|
||||
|
||||
def test_second_stall_raises_download_stall_error(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
|
||||
)
|
||||
fake = _install(monkeypatch, [("stall", None), ("stall", None)])
|
||||
with pytest.raises(xf.DownloadStallError):
|
||||
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
|
||||
assert len(fake.calls) == 2
|
||||
def test_shim_snapshot_injects_studio_prepare(monkeypatch):
|
||||
"""The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper."""
|
||||
captured = {}
|
||||
|
||||
def fake_snapshot(repo_id, **kwargs):
|
||||
captured["repo_id"] = repo_id
|
||||
captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn")
|
||||
return "/tmp/snap-dir"
|
||||
|
||||
monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot)
|
||||
out = xf.snapshot_download_with_xet_fallback("org/model")
|
||||
assert out == "/tmp/snap-dir"
|
||||
assert captured["repo_id"] == "org/model"
|
||||
assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http
|
||||
|
||||
|
||||
def test_cancelled_midattempt_raises_no_fallback(monkeypatch):
|
||||
fake = _install(monkeypatch, [("cancelled", None)])
|
||||
with pytest.raises(RuntimeError, match = "Cancelled"):
|
||||
xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
|
||||
assert len(fake.calls) == 1
|
||||
def test_degrades_gracefully_without_shared_helper(monkeypatch):
|
||||
"""On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio
|
||||
boots) and exposes stub API doing plain HF downloads with the watchdog disabled."""
|
||||
import importlib
|
||||
|
||||
class _BlockShared:
|
||||
def find_spec(
|
||||
self,
|
||||
name,
|
||||
path = None,
|
||||
target = None,
|
||||
):
|
||||
if name == "unsloth_zoo.hf_xet_fallback":
|
||||
raise ModuleNotFoundError(f"No module named '{name}'", name = name)
|
||||
return None
|
||||
|
||||
finder = _BlockShared()
|
||||
saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None)
|
||||
saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
sys.meta_path.insert(0, finder)
|
||||
try:
|
||||
degraded = importlib.import_module("utils.hf_xet_fallback")
|
||||
|
||||
# Boots without raising and mirrors the shared API surface.
|
||||
assert issubclass(degraded.DownloadStallError, RuntimeError)
|
||||
assert degraded.child_should_disable_xet({"disable_xet": True}) is True
|
||||
assert degraded.get_hf_download_state(["x"]) is None # unmeasurable
|
||||
event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None)
|
||||
assert hasattr(event, "set") and not event.is_set() # never fires
|
||||
|
||||
# Degraded mode still emits heartbeats so the inactivity deadline is not tripped.
|
||||
import time as _time
|
||||
|
||||
beats = []
|
||||
hb_stop = degraded.start_watchdog(
|
||||
repo_ids = ["x"],
|
||||
on_stall = lambda m: None,
|
||||
on_heartbeat = beats.append,
|
||||
interval = 0.02,
|
||||
)
|
||||
try:
|
||||
deadline = _time.monotonic() + 2.0
|
||||
while not beats and _time.monotonic() < deadline:
|
||||
_time.sleep(0.02)
|
||||
assert beats, "degraded watchdog emitted no heartbeat"
|
||||
finally:
|
||||
hb_stop.set()
|
||||
|
||||
# Downloads fall back to plain huggingface_hub (no watchdog, no crash).
|
||||
called = {}
|
||||
|
||||
def _fake_snapshot(repo_id, **kwargs):
|
||||
called["repo_id"] = repo_id
|
||||
return "/snap-dir"
|
||||
|
||||
monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot)
|
||||
assert degraded.snapshot_download_with_xet_fallback("org/model") == "/snap-dir"
|
||||
assert called["repo_id"] == "org/model"
|
||||
|
||||
# Cancellation still holds: an already-set cancel_event aborts before the HF download.
|
||||
import threading as _threading
|
||||
|
||||
cancelled = _threading.Event()
|
||||
cancelled.set()
|
||||
called.clear()
|
||||
with pytest.raises(RuntimeError, match = "Cancelled"):
|
||||
degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled)
|
||||
assert "repo_id" not in called, "degraded download ran despite cancellation"
|
||||
finally:
|
||||
sys.meta_path.remove(finder)
|
||||
sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
if saved_shared is not None:
|
||||
sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared
|
||||
if saved_shim is not None:
|
||||
sys.modules["utils.hf_xet_fallback"] = saved_shim
|
||||
|
||||
|
||||
def test_per_file_independent_fallback(monkeypatch):
|
||||
"""A stalled shard falls back; a sibling shard that succeeds does not."""
|
||||
monkeypatch.setattr(
|
||||
"hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
|
||||
)
|
||||
fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")])
|
||||
assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a"
|
||||
assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b"
|
||||
assert [c.disable_xet for c in fake.calls] == [False, False, True]
|
||||
def test_degrades_when_unsloth_zoo_entirely_absent():
|
||||
"""When unsloth_zoo is absent entirely, the import raises
|
||||
ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still
|
||||
degrades and does not re-raise, breaking every Studio import that pulls it in."""
|
||||
import importlib
|
||||
|
||||
class _BlockZoo:
|
||||
def find_spec(
|
||||
self,
|
||||
name,
|
||||
path = None,
|
||||
target = None,
|
||||
):
|
||||
# Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'.
|
||||
if name == "unsloth_zoo" or name.startswith("unsloth_zoo."):
|
||||
raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo")
|
||||
return None
|
||||
|
||||
finder = _BlockZoo()
|
||||
saved = {
|
||||
k: v
|
||||
for k, v in list(sys.modules.items())
|
||||
if k == "unsloth_zoo" or k.startswith("unsloth_zoo.")
|
||||
}
|
||||
for k in saved:
|
||||
del sys.modules[k]
|
||||
saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
sys.meta_path.insert(0, finder)
|
||||
try:
|
||||
degraded = importlib.import_module("utils.hf_xet_fallback")
|
||||
# Boots without raising and exposes the stub API.
|
||||
assert issubclass(degraded.DownloadStallError, RuntimeError)
|
||||
assert degraded.get_hf_download_state(["x"]) is None
|
||||
event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None)
|
||||
assert hasattr(event, "set") and not event.is_set()
|
||||
finally:
|
||||
sys.meta_path.remove(finder)
|
||||
sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
sys.modules.update(saved)
|
||||
if saved_shim is not None:
|
||||
sys.modules["utils.hf_xet_fallback"] = saved_shim
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect
|
||||
# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _safe_path() -> str:
|
||||
def test_degrades_when_shared_helper_import_raises_importerror():
|
||||
"""unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only
|
||||
Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too."""
|
||||
import importlib
|
||||
|
||||
class _BlockWithImportError:
|
||||
def find_spec(
|
||||
self,
|
||||
name,
|
||||
path = None,
|
||||
target = None,
|
||||
):
|
||||
if name == "unsloth_zoo.hf_xet_fallback":
|
||||
# Mirror a torch-less install: a plain ImportError with no .name.
|
||||
raise ImportError("Unsloth: Pytorch is not installed.")
|
||||
return None
|
||||
|
||||
finder = _BlockWithImportError()
|
||||
saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None)
|
||||
saved_zoo = sys.modules.pop("unsloth_zoo", None)
|
||||
saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
sys.meta_path.insert(0, finder)
|
||||
try:
|
||||
degraded = importlib.import_module("utils.hf_xet_fallback")
|
||||
assert issubclass(degraded.DownloadStallError, RuntimeError)
|
||||
assert degraded.get_hf_download_state(["x"]) is None
|
||||
event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None)
|
||||
assert hasattr(event, "set") and not event.is_set()
|
||||
finally:
|
||||
sys.meta_path.remove(finder)
|
||||
sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
if saved_shared is not None:
|
||||
sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared
|
||||
if saved_zoo is not None:
|
||||
sys.modules["unsloth_zoo"] = saved_zoo
|
||||
if saved_shim is not None:
|
||||
sys.modules["utils.hf_xet_fallback"] = saved_shim
|
||||
|
||||
|
||||
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.
|
||||
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
|
||||
return os.environ.get("PATH", "")
|
||||
|
||||
monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False)
|
||||
seen_env = []
|
||||
|
||||
class _GpuGatedBlocker:
|
||||
def find_spec(
|
||||
self,
|
||||
name,
|
||||
path = None,
|
||||
target = None,
|
||||
):
|
||||
# Crash is in unsloth_zoo's __init__, so intercept "unsloth_zoo" itself (the parent).
|
||||
if name == "unsloth_zoo":
|
||||
# Record the env each attempt sees; raise the no-GPU error both times so the shim
|
||||
# degrades.
|
||||
seen_env.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT"))
|
||||
raise NotImplementedError("Unsloth cannot find any torch accelerator")
|
||||
return None
|
||||
|
||||
finder = _GpuGatedBlocker()
|
||||
saved = {
|
||||
k: v
|
||||
for k, v in list(sys.modules.items())
|
||||
if k == "unsloth_zoo" or k.startswith("unsloth_zoo.")
|
||||
}
|
||||
for k in saved:
|
||||
del sys.modules[k]
|
||||
saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
sys.meta_path.insert(0, finder)
|
||||
try:
|
||||
degraded = importlib.import_module("utils.hf_xet_fallback")
|
||||
# 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(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)
|
||||
sys.modules.pop("utils.hf_xet_fallback", None)
|
||||
sys.modules.update(saved)
|
||||
if saved_shim is not None:
|
||||
sys.modules["utils.hf_xet_fallback"] = saved_shim
|
||||
|
||||
|
||||
def test_disable_xet_constant_set_in_fresh_interpreter():
|
||||
code = (
|
||||
"from huggingface_hub import constants as c; "
|
||||
"import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)"
|
||||
)
|
||||
proc = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()},
|
||||
capture_output = True,
|
||||
text = True,
|
||||
)
|
||||
assert proc.returncode == 0, (
|
||||
f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True "
|
||||
f"(rc={proc.returncode}): {proc.stderr}"
|
||||
)
|
||||
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)
|
||||
|
||||
def test_default_leaves_xet_enabled():
|
||||
code = (
|
||||
"from huggingface_hub import constants as c; "
|
||||
"import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)"
|
||||
)
|
||||
proc = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET
|
||||
capture_output = True,
|
||||
text = True,
|
||||
)
|
||||
assert proc.returncode == 0, (
|
||||
f"without the env var, constants.HF_HUB_DISABLE_XET was not False "
|
||||
f"(rc={proc.returncode}): {proc.stderr}"
|
||||
)
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,42 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Default Chat model metadata must not block on remote Hugging Face discovery."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND = Path(__file__).resolve().parent.parent
|
||||
if str(_BACKEND) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND))
|
||||
|
||||
from core.inference.orchestrator import InferenceOrchestrator # noqa: E402
|
||||
|
||||
|
||||
def test_default_models_returns_static_defaults_before_top_fetch(monkeypatch):
|
||||
sleep_seconds = 2.0
|
||||
|
||||
def _slow_fetch(self: InferenceOrchestrator) -> None:
|
||||
time.sleep(sleep_seconds)
|
||||
self._top_gguf_cache = ["unsloth/slow-GGUF"]
|
||||
self._top_models_ready.set()
|
||||
|
||||
monkeypatch.setattr(InferenceOrchestrator, "_fetch_top_models", _slow_fetch)
|
||||
|
||||
orchestrator = InferenceOrchestrator()
|
||||
started = time.monotonic()
|
||||
defaults = orchestrator.default_models
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert elapsed < 0.5, f"default_models blocked for {elapsed:.2f}s"
|
||||
assert defaults == orchestrator._static_models
|
||||
assert "unsloth/slow-GGUF" not in defaults
|
||||
|
||||
deadline = time.monotonic() + sleep_seconds + 5
|
||||
while not orchestrator._top_models_ready.is_set() and time.monotonic() < deadline:
|
||||
time.sleep(0.05)
|
||||
|
||||
assert "unsloth/slow-GGUF" in orchestrator.default_models
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
320
studio/backend/tests/test_llama_admission.py
Normal file
320
studio/backend/tests/test_llama_admission.py
Normal file
|
|
@ -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())
|
||||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
81
studio/backend/tests/test_llama_cpp_stream_cancel.py
Normal file
81
studio/backend/tests/test_llama_cpp_stream_cancel.py
Normal file
|
|
@ -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 == []
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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,
|
||||
|
|
|
|||
193
studio/backend/tests/test_llama_cpp_vulkan_probe.py
Normal file
193
studio/backend/tests/test_llama_cpp_vulkan_probe.py
Normal file
|
|
@ -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<i>``.
|
||||
|
||||
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"]))
|
||||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -587,10 +587,12 @@ def test_tool_xml_strip_handles_hyphenated_function_names():
|
|||
import re as _re
|
||||
from pathlib import Path
|
||||
|
||||
from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
|
||||
|
||||
src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text()
|
||||
m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL)
|
||||
assert m, "could not extract _TOOL_XML_RE"
|
||||
ns: dict = {"_re": _re}
|
||||
ns: dict = {"_re": _re, "_DS_OPEN_SRC": _DS_OPEN_SRC}
|
||||
exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns)
|
||||
rx = ns["_TOOL_XML_RE"]
|
||||
stripped = rx.sub(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ import sys
|
|||
import types
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class _DummyMetal:
|
||||
@staticmethod
|
||||
|
|
@ -100,6 +102,32 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):
|
|||
]
|
||||
assert backend._is_vlm is False
|
||||
assert isinstance(backend._tokenizer, _DummyTokenizer)
|
||||
# Non-LoRA text model: no base_model on the record.
|
||||
assert backend.models["fake/text"]["base_model"] is None
|
||||
|
||||
|
||||
def test_mlx_text_lora_record_keeps_base_model_for_native_template(monkeypatch):
|
||||
# A LoRA adapter's own tokenizer often ships no chat template; the native tool-calling template
|
||||
# lives on the base model.
|
||||
_install_fake_mlx(monkeypatch)
|
||||
calls = []
|
||||
_install_fake_fast_mlx(monkeypatch, calls)
|
||||
|
||||
from core.inference.mlx_inference import MLXInferenceBackend
|
||||
|
||||
backend = MLXInferenceBackend()
|
||||
config = SimpleNamespace(
|
||||
identifier = "fake/text-adapter",
|
||||
is_vision = False,
|
||||
is_lora = True,
|
||||
base_model = "fake/text-base",
|
||||
)
|
||||
|
||||
assert backend.load_model(config, max_seq_length = 4096, hf_token = "hf-token")
|
||||
|
||||
record = backend.models["fake/text-adapter"]
|
||||
assert record["is_lora"] is True
|
||||
assert record["base_model"] == "fake/text-base"
|
||||
|
||||
|
||||
def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite(
|
||||
|
|
@ -159,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
|
||||
|
|
@ -188,12 +339,12 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
|
|||
_install_fake_mlx(monkeypatch)
|
||||
from core.inference.mlx_inference import MLXInferenceBackend
|
||||
|
||||
captured = {}
|
||||
# The text path renders once with tools, then the native-template fallback makes a second no-
|
||||
# tools probe call (tools=None) to detect whether the template dropped the schema.
|
||||
captured_calls = []
|
||||
|
||||
def _fake_apply(tokenizer, messages, **kwargs):
|
||||
captured["tokenizer"] = tokenizer
|
||||
captured["messages"] = messages
|
||||
captured["kwargs"] = kwargs
|
||||
captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs})
|
||||
return "<rendered prompt>"
|
||||
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -248,8 +399,15 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
|
|||
)
|
||||
)
|
||||
assert out == ["hi"]
|
||||
# The toggled kwargs must reach the chat-template helper.
|
||||
assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}]
|
||||
assert captured["kwargs"]["enable_thinking"] is True
|
||||
assert captured["kwargs"]["reasoning_effort"] == "medium"
|
||||
assert captured["kwargs"]["preserve_thinking"] is True
|
||||
# The toggled kwargs must reach the chat-template helper on the real render
|
||||
# (one of the calls carries the tools; the fallback probe passes tools=None).
|
||||
tool_renders = [
|
||||
c
|
||||
for c in captured_calls
|
||||
if c["kwargs"].get("tools") == [{"function": {"name": "web_search"}}]
|
||||
]
|
||||
assert tool_renders, captured_calls
|
||||
render = tool_renders[0]
|
||||
assert render["kwargs"]["enable_thinking"] is True
|
||||
assert render["kwargs"]["reasoning_effort"] == "medium"
|
||||
assert render["kwargs"]["preserve_thinking"] is True
|
||||
|
|
|
|||
|
|
@ -271,6 +271,29 @@ def test_stack_available_requires_runtime_imports_and_versions(monkeypatch):
|
|||
assert imported == list(mr._MLX_RUNTIME_IMPORTS)
|
||||
|
||||
|
||||
def test_mlx_packages_exclude_known_bad_mlx_lm():
|
||||
# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5); the install spec
|
||||
# must exclude it so the resolver picks 0.31.2 or >=0.31.4. See mlx-lm #1242.
|
||||
(mlx_lm_spec,) = [p for p in mr.MLX_PACKAGES if p.startswith("mlx-lm")]
|
||||
assert mlx_lm_spec == "mlx-lm>=0.22.0,!=0.31.3"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_form", ["0.31.3", "0.31.3.0"])
|
||||
def test_known_bad_installed_mlx_lm_triggers_repair(monkeypatch, bad_form):
|
||||
# An installed 0.31.3 counts as unsatisfied so the self-heal replaces it;
|
||||
# parsed-Version compare also catches the trailing-zero form 0.31.3.0.
|
||||
import importlib.metadata as metadata
|
||||
|
||||
def _version(name):
|
||||
return bad_form if name == "mlx-lm" else mr._MLX_MIN_VERSIONS[name]
|
||||
|
||||
monkeypatch.setattr(metadata, "version", _version)
|
||||
monkeypatch.setattr(
|
||||
mr.importlib, "import_module", lambda _n: pytest.fail("versions must gate imports")
|
||||
)
|
||||
assert mr.mlx_stack_available() is False
|
||||
|
||||
|
||||
def test_no_op_off_apple_silicon(monkeypatch):
|
||||
monkeypatch.setattr(mr, "is_apple_silicon", lambda: False)
|
||||
called = {"n": 0}
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@
|
|||
Covers:
|
||||
* GGUF variant listing computes update_available from the already-fetched
|
||||
sibling metadata instead of a second Hub call.
|
||||
* hf_hub_download_with_xet_fallback(force_download=True) bypasses the
|
||||
try_to_load_from_cache cache-first early-return.
|
||||
* hf_hub_download_with_xet_fallback forwards force_download through the shim to the
|
||||
shared unsloth_zoo helper (which owns the cache-first early-return and its bypass).
|
||||
|
||||
The cache "Update" action now runs through the download manager as a normal
|
||||
managed download (so it shows in the Downloads panel with progress + cancel),
|
||||
|
|
@ -341,44 +341,26 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
|
|||
# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ───
|
||||
|
||||
|
||||
def test_force_download_bypasses_cache_first_early_return(monkeypatch):
|
||||
"""force_download=True skips the try_to_load_from_cache early-return and
|
||||
proceeds to the real download path; force_download=False returns the cached
|
||||
path without ever attempting a download (X2/F2)."""
|
||||
import huggingface_hub as hf
|
||||
def test_force_download_is_forwarded_through_the_shim(monkeypatch):
|
||||
"""The shim's contract is to forward force_download unchanged to the shared helper (which owns the
|
||||
cache-first early-return and bypass). Verify both False and True reach it (X2/F2)."""
|
||||
import utils.hf_xet_fallback as X
|
||||
|
||||
cached_path = "/cache/blob/cached.gguf"
|
||||
seen = []
|
||||
|
||||
# Pretend the blob IS cached on disk (try_to_load_from_cache is imported
|
||||
# inside the function from huggingface_hub, and os.path.exists must agree).
|
||||
monkeypatch.setattr(hf, "try_to_load_from_cache", lambda *a, **k: cached_path, raising = False)
|
||||
monkeypatch.setattr(X.os.path, "exists", lambda p: True, raising = False)
|
||||
def fake_shared(repo_id, filename, token, **kwargs):
|
||||
seen.append(kwargs.get("force_download"))
|
||||
return "/downloaded/path"
|
||||
|
||||
attempts = []
|
||||
monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True)
|
||||
|
||||
def fake_attempt(repo_id, filename, token, **kwargs):
|
||||
attempts.append(
|
||||
{"repo_id": repo_id, "filename": filename, "force": kwargs.get("force_download")}
|
||||
)
|
||||
return ("ok", "/freshly/downloaded/path")
|
||||
|
||||
monkeypatch.setattr(X, "_run_download_attempt", fake_attempt, raising = True)
|
||||
|
||||
# force_download=False: cache-first early-return, no download attempt.
|
||||
out = X.hf_hub_download_with_xet_fallback(
|
||||
X.hf_hub_download_with_xet_fallback(
|
||||
"unsloth/repo", "model.gguf", token = None, force_download = False
|
||||
)
|
||||
assert out == cached_path
|
||||
assert attempts == [] # never reached the real download
|
||||
|
||||
# force_download=True: bypass the early-return, run the real download.
|
||||
out2 = X.hf_hub_download_with_xet_fallback(
|
||||
X.hf_hub_download_with_xet_fallback(
|
||||
"unsloth/repo", "model.gguf", token = None, force_download = True
|
||||
)
|
||||
assert out2 == "/freshly/downloaded/path"
|
||||
assert len(attempts) == 1
|
||||
assert attempts[0]["force"] is True
|
||||
assert seen == [False, True] # the shim forwards force_download to the shared helper unchanged
|
||||
|
||||
|
||||
# ── multi-revision GGUF blob comparison and update reclaim ──
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
176
studio/backend/tests/test_native_template_trust_remote_code.py
Normal file
176
studio/backend/tests/test_native_template_trust_remote_code.py
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Regression tests for trust_remote_code in the native-template fallback.
|
||||
|
||||
``render_native_template`` re-fetches a model's native chat template from its
|
||||
repo when an Unsloth override template (mistral, gemma-4) dropped the tools
|
||||
schema. For a model loaded with ``trust_remote_code=True`` whose tokenizer repo
|
||||
carries custom code, the secondary ``AutoTokenizer.from_pretrained`` must re-use
|
||||
that same consent or transformers raises (it requires ``trust_remote_code`` to
|
||||
instantiate a custom tokenizer class), the ``except`` swallows it, and the
|
||||
request silently keeps the tool-dropping prompt even though the user already
|
||||
consented to remote code for the model load.
|
||||
|
||||
These tests pin that the stored ``trust_remote_code`` is threaded to the reload,
|
||||
that the reload is skipped (returns ``None`` without executing code) when no
|
||||
consent is stored, and that both backend ``model_info`` dicts persist the flag at
|
||||
load time so the read lands on a value ``load_model`` actually set.
|
||||
"""
|
||||
|
||||
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)
|
||||
|
||||
# ``chat_template_helpers`` is dependency-light (copy / logging / typing, with the
|
||||
# transformers import deferred inside the function). Load it directly so the test
|
||||
# runs without importing the heavy ``core.inference`` package (unsloth / torch).
|
||||
_HELPERS_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_template_helpers.py"
|
||||
_spec = importlib.util.spec_from_file_location("_native_tpl_trc_test", _HELPERS_PATH)
|
||||
chat_template_helpers = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(chat_template_helpers)
|
||||
|
||||
render_native_template = chat_template_helpers.render_native_template
|
||||
|
||||
|
||||
# A native template that emits a tools section only when tools are provided, so the
|
||||
# with-tools vs no-tools render differs and ``render_native_template`` accepts it.
|
||||
_NATIVE_TEMPLATE = (
|
||||
"{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}"
|
||||
"{% if tools %}[AVAILABLE_TOOLS]{{ tools }}[/AVAILABLE_TOOLS]\n{% endif %}"
|
||||
"{% if add_generation_prompt %}assistant:{% endif %}"
|
||||
)
|
||||
|
||||
_MESSAGES = [{"role": "user", "content": "what is the weather"}]
|
||||
_TOOLS = [{"type": "function", "function": {"name": "get_weather"}}]
|
||||
|
||||
|
||||
class _JinjaTokenizer:
|
||||
"""Minimal tokenizer whose ``apply_chat_template`` renders ``self.chat_template``.
|
||||
|
||||
Stands in for the live model tokenizer that ``render_native_template`` shallow-
|
||||
copies and re-points at the native template before rendering.
|
||||
"""
|
||||
|
||||
def __init__(self, chat_template):
|
||||
self.chat_template = chat_template
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
tools = None,
|
||||
**kwargs,
|
||||
):
|
||||
from jinja2 import BaseLoader, Environment
|
||||
env = Environment(loader = BaseLoader())
|
||||
return env.from_string(self.chat_template).render(
|
||||
messages = messages,
|
||||
tools = tools,
|
||||
add_generation_prompt = add_generation_prompt,
|
||||
)
|
||||
|
||||
|
||||
def _install_custom_code_tokenizer(monkeypatch):
|
||||
"""Patch ``AutoTokenizer.from_pretrained`` to mimic a custom-code repo: raise
|
||||
unless ``trust_remote_code`` is truthy, else return a tokenizer carrying the
|
||||
native template. Records the ``trust_remote_code`` it was called with."""
|
||||
pytest.importorskip("jinja2")
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
calls = {}
|
||||
|
||||
def fake_from_pretrained(
|
||||
model_id,
|
||||
*args,
|
||||
trust_remote_code = False,
|
||||
token = None,
|
||||
**kwargs,
|
||||
):
|
||||
calls["trust_remote_code"] = trust_remote_code
|
||||
calls["model_id"] = model_id
|
||||
calls["token"] = token
|
||||
if not trust_remote_code:
|
||||
# Mirrors transformers.dynamic_module_utils.resolve_trust_remote_code:
|
||||
# has_remote_code and not has_local_code and not trust_remote_code -> ValueError.
|
||||
raise ValueError(
|
||||
f"The repository {model_id} contains custom code which must be executed "
|
||||
"to correctly load the model. Please pass the argument "
|
||||
"`trust_remote_code=True` to allow custom code to be run."
|
||||
)
|
||||
return _JinjaTokenizer(_NATIVE_TEMPLATE)
|
||||
|
||||
monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained))
|
||||
return calls
|
||||
|
||||
|
||||
def _model_info(trust_remote_code):
|
||||
return {
|
||||
"native_chat_template": None, # force the repo reload path
|
||||
"base_model": None, # non-LoRA: template_source == active_model_name
|
||||
"trust_remote_code": trust_remote_code,
|
||||
# Live tokenizer that gets shallow-copied + re-pointed at the native template.
|
||||
"tokenizer": _JinjaTokenizer("OVERRIDE-THAT-DROPS-TOOLS"),
|
||||
}
|
||||
|
||||
|
||||
def test_native_reload_passes_stored_trust_remote_code(monkeypatch):
|
||||
"""With ``trust_remote_code`` stored on ``model_info`` the custom-code reload
|
||||
succeeds and the tools-advertising native prompt is returned. This FAILS before
|
||||
the fix (reload omits the flag, raises, is swallowed, returns None)."""
|
||||
calls = _install_custom_code_tokenizer(monkeypatch)
|
||||
model_info = _model_info(trust_remote_code = True)
|
||||
|
||||
out = render_native_template(
|
||||
model_info = model_info,
|
||||
active_model_name = "acme/custom-tokenizer-model",
|
||||
messages = _MESSAGES,
|
||||
tools = _TOOLS,
|
||||
)
|
||||
|
||||
assert out is not None, "native fallback should render the tools prompt with consent"
|
||||
assert "[AVAILABLE_TOOLS]" in out
|
||||
assert "get_weather" in out
|
||||
assert calls["trust_remote_code"] is True # the stored consent was threaded through
|
||||
# A successful fetch is cached so the next tool turn skips the reload.
|
||||
assert model_info["native_chat_template"] == _NATIVE_TEMPLATE
|
||||
|
||||
|
||||
def test_native_reload_without_consent_returns_none(monkeypatch):
|
||||
"""Without stored consent the custom-code reload raises, is swallowed, and
|
||||
``render_native_template`` returns None (no unconsented code execution). Proves
|
||||
the stored flag -- not a hard-coded True -- drives the reload."""
|
||||
calls = _install_custom_code_tokenizer(monkeypatch)
|
||||
model_info = _model_info(trust_remote_code = False)
|
||||
|
||||
out = render_native_template(
|
||||
model_info = model_info,
|
||||
active_model_name = "acme/custom-tokenizer-model",
|
||||
messages = _MESSAGES,
|
||||
tools = _TOOLS,
|
||||
)
|
||||
|
||||
assert out is None
|
||||
assert calls["trust_remote_code"] is False
|
||||
# A failed fetch must not be cached as "no template" (would pin the tool drop).
|
||||
assert model_info["native_chat_template"] is None
|
||||
|
||||
|
||||
def test_backend_model_info_persists_trust_remote_code():
|
||||
"""Both backends must store ``trust_remote_code`` on their per-model info dict so
|
||||
``render_native_template`` can source the consent value. Guards against the read
|
||||
landing on a key ``load_model`` never sets (which would silently no-op the fix)."""
|
||||
inf = (Path(_BACKEND_DIR) / "core" / "inference" / "inference.py").read_text()
|
||||
mlx = (Path(_BACKEND_DIR) / "core" / "inference" / "mlx_inference.py").read_text()
|
||||
assert '"trust_remote_code": trust_remote_code,' in inf
|
||||
assert '"trust_remote_code": trust_remote_code,' in mlx
|
||||
99
studio/backend/tests/test_nudge_tool_calls_wiring.py
Normal file
99
studio/backend/tests/test_nudge_tool_calls_wiring.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Wiring guard for the plan-without-action ``nudge_tool_calls`` policy.
|
||||
|
||||
Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths
|
||||
(safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat +
|
||||
Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off).
|
||||
|
||||
Mechanism (verified here without loading a model):
|
||||
|
||||
* every backend tool-loop entry point accepts and forwards ``nudge_tool_calls``
|
||||
(safetensors -> ``InferenceBackend``; MLX -> ``InferenceOrchestrator``; both
|
||||
call the shared ``run_safetensors_tool_loop``; GGUF -> ``LlamaCppBackend``);
|
||||
* the safetensors/MLX loop gates the retry on a truthy flag (new retry ->
|
||||
opt-in), while the GGUF loop keeps its pre-existing default-on behaviour
|
||||
(``None`` keeps nudging) so an omitted flag never disables GGUF;
|
||||
* the API request models default the flag to ``None`` (opt-in / off);
|
||||
* the Studio-facing routes forward the request's flag, and the Studio frontend
|
||||
sends ``nudge_tool_calls: true`` -- exercised behaviourally in
|
||||
``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
from core.inference.orchestrator import InferenceOrchestrator
|
||||
from core.inference.safetensors_agentic import run_safetensors_tool_loop
|
||||
|
||||
try:
|
||||
# core.inference.inference imports unsloth at module scope, which requires
|
||||
# unsloth_zoo. The dependency-light backend CI matrix job does not install
|
||||
# it, so the safetensors InferenceBackend is folded into the checks below
|
||||
# only when the unsloth stack is importable (local runs / full CI); the
|
||||
# other entry points are always checked.
|
||||
from core.inference.inference import InferenceBackend
|
||||
except ImportError:
|
||||
InferenceBackend = None
|
||||
|
||||
|
||||
def _params(fn):
|
||||
return inspect.signature(fn).parameters
|
||||
|
||||
|
||||
def test_shared_loop_accepts_nudge_flag():
|
||||
assert "nudge_tool_calls" in _params(run_safetensors_tool_loop)
|
||||
|
||||
|
||||
def test_backends_accept_the_flag():
|
||||
methods = [
|
||||
InferenceOrchestrator.generate_chat_completion_with_tools,
|
||||
LlamaCppBackend.generate_chat_completion_with_tools,
|
||||
]
|
||||
if InferenceBackend is not None: # safetensors path; needs the unsloth stack
|
||||
methods.append(InferenceBackend.generate_chat_completion_with_tools)
|
||||
for method in methods:
|
||||
assert "nudge_tool_calls" in _params(method), method.__qualname__
|
||||
|
||||
|
||||
def test_delegating_backends_forward_the_flag_to_the_shared_loop():
|
||||
# safetensors (in-process transformers) and MLX (parent-process orchestrator)
|
||||
# both delegate to run_safetensors_tool_loop; GGUF runs its own in-file loop
|
||||
# and consumes the flag directly (asserted separately by the gate test).
|
||||
methods = [InferenceOrchestrator.generate_chat_completion_with_tools]
|
||||
if InferenceBackend is not None: # safetensors path; needs the unsloth stack
|
||||
methods.append(InferenceBackend.generate_chat_completion_with_tools)
|
||||
for method in methods:
|
||||
src = inspect.getsource(method)
|
||||
assert "nudge_tool_calls = nudge_tool_calls" in src, method.__qualname__
|
||||
|
||||
|
||||
def test_safetensors_loop_is_opt_in_while_gguf_stays_default_on():
|
||||
# Safetensors/MLX: the retry is new here, so it requires a truthy flag.
|
||||
sf_src = inspect.getsource(run_safetensors_tool_loop)
|
||||
assert "and nudge_tool_calls" in sf_src
|
||||
# GGUF: pre-existing nudge must not be accidentally disabled -- an omitted
|
||||
# (None) flag keeps nudging; only an explicit False turns it off.
|
||||
gguf_src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools)
|
||||
assert "nudge_tool_calls is None or nudge_tool_calls" in gguf_src
|
||||
|
||||
|
||||
def test_api_request_models_default_the_flag_off():
|
||||
from models.inference import AnthropicMessagesRequest, ChatCompletionRequest
|
||||
for model in (ChatCompletionRequest, AnthropicMessagesRequest):
|
||||
field = model.model_fields["nudge_tool_calls"]
|
||||
assert field.default is None, model.__name__
|
||||
|
||||
|
||||
def test_studio_routes_forward_the_request_flag():
|
||||
# The Studio chat frontend posts to /v1/chat/completions and /v1/messages
|
||||
# with nudge_tool_calls=true; the route handlers forward the request value
|
||||
# (external API clients that omit it fall back to the opt-in default).
|
||||
from routes import inference as routes_inference
|
||||
for handler in (
|
||||
routes_inference.openai_chat_completions,
|
||||
routes_inference.anthropic_messages,
|
||||
):
|
||||
src = inspect.getsource(handler)
|
||||
assert "nudge_tool_calls = payload.nudge_tool_calls" in src, handler.__name__
|
||||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
1622
studio/backend/tests/test_orchestrator_unload_cancel.py
Normal file
1622
studio/backend/tests/test_orchestrator_unload_cancel.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -280,6 +280,56 @@ class TestStreamHealer:
|
|||
assert [c["id"] for c in calls] == ["call_0", "call_1"]
|
||||
assert _events_text(events).strip() == "then"
|
||||
|
||||
def test_mistral_array_multiple_calls_all_promoted_in_stream(self):
|
||||
# A canonical Mistral [TOOL_CALLS] array carries several calls under a
|
||||
# SINGLE signal. Draining only the first call would leave the residue
|
||||
# starting at ",{...}]" (no signal), so later calls in the same array
|
||||
# must be promoted in the same pass, not flushed as raw text.
|
||||
healer = StreamToolCallHealer({"get_weather", "get_time"})
|
||||
array = (
|
||||
'[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},'
|
||||
'{"name":"get_time","arguments":{"tz":"UTC"}}]'
|
||||
)
|
||||
events = healer.feed(array) + healer.finalize()
|
||||
calls = _events_calls(events)
|
||||
assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"]
|
||||
assert [c["id"] for c in calls] == ["call_0", "call_1"]
|
||||
assert _events_text(events) == ""
|
||||
|
||||
def test_mistral_array_multiple_calls_promoted_char_by_char(self):
|
||||
healer = StreamToolCallHealer({"get_weather", "get_time"})
|
||||
array = (
|
||||
'[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},'
|
||||
'{"name":"get_time","arguments":{"tz":"UTC"}}]'
|
||||
)
|
||||
events = []
|
||||
for ch in array:
|
||||
events += healer.feed(ch)
|
||||
events += healer.finalize()
|
||||
calls = _events_calls(events)
|
||||
assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"]
|
||||
assert _events_text(events) == ""
|
||||
|
||||
def test_mistral_array_undeclared_middle_kept_as_text_others_promoted(self):
|
||||
# A mid-array element for a tool that is not declared must survive as
|
||||
# text while the declared neighbours on either side still promote in
|
||||
# document order.
|
||||
healer = StreamToolCallHealer({"a", "c"})
|
||||
array = (
|
||||
'[TOOL_CALLS][{"name":"a","arguments":{}},'
|
||||
'{"name":"b","arguments":{}},{"name":"c","arguments":{}}]'
|
||||
)
|
||||
events = healer.feed(array) + healer.finalize()
|
||||
assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "c"]
|
||||
assert '"b"' in _events_text(events)
|
||||
|
||||
def test_mistral_array_then_trailing_prose(self):
|
||||
healer = StreamToolCallHealer({"a", "b"})
|
||||
array = '[TOOL_CALLS][{"name":"a","arguments":{}},{"name":"b","arguments":{}}]'
|
||||
events = healer.feed(f"{array} all done") + healer.finalize()
|
||||
assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "b"]
|
||||
assert "all done" in _events_text(events)
|
||||
|
||||
def test_incomplete_call_healed_at_finalize(self):
|
||||
healer = StreamToolCallHealer({"Bash"})
|
||||
events = healer.feed('<tool_call>{"name":"Bash","arguments":{"cmd":"ls"}}')
|
||||
|
|
@ -465,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)])
|
||||
|
|
@ -1356,3 +1407,42 @@ class TestOpenaiStreamingRoute:
|
|||
assert chunks[0] == line + "\n\n" # byte-for-byte relay
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
class TestHealerSignalAlignment:
|
||||
"""The passthrough healer buffers only formats its parser can promote.
|
||||
The loops' bare [ARGS] rehearsal signal is gated on active tool names
|
||||
there; ungated in the healer it would stall legitimate prose until
|
||||
finalization without ever producing a promotable call."""
|
||||
|
||||
def test_heal_signals_are_promotable_formats_only(self):
|
||||
from core.inference.passthrough_healing import _HEAL_SIGNALS
|
||||
assert set(_HEAL_SIGNALS) == {
|
||||
"<tool_call>",
|
||||
"<|tool_call>",
|
||||
"<function=",
|
||||
"[TOOL_CALLS]",
|
||||
}
|
||||
|
||||
def test_prose_with_bare_args_marker_streams_through(self):
|
||||
healer = StreamToolCallHealer({"Bash"})
|
||||
chunks = [
|
||||
"Use the pattern foo",
|
||||
"[ARGS] in templates when calling tools, ",
|
||||
"and remember to close it.",
|
||||
]
|
||||
streamed = ""
|
||||
for chunk in chunks:
|
||||
streamed += _events_text(healer.feed(chunk))
|
||||
# Incremental relay: nothing withheld for finalize.
|
||||
assert streamed == "".join(chunks)
|
||||
final = healer.finalize()
|
||||
assert not _events_calls(final)
|
||||
assert not healer.healed
|
||||
|
||||
def test_bracket_tool_calls_still_promote_in_stream(self):
|
||||
healer = StreamToolCallHealer({"web_search"})
|
||||
events = healer.feed('[TOOL_CALLS]web_search{"query": "unsloth docs"}') + healer.finalize()
|
||||
(call,) = _events_calls(events)
|
||||
assert call["function"]["name"] == "web_search"
|
||||
assert healer.healed
|
||||
|
|
|
|||
1011
studio/backend/tests/test_pr5624_regressions.py
Normal file
1011
studio/backend/tests/test_pr5624_regressions.py
Normal file
File diff suppressed because it is too large
Load diff
252
studio/backend/tests/test_presence_penalty.py
Normal file
252
studio/backend/tests/test_presence_penalty.py
Normal file
|
|
@ -0,0 +1,252 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
"""Presence-penalty parity between the GGUF path and the safetensors/MLX paths.
|
||||
|
||||
The safetensors path historically dropped ``presence_penalty``, so the SAME model
|
||||
looked worse served as safetensors. These tests pin the processor semantics
|
||||
(subtract once per distinct completion token, prompt excluded, presence not
|
||||
frequency, zero a no-op, negatives raise) plus a param-propagation regression
|
||||
over route -> orchestrator cmd -> worker gen_kwargs.
|
||||
"""
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from core.inference.presence_penalty import (
|
||||
apply_presence_penalty,
|
||||
_make_presence_penalty_processor,
|
||||
)
|
||||
|
||||
|
||||
def test_seen_token_gets_exactly_minus_penalty_unseen_unchanged():
|
||||
input_ids = torch.tensor([[0, 1, 3]]) # prompt [0, 1], completion [3]
|
||||
scores = torch.zeros(1, 5)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 2)
|
||||
assert out[0, 3].item() == pytest.approx(-1.5)
|
||||
for tok in (0, 1, 2, 4):
|
||||
assert out[0, tok].item() == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_multiplicity_ignored_presence_not_frequency():
|
||||
# Token 3 emitted three times -> still a single -penalty (presence, not freq).
|
||||
input_ids = torch.tensor([[0, 3, 3, 3]])
|
||||
scores = torch.zeros(1, 5)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 2.0, prompt_len = 1)
|
||||
assert out[0, 3].item() == pytest.approx(-2.0)
|
||||
|
||||
|
||||
def test_negative_penalty_raises_seen_logits():
|
||||
input_ids = torch.tensor([[0, 2]])
|
||||
scores = torch.zeros(1, 4)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = -0.5, prompt_len = 1)
|
||||
assert out[0, 2].item() == pytest.approx(0.5)
|
||||
|
||||
|
||||
def test_prompt_tokens_excluded():
|
||||
# Token 7 is prompt-only (untouched); token 4 in the completion is penalized.
|
||||
input_ids = torch.tensor([[7, 4, 4]])
|
||||
scores = torch.zeros(1, 8)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
|
||||
assert out[0, 7].item() == pytest.approx(0.0)
|
||||
assert out[0, 4].item() == pytest.approx(-1.0)
|
||||
|
||||
|
||||
def test_batch_rows_isolated():
|
||||
input_ids = torch.tensor([[0, 1], [0, 2]]) # row completions [1] and [2]
|
||||
scores = torch.zeros(2, 4)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
|
||||
assert out[0, 1].item() == pytest.approx(-1.0)
|
||||
assert out[0, 2].item() == pytest.approx(0.0)
|
||||
assert out[1, 2].item() == pytest.approx(-1.0)
|
||||
assert out[1, 1].item() == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_zero_penalty_is_noop():
|
||||
input_ids = torch.tensor([[0, 1, 2]])
|
||||
scores = torch.randn(1, 5)
|
||||
original = scores.clone()
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 0.0, prompt_len = 1)
|
||||
assert torch.equal(out, original)
|
||||
|
||||
|
||||
def test_empty_completion_is_noop():
|
||||
# prompt_len covers the whole sequence -> nothing generated yet.
|
||||
input_ids = torch.tensor([[0, 1, 2]])
|
||||
scores = torch.randn(1, 5)
|
||||
original = scores.clone()
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 3)
|
||||
assert torch.equal(out, original)
|
||||
|
||||
|
||||
def test_out_of_vocab_id_ignored():
|
||||
# A generated id >= vocab_size (defensive) must not index out of bounds.
|
||||
input_ids = torch.tensor([[0, 9]])
|
||||
scores = torch.zeros(1, 5) # vocab 5, token 9 is out of range
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
|
||||
assert torch.equal(out, torch.zeros(1, 5))
|
||||
|
||||
|
||||
def test_negative_generated_id_ignored():
|
||||
# A negative generated id (defensive) must be dropped, not wrap to scores[-1].
|
||||
input_ids = torch.tensor([[0, -1]])
|
||||
scores = torch.zeros(1, 5)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
|
||||
# Nothing penalized; in particular the last row (the numpy/torch wrap target
|
||||
# for id -1) is untouched.
|
||||
assert torch.equal(out, torch.zeros(1, 5))
|
||||
|
||||
|
||||
def test_mixed_oob_negative_and_valid_ids_only_in_range_penalized():
|
||||
# Completion mixes a valid id (1), an out-of-vocab id (9 >= vocab 5) and a
|
||||
# negative id (-1). Only the in-range distinct id is penalized; OOB/negative
|
||||
# ids are ignored with no crash and no wrong-index wrap. This fails under the
|
||||
# old ``seen[seen < vocab_size]`` filter (id -1 wraps to the last row) and
|
||||
# passes only with the both-ends bound.
|
||||
input_ids = torch.tensor([[0, 1, 9, -1, 1]]) # prompt [0], completion [1, 9, -1, 1]
|
||||
scores = torch.zeros(1, 5)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
|
||||
expected = torch.zeros(1, 5)
|
||||
expected[0, 1] = -1.0 # once per distinct in-range id (multiplicity ignored)
|
||||
assert torch.equal(out, expected)
|
||||
assert out[0, 4].item() == pytest.approx(0.0) # id -1 did not wrap to the last row
|
||||
|
||||
|
||||
def test_dtype_and_device_preserved():
|
||||
input_ids = torch.tensor([[0, 1]])
|
||||
scores = torch.zeros(1, 4, dtype = torch.float16)
|
||||
out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1)
|
||||
assert out.dtype == torch.float16
|
||||
assert out.device == scores.device
|
||||
|
||||
|
||||
def test_processor_none_when_zero():
|
||||
assert _make_presence_penalty_processor(0.0, prompt_len = 0) is None
|
||||
|
||||
|
||||
def test_processor_applies_penalty():
|
||||
proc = _make_presence_penalty_processor(1.5, prompt_len = 2)
|
||||
assert proc is not None
|
||||
input_ids = torch.tensor([[0, 1, 3]])
|
||||
scores = torch.zeros(1, 5)
|
||||
out = proc(input_ids, scores)
|
||||
assert out[0, 3].item() == pytest.approx(-1.5)
|
||||
|
||||
|
||||
def test_processor_composes_with_other_processors():
|
||||
# LogitsProcessorList must run our processor alongside a pre-existing one.
|
||||
from transformers import LogitsProcessor, LogitsProcessorList
|
||||
|
||||
class _AddToTokenZero(LogitsProcessor):
|
||||
def __call__(self, input_ids, scores):
|
||||
scores[:, 0] = scores[:, 0] + 100.0
|
||||
return scores
|
||||
|
||||
presence = _make_presence_penalty_processor(1.0, prompt_len = 1)
|
||||
combined = LogitsProcessorList([_AddToTokenZero(), *presence])
|
||||
input_ids = torch.tensor([[5, 2]]) # completion = [2]
|
||||
scores = torch.zeros(1, 6)
|
||||
out = combined(input_ids, scores)
|
||||
assert out[0, 0].item() == pytest.approx(100.0) # other processor ran
|
||||
assert out[0, 2].item() == pytest.approx(-1.0) # presence ran
|
||||
|
||||
|
||||
def test_mlx_presence_penalty_callable():
|
||||
mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS")
|
||||
from core.inference.mlx_inference import _make_mlx_presence_penalty_processor
|
||||
|
||||
proc = _make_mlx_presence_penalty_processor(1.5)
|
||||
# First call = prompt only (latches prompt_len, penalizes nothing).
|
||||
prompt = mx.array([10, 11])
|
||||
logits0 = mx.zeros((1, 20))
|
||||
out0 = proc(prompt, logits0)
|
||||
assert float(out0[0, 10]) == pytest.approx(0.0)
|
||||
# Second call: one completion token (5) appended -> penalized once.
|
||||
seq = mx.array([10, 11, 5])
|
||||
logits1 = mx.zeros((1, 20))
|
||||
out1 = proc(seq, logits1)
|
||||
assert float(out1[0, 5]) == pytest.approx(-1.5)
|
||||
assert float(out1[0, 10]) == pytest.approx(0.0) # prompt token untouched
|
||||
|
||||
|
||||
def test_mlx_presence_penalty_bounds_out_of_range_ids():
|
||||
# Documents (and, on Apple Silicon CI, enforces) the intended MLX bound:
|
||||
# out-of-vocab and negative completion ids must be ignored. MLX does no
|
||||
# bounds checking and OOB indexing is undefined behavior (crash / memory
|
||||
# corruption), so the processor routes stray ids to a discarded scratch slot
|
||||
# and penalizes only in-range distinct ids -- matching the torch filter
|
||||
# seen[(seen >= 0) & (seen < vocab)]. Skips off arm64 macOS where MLX is absent.
|
||||
mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS")
|
||||
from core.inference.mlx_inference import _make_mlx_presence_penalty_processor
|
||||
|
||||
proc = _make_mlx_presence_penalty_processor(1.0)
|
||||
proc(mx.array([10, 11]), mx.zeros((1, 8))) # first call latches prompt_len = 2
|
||||
# Completion appends a valid id (3), an out-of-vocab id (99 >= vocab 8) and a
|
||||
# negative id (-1); only the in-range id is penalized and nothing crashes.
|
||||
seq = mx.array([10, 11, 3, 99, -1])
|
||||
out = proc(seq, mx.zeros((1, 8)))
|
||||
assert float(out[0, 3]) == pytest.approx(-1.0)
|
||||
for tok in range(8):
|
||||
if tok != 3:
|
||||
assert float(out[0, tok]) == pytest.approx(0.0)
|
||||
|
||||
|
||||
# Param propagation: route payload -> orchestrator cmd -> worker gen_kwargs
|
||||
_SAMPLING = {
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.8,
|
||||
"top_k": 20,
|
||||
"min_p": 0.05,
|
||||
"repetition_penalty": 1.1,
|
||||
"presence_penalty": 1.5,
|
||||
}
|
||||
|
||||
|
||||
def test_orchestrator_cmd_carries_all_sampling_params():
|
||||
from core.inference.orchestrator import InferenceOrchestrator
|
||||
|
||||
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
|
||||
cmd = o._build_generate_cmd(
|
||||
"req1",
|
||||
None,
|
||||
messages = [{"role": "user", "content": "hi"}],
|
||||
max_new_tokens = 128,
|
||||
**_SAMPLING,
|
||||
)
|
||||
for key, val in _SAMPLING.items():
|
||||
assert cmd[key] == val, f"{key} dropped/altered in orchestrator cmd"
|
||||
|
||||
|
||||
def test_worker_forwards_all_sampling_params_to_backend():
|
||||
from core.inference.worker import _handle_generate
|
||||
|
||||
class _RecordingBackend:
|
||||
last_generation_stats = None
|
||||
|
||||
def __init__(self):
|
||||
self.received = None
|
||||
|
||||
def generate_chat_response(self, **kwargs):
|
||||
self.received = kwargs
|
||||
return iter(()) # empty stream -> loop exits, gen_done is sent
|
||||
|
||||
class _FakeQueue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
cmd = {
|
||||
"type": "generate",
|
||||
"request_id": "r",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_new_tokens": 128,
|
||||
**_SAMPLING,
|
||||
}
|
||||
backend = _RecordingBackend()
|
||||
_handle_generate(backend, cmd, _FakeQueue(), threading.Event())
|
||||
|
||||
assert backend.received is not None
|
||||
for key, val in _SAMPLING.items():
|
||||
assert backend.received[key] == val, f"{key} dropped/altered in worker gen_kwargs"
|
||||
216
studio/backend/tests/test_response_template_markers.py
Normal file
216
studio/backend/tests/test_response_template_markers.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""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
|
||||
<s> as bare "[" on transformers 5.x while the standalone
|
||||
encoding gives "▁[", so the marker must anchor on <s>.
|
||||
starling - trailing space after "GPT4 Correct Assistant:" folds
|
||||
into the next content token ("▁Hello").
|
||||
glm - "[gMASK]<sop>" renders once at text start, never before
|
||||
later user turns; "<think>" is generation scaffolding
|
||||
that non-final turns render as a lone "</think>".
|
||||
qwen3-thinking - "<think>" 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
|
||||
"</s>\\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": "<s>[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": "<start_of_turn>user\n", "response": "<start_of_turn>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"]))
|
||||
|
|
@ -59,6 +59,7 @@ from models.inference import (
|
|||
ResponsesUsage,
|
||||
)
|
||||
from routes.inference import (
|
||||
_ResponsesReasoningExtractor,
|
||||
_SameTaskStreamingResponse,
|
||||
_build_chat_request,
|
||||
_chat_tool_calls_to_responses_output,
|
||||
|
|
@ -795,6 +796,7 @@ class TestResponsesNonStreamingAdapter:
|
|||
|
||||
def test_monitor_records_translated_visible_text(self, monkeypatch):
|
||||
import routes.inference as inf_mod
|
||||
import routes.inference as inf_mod
|
||||
|
||||
async def fake_chat_completions(chat_req, request):
|
||||
assert request.state.skip_api_monitor is True
|
||||
|
|
@ -1988,6 +1990,123 @@ class TestTranslatedMessagesValidate:
|
|||
ChatMessage(**m.model_dump(exclude_none = True))
|
||||
|
||||
|
||||
# reasoning_prefilled: enable_thinking templates prefill an unclosed <think>, so
|
||||
# generation begins inside the block; the extractor must start in reasoning.
|
||||
class TestReasoningPrefilledExtractor:
|
||||
def test_prefilled_single_feed_splits_lone_close(self):
|
||||
# T1: reasoning...</think>answer with a prefilled (unseen) open tag.
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"plan</think>answer",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = True,
|
||||
)
|
||||
assert reasoning == "plan"
|
||||
assert visible == "answer"
|
||||
|
||||
def test_prefilled_never_closed_is_all_reasoning(self):
|
||||
# T2: truncated mid-thought (no </think>) -> all reasoning (GGUF parity).
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"still thinking with no close",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = True,
|
||||
)
|
||||
assert reasoning == "still thinking with no close"
|
||||
assert visible == ""
|
||||
|
||||
def test_prefilled_close_split_across_feeds(self):
|
||||
# T3: </think> straddles two feed() calls; holdback resolves it.
|
||||
ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True)
|
||||
r1, v1 = ex.feed("plan</th")
|
||||
r2, v2 = ex.feed("ink>ans")
|
||||
fr, fv = ex.finish()
|
||||
assert (r1 + r2 + fr) == "plan"
|
||||
assert (v1 + v2 + fv) == "ans"
|
||||
|
||||
def test_prefilled_close_split_one_char_per_feed(self):
|
||||
# T4: every char in its own feed still splits correctly.
|
||||
ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True)
|
||||
reasoning, visible = "", ""
|
||||
for ch in "plan</think>x":
|
||||
r, v = ex.feed(ch)
|
||||
reasoning += r
|
||||
visible += v
|
||||
fr, fv = ex.finish()
|
||||
assert (reasoning + fr) == "plan"
|
||||
assert (visible + fv) == "x"
|
||||
|
||||
def test_prefilled_empty_generation(self):
|
||||
# T5: nothing generated.
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = True,
|
||||
)
|
||||
assert reasoning == ""
|
||||
assert visible == ""
|
||||
|
||||
def test_prefilled_whitespace_after_close_is_visible(self):
|
||||
# T6: Qwen commonly emits </think>\n\n before the answer.
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"plan</think>\n\nanswer",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = True,
|
||||
)
|
||||
assert reasoning == "plan"
|
||||
assert visible == "\n\nanswer"
|
||||
|
||||
def test_prefilled_stray_open_tag_is_suppressed(self):
|
||||
# T7: a re-emitted literal <think> inside prefilled reasoning is dropped,
|
||||
# not leaked into the drawer (covers enable_thinking_effort full-tag output).
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"a<think>b</think>c",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = True,
|
||||
)
|
||||
assert reasoning == "ab"
|
||||
assert visible == "c"
|
||||
assert "<think>" not in reasoning
|
||||
|
||||
def test_prefilled_close_at_start_empty_reasoning(self):
|
||||
# T8: model closed immediately (empty reasoning) then answered.
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"</think>hi",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = True,
|
||||
)
|
||||
assert reasoning == ""
|
||||
assert visible == "hi"
|
||||
|
||||
def test_not_prefilled_lone_close_preserves_current_behavior(self):
|
||||
# T9: without prefilled, a lone close tag keeps the pre-fix behavior (parity guard).
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"reasoning</think>ans",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = False,
|
||||
)
|
||||
assert reasoning == ""
|
||||
assert visible == "reasoningans"
|
||||
|
||||
def test_not_prefilled_full_pair_still_splits(self):
|
||||
# T10: normal explicit <think>..</think> (GGUF / Harmony) unchanged.
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"<think>r</think>v",
|
||||
parse_think_markers = True,
|
||||
reasoning_prefilled = False,
|
||||
)
|
||||
assert reasoning == "r"
|
||||
assert visible == "v"
|
||||
|
||||
def test_prefilled_ignored_when_markers_not_parsed(self):
|
||||
# T11: a non-reasoning model passes text through even with reasoning_prefilled False.
|
||||
reasoning, visible = _extract_responses_reasoning(
|
||||
"just an answer",
|
||||
parse_think_markers = False,
|
||||
reasoning_prefilled = False,
|
||||
)
|
||||
assert reasoning == ""
|
||||
assert visible == "just an answer"
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Streaming passthrough healing — text-form calls promoted in order
|
||||
# =====================================================================
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ from pathlib import Path
|
|||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
_backend_root = Path(__file__).resolve().parent.parent
|
||||
if str(_backend_root) not in sys.path:
|
||||
sys.path.insert(0, str(_backend_root))
|
||||
|
|
@ -46,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }}
|
|||
"""
|
||||
|
||||
|
||||
# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort
|
||||
# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders
|
||||
# identically to thinking-on-without-the-preamble), so the literal scan alone
|
||||
# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to
|
||||
# expose the encoder's full none/high/max ladder.
|
||||
DEEPSEEK_V4_TEMPLATE = (
|
||||
"{%- if not thinking is defined %}"
|
||||
"{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}"
|
||||
"{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n"
|
||||
"{%- if thinking and reasoning_effort == 'max' %}"
|
||||
"{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n"
|
||||
"{%- for message in messages %}{{- message.content }}{%- endfor %}"
|
||||
)
|
||||
|
||||
|
||||
PLAIN_TEMPLATE = """
|
||||
{%- for message in messages %}
|
||||
{{- message.role + ': ' + message.content + '\\n' }}
|
||||
|
|
@ -88,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false():
|
|||
assert flags["reasoning_style"] == "enable_thinking"
|
||||
|
||||
|
||||
def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max():
|
||||
"""DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble.
|
||||
Classified as the hybrid style with the full none/high/max ladder even
|
||||
though the template only branches on 'max'."""
|
||||
from core.inference.llama_cpp import detect_reasoning_flags
|
||||
|
||||
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF")
|
||||
assert flags["supports_reasoning"] is True
|
||||
assert flags["reasoning_style"] == "enable_thinking_effort"
|
||||
assert flags["reasoning_effort_levels"] == ["high", "max"]
|
||||
assert flags["reasoning_always_on"] is False
|
||||
|
||||
|
||||
def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected():
|
||||
"""The 'high' injection is scoped to deepseek-v4: a different model whose
|
||||
template only branches on 'max' keeps ['max'] (no phantom 'high')."""
|
||||
from core.inference.llama_cpp import detect_reasoning_flags
|
||||
|
||||
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF")
|
||||
assert flags["reasoning_style"] == "enable_thinking_effort"
|
||||
assert flags["reasoning_effort_levels"] == ["max"]
|
||||
|
||||
|
||||
def test_detect_safetensors_features_passes_template_through_to_classifier():
|
||||
"""Route wrapper forwards a real template to the inner classifier."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
|
@ -127,9 +167,8 @@ def test_detect_safetensors_features_gptoss_disables_tools():
|
|||
assert flags["supports_tools"] is False
|
||||
|
||||
|
||||
# Llama-3 / Mistral advertise tools but emit <|python_tag|> / [TOOL_CALLS],
|
||||
# which our parser can't read. The route helper must not flip supports_tools=True
|
||||
# for them, else the UI enables a pill the agentic loop can't honour.
|
||||
# Llama-3 / Mistral / Gemma 4 tool-call formats are now parser-supported, so supports_tools=True
|
||||
# must hold for all of them; only templates matching none of the five known markers are suppressed.
|
||||
|
||||
LLAMA3_TEMPLATE = """
|
||||
{%- if tools %}
|
||||
|
|
@ -161,27 +200,188 @@ MISTRAL_TEMPLATE = """
|
|||
{%- endfor %}
|
||||
"""
|
||||
|
||||
GEMMA4_TEMPLATE = """
|
||||
{%- if tools %}
|
||||
{{- 'Tools available. Emit calls as ' }}
|
||||
{{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}<tool_call|>' }}
|
||||
{%- for tool in tools %}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
"""
|
||||
|
||||
def test_detect_safetensors_features_llama3_template_suppresses_tools():
|
||||
"""Llama-3 emits <|python_tag|>; safetensors loop cannot parse it."""
|
||||
|
||||
def test_detect_safetensors_features_llama3_template_keeps_tools_on():
|
||||
"""Llama-3 emits <|python_tag|>; parser now supports it."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
|
||||
flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE)
|
||||
assert flags["supports_tools"] is False
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
def test_detect_safetensors_features_mistral_template_suppresses_tools():
|
||||
"""Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it."""
|
||||
def test_detect_safetensors_features_mistral_template_keeps_tools_on():
|
||||
"""Mistral emits [TOOL_CALLS]name{json}, which the safetensors loop now parses
|
||||
(the shared bracket-tag parser). The gate must no longer suppress it, or the
|
||||
PR's Mistral tool support is unreachable through normal capability detection."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3")
|
||||
flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
def test_detect_safetensors_features_gemma4_template_keeps_tools_on():
|
||||
"""Gemma 4 emits <|tool_call>; parser now supports it."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit")
|
||||
flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
# DeepSeek V3 / V3.1 / R1 emit ``<｜tool▁calls▁begin｜>...`` blocks.
|
||||
# Note the full-width pipe (U+FF5C) and lower-1/8-block (U+2581).
|
||||
DEEPSEEK_TEMPLATE = """
|
||||
{%- if tools %}
|
||||
{%- for tool in tools %}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- for message in messages %}
|
||||
{%- if message.role == 'assistant' and message.tool_calls %}
|
||||
{%- for tc in message.tool_calls %}
|
||||
{{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tc.function.name +
|
||||
'<｜tool▁sep｜>' + tc.function.arguments + '<｜tool▁call▁end｜>' }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"""
|
||||
|
||||
|
||||
def test_detect_safetensors_features_deepseek_template_keeps_tools_on():
|
||||
"""DeepSeek emits ``<｜tool▁calls▁begin｜>...``; parser now supports it."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1")
|
||||
flags = _detect_safetensors_features(backend, DEEPSEEK_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
# GLM 4.5 / 4.6 / 4.7 emit ``<tool_call>NAME\n<arg_key>...<arg_value>...
|
||||
GLM_TEMPLATE = """
|
||||
{%- if tools %}
|
||||
For each function call, output the function name and arguments within
|
||||
the following XML format:
|
||||
<tool_call>{function-name}
|
||||
<arg_key>{arg-key}</arg_key>
|
||||
<arg_value>{arg-value}</arg_value>
|
||||
</tool_call>
|
||||
{%- for tool in tools %}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
"""
|
||||
|
||||
|
||||
def test_detect_safetensors_features_glm_template_keeps_tools_on():
|
||||
"""GLM 4.x emits ``<tool_call>NAME\\n<arg_key>...``; parser handles it."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/GLM-4.6")
|
||||
flags = _detect_safetensors_features(backend, GLM_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
# Kimi K2 / Moonshot uses ``<|tool_calls_section_begin|>...`` blocks
|
||||
# with ``functions.NAME:IDX`` as the per-call id.
|
||||
KIMI_TEMPLATE = """
|
||||
{%- if tools %}
|
||||
<|im_system|>tool_declare<|im_middle|>{{ tools | tojson }}<|im_end|>
|
||||
{%- endif %}
|
||||
{%- for message in messages %}
|
||||
{%- if message.role == 'assistant' and message.tool_calls %}
|
||||
<|tool_calls_section_begin|>
|
||||
{%- for tc in message.tool_calls %}
|
||||
<|tool_call_begin|>{{ tc.id }}<|tool_call_argument_begin|>{{ tc.function.arguments | tojson }}<|tool_call_end|>
|
||||
{%- endfor %}
|
||||
<|tool_calls_section_end|>
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"""
|
||||
|
||||
|
||||
def test_detect_safetensors_features_kimi_template_keeps_tools_on():
|
||||
"""Kimi K2 emits ``<|tool_calls_section_begin|>...``; parser handles it."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Kimi-K2-Instruct")
|
||||
flags = _detect_safetensors_features(backend, KIMI_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
LLAMA3_2_BARE_JSON_TEMPLATE = """
|
||||
{%- if tools %}
|
||||
{{- 'Given the following functions, respond with JSON for a function call.' }}
|
||||
{{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }}
|
||||
{%- for tool in tools %}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{%- for message in messages %}
|
||||
{%- if 'tool_calls' in message %}
|
||||
{{- '{"name": "' + message.tool_calls[0].function.name + '", '}}
|
||||
{{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
"""
|
||||
|
||||
|
||||
def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on():
|
||||
"""Llama-3.2 bare JSON is supported, so the pill stays enabled."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
|
||||
flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
MINICPM5_ATTRIBUTE_TEMPLATE = """
|
||||
{%- if tools %}
|
||||
{{- 'Available tools. Emit calls as ' }}
|
||||
{{- '<function name="NAME"><parameter name="key">value</parameter></function>' }}
|
||||
{%- for tool in tools %}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
"""
|
||||
|
||||
|
||||
def test_detect_safetensors_features_attribute_function_form_keeps_tools_on():
|
||||
"""The attribute form ``<function name="...">`` must be whitelisted or the pill is wrongly suppressed."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "openbmb/MiniCPM-5")
|
||||
flags = _detect_safetensors_features(backend, MINICPM5_ATTRIBUTE_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
def test_detect_safetensors_features_unknown_format_suppresses_tools():
|
||||
"""Tools advertised with no known marker must be suppressed."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
tpl = (
|
||||
"{%- if tools %}<|im_start|>system\n"
|
||||
"Emit tool calls as JSON-RPC notifications inside the response."
|
||||
"<|im_end|>{%- endif %}"
|
||||
)
|
||||
backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format")
|
||||
flags = _detect_safetensors_features(backend, tpl)
|
||||
assert flags["supports_tools"] is False
|
||||
|
||||
|
||||
def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on():
|
||||
"""Sanity check: gate only suppresses non-Qwen formats."""
|
||||
"""Sanity check: Qwen <tool_call> marker still flips supports_tools."""
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B")
|
||||
|
|
@ -454,3 +654,184 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors():
|
|||
assert flags["supports_tools"] is True
|
||||
assert flags["supports_reasoning"] is True
|
||||
assert flags["supports_preserve_thinking"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"opener",
|
||||
[
|
||||
"<｜tool▁calls▁begin｜>", # canonical
|
||||
"<｜tool_calls_begin｜>", # ASCII underscores
|
||||
"<｜tool▁calls｜>", # short form
|
||||
"<｜tool calls begin｜>", # spaces
|
||||
"<｜tool\\_calls\\_begin｜>", # escaped underscores
|
||||
],
|
||||
)
|
||||
def test_detect_safetensors_features_deepseek_opener_variants_keep_tools_on(opener):
|
||||
# Every DeepSeek opener the parser accepts must keep supports_tools on; the route gate derives
|
||||
# its markers from the parser's TOOL_XML_SIGNALS so it can no longer drift behind the parser ...
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
tpl = (
|
||||
"{%- if tools %}tools{%- endif %}"
|
||||
+ opener
|
||||
+ "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time{}"
|
||||
"<｜tool▁call▁end｜><｜tool▁calls▁end｜>"
|
||||
)
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1")
|
||||
flags = _detect_safetensors_features(backend, tpl)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
# Templates that advertise tools ({%- if tools %}) and prompt the bare-JSON
|
||||
# call form, but whose ``{"name":`` example is pretty-printed or JSON-escaped.
|
||||
_WHITESPACE_BARE_JSON_TEMPLATE = (
|
||||
"{%- if tools %}\n"
|
||||
"To call a tool, output JSON of the form:\n"
|
||||
'{ "name" : "function_name", "parameters": { } }\n'
|
||||
"{%- endif %}\n"
|
||||
"{{ messages }}"
|
||||
)
|
||||
_ESCAPED_BARE_JSON_TEMPLATE = (
|
||||
"{%- if tools %}\n"
|
||||
'Respond with {\\"name\\": \\"fn\\", \\"parameters\\": {}}\n'
|
||||
"{%- endif %}\n"
|
||||
"{{ messages }}"
|
||||
)
|
||||
_TOOLS_ADVERTISED_NO_PARSEABLE_FORM = (
|
||||
"{%- if tools %}\nYou may use the available tools.\n{%- endif %}\n{{ messages }}"
|
||||
)
|
||||
|
||||
|
||||
def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json():
|
||||
# A pretty-printed bare-JSON example (``{ "name" :``) must keep supports_tools since the parser
|
||||
# accepts that whitespace via raw_decode.
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
|
||||
flags = _detect_safetensors_features(backend, _WHITESPACE_BARE_JSON_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json():
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
|
||||
flags = _detect_safetensors_features(backend, _ESCAPED_BARE_JSON_TEMPLATE)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
def test_detect_safetensors_features_drops_tools_when_no_parseable_form():
|
||||
# Negative control: tools advertised but no parser-recognised emission form at
|
||||
# all -> the pill is still dropped (the gate is not now matching everything).
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
|
||||
flags = _detect_safetensors_features(backend, _TOOLS_ADVERTISED_NO_PARSEABLE_FORM)
|
||||
assert flags["supports_tools"] is False
|
||||
|
||||
|
||||
def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json():
|
||||
# A template documenting the parser-supported {"function":...} bare-JSON alias
|
||||
# must keep supports_tools, mirroring the {"name":...} form.
|
||||
from routes.inference import _detect_safetensors_features
|
||||
|
||||
tpl = (
|
||||
"{%- if tools %}\n"
|
||||
'Respond with {"function": "fn", "parameters": {}}\n'
|
||||
"{%- endif %}\n"
|
||||
"{{ messages }}"
|
||||
)
|
||||
backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
|
||||
flags = _detect_safetensors_features(backend, tpl)
|
||||
assert flags["supports_tools"] is True
|
||||
|
||||
|
||||
# _sf_reasoning_prefill_mode gates the prefilled-<think> extractor (GGUF reasoning parity).
|
||||
class TestSafetensorsReasoningPrefillGate:
|
||||
# A minimal Qwen3-style template with the standard <think>/</think> markers.
|
||||
_QWEN_TPL = "{% if enable_thinking %}<think>{% endif %}...</think>..."
|
||||
# gemma-style bespoke reasoning channel -- no standard markers.
|
||||
_GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought<channel|>"
|
||||
# always-on template whose GENERATION PROMPT opens an unclosed <think> (DeepSeek-R1 / QwQ /
|
||||
# Qwen3-Thinking shape): the model emits only the closing </think>, so prefill.
|
||||
_ALWAYS_ON_OPEN_TPL = (
|
||||
"{% for m in messages %}{{ m['content'] }}{% endfor %}"
|
||||
"{% if add_generation_prompt %}<|assistant|><think>\n{% endif %}"
|
||||
)
|
||||
# always-on template that renders PAST assistant <think>...</think> history but leaves the
|
||||
# generation prompt open with no <think> (Kimi-K2-Thinking shape): the model self-emits its
|
||||
# own block, so prefill mode would blank a normal answer.
|
||||
_ALWAYS_ON_HISTORY_TPL = (
|
||||
"{% for m in messages %}"
|
||||
"{% if m['role'] == 'assistant' %}<think>{{ m.get('reasoning_content', '') }}</think>"
|
||||
"{{ m['content'] }}{% endif %}"
|
||||
"{% endfor %}"
|
||||
"{% if add_generation_prompt %}<|im_assistant|>assistant<|im_middle|>{% endif %}"
|
||||
)
|
||||
|
||||
def _features(self, **over):
|
||||
base = {
|
||||
"supports_reasoning": True,
|
||||
"reasoning_always_on": False,
|
||||
"reasoning_style": "enable_thinking",
|
||||
}
|
||||
base.update(over)
|
||||
return base
|
||||
|
||||
def test_g1_enable_thinking_true(self):
|
||||
# G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True
|
||||
|
||||
def test_g2_enable_thinking_none_defaults_on(self):
|
||||
# G2: default request (None) -> prefilled (Qwen3/GLM templates default on).
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True
|
||||
|
||||
def test_g3_enable_thinking_false(self):
|
||||
# G3: thinking explicitly off -> not prefilled.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False
|
||||
|
||||
def test_g4_gpt_oss_reasoning_effort_excluded(self):
|
||||
# G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
feats = self._features(reasoning_style = "reasoning_effort")
|
||||
assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False
|
||||
|
||||
def test_g5_enable_thinking_effort_included(self):
|
||||
# G5: GLM-style enable_thinking_effort also prefills.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
feats = self._features(reasoning_style = "enable_thinking_effort")
|
||||
assert _sf_reasoning_prefill_mode(feats, None, self._QWEN_TPL) is True
|
||||
|
||||
def test_g6_non_reasoning_model(self):
|
||||
# G6: no reasoning capability -> never prefilled.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
feats = self._features(supports_reasoning = False, reasoning_style = None)
|
||||
assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False
|
||||
|
||||
def test_g7_reasoning_always_on_prompt_opens_think(self):
|
||||
# G7: always-on template whose generation prompt opens <think> -> prefilled regardless of the flag.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
feats = self._features(reasoning_always_on = True)
|
||||
assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True
|
||||
|
||||
def test_g7b_reasoning_always_on_history_only_not_prefilled(self):
|
||||
# G7b (#5704): always-on classification from rendered assistant HISTORY <think></think>
|
||||
# (Kimi-K2-Thinking) whose generation prompt opens no <think>. Prefill mode would capture a
|
||||
# normal answer entirely as reasoning_content and blank the visible answer, so it must be off.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
feats = self._features(reasoning_always_on = True)
|
||||
assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False
|
||||
|
||||
def test_g8_gemma_bespoke_channel_excluded(self):
|
||||
# G8: gemma's <|think|>/<|channel> format has no </think> -> NOT prefilled
|
||||
# (would otherwise swallow the whole answer as reasoning). Regression guard.
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False
|
||||
|
||||
def test_g9_missing_template_not_prefilled(self):
|
||||
# G9: no template available -> conservative (not prefilled).
|
||||
from routes.inference import _sf_reasoning_prefill_mode
|
||||
assert _sf_reasoning_prefill_mode(self._features(), True, None) is False
|
||||
|
|
|
|||
217
studio/backend/tests/test_safetensors_reasoning_stream.py
Normal file
217
studio/backend/tests/test_safetensors_reasoning_stream.py
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Safetensors/MLX reasoning-block parity with GGUF.
|
||||
|
||||
enable_thinking templates (Qwen3/GLM) prefill an unclosed ``<think>`` so the model
|
||||
emits only the closing ``</think>`` then the answer; the safetensors stream must
|
||||
split the leading text into ``reasoning_content`` deltas (plain stream and tool
|
||||
loop), resetting per turn and appending only visible text to the monitor. Replays a
|
||||
copy of ``sf_tool_stream``'s reasoning loop against synthetic events.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
from routes.inference import (
|
||||
_ResponsesReasoningExtractor,
|
||||
_sf_reasoning_prefill_mode,
|
||||
_strip_tool_xml_for_display,
|
||||
)
|
||||
|
||||
|
||||
_THINK_TPL = "...<think>...</think>..."
|
||||
_ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True}
|
||||
_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True}
|
||||
|
||||
|
||||
def test_prefill_mode_on_for_enable_thinking_default():
|
||||
assert _sf_reasoning_prefill_mode(_ETHINK, None, _THINK_TPL) is True
|
||||
|
||||
|
||||
def test_prefill_mode_off_when_thinking_disabled():
|
||||
assert _sf_reasoning_prefill_mode(_ETHINK, False, _THINK_TPL) is False
|
||||
|
||||
|
||||
def test_prefill_mode_off_for_reasoning_effort_none():
|
||||
# enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode
|
||||
# would capture the whole answer as reasoning_content.
|
||||
assert (
|
||||
_sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none")
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
_sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high")
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_prefill_mode_off_without_think_markers():
|
||||
assert _sf_reasoning_prefill_mode(_ETHINK, None, "no markers here") is False
|
||||
|
||||
|
||||
def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict:
|
||||
"""Mirror sf_tool_stream's reasoning loop: diff each cumulative ``content``
|
||||
snapshot, feed the delta through the extractor, and reset (flushing first) on
|
||||
``tool_start`` / empty ``status`` so each turn splits independently."""
|
||||
prev_text = ""
|
||||
extractor = _ResponsesReasoningExtractor(
|
||||
parse_think_markers = True, reasoning_prefilled = prefilled
|
||||
)
|
||||
reasoning_deltas: list[str] = []
|
||||
visible_deltas: list[str] = []
|
||||
monitor: list[str] = []
|
||||
tool_starts: list[dict] = []
|
||||
order: list[str] = [] # sequence of ("reasoning"|"visible"|"tool_start") events
|
||||
|
||||
def _flush():
|
||||
fr, fv = extractor.finish()
|
||||
if fr:
|
||||
reasoning_deltas.append(fr)
|
||||
order.append("reasoning")
|
||||
if fv:
|
||||
visible_deltas.append(fv)
|
||||
monitor.append(fv)
|
||||
order.append("visible")
|
||||
|
||||
for event in events:
|
||||
etype = event["type"]
|
||||
if etype == "status":
|
||||
if not event["text"]:
|
||||
_flush()
|
||||
prev_text = ""
|
||||
extractor = _ResponsesReasoningExtractor(
|
||||
parse_think_markers = True, reasoning_prefilled = prefilled
|
||||
)
|
||||
continue
|
||||
if etype in ("tool_start", "tool_end"):
|
||||
if etype == "tool_start":
|
||||
_flush()
|
||||
prev_text = ""
|
||||
extractor = _ResponsesReasoningExtractor(
|
||||
parse_think_markers = True, reasoning_prefilled = prefilled
|
||||
)
|
||||
tool_starts.append(event)
|
||||
order.append("tool_start")
|
||||
continue
|
||||
clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True)
|
||||
new_text = clean[len(prev_text) :]
|
||||
prev_text = clean
|
||||
if not new_text:
|
||||
continue
|
||||
r, v = extractor.feed(new_text)
|
||||
if r:
|
||||
reasoning_deltas.append(r)
|
||||
order.append("reasoning")
|
||||
if v:
|
||||
visible_deltas.append(v)
|
||||
monitor.append(v)
|
||||
order.append("visible")
|
||||
_flush()
|
||||
return {
|
||||
"reasoning": "".join(reasoning_deltas),
|
||||
"visible": "".join(visible_deltas),
|
||||
"monitor": "".join(monitor),
|
||||
"tool_starts": tool_starts,
|
||||
"order": order,
|
||||
}
|
||||
|
||||
|
||||
def test_s1_plain_stream_splits_prefilled_reasoning():
|
||||
# S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only.
|
||||
events = [
|
||||
{"type": "content", "text": "Let me compute 17*23"},
|
||||
{"type": "content", "text": "Let me compute 17*23 = 391</think>The answer is 391."},
|
||||
]
|
||||
out = _replay_sf_reasoning_stream(events, prefilled = True)
|
||||
assert out["reasoning"] == "Let me compute 17*23 = 391"
|
||||
assert out["visible"] == "The answer is 391."
|
||||
assert out["monitor"] == "The answer is 391."
|
||||
assert "<think>" not in out["reasoning"] and "</think>" not in out["visible"]
|
||||
|
||||
|
||||
def test_s2_reasoning_flushed_before_tool_start():
|
||||
# S2: reasoning streamed as reasoning_content, then flushed BEFORE tool_start.
|
||||
events = [
|
||||
{"type": "content", "text": "I should search"},
|
||||
{"type": "content", "text": "I should search Sydney weather</think>"},
|
||||
{"type": "tool_start", "tool_name": "web_search", "tool_call_id": "c0"},
|
||||
{"type": "tool_end", "tool_name": "web_search", "tool_call_id": "c0"},
|
||||
{"type": "status", "text": ""},
|
||||
{"type": "content", "text": "Found it</think>Sydney is 21C today."},
|
||||
]
|
||||
out = _replay_sf_reasoning_stream(events, prefilled = True)
|
||||
# Both turns' reasoning surfaced, answer only from turn 2.
|
||||
assert "I should search Sydney weather" in out["reasoning"]
|
||||
assert "Found it" in out["reasoning"]
|
||||
assert out["visible"] == "Sydney is 21C today."
|
||||
assert out["monitor"] == "Sydney is 21C today."
|
||||
# Ordering: the pre-tool reasoning is emitted before the tool_start.
|
||||
assert out["order"].index("reasoning") < out["order"].index("tool_start")
|
||||
|
||||
|
||||
def test_s3_extractor_resets_each_turn():
|
||||
# S3: multi-turn -> the two turns' reasoning are distinct (fresh extractor each).
|
||||
events = [
|
||||
{"type": "content", "text": "turn1 thoughts</think>partial"},
|
||||
{"type": "status", "text": ""},
|
||||
{"type": "content", "text": "turn2 thoughts</think>final answer"},
|
||||
]
|
||||
out = _replay_sf_reasoning_stream(events, prefilled = True)
|
||||
assert out["reasoning"] == "turn1 thoughtsturn2 thoughts"
|
||||
assert out["visible"] == "partialfinal answer"
|
||||
|
||||
|
||||
def test_s4_harmony_full_tags_normal_mode():
|
||||
# S4: gpt-oss / explicit-tag models use normal mode (prefilled=False).
|
||||
events = [{"type": "content", "text": "<think>reasoning here</think>visible answer"}]
|
||||
out = _replay_sf_reasoning_stream(events, prefilled = False)
|
||||
assert out["reasoning"] == "reasoning here"
|
||||
assert out["visible"] == "visible answer"
|
||||
|
||||
|
||||
def test_s5_thinking_off_no_reasoning_deltas():
|
||||
# S5: thinking disabled -> not prefilled, no </think>, all content is visible.
|
||||
events = [{"type": "content", "text": "Just the plain answer, no thinking."}]
|
||||
out = _replay_sf_reasoning_stream(events, prefilled = False)
|
||||
assert out["reasoning"] == ""
|
||||
assert out["visible"] == "Just the plain answer, no thinking."
|
||||
assert out["monitor"] == "Just the plain answer, no thinking."
|
||||
|
||||
|
||||
def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort():
|
||||
# GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and
|
||||
# enable_thinking omitted) disables thinking exactly like enable_thinking=False, so
|
||||
# prefilled mode must be OFF. Otherwise the model emits no </think> and a plain
|
||||
# answer is swallowed whole into reasoning_content, leaving the visible response
|
||||
# empty (the exact bug: prefilled=True below eats the whole answer).
|
||||
feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True}
|
||||
assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False
|
||||
# Thinking on (effort level or default) still prefills.
|
||||
assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "high") is True
|
||||
assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, None) is True
|
||||
# An explicit enable_thinking=False also disables (unchanged).
|
||||
assert _sf_reasoning_prefill_mode(feats, False, _THINK_TPL, "high") is False
|
||||
# reasoning_always_on wins regardless of reasoning_effort.
|
||||
always = {**feats, "reasoning_always_on": True}
|
||||
assert _sf_reasoning_prefill_mode(always, None, _THINK_TPL, "none") is True
|
||||
# Plain enable_thinking models (Qwen) have no "none" sentinel; unaffected.
|
||||
plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True}
|
||||
assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True
|
||||
|
||||
# End-to-end: with the corrected prefilled=False, a plain no-</think> answer is
|
||||
# emitted as visible content rather than swallowed into the thinking drawer.
|
||||
events = [{"type": "content", "text": "The capital of France is Paris."}]
|
||||
out = _replay_sf_reasoning_stream(events, prefilled = False)
|
||||
assert out["visible"] == "The capital of France is Paris."
|
||||
assert out["reasoning"] == ""
|
||||
# The buggy prefilled=True path is what swallowed the whole answer (guard the delta).
|
||||
swallowed = _replay_sf_reasoning_stream(events, prefilled = True)
|
||||
assert swallowed["visible"] == ""
|
||||
assert swallowed["reasoning"] == "The capital of France is Paris."
|
||||
File diff suppressed because it is too large
Load diff
179
studio/backend/tests/test_safetensors_toolcall_wiring.py
Normal file
179
studio/backend/tests/test_safetensors_toolcall_wiring.py
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path.
|
||||
|
||||
The parser and the cumulative-text state machine are already covered exhaustively by
|
||||
``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the
|
||||
*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and
|
||||
``MLXInferenceBackend`` render the prompt through the shared
|
||||
``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the
|
||||
shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and
|
||||
``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU
|
||||
test of that seam covers the macOS MLX path too).
|
||||
|
||||
This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the
|
||||
``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full
|
||||
agentic chain end to end:
|
||||
|
||||
tools injected into the template -> loop parses the call -> tool dispatched once ->
|
||||
tool result fed back -> generation re-entered -> final answer streamed.
|
||||
|
||||
It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling
|
||||
end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the
|
||||
tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704).
|
||||
"""
|
||||
|
||||
from core.inference.chat_template_helpers import apply_chat_template_for_generation
|
||||
from core.inference.safetensors_agentic import run_safetensors_tool_loop
|
||||
|
||||
TOOL_NAME = "get_weather"
|
||||
TOOL_ARGS = {"city": "Paris"}
|
||||
FAKE_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": TOOL_NAME,
|
||||
"description": "Get the current weather for a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
# Full parser matrix lives in test_safetensors_tool_loop.py.
|
||||
TOOL_CALL_TEXT = '<tool_call>{"name": "get_weather", "arguments": {"city": "Paris"}}</tool_call>'
|
||||
FINAL_ANSWER = "The weather in Paris is sunny and 22C."
|
||||
TOOL_RESULT = "Paris: sunny, 22C"
|
||||
|
||||
|
||||
class RecordingTokenizer:
|
||||
"""Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``.
|
||||
|
||||
Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the
|
||||
real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually
|
||||
forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen``
|
||||
holding ``None``.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.tools_seen: list = []
|
||||
self.call_count = 0
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
*,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kwargs,
|
||||
):
|
||||
self.call_count += 1
|
||||
self.tools_seen.append(kwargs.get("tools"))
|
||||
return "PROMPT"
|
||||
|
||||
|
||||
class StubExecutor:
|
||||
"""Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result.
|
||||
|
||||
A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run.
|
||||
"""
|
||||
|
||||
def __init__(self, result: str):
|
||||
self.result = result
|
||||
self.calls: list[tuple[str, dict]] = []
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
name,
|
||||
arguments,
|
||||
*,
|
||||
cancel_event = None,
|
||||
timeout = None,
|
||||
session_id = None,
|
||||
rag_scope = None,
|
||||
disable_sandbox = False,
|
||||
):
|
||||
self.calls.append((name, arguments))
|
||||
return self.result
|
||||
|
||||
|
||||
def _collect(generator, max_events = 200):
|
||||
events = []
|
||||
for ev in generator:
|
||||
events.append(ev)
|
||||
if len(events) >= max_events:
|
||||
break
|
||||
return events
|
||||
|
||||
|
||||
def _tool_names(tools):
|
||||
return [(t.get("function") or {}).get("name") for t in (tools or [])]
|
||||
|
||||
|
||||
def test_backend_seam_injects_tools_and_drives_full_tool_loop():
|
||||
"""The shared backend seam forwards tools into the chat template, and the loop parses the call,
|
||||
dispatches it once, feeds the result back, and re-enters generation for the final answer."""
|
||||
tok = RecordingTokenizer()
|
||||
executor = StubExecutor(TOOL_RESULT)
|
||||
turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER])
|
||||
active_tools_seen: list = []
|
||||
conversations_seen: list = []
|
||||
|
||||
def single_turn(conversation, *, active_tools = None):
|
||||
# Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots.
|
||||
active_tools_seen.append(active_tools)
|
||||
conversations_seen.append([dict(m) for m in conversation])
|
||||
apply_chat_template_for_generation(tok, conversation, tools = active_tools)
|
||||
text = next(turns)
|
||||
mid = len(text) // 2
|
||||
acc = ""
|
||||
for chunk in (text[:mid], text[mid:]):
|
||||
acc += chunk
|
||||
yield acc
|
||||
|
||||
events = _collect(
|
||||
run_safetensors_tool_loop(
|
||||
single_turn = single_turn,
|
||||
messages = [{"role": "user", "content": "What is the weather in Paris?"}],
|
||||
tools = [FAKE_TOOL],
|
||||
execute_tool = executor,
|
||||
max_tool_iterations = 3,
|
||||
)
|
||||
)
|
||||
|
||||
# 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools).
|
||||
assert tok.tools_seen, "tokenizer.apply_chat_template was never called"
|
||||
assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer"
|
||||
assert TOOL_NAME in _tool_names(tok.tools_seen[0])
|
||||
|
||||
# 2. Loop offered the tool to the first generation turn.
|
||||
assert active_tools_seen and active_tools_seen[0] is not None
|
||||
assert TOOL_NAME in _tool_names(active_tools_seen[0])
|
||||
|
||||
# 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result.
|
||||
tool_starts = [e for e in events if e["type"] == "tool_start"]
|
||||
tool_ends = [e for e in events if e["type"] == "tool_end"]
|
||||
assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME
|
||||
assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls
|
||||
assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT
|
||||
|
||||
# 6. Final answer streams after the tool result: loop appended it and re-entered generation.
|
||||
contents = [e for e in events if e["type"] == "content"]
|
||||
assert contents and FINAL_ANSWER in contents[-1]["text"]
|
||||
last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end")
|
||||
last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content")
|
||||
assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result"
|
||||
|
||||
# 6b. Tool result fed back into the conversation before the final turn (6 alone misses this:
|
||||
# the fake generation ignores the conversation).
|
||||
assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call"
|
||||
final_turn_convo = conversations_seen[1]
|
||||
assert any(
|
||||
TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo
|
||||
), "tool result was not fed back into the conversation before the final generation turn"
|
||||
|
||||
# 7. Guard: raw tool-call markup never leaked to the client as content.
|
||||
for e in contents:
|
||||
assert "<tool_call>" not in e["text"]
|
||||
assert TOOL_NAME not in e["text"]
|
||||
|
|
@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base():
|
|||
if runs_gate and not resolves_base:
|
||||
offenders.append(f"{rel} runs a load gate but never resolves the LoRA base")
|
||||
assert not offenders, "\n".join(offenders)
|
||||
|
||||
|
||||
def test_rag_embedding_path_runs_the_malware_gate():
|
||||
"""The RAG embedding model is set through /settings and later loaded by
|
||||
SentenceTransformer, which deserializes pickles; both sites must run the malware gate
|
||||
or a flagged repo loads unscanned (bypassing the normal model-load protections)."""
|
||||
offenders = []
|
||||
for rel in ("routes/settings.py", "core/rag/embeddings.py"):
|
||||
if "evaluate_file_security(" not in (_BACKEND / rel).read_text():
|
||||
offenders.append(
|
||||
f"{rel} loads/persists an embedding model without evaluate_file_security"
|
||||
)
|
||||
assert not offenders, "\n".join(offenders)
|
||||
|
|
|
|||
786
studio/backend/tests/test_sf_client_tools_passthrough.py
Normal file
786
studio/backend/tests/test_sf_client_tools_passthrough.py
Normal file
|
|
@ -0,0 +1,786 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Client-tools passthrough healing for the safetensors/MLX backend.
|
||||
|
||||
Parity for #6801: when a NON-GGUF model is loaded and the request declares its
|
||||
own ``tools`` with server-side tools OFF, text-form tool calls are promoted back
|
||||
into structured ``tool_calls`` (declared tools only) via the shared healer. MLX
|
||||
rides the same orchestrator path, so a single scripted backend covers both.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
from models.inference import ChatCompletionRequest, ChatMessage
|
||||
from routes.inference import openai_chat_completions
|
||||
from core.inference.api_monitor import ApiMonitor
|
||||
|
||||
|
||||
LOOKUP_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "Look something up",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"q": {"type": "string"}},
|
||||
"required": ["q"],
|
||||
},
|
||||
},
|
||||
}
|
||||
SEARCH_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
_CALL_XML = '<tool_call>{"name": "lookup", "arguments": {"q": "cats"}}</tool_call>'
|
||||
_SEARCH_XML = '<tool_call>{"name": "search", "arguments": {"query": "dogs"}}</tool_call>'
|
||||
|
||||
|
||||
class _Request:
|
||||
state = SimpleNamespace()
|
||||
url = SimpleNamespace(path = "/v1/chat/completions")
|
||||
method = "POST"
|
||||
scope: dict = {}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
|
||||
class _ScriptedBackend:
|
||||
"""Non-GGUF backend: ``generate_chat_response`` replays scripted
|
||||
CUMULATIVE snapshots. ``responder(messages, tools)`` returns the snapshot
|
||||
list for one generation, so nudge tests can vary output across turns."""
|
||||
|
||||
active_model_name = "sf-model"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
responder,
|
||||
*,
|
||||
stats = None,
|
||||
):
|
||||
self.models = {
|
||||
"sf-model": {
|
||||
"chat_template_info": {"template": "<tool_call> chatml"},
|
||||
"context_length": 2048,
|
||||
}
|
||||
}
|
||||
self._responder = responder
|
||||
self._stats = stats
|
||||
self.calls: list = []
|
||||
self.reset_count = 0
|
||||
|
||||
def generate_chat_response(
|
||||
self,
|
||||
*,
|
||||
messages,
|
||||
tools = None,
|
||||
stats_holder = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.calls.append({"messages": messages, "tools": tools, **kwargs})
|
||||
snapshots = self._responder(messages, tools)
|
||||
if stats_holder is not None and self._stats is not None:
|
||||
stats_holder["stats"] = self._stats
|
||||
for snap in snapshots:
|
||||
yield snap
|
||||
|
||||
def reset_generation_state(self):
|
||||
self.reset_count += 1
|
||||
|
||||
|
||||
def _fixed(*snapshots):
|
||||
"""Responder that always replays the given cumulative snapshots."""
|
||||
return lambda messages, tools: list(snapshots)
|
||||
|
||||
|
||||
def _llama_stub():
|
||||
return SimpleNamespace(
|
||||
is_loaded = False,
|
||||
supports_tools = False,
|
||||
is_vision = False,
|
||||
context_length = None,
|
||||
)
|
||||
|
||||
|
||||
def _install(
|
||||
monkeypatch,
|
||||
backend,
|
||||
*,
|
||||
supports_tools = True,
|
||||
):
|
||||
import routes.inference as inf
|
||||
from state.tool_policy import reset_tool_policy
|
||||
|
||||
reset_tool_policy()
|
||||
monitor = ApiMonitor(max_entries = 8)
|
||||
monkeypatch.setattr(inf, "api_monitor", monitor)
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _llama_stub())
|
||||
monkeypatch.setattr(inf, "get_inference_backend", lambda: backend)
|
||||
monkeypatch.setattr(
|
||||
inf,
|
||||
"_detect_safetensors_features",
|
||||
lambda *a, **k: {"supports_tools": supports_tools},
|
||||
)
|
||||
return monitor
|
||||
|
||||
|
||||
def _request(**kwargs):
|
||||
base = dict(model = "default", messages = [ChatMessage(role = "user", content = "hi")])
|
||||
base.update(kwargs)
|
||||
return ChatCompletionRequest(**base)
|
||||
|
||||
|
||||
def _call(payload, monkeypatch, backend, **install_kwargs):
|
||||
_install(monkeypatch, backend, **install_kwargs)
|
||||
|
||||
async def _run():
|
||||
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
|
||||
|
||||
return asyncio.run(_run())
|
||||
|
||||
|
||||
def _json_body(response):
|
||||
return json.loads(response.body if hasattr(response, "body") else response.content)
|
||||
|
||||
|
||||
def _collect_sse(response):
|
||||
async def _run():
|
||||
return [c async for c in response.body_iterator]
|
||||
|
||||
return asyncio.run(_run())
|
||||
|
||||
|
||||
def _sse_objects(chunks):
|
||||
out = []
|
||||
for chunk in chunks:
|
||||
if isinstance(chunk, bytes):
|
||||
chunk = chunk.decode()
|
||||
for line in str(chunk).splitlines():
|
||||
if line.startswith("data: "):
|
||||
data = line.removeprefix("data: ")
|
||||
if data != "[DONE]":
|
||||
out.append(json.loads(data))
|
||||
return out
|
||||
|
||||
|
||||
# ── Non-streaming ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_xml_healed_to_tool_calls_non_streaming(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "tool_calls"
|
||||
assert choice["message"]["content"] is None
|
||||
calls = choice["message"]["tool_calls"]
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["function"]["name"] == "lookup"
|
||||
assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"}
|
||||
# The client tools reached the generator (template injection).
|
||||
assert backend.calls[0]["tools"] == [LOOKUP_TOOL]
|
||||
|
||||
|
||||
def test_undeclared_call_stays_text(monkeypatch):
|
||||
xml = '<tool_call>{"name": "other", "arguments": {}}</tool_call>'
|
||||
backend = _ScriptedBackend(_fixed(xml))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "stop"
|
||||
assert choice["message"].get("tool_calls") is None
|
||||
assert choice["message"]["content"] == xml
|
||||
|
||||
|
||||
def test_opt_out_relays_verbatim(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False, auto_heal_tool_calls = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "stop"
|
||||
assert choice["message"].get("tool_calls") is None
|
||||
assert choice["message"]["content"] == _CALL_XML
|
||||
|
||||
|
||||
def test_env_kill_switch_relays_verbatim(monkeypatch):
|
||||
import core.inference.passthrough_healing as ph
|
||||
|
||||
monkeypatch.setattr(ph, "_HEALING_DISABLED", True)
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "stop"
|
||||
assert choice["message"].get("tool_calls") is None
|
||||
assert choice["message"]["content"] == _CALL_XML
|
||||
|
||||
|
||||
def test_no_tools_request_untouched(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed("just a plain answer"))
|
||||
payload = _request(stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
# No tools and no tool messages -> plain path, normal ChatCompletion.
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "stop"
|
||||
assert choice["message"]["content"] == "just a plain answer"
|
||||
assert choice["message"].get("tool_calls") is None
|
||||
|
||||
|
||||
def test_prose_around_call_retained(monkeypatch):
|
||||
text = "Let me look:\n" + _CALL_XML + "\ndone"
|
||||
backend = _ScriptedBackend(_fixed(text))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "tool_calls"
|
||||
assert choice["message"]["content"] == "Let me look:\n\ndone"
|
||||
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
|
||||
|
||||
|
||||
def test_empty_output_is_valid_stop(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed(""))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "stop"
|
||||
assert choice["message"]["content"] in ("", None)
|
||||
assert choice["message"].get("tool_calls") is None
|
||||
|
||||
|
||||
def test_tool_role_follow_up_turn_preserves_history(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed("The weather is sunny."))
|
||||
payload = _request(
|
||||
tools = [LOOKUP_TOOL],
|
||||
stream = False,
|
||||
messages = [
|
||||
ChatMessage(role = "user", content = "weather?"),
|
||||
ChatMessage(
|
||||
role = "assistant",
|
||||
content = None,
|
||||
tool_calls = [
|
||||
{
|
||||
"id": "call_0",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": '{"q": "weather"}'},
|
||||
}
|
||||
],
|
||||
),
|
||||
ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"),
|
||||
],
|
||||
)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert body["choices"][0]["message"]["content"] == "The weather is sunny."
|
||||
# The tool history reached the generator intact (role=tool + assistant.tool_calls).
|
||||
sent = backend.calls[0]["messages"]
|
||||
roles = [m["role"] for m in sent]
|
||||
assert "tool" in roles
|
||||
assistant = next(m for m in sent if m["role"] == "assistant")
|
||||
assert assistant.get("tool_calls")
|
||||
|
||||
|
||||
def test_dict_arguments_history_does_not_crash(monkeypatch):
|
||||
# Non-spec client: assistant tool_calls[].function.arguments as a dict.
|
||||
backend = _ScriptedBackend(_fixed("ok"))
|
||||
payload = _request(
|
||||
tools = [LOOKUP_TOOL],
|
||||
stream = False,
|
||||
messages = [
|
||||
ChatMessage(role = "user", content = "hi"),
|
||||
ChatMessage(
|
||||
role = "assistant",
|
||||
content = None,
|
||||
tool_calls = [
|
||||
{
|
||||
"id": "call_0",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": {"q": "x"}},
|
||||
}
|
||||
],
|
||||
),
|
||||
ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"),
|
||||
],
|
||||
)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert body["choices"][0]["message"]["content"] == "ok"
|
||||
|
||||
|
||||
def test_forced_tool_choice_narrows_promotion(monkeypatch):
|
||||
# tool_choice forces `search`; a `lookup` text call must NOT promote.
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(
|
||||
tools = [LOOKUP_TOOL, SEARCH_TOOL],
|
||||
stream = False,
|
||||
tool_choice = {"type": "function", "function": {"name": "search"}},
|
||||
)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "stop"
|
||||
assert choice["message"].get("tool_calls") is None
|
||||
|
||||
|
||||
def test_parallel_cap_non_streaming(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
calls = body["choices"][0]["message"]["tool_calls"]
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["function"]["name"] == "lookup"
|
||||
|
||||
|
||||
def test_usage_recorded_when_stats_present(monkeypatch):
|
||||
stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}}
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML), stats = stats)
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
monitor = _install(monkeypatch, backend)
|
||||
|
||||
async def _run():
|
||||
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
|
||||
|
||||
asyncio.run(_run())
|
||||
[entry] = monitor.snapshot()
|
||||
assert entry["prompt_tokens"] == 7
|
||||
assert entry["completion_tokens"] == 3
|
||||
|
||||
|
||||
# ── Nudge ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_nudge_default_off_single_generation(monkeypatch):
|
||||
# Signal present but unparseable; without opt-in, no retry.
|
||||
truncated = '<tool_call>{"name": "lookup"'
|
||||
backend = _ScriptedBackend(_fixed(truncated))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
_call(payload, monkeypatch, backend)
|
||||
assert len(backend.calls) == 1
|
||||
|
||||
|
||||
def test_nudge_opt_in_retry_recovers(monkeypatch):
|
||||
truncated = '<tool_call>{"name": "lookup"'
|
||||
|
||||
def responder(messages, tools):
|
||||
nudged = any(
|
||||
"native tool-call format" in (m.get("content") or "")
|
||||
for m in messages
|
||||
if m.get("role") == "user"
|
||||
)
|
||||
return [_CALL_XML] if nudged else [truncated]
|
||||
|
||||
backend = _ScriptedBackend(responder)
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert len(backend.calls) == 2
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "tool_calls"
|
||||
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
|
||||
|
||||
|
||||
def test_nudge_double_failure_relays_original(monkeypatch):
|
||||
truncated = '<tool_call>{"name": "lookup"'
|
||||
backend = _ScriptedBackend(_fixed(truncated))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert len(backend.calls) == 2 # exactly one retry
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "stop"
|
||||
assert choice["message"]["content"] == truncated
|
||||
|
||||
|
||||
# ── Streaming ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_streaming_heals_split_call_into_one_delta(monkeypatch):
|
||||
# Cumulative snapshots that build the call across many increments.
|
||||
pieces = ["<tool", '<tool_call>{"name": "loo', '<tool_call>{"name": "lookup", "argum']
|
||||
cumulative = pieces + [_CALL_XML]
|
||||
backend = _ScriptedBackend(_fixed(*cumulative))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = True)
|
||||
response = _call(payload, monkeypatch, backend)
|
||||
objs = _sse_objects(_collect_sse(response))
|
||||
tool_deltas = [
|
||||
tc
|
||||
for o in objs
|
||||
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
|
||||
]
|
||||
assert len(tool_deltas) == 1
|
||||
assert tool_deltas[0]["function"]["name"] == "lookup"
|
||||
finishes = [
|
||||
o["choices"][0]["finish_reason"]
|
||||
for o in objs
|
||||
if o["choices"] and o["choices"][0].get("finish_reason")
|
||||
]
|
||||
assert finishes == ["tool_calls"]
|
||||
|
||||
|
||||
def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch):
|
||||
# A stream cancelled via the registry ("Stop") must NOT promote the
|
||||
# buffered-but-unclosed tool markup at finalize, else it executes a tool
|
||||
# the user just cancelled. Guarded on cancel_event at the finalize step.
|
||||
import routes.inference as inf
|
||||
|
||||
cancel_id = "cancel-me-6870"
|
||||
# Balanced JSON but no closing </tool_call> -> healer HOLDS it until finalize.
|
||||
held = '<tool_call>{"name": "lookup", "arguments": {"q": "cats"}}'
|
||||
|
||||
class _CancelMidStream(_ScriptedBackend):
|
||||
def __init__(self):
|
||||
super().__init__(_fixed(held))
|
||||
|
||||
def generate_chat_response(
|
||||
self,
|
||||
*,
|
||||
messages,
|
||||
tools = None,
|
||||
stats_holder = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.calls.append({"messages": messages, "tools": tools, **kwargs})
|
||||
yield held # healer holds the unclosed call
|
||||
inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF
|
||||
|
||||
backend = _CancelMidStream()
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = True, cancel_id = cancel_id)
|
||||
response = _call(payload, monkeypatch, backend)
|
||||
objs = _sse_objects(_collect_sse(response))
|
||||
tool_deltas = [
|
||||
tc
|
||||
for o in objs
|
||||
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
|
||||
]
|
||||
assert tool_deltas == [] # no tool promoted after cancel
|
||||
finishes = [
|
||||
o["choices"][0]["finish_reason"]
|
||||
for o in objs
|
||||
if o["choices"] and o["choices"][0].get("finish_reason")
|
||||
]
|
||||
assert "tool_calls" not in finishes # ends with finish_reason=stop, not tool_calls
|
||||
|
||||
|
||||
def test_streaming_no_tools_verbatim(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed("hello ", "hello world"))
|
||||
payload = _request(stream = True)
|
||||
response = _call(payload, monkeypatch, backend)
|
||||
objs = _sse_objects(_collect_sse(response))
|
||||
text = "".join(
|
||||
(o["choices"][0]["delta"].get("content") or "")
|
||||
for o in objs
|
||||
if o["choices"] and "delta" in o["choices"][0]
|
||||
)
|
||||
assert text == "hello world"
|
||||
finishes = [
|
||||
o["choices"][0]["finish_reason"]
|
||||
for o in objs
|
||||
if o["choices"] and o["choices"][0].get("finish_reason")
|
||||
]
|
||||
assert finishes == ["stop"]
|
||||
|
||||
|
||||
def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch):
|
||||
# Repeated then shrunk cumulative snapshots must not double-heal.
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = True)
|
||||
response = _call(payload, monkeypatch, backend)
|
||||
objs = _sse_objects(_collect_sse(response))
|
||||
tool_deltas = [
|
||||
tc
|
||||
for o in objs
|
||||
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
|
||||
]
|
||||
assert len(tool_deltas) == 1
|
||||
|
||||
|
||||
def test_streaming_parallel_cap(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False)
|
||||
response = _call(payload, monkeypatch, backend)
|
||||
objs = _sse_objects(_collect_sse(response))
|
||||
tool_deltas = [
|
||||
tc
|
||||
for o in objs
|
||||
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
|
||||
]
|
||||
assert len(tool_deltas) == 1
|
||||
assert tool_deltas[0]["function"]["name"] == "lookup"
|
||||
|
||||
|
||||
def test_streaming_generator_error_closes_cleanly(monkeypatch):
|
||||
def responder(messages, tools):
|
||||
raise RuntimeError("boom /secret/path")
|
||||
|
||||
backend = _ScriptedBackend(responder)
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = True)
|
||||
response = _call(payload, monkeypatch, backend)
|
||||
chunks = _collect_sse(response)
|
||||
joined = "".join(c.decode() if isinstance(c, bytes) else c for c in chunks)
|
||||
assert "An internal error occurred" in joined
|
||||
assert "secret/path" not in joined # CWE-209: no path leak
|
||||
assert backend.reset_count >= 1
|
||||
|
||||
|
||||
def test_streaming_disconnect_resets_once(monkeypatch):
|
||||
class _DisconnectRequest(_Request):
|
||||
async def is_disconnected(self):
|
||||
return True
|
||||
|
||||
backend = _ScriptedBackend(_fixed("a", "ab", "abc"))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = True)
|
||||
_install(monkeypatch, backend)
|
||||
|
||||
async def _run():
|
||||
resp = await openai_chat_completions(
|
||||
payload, request = _DisconnectRequest(), current_subject = "u"
|
||||
)
|
||||
return [c async for c in resp.body_iterator]
|
||||
|
||||
asyncio.run(_run())
|
||||
assert backend.reset_count == 1
|
||||
|
||||
|
||||
def test_mlx_uses_same_path(monkeypatch):
|
||||
# MLX and safetensors share get_inference_backend(); one scripted backend covers both.
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert body["choices"][0]["finish_reason"] == "tool_calls"
|
||||
|
||||
|
||||
def test_tool_choice_none_does_not_advertise_tools(monkeypatch):
|
||||
# tool_choice="none": no tools rendered into the template; history templating still applies.
|
||||
backend = _ScriptedBackend(_fixed("plain answer"))
|
||||
payload = _request(tools = [LOOKUP_TOOL], tool_choice = "none", stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert body["choices"][0]["message"]["content"] == "plain answer"
|
||||
assert backend.calls[0]["tools"] is None
|
||||
|
||||
|
||||
def test_developer_message_folded_into_system_prompt(monkeypatch):
|
||||
# The "developer" role folds into one leading system message (local templates reject it).
|
||||
backend = _ScriptedBackend(_fixed("ok"))
|
||||
payload = _request(
|
||||
messages = [
|
||||
ChatMessage(role = "developer", content = "always be terse"),
|
||||
ChatMessage(role = "user", content = "hi"),
|
||||
],
|
||||
tools = [LOOKUP_TOOL],
|
||||
stream = False,
|
||||
)
|
||||
_call(payload, monkeypatch, backend)
|
||||
sent = backend.calls[0]["messages"]
|
||||
assert sent[0]["role"] == "system"
|
||||
assert "always be terse" in sent[0]["content"]
|
||||
assert all(m.get("role") != "developer" for m in sent)
|
||||
|
||||
|
||||
def test_failed_nudge_retry_keeps_original_response(monkeypatch):
|
||||
# A raising retry must not 500; the first response is returned.
|
||||
state = {"n": 0}
|
||||
|
||||
def responder(messages, tools):
|
||||
state["n"] += 1
|
||||
if state["n"] == 1:
|
||||
return ['<tool_call>{"name":"lookup"'] # unhealable signal
|
||||
raise RuntimeError("retry blew up")
|
||||
|
||||
backend = _ScriptedBackend(responder)
|
||||
payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert state["n"] == 2
|
||||
assert body["choices"][0]["finish_reason"] == "stop"
|
||||
assert body["choices"][0]["message"]["content"] == '<tool_call>{"name":"lookup"'
|
||||
|
||||
|
||||
def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch):
|
||||
# Double-failure nudge: the first response is delivered, but the retry's
|
||||
# generate() overwrites stats_holder. The monitor must record the FIRST
|
||||
# attempt's usage, not the discarded retry's.
|
||||
first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}}
|
||||
retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}}
|
||||
|
||||
class _PerCallStatsBackend(_ScriptedBackend):
|
||||
def __init__(self):
|
||||
# Unhealable truncated markup on both attempts -> retry is discarded.
|
||||
super().__init__(lambda m, t: ['<tool_call>{"name":"lookup"'])
|
||||
self._stats_seq = [first_stats, retry_stats]
|
||||
|
||||
def generate_chat_response(
|
||||
self,
|
||||
*,
|
||||
messages,
|
||||
tools = None,
|
||||
stats_holder = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.calls.append({"messages": messages, "tools": tools, **kwargs})
|
||||
stats = self._stats_seq[min(len(self.calls) - 1, len(self._stats_seq) - 1)]
|
||||
if stats_holder is not None:
|
||||
stats_holder["stats"] = stats
|
||||
for snap in self._responder(messages, tools):
|
||||
yield snap
|
||||
|
||||
backend = _PerCallStatsBackend()
|
||||
payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False)
|
||||
monitor = _install(monkeypatch, backend)
|
||||
|
||||
async def _run():
|
||||
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
|
||||
|
||||
asyncio.run(_run())
|
||||
assert len(backend.calls) == 2 # first attempt + one discarded retry
|
||||
[entry] = monitor.snapshot()
|
||||
# The delivered response is the first attempt, so its usage must be reported.
|
||||
assert entry["prompt_tokens"] == 7
|
||||
assert entry["completion_tokens"] == 3
|
||||
|
||||
|
||||
def test_monitor_records_healed_call_not_raw_xml(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False)
|
||||
monitor = _install(monkeypatch, backend)
|
||||
|
||||
async def _run():
|
||||
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
|
||||
|
||||
asyncio.run(_run())
|
||||
snap = monitor.snapshot(include_details = True)
|
||||
replies = json.dumps(snap)
|
||||
assert "<tool_call>" not in replies
|
||||
assert "lookup" in replies
|
||||
|
||||
|
||||
def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch):
|
||||
# Monitor mirrors what the client received, never the healed-away raw markup.
|
||||
backend = _ScriptedBackend(
|
||||
_fixed("Sure. ", 'Sure. <tool_call>{"name": "loo', "Sure. " + _CALL_XML)
|
||||
)
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = True)
|
||||
monitor = _install(monkeypatch, backend)
|
||||
|
||||
async def _run():
|
||||
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
|
||||
|
||||
response = asyncio.run(_run())
|
||||
_collect_sse(response)
|
||||
replies = json.dumps(monitor.snapshot(include_details = True))
|
||||
assert "<tool_call>" not in replies
|
||||
assert "Sure. " in replies
|
||||
assert "[tool_calls] lookup(" in replies
|
||||
|
||||
|
||||
def test_forced_tool_choice_narrows_templated_tools(monkeypatch):
|
||||
# A forced function is the only schema rendered into the template.
|
||||
backend = _ScriptedBackend(_fixed(_SEARCH_XML))
|
||||
payload = _request(
|
||||
tools = [LOOKUP_TOOL, SEARCH_TOOL],
|
||||
stream = False,
|
||||
tool_choice = {"type": "function", "function": {"name": "search"}},
|
||||
)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
templated = backend.calls[0]["tools"]
|
||||
assert [t["function"]["name"] for t in templated] == ["search"]
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "tool_calls"
|
||||
assert choice["message"]["tool_calls"][0]["function"]["name"] == "search"
|
||||
|
||||
|
||||
def test_multimodal_content_parts_flattened_for_local_template(monkeypatch):
|
||||
# Remote image URLs leave image=None, so content arrives as a part LIST:
|
||||
# text parts are kept, the image part dropped.
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(
|
||||
messages = [
|
||||
ChatMessage(
|
||||
role = "user",
|
||||
content = [
|
||||
{"type": "text", "text": "what is this?"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "https://example.com/cat.png"},
|
||||
},
|
||||
],
|
||||
)
|
||||
],
|
||||
tools = [LOOKUP_TOOL],
|
||||
stream = False,
|
||||
)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
templated = backend.calls[0]["messages"]
|
||||
assert all(isinstance(m.get("content"), str) for m in templated)
|
||||
assert any(m["content"] == "what is this?" for m in templated)
|
||||
assert body["choices"][0]["finish_reason"] == "tool_calls"
|
||||
|
||||
|
||||
def test_string_arguments_history_deserialized_for_template(monkeypatch):
|
||||
# JSON-string tool_calls arguments become dicts in the templated copy;
|
||||
# the HTTP response stays OpenAI-shaped.
|
||||
backend = _ScriptedBackend(_fixed("done"))
|
||||
payload = _request(
|
||||
tools = [LOOKUP_TOOL],
|
||||
stream = False,
|
||||
messages = [
|
||||
ChatMessage(role = "user", content = "weather?"),
|
||||
ChatMessage(
|
||||
role = "assistant",
|
||||
content = None,
|
||||
tool_calls = [
|
||||
{
|
||||
"id": "call_0",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": '{"q": "weather"}'},
|
||||
}
|
||||
],
|
||||
),
|
||||
ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"),
|
||||
],
|
||||
)
|
||||
_json_body(_call(payload, monkeypatch, backend))
|
||||
assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant")
|
||||
assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"}
|
||||
|
||||
|
||||
def test_unparseable_arguments_string_left_untouched(monkeypatch):
|
||||
backend = _ScriptedBackend(_fixed("ok"))
|
||||
payload = _request(
|
||||
tools = [LOOKUP_TOOL],
|
||||
stream = False,
|
||||
messages = [
|
||||
ChatMessage(role = "user", content = "hi"),
|
||||
ChatMessage(
|
||||
role = "assistant",
|
||||
content = None,
|
||||
tool_calls = [
|
||||
{
|
||||
"id": "call_0",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "not json {"},
|
||||
}
|
||||
],
|
||||
),
|
||||
ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"),
|
||||
],
|
||||
)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
assert body["choices"][0]["message"]["content"] == "ok"
|
||||
assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant")
|
||||
assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {"
|
||||
|
||||
|
||||
def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch):
|
||||
# mcp_enabled=true with an empty registry must not silently drop the
|
||||
# declared tools; the gate keys on the server-side path claiming the request.
|
||||
backend = _ScriptedBackend(_fixed(_CALL_XML))
|
||||
payload = _request(tools = [LOOKUP_TOOL], stream = False, mcp_enabled = True)
|
||||
body = _json_body(_call(payload, monkeypatch, backend))
|
||||
choice = body["choices"][0]
|
||||
assert choice["finish_reason"] == "tool_calls"
|
||||
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
|
||||
assert backend.calls[0]["tools"] == [LOOKUP_TOOL]
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue