Merge remote-tracking branch 'upstream/main' into feat/model-picker-per-model-config
# Conflicts: # studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
This commit is contained in:
commit
a615a40284
84 changed files with 13015 additions and 2157 deletions
148
.github/scripts/agent-guides-drive.sh
vendored
148
.github/scripts/agent-guides-drive.sh
vendored
|
|
@ -527,6 +527,154 @@ case "$MODE" in
|
|||
echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)"
|
||||
;;
|
||||
|
||||
# ── resume: does a launched agent's session survive exit and resume? ────
|
||||
# Unlike the other modes, this drives the real LAUNCH path (`unsloth start
|
||||
# <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
|
||||
|
|
|
|||
2
.github/workflows/consolidated-tests-ci.yml
vendored
2
.github/workflows/consolidated-tests-ci.yml
vendored
|
|
@ -364,7 +364,9 @@ jobs:
|
|||
tests/utils/test_attention_masks.py \
|
||||
tests/utils/test_trunc_normal_patch.py \
|
||||
tests/python/test_fast_language_model_text_only.py \
|
||||
tests/test_bad_mappings_redirect.py \
|
||||
tests/test_prefetch_snapshot_scope.py \
|
||||
tests/test_gemma_2b_mapper_key.py \
|
||||
--deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap'
|
||||
# The deselected test monkeypatches flash_attn_varlen_func, which is
|
||||
# only bound on the module when `flash_attn` is importable. flash_attn
|
||||
|
|
|
|||
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
|
||||
|
|
|
|||
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
|
||||
|
|
|
|||
205
.github/workflows/studio-windows-inference-smoke.yml
vendored
205
.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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
|
|
|||
24
install.ps1
24
install.ps1
|
|
@ -469,6 +469,17 @@ function Install-UnslothStudio {
|
|||
param(
|
||||
[Parameter(Mandatory = $true)][ScriptBlock]$Command
|
||||
)
|
||||
# Installer-pinned index installs (torch) must beat an inherited uv mirror
|
||||
# (#6898): when the command pins an index, clear every uv index env var so
|
||||
# it wins, then restore in finally. Other installs keep the user's mirror.
|
||||
$savedUvIndex = $null
|
||||
if ($Command.ToString() -match '--default-index') {
|
||||
$savedUvIndex = @{}
|
||||
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
|
||||
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
|
||||
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
|
||||
}
|
||||
}
|
||||
$prevEap = $ErrorActionPreference
|
||||
$ErrorActionPreference = "Continue"
|
||||
try {
|
||||
|
|
@ -488,6 +499,7 @@ function Install-UnslothStudio {
|
|||
return [int]$LASTEXITCODE
|
||||
} finally {
|
||||
$ErrorActionPreference = $prevEap
|
||||
if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -2200,7 +2212,7 @@ exit 0
|
|||
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
|
||||
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
|
||||
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
# Transient AMD-index failure: fall back to a CPU base so the install
|
||||
# still completes; Studio setup retries ROCm afterwards.
|
||||
|
|
@ -2209,7 +2221,7 @@ exit 0
|
|||
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
|
||||
# torch>= range, so without it uv would keep the ROCm build and only swap
|
||||
# the companions -- a mismatched venv the flavor-repair block won't fix.
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
|
||||
|
|
@ -2223,7 +2235,7 @@ exit 0
|
|||
} else {
|
||||
Write-TauriLog "STEP" "Installing PyTorch"
|
||||
substep "installing PyTorch ($TorchIndexUrl)..."
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
|
||||
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
|
||||
|
|
@ -2306,7 +2318,7 @@ exit 0
|
|||
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
|
||||
# "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
|
||||
# expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
|
||||
# is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
|
||||
# is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install
|
||||
# above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
|
||||
if (-not $SkipTorch) {
|
||||
$expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
|
||||
|
|
@ -2322,7 +2334,7 @@ exit 0
|
|||
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
|
||||
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
|
||||
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
|
||||
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
|
||||
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
|
||||
if ($torchFixExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
|
||||
|
|
@ -2331,7 +2343,7 @@ exit 0
|
|||
} elseif ($expectedTorchTag -ne 'rocm') {
|
||||
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
|
||||
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
|
||||
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
|
||||
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
|
||||
if ($torchFixExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
|
||||
|
|
|
|||
28
install.sh
28
install.sh
|
|
@ -159,6 +159,12 @@ run_maybe_quiet() {
|
|||
run_install_cmd() {
|
||||
_label="$1"
|
||||
shift
|
||||
# Installer-pinned index installs (torch) must beat an inherited uv mirror
|
||||
# (#6898): when we pass --default-index, neutralize every uv index env var so
|
||||
# the pinned index wins. Other installs keep the user's mirror.
|
||||
case " $* " in
|
||||
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
|
||||
esac
|
||||
if _is_verbose; then
|
||||
"$@" && return 0
|
||||
_rc=$?
|
||||
|
|
@ -2190,9 +2196,9 @@ _expected_torch_flavor_tag() {
|
|||
esac
|
||||
}
|
||||
|
||||
# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
|
||||
# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX /
|
||||
# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
|
||||
# resolves (torch + every transitive dep) via --index-url -- the same URLs the
|
||||
# resolves (torch + every transitive dep) via --default-index -- the same URLs the
|
||||
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
|
||||
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
|
||||
_torch_index_repairable() {
|
||||
|
|
@ -2744,7 +2750,7 @@ if [ "$_MIGRATED" = true ]; then
|
|||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
fi
|
||||
;;
|
||||
|
|
@ -2870,7 +2876,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
|
||||
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
--default-index "$TORCH_INDEX_URL"
|
||||
else
|
||||
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
|
||||
# Pass explicit wheel URLs so the matched trio is
|
||||
|
|
@ -2893,18 +2899,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
|
||||
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
--default-index "$TORCH_INDEX_URL"
|
||||
fi
|
||||
else
|
||||
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
|
||||
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
--default-index "$TORCH_INDEX_URL"
|
||||
fi
|
||||
else
|
||||
substep "installing PyTorch ($TORCH_INDEX_URL)..."
|
||||
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
--default-index "$TORCH_INDEX_URL"
|
||||
fi
|
||||
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
|
||||
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
|
||||
|
|
@ -2964,7 +2970,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--force-reinstall
|
||||
fi
|
||||
;;
|
||||
|
|
@ -2999,14 +3005,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
|
|||
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
|
||||
_installed_torch_tag=""
|
||||
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
|
||||
# Repair when flavor is wrong AND the index is plain --index-url reinstallable
|
||||
# Repair when flavor is wrong AND the index is plain --default-index reinstallable
|
||||
# (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
|
||||
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
|
||||
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
|
||||
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
|
||||
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
|
||||
"$TORCH_CONSTRAINT" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL" \
|
||||
--default-index "$TORCH_INDEX_URL" \
|
||||
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
|
||||
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
|
||||
_installed_torch_tag=""
|
||||
|
|
@ -3017,7 +3023,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
|
|||
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
|
||||
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
|
||||
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
|
||||
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
|
||||
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
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())
|
||||
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()
|
||||
|
|
@ -100,6 +100,10 @@ class LlamaServerNotFoundError(RuntimeError):
|
|||
Subclasses RuntimeError so existing handlers still catch it."""
|
||||
|
||||
|
||||
class _LlamaStreamCancelled(Exception):
|
||||
"""Internal signal for an expected client/request cancellation."""
|
||||
|
||||
|
||||
# Shared so the from_identifier preflight and the load-time raise stay in sync.
|
||||
LLAMA_SERVER_NOT_FOUND_DETAIL = (
|
||||
"This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
|
||||
|
|
@ -1436,6 +1440,50 @@ def _backfill_usage_from_timings(usage, timings):
|
|||
return out
|
||||
|
||||
|
||||
def _vulkan_lib_filename() -> str:
|
||||
return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so"
|
||||
|
||||
|
||||
# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit
|
||||
# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared
|
||||
# system RAM, so hold back the same margin rather than inventing a larger one.
|
||||
_IGPU_HOST_RESERVE_MIB = 1024
|
||||
|
||||
|
||||
def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int:
|
||||
"""Reserve host headroom on an integrated (shared-memory) Vulkan GPU.
|
||||
|
||||
An iGPU's reported free "VRAM" is really free system RAM, so sizing
|
||||
context/offload against all of it would push the host into swap or the OOM
|
||||
killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from
|
||||
ggml's device type, so a discrete card is never touched; only ever reduces.
|
||||
"""
|
||||
if not is_igpu:
|
||||
return free_mib
|
||||
return max(0, free_mib - _IGPU_HOST_RESERVE_MIB)
|
||||
|
||||
|
||||
def _llama_lib_dir(binary: str) -> Path:
|
||||
# The installer exposes llama-server as a top-level entrypoint into build/bin/,
|
||||
# where the ggml backend libs live, so callers looking for sibling libs (Vulkan
|
||||
# detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a
|
||||
# symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to
|
||||
# a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it
|
||||
# cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's
|
||||
# exec target too, so a wrapper-based install still finds build/bin.
|
||||
resolved = Path(binary).resolve()
|
||||
try:
|
||||
with open(resolved, "rb") as _f:
|
||||
_head = _f.read(256)
|
||||
if _head.startswith(b"#!"):
|
||||
_m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore"))
|
||||
if _m:
|
||||
return (resolved.parent / _m.group(1)).resolve().parent
|
||||
except OSError:
|
||||
pass
|
||||
return resolved.parent
|
||||
|
||||
|
||||
def _is_external_link(path: Path) -> bool:
|
||||
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
|
||||
or a Windows directory junction / reparse point. Such a link resolves into
|
||||
|
|
@ -1493,6 +1541,7 @@ class LlamaCppBackend:
|
|||
self._context_length: Optional[int] = None
|
||||
self._effective_context_length: Optional[int] = None
|
||||
self._max_context_length: Optional[int] = None
|
||||
self._effective_parallel_slots: int = 1
|
||||
self._chat_template: Optional[str] = None
|
||||
self._chat_template_override: Optional[str] = None
|
||||
self._supports_reasoning: bool = False
|
||||
|
|
@ -1678,6 +1727,15 @@ class LlamaCppBackend:
|
|||
"""Return the effective context length the server is running at."""
|
||||
return self._effective_context_length or self._context_length
|
||||
|
||||
@property
|
||||
def effective_parallel_slots(self) -> int:
|
||||
"""Return the serving-slot count the active llama-server actually uses."""
|
||||
try:
|
||||
slots = int(getattr(self, "_effective_parallel_slots", 1))
|
||||
except (TypeError, ValueError):
|
||||
slots = 1
|
||||
return max(1, slots)
|
||||
|
||||
@property
|
||||
def max_context_length(self) -> Optional[int]:
|
||||
"""Return the largest context that fits on this hardware at load time.
|
||||
|
|
@ -1694,6 +1752,16 @@ class LlamaCppBackend:
|
|||
"""Return the model's native context length from GGUF metadata."""
|
||||
return self._context_length
|
||||
|
||||
def _commit_effective_parallel_slots(self, n_parallel: int) -> None:
|
||||
try:
|
||||
slots = int(n_parallel)
|
||||
except (TypeError, ValueError):
|
||||
slots = 1
|
||||
self._effective_parallel_slots = max(1, slots)
|
||||
|
||||
def _reset_effective_parallel_slots(self) -> None:
|
||||
self._effective_parallel_slots = 1
|
||||
|
||||
@staticmethod
|
||||
def _read_rss_bytes(pid: int) -> Optional[int]:
|
||||
"""Resident set size of ``pid`` in bytes, from /proc/<pid>/status (Linux).
|
||||
|
|
@ -2278,6 +2346,30 @@ class LlamaCppBackend:
|
|||
|
||||
return total
|
||||
|
||||
@staticmethod
|
||||
def _is_vulkan_backend(binary: Optional[str] = None) -> bool:
|
||||
"""True if the installed llama.cpp build is Vulkan-only.
|
||||
|
||||
The official prebuilts are single-backend, so the Vulkan ggml lib next
|
||||
to llama-server identifies a Vulkan build. Keeps the free-memory probe
|
||||
and GPU pin in ggml's Vulkan device-index space. For a custom
|
||||
multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer
|
||||
to that backend (torch-usable, better-understood probe/pin).
|
||||
"""
|
||||
binary = binary or LlamaCppBackend._find_llama_server_binary()
|
||||
if not binary:
|
||||
return False
|
||||
lib_dir = _llama_lib_dir(binary)
|
||||
if not (lib_dir / _vulkan_lib_filename()).is_file():
|
||||
return False
|
||||
for _backend in ("cuda", "hip"):
|
||||
sibling = (
|
||||
f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so"
|
||||
)
|
||||
if (lib_dir / sibling).is_file():
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _resolve_visible_physical_ids() -> Optional[list[int]]:
|
||||
"""Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on
|
||||
|
|
@ -2440,11 +2532,42 @@ class LlamaCppBackend:
|
|||
return True
|
||||
|
||||
@staticmethod
|
||||
def _get_gpu_free_memory() -> list[tuple[int, int]]:
|
||||
def _visible_devices_mask(env_name: str) -> Optional[set[int]]:
|
||||
"""Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset.
|
||||
|
||||
``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask
|
||||
("") yields an empty set (all devices hidden), distinct from an unset
|
||||
var (None, no mask). Used by the nvidia-smi probe.
|
||||
"""
|
||||
raw = os.environ.get(env_name)
|
||||
if raw is None:
|
||||
return None
|
||||
try:
|
||||
return set(int(x.strip()) for x in raw.split(",") if x.strip())
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]:
|
||||
"""``--device Vulkan<i>,...`` to pin a Vulkan launch to selected GPUs.
|
||||
|
||||
The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory
|
||||
reports and the registry names ``Vulkan<i>``). Pin by that name, NOT via
|
||||
GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw
|
||||
vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices
|
||||
and deduplicating ICDs), so a compact ordinal there could select a
|
||||
different physical device or the CPU rasterizer.
|
||||
"""
|
||||
if not gpu_indices:
|
||||
return []
|
||||
return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)]
|
||||
|
||||
@staticmethod
|
||||
def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]:
|
||||
"""Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by
|
||||
index; empty if no supported GPU is reachable. Thin wrapper over
|
||||
``_get_gpu_memory`` for callers that only need free VRAM."""
|
||||
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()]
|
||||
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)]
|
||||
|
||||
@staticmethod
|
||||
def _apple_metal_memory_budget_bytes() -> int:
|
||||
|
|
@ -2475,7 +2598,7 @@ class LlamaCppBackend:
|
|||
return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION)
|
||||
|
||||
@staticmethod
|
||||
def _get_gpu_memory() -> list[tuple[int, int, int]]:
|
||||
def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
|
||||
"""Query free AND total memory per GPU.
|
||||
|
||||
Order:
|
||||
|
|
@ -2487,9 +2610,18 @@ class LlamaCppBackend:
|
|||
probe returned [] on AMD) and NVIDIA hosts missing
|
||||
``nvidia-smi`` from PATH.
|
||||
|
||||
On a Vulkan build the ggml Vulkan probe is authoritative, so the indices
|
||||
are ggml's compact Vulkan ordinals (the space the pin selects via
|
||||
``--device Vulkan<i>``). It reports ``total`` for discrete cards and 0
|
||||
for an iGPU (shared RAM) so the fit falls back to free*frac there.
|
||||
Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm.
|
||||
|
||||
Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no
|
||||
supported GPU is reachable. ``total`` lets the fit reserve absolute headroom.
|
||||
supported GPU is reachable.
|
||||
"""
|
||||
binary = binary or LlamaCppBackend._find_llama_server_binary()
|
||||
if LlamaCppBackend._is_vulkan_backend(binary):
|
||||
return LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
|
||||
# ── NVIDIA via nvidia-smi ────────────────────────────────────
|
||||
try:
|
||||
result = subprocess.run(
|
||||
|
|
@ -2505,16 +2637,7 @@ class LlamaCppBackend:
|
|||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode == 0:
|
||||
allowed: Optional[set[int]] = None
|
||||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if cvd is not None:
|
||||
try:
|
||||
# `if x.strip()` filters trailing-comma masks ("0,1,").
|
||||
# Empty mask (CVD="") yields an empty set -> all GPUs
|
||||
# filtered out, per codebase convention.
|
||||
allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip())
|
||||
except ValueError:
|
||||
pass
|
||||
allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES")
|
||||
gpus: list[tuple[int, int, int]] = []
|
||||
for line in result.stdout.strip().splitlines():
|
||||
parts = [p.strip() for p in line.split(",")]
|
||||
|
|
@ -2579,6 +2702,91 @@ class LlamaCppBackend:
|
|||
logger.debug(f"torch GPU probe failed: {e}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
|
||||
"""Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
|
||||
|
||||
Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance
|
||||
in this process) and returns (device_index, free_mib, total_mib) sorted
|
||||
by index. The index is ggml's compact Vulkan ordinal -- the one the
|
||||
registry names ``Vulkan<index>`` and load_model pins with ``--device``,
|
||||
NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set
|
||||
``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the
|
||||
list already reflects it. iGPUs leave a host-RAM margin (see
|
||||
``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
|
||||
their real total through. [] when no Vulkan build or device is reachable.
|
||||
"""
|
||||
binary = binary or LlamaCppBackend._find_llama_server_binary()
|
||||
if not binary:
|
||||
return []
|
||||
binary_dir = _llama_lib_dir(binary)
|
||||
if not (binary_dir / _vulkan_lib_filename()).is_file():
|
||||
return []
|
||||
|
||||
env = child_env_without_native_path_secret()
|
||||
# Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so
|
||||
# the probe enumerates the same device list the launch will, named
|
||||
# Vulkan0..N in the compact order reported here and pinned by that name
|
||||
# via --device -- probe, mask, and pin stay in one index space. Do NOT
|
||||
# filter the mask in Python: ggml parses the env var in raw
|
||||
# vkEnumeratePhysicalDevices space while this probe reports the compact
|
||||
# post-filter ordinal, so a Python filter would compare mismatched spaces.
|
||||
if sys.platform != "win32":
|
||||
# Let the loader resolve sibling ggml libs next to the binary.
|
||||
existing_ld = env.get("LD_LIBRARY_PATH", "")
|
||||
env["LD_LIBRARY_PATH"] = (
|
||||
f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir)
|
||||
)
|
||||
probe_script = Path(__file__).with_name("_vulkan_probe.py")
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(probe_script), str(binary_dir)],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 15,
|
||||
env = env,
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.debug(
|
||||
f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}"
|
||||
)
|
||||
return []
|
||||
except Exception as e:
|
||||
logger.debug(f"vulkan GPU probe failed: {e}")
|
||||
return []
|
||||
|
||||
gpus: list[tuple[int, int, int]] = []
|
||||
for line in result.stdout.strip().splitlines():
|
||||
parts = line.split("\t")
|
||||
if len(parts) != 4:
|
||||
continue
|
||||
try:
|
||||
idx = int(parts[0])
|
||||
free_mib = int(parts[1]) // (1024 * 1024)
|
||||
is_igpu = parts[2] == "1"
|
||||
# iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
|
||||
# fit stays on free*frac (the host reserve below is its
|
||||
# headroom); a discrete card passes its real total through.
|
||||
total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024)
|
||||
except ValueError:
|
||||
continue
|
||||
capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu)
|
||||
if capped < free_mib:
|
||||
logger.info(
|
||||
f"Vulkan device VK{idx} is an integrated GPU sharing system "
|
||||
f"RAM; reserving {free_mib - capped}MiB host headroom "
|
||||
f"({free_mib}->{capped}MiB usable)"
|
||||
)
|
||||
gpus.append((idx, capped, total_mib))
|
||||
gpus.sort(key = lambda g: g[0])
|
||||
if gpus:
|
||||
logger.info(
|
||||
"Vulkan GPU memory detected: "
|
||||
+ ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus)
|
||||
)
|
||||
return gpus
|
||||
|
||||
@staticmethod
|
||||
def _available_system_memory_mib() -> Optional[int]:
|
||||
"""Available system RAM in MiB (psutil, then /proc/meminfo), or None if
|
||||
|
|
@ -2807,7 +3015,8 @@ class LlamaCppBackend:
|
|||
def _llama_server_env_for_binary(binary: str) -> dict[str, str]:
|
||||
"""Build a subprocess env that lets llama-server resolve native libs."""
|
||||
env = child_env_without_native_path_secret()
|
||||
binary_dir = str(Path(binary).parent)
|
||||
# _llama_lib_dir resolves the llama-server symlink to the real build/bin.
|
||||
binary_dir = str(_llama_lib_dir(binary))
|
||||
|
||||
if sys.platform == "win32":
|
||||
# Ordering: see _build_windows_path_dirs. #5106.
|
||||
|
|
@ -4488,6 +4697,29 @@ class LlamaCppBackend:
|
|||
cancel_event = cancel_event,
|
||||
)
|
||||
|
||||
def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
|
||||
"""A drafter already in this repo's local HF cache, reused offline when a
|
||||
fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all
|
||||
cached snapshots; else an existing ``MTP/`` copy (any precision -- the
|
||||
target verifies every drafted token). None if none is cached."""
|
||||
try:
|
||||
from utils.models.model_config import _iter_hf_cache_snapshots
|
||||
|
||||
roots: list[Path] = []
|
||||
subdirs: list[Path] = []
|
||||
for snap in _iter_hf_cache_snapshots(hf_repo): # newest first
|
||||
for f in sorted(_gguf_snapshot_files(snap)):
|
||||
if _is_companion_gguf_path(f) and "mmproj" not in f.lower():
|
||||
(roots if "/" not in f else subdirs).append(snap / f)
|
||||
# Keep snapshot order (newest first), root before any MTP/ copy, so a
|
||||
# newer main GGUF pairs with the newest cached drafter, not a stale one.
|
||||
for cand in roots + subdirs:
|
||||
if cand.is_file():
|
||||
return str(cand)
|
||||
except Exception as e:
|
||||
logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e)
|
||||
return None
|
||||
|
||||
def _download_mtp(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -4504,11 +4736,25 @@ class LlamaCppBackend:
|
|||
are intentionally skipped. Returns the local path, or None.
|
||||
"""
|
||||
|
||||
# Offline, reuse any drafter already on disk (a fresh copy can't be
|
||||
# fetched). Online, _download_companion_gguf/hf_hub_download reuse the
|
||||
# current cached file and refetch a changed one, so skip the probe here
|
||||
# rather than pair new weights with a stale draft.
|
||||
if _hf_env_offline():
|
||||
cached = self._cached_repo_mtp_drafter(hf_repo)
|
||||
if cached:
|
||||
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
|
||||
return cached
|
||||
|
||||
def _pick_mtp(candidates: list[str]) -> Optional[str]:
|
||||
# Root-level only: MTP/ subdir copies now share the mtp- prefix but
|
||||
# are explicit-selection, not auto-fetch (they'd sort ahead of root).
|
||||
mtp_files = sorted(
|
||||
f
|
||||
for f in candidates
|
||||
if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-")
|
||||
if f.lower().endswith(".gguf")
|
||||
and "/" not in f
|
||||
and Path(f).name.lower().startswith("mtp-")
|
||||
)
|
||||
return mtp_files[0] if mtp_files else None
|
||||
|
||||
|
|
@ -5210,6 +5456,7 @@ class LlamaCppBackend:
|
|||
# Resolve llama-server now but defer a not-found error: a block-diffusion
|
||||
# GGUF uses the diffusion runner, and its arch is only known after the header.
|
||||
binary = self._find_llama_server_binary()
|
||||
is_vulkan_backend = self._is_vulkan_backend(binary)
|
||||
|
||||
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
|
||||
# mtp_draft_path arrives set for local Gemma loads (detected
|
||||
|
|
@ -5449,7 +5696,8 @@ class LlamaCppBackend:
|
|||
model_size = gguf_size + mmproj_size
|
||||
# 2-tuple gpus for existing logic + a total map for the absolute
|
||||
# per-GPU headroom (correct when the GPU is already partly used).
|
||||
_gpu_mem = self._get_gpu_memory()
|
||||
# Pass binary so a Vulkan build probes ggml's Vulkan ordinals.
|
||||
_gpu_mem = self._get_gpu_memory(binary)
|
||||
gpus = [(idx, free) for idx, free, _t in _gpu_mem]
|
||||
total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
|
||||
|
||||
|
|
@ -6222,7 +6470,12 @@ class LlamaCppBackend:
|
|||
# cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
|
||||
# oversize load the OS would otherwise kill mid-flight. Base model
|
||||
# only: an optional MTP drafter is dropped by the MTP-drop fallback.
|
||||
if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices):
|
||||
# CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals.
|
||||
if (
|
||||
model_size is not None
|
||||
and not is_vulkan_backend
|
||||
and self._amd_apu_wants_unified_memory(gpu_indices)
|
||||
):
|
||||
_ram_msg = self._apu_ram_shortfall_message(
|
||||
model_size, self._available_system_memory_mib()
|
||||
)
|
||||
|
|
@ -6485,6 +6738,12 @@ class LlamaCppBackend:
|
|||
", ".join(unsupported_cache_flags),
|
||||
)
|
||||
|
||||
# Vulkan pins via --device (a cmd arg, unlike the env-based
|
||||
# CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's
|
||||
# last-wins parsing lets a user --device override Studio's pick.
|
||||
if is_vulkan_backend and gpu_indices is not None:
|
||||
cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices)
|
||||
|
||||
# User pass-through args go last so llama.cpp's last-wins parsing
|
||||
# lets the user override Studio's auto-set flags. Already
|
||||
# validated by the route via validate_extra_args().
|
||||
|
|
@ -6536,23 +6795,25 @@ class LlamaCppBackend:
|
|||
env.setdefault("OMP_NUM_THREADS", "2")
|
||||
|
||||
# AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
|
||||
# shared system RAM. setdefault so a user value wins.
|
||||
if self._amd_apu_wants_unified_memory(gpu_indices):
|
||||
# shared system RAM. setdefault so a user value wins. Not on Vulkan
|
||||
# (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids.
|
||||
if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices):
|
||||
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
|
||||
logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1")
|
||||
|
||||
# DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU).
|
||||
# See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1.
|
||||
if self._apply_datacenter_env(env, gpu_indices):
|
||||
if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices):
|
||||
multi_gpu = self._effective_gpu_count(gpu_indices) > 1
|
||||
logger.info(
|
||||
f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})"
|
||||
)
|
||||
|
||||
# Pin to selected GPU(s). On ROCm, narrowing only
|
||||
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
|
||||
# set, so set HIP_VISIBLE_DEVICES too.
|
||||
if gpu_indices is not None:
|
||||
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so
|
||||
# set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device
|
||||
# (above), not here.
|
||||
if gpu_indices is not None and not is_vulkan_backend:
|
||||
pinned = ",".join(str(i) for i in gpu_indices)
|
||||
env["CUDA_VISIBLE_DEVICES"] = pinned
|
||||
try:
|
||||
|
|
@ -6946,6 +7207,7 @@ class LlamaCppBackend:
|
|||
)
|
||||
|
||||
self._healthy = True
|
||||
self._commit_effective_parallel_slots(n_parallel)
|
||||
|
||||
# Commit caller intent only after _healthy=True so a failed start
|
||||
# can't poison the next inheritance check. None keeps prior, []
|
||||
|
|
@ -7483,6 +7745,7 @@ class LlamaCppBackend:
|
|||
self._context_length = None
|
||||
self._effective_context_length = None
|
||||
self._max_context_length = None
|
||||
self._reset_effective_parallel_slots()
|
||||
self._chat_template = None
|
||||
self._chat_template_override = None
|
||||
self._supports_reasoning = False
|
||||
|
|
@ -7538,6 +7801,7 @@ class LlamaCppBackend:
|
|||
# Stop the watchdog before a deliberate kill so a planned reload/unload
|
||||
# isn't seen as a crash; a real crash never routes through here.
|
||||
self._stop_mtp_crash_watchdog()
|
||||
self._reset_effective_parallel_slots()
|
||||
if self._process is None:
|
||||
return
|
||||
try:
|
||||
|
|
@ -8379,7 +8643,7 @@ class LlamaCppBackend:
|
|||
):
|
||||
"""Open one streaming POST and let cancel interrupt prefill or reads."""
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise GeneratorExit
|
||||
raise _LlamaStreamCancelled
|
||||
|
||||
_cancel_closed = threading.Event()
|
||||
_response_ref: list = [None]
|
||||
|
|
@ -8424,13 +8688,13 @@ class LlamaCppBackend:
|
|||
) as response:
|
||||
_response_ref[0] = response
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise GeneratorExit
|
||||
raise _LlamaStreamCancelled
|
||||
yield response
|
||||
return
|
||||
except (httpx.RequestError, RuntimeError):
|
||||
# Response was closed by the cancel watcher
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise GeneratorExit
|
||||
raise _LlamaStreamCancelled
|
||||
raise
|
||||
finally:
|
||||
_cancel_closed.set()
|
||||
|
|
@ -8633,6 +8897,8 @@ class LlamaCppBackend:
|
|||
"finish_reason": _metadata_finish_reason,
|
||||
}
|
||||
|
||||
except _LlamaStreamCancelled:
|
||||
return
|
||||
except httpx.ConnectError as e:
|
||||
# Server already down. If this was an MTP+tensor crash, recover by
|
||||
# reloading without MTP (scheduled in the background) and fail this
|
||||
|
|
@ -9757,6 +10023,8 @@ class LlamaCppBackend:
|
|||
break
|
||||
continue
|
||||
|
||||
except _LlamaStreamCancelled:
|
||||
return
|
||||
except httpx.ConnectError:
|
||||
# Mark unresolved provisional cards as failed before raising.
|
||||
for _pid, _pname in provisional_started_tool_calls.items():
|
||||
|
|
@ -9939,6 +10207,8 @@ class LlamaCppBackend:
|
|||
if _meta is not None:
|
||||
yield _meta
|
||||
|
||||
except _LlamaStreamCancelled:
|
||||
return
|
||||
except httpx.ConnectError:
|
||||
raise RuntimeError("Lost connection to llama-server")
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1533,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel):
|
|||
tool_use_id: str
|
||||
content: Union[str, list] = ""
|
||||
|
||||
@field_validator("content", mode = "before")
|
||||
@classmethod
|
||||
def _coerce_null_content(cls, v):
|
||||
# Some clients send null content for an empty tool result; the str|list
|
||||
# union would 400 on it, so treat null as "".
|
||||
return "" if v is None else v
|
||||
|
||||
|
||||
# Block types the converter translates explicitly. Anything else (thinking /
|
||||
# redacted_thinking, a provider block a resumed session replays, or a future type)
|
||||
# is accepted as an unknown block and dropped by the converter, rather than 400-ing
|
||||
# the whole request on strict validation.
|
||||
_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"})
|
||||
|
||||
|
||||
class AnthropicUnknownBlock(BaseModel):
|
||||
type: str
|
||||
model_config = {"extra": "allow"}
|
||||
|
||||
@field_validator("type")
|
||||
@classmethod
|
||||
def _only_unknown_types(cls, v):
|
||||
# Known types parse as their typed models above (so a malformed known block
|
||||
# still fails cleanly); this fallback only catches the rest.
|
||||
if v in _KNOWN_ANTHROPIC_BLOCK_TYPES:
|
||||
raise ValueError("known block type handled by its typed model")
|
||||
return v
|
||||
|
||||
|
||||
AnthropicContentBlock = Union[
|
||||
AnthropicTextBlock,
|
||||
AnthropicImageBlock,
|
||||
AnthropicToolUseBlock,
|
||||
AnthropicToolResultBlock,
|
||||
AnthropicUnknownBlock,
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -1583,6 +1612,40 @@ class AnthropicMessage(BaseModel):
|
|||
role: Literal["user", "assistant"]
|
||||
content: Union[str, list[AnthropicContentBlock]]
|
||||
|
||||
@model_validator(mode = "before")
|
||||
@classmethod
|
||||
def _normalize_content(cls, data):
|
||||
# Role-aware leniency that never silently drops real user input:
|
||||
# - assistant: a resumed tool-only turn's null content -> "" (str|list would
|
||||
# 400 on null; "" keeps the converter's `for block in content` safe).
|
||||
# Unknown blocks (thinking / future types) validate via
|
||||
# AnthropicUnknownBlock and are dropped by the converter.
|
||||
# - user: keep strict. Null user content stays None so str|list rejects it
|
||||
# (400) rather than forwarding an empty prompt; and reject block types the
|
||||
# converter cannot translate, since it silently skips unknown user blocks
|
||||
# -- a user turn made only of them would validate yet send no content
|
||||
# (silent data loss).
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
content = data.get("content")
|
||||
if data.get("role") == "assistant":
|
||||
# Coerce only an explicit null (resumed tool-only turn). A missing
|
||||
# content key stays malformed so the required-field check still 400s.
|
||||
if "content" in data and content is None:
|
||||
return {**data, "content": ""}
|
||||
return data
|
||||
if isinstance(content, list):
|
||||
for block in content:
|
||||
btype = (
|
||||
block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
|
||||
)
|
||||
# Guard the value: a non-string type is unsupported too, and a
|
||||
# membership test on an unhashable value would raise TypeError
|
||||
# (escaping as a 500 instead of a clean 400).
|
||||
if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES:
|
||||
raise ValueError(f"unsupported content block type {btype!r} in a user message")
|
||||
return data
|
||||
|
||||
|
||||
class AnthropicTool(BaseModel):
|
||||
# Client tools have input_schema; server tools may only have type/name.
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -1770,3 +1770,187 @@ class TestAnthropicMessagesToolRouting:
|
|||
|
||||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
assert backend.calls[0][0] == "plain"
|
||||
|
||||
|
||||
def test_resumed_session_thinking_and_null_content_do_not_400():
|
||||
# A resumed session replays assistant turns with `thinking` (and sometimes null)
|
||||
# content. Those must be accepted (thinking dropped by the converter), not 400ed.
|
||||
from pydantic import ValidationError
|
||||
|
||||
req = AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "secret reasoning", "signature": "s"},
|
||||
{"type": "text", "text": "the answer"},
|
||||
{"type": "tool_use", "id": "t1", "name": "f", "input": {}},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": None}, # tool-only turn serialized as null
|
||||
],
|
||||
)
|
||||
# Known blocks still parse as their typed models; only the unknown one is loose.
|
||||
assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock"
|
||||
assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock"
|
||||
assert req.messages[2].content == "" # null coerced
|
||||
|
||||
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
|
||||
assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content"))
|
||||
assert assistant["content"] == "the answer"
|
||||
assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded
|
||||
|
||||
# A malformed KNOWN block still fails cleanly instead of being swallowed.
|
||||
with pytest.raises(ValidationError):
|
||||
AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}],
|
||||
)
|
||||
|
||||
|
||||
def test_user_null_content_rejected():
|
||||
# The null->"" leniency is assistant-only; a null user content must be rejected
|
||||
# at the boundary, not coerced into an empty prompt and forwarded to the model.
|
||||
from pydantic import ValidationError
|
||||
with pytest.raises(ValidationError):
|
||||
AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [{"role": "user", "content": None}],
|
||||
)
|
||||
|
||||
|
||||
def test_user_unknown_block_rejected_not_silently_dropped():
|
||||
# The converter skips user blocks it cannot translate, so a user turn whose only
|
||||
# block is unknown would validate yet forward no content. Reject at the boundary
|
||||
# to avoid that silent data loss (the assistant fallback is unaffected).
|
||||
from pydantic import ValidationError
|
||||
with pytest.raises(ValidationError):
|
||||
AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [
|
||||
{"role": "user", "content": [{"type": "document", "source": {}}]},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_user_translatable_blocks_still_accepted():
|
||||
# text / image / tool_result are translatable, so a real user message built from
|
||||
# them must still pass; the unknown-block guard only trips on other types.
|
||||
req = AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is this?"},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {"type": "base64", "media_type": "image/png", "data": "AA"},
|
||||
},
|
||||
{"type": "tool_result", "tool_use_id": "t1", "content": "ok"},
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
assert [type(b).__name__ for b in req.messages[0].content] == [
|
||||
"AnthropicTextBlock",
|
||||
"AnthropicImageBlock",
|
||||
"AnthropicToolResultBlock",
|
||||
]
|
||||
|
||||
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
|
||||
assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai)
|
||||
|
||||
|
||||
def test_user_malformed_known_block_still_rejected():
|
||||
# The guard only allow-lists a user block's *type*; the union still validates its
|
||||
# shape, so a known-but-malformed block (tool_result without tool_use_id) fails.
|
||||
from pydantic import ValidationError
|
||||
with pytest.raises(ValidationError):
|
||||
AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [
|
||||
{"role": "user", "content": [{"type": "tool_result", "content": "x"}]},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_user_content_block_non_string_type_rejected_cleanly():
|
||||
# A user block whose `type` is a non-string (unhashable list / dict, or a stray
|
||||
# int) must fail as a clean validation error, not raise TypeError from the
|
||||
# frozenset membership test and escape as a 500.
|
||||
from pydantic import ValidationError
|
||||
for bad_type in ([], {}, 5):
|
||||
with pytest.raises(ValidationError):
|
||||
AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [{"role": "user", "content": [{"type": bad_type}]}],
|
||||
)
|
||||
|
||||
|
||||
def test_assistant_missing_content_key_still_rejected():
|
||||
# The null -> "" leniency is only for an EXPLICIT null. An assistant message that
|
||||
# omits content entirely stays malformed and must fail required-field validation.
|
||||
from pydantic import ValidationError
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [{"role": "assistant"}],
|
||||
)
|
||||
# An explicit null is still accepted and coerced (regression guard).
|
||||
req = AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": None},
|
||||
],
|
||||
)
|
||||
assert req.messages[1].content == ""
|
||||
|
||||
|
||||
def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch):
|
||||
# user -> assistant(null) -> user is now accepted: the null assistant turn coerces
|
||||
# to "" and is dropped. The route must then coalesce the two remaining user turns
|
||||
# so a strict GGUF chat template does not 400 on non-alternating roles.
|
||||
backend = _mock_backend(monkeypatch, context_length = 2048)
|
||||
|
||||
class _Req:
|
||||
state = SimpleNamespace()
|
||||
url = SimpleNamespace(path = "/v1/messages")
|
||||
method = "POST"
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
payload = AnthropicMessagesRequest(
|
||||
model = "x",
|
||||
max_tokens = 16,
|
||||
messages = [
|
||||
{"role": "user", "content": "first question"},
|
||||
{"role": "assistant", "content": None},
|
||||
{"role": "user", "content": "please continue"},
|
||||
],
|
||||
)
|
||||
|
||||
response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t"))
|
||||
assert response.status_code == 200
|
||||
|
||||
[(_path, kwargs)] = backend.calls
|
||||
user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"]
|
||||
assert len(user_turns) == 1 # the two user turns were merged, not left adjacent
|
||||
merged = user_turns[0]["content"]
|
||||
if isinstance(merged, list):
|
||||
merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict))
|
||||
assert "first question" in merged and "please continue" in merged
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -55,6 +55,27 @@ def _host(**kw):
|
|||
return ilp.HostInfo(**base)
|
||||
|
||||
|
||||
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():
|
||||
pre26 = _host(
|
||||
system = "Darwin",
|
||||
|
|
@ -313,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
|
||||
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 == []
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -515,6 +515,7 @@ class ScriptedClient:
|
|||
_url,
|
||||
json = None,
|
||||
timeout = None,
|
||||
headers = None,
|
||||
):
|
||||
self.posts.append(json)
|
||||
return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])
|
||||
|
|
|
|||
|
|
@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
|
|||
logger.info("llama update: installing", cmd = " ".join(cmd))
|
||||
# Stream progress lines into job["progress"].
|
||||
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
|
||||
# Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm
|
||||
# box would otherwise re-route and silently replace the Vulkan build.
|
||||
# Re-assert it via the same env flag setup uses (mirrors
|
||||
# _rocm_install_args).
|
||||
if asset and "vulkan" in asset.lower():
|
||||
env["UNSLOTH_FORCE_VULKAN"] = "1"
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout = subprocess.PIPE,
|
||||
|
|
|
|||
|
|
@ -81,7 +81,6 @@ import {
|
|||
TestTube01Icon,
|
||||
ZapIcon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage";
|
||||
import {
|
||||
Tooltip,
|
||||
TooltipContent,
|
||||
|
|
@ -97,6 +96,7 @@ import {
|
|||
createChatProject,
|
||||
deleteChatProject,
|
||||
deleteChatItem,
|
||||
listStoredChatThreads,
|
||||
moveChatItemToProject,
|
||||
renameChatItem,
|
||||
renameChatProject,
|
||||
|
|
@ -582,7 +582,14 @@ export function AppSidebar() {
|
|||
useEffect(() => {
|
||||
if (!pendingRename) return;
|
||||
const match = allChatItems.find((i) => i.id === pendingRename.id);
|
||||
if (match && match.title === pendingRename.title) setPendingRename(null);
|
||||
if (!match || match.title !== pendingRename.title) return;
|
||||
queueMicrotask(() => {
|
||||
setPendingRename((current) =>
|
||||
current?.id === pendingRename.id && current.title === pendingRename.title
|
||||
? null
|
||||
: current,
|
||||
);
|
||||
});
|
||||
}, [allChatItems, pendingRename]);
|
||||
const [creatingProject, setCreatingProject] = useState(false);
|
||||
const [projectNameDraft, setProjectNameDraft] = useState("");
|
||||
|
|
@ -680,12 +687,6 @@ export function AppSidebar() {
|
|||
useState<DeleteTarget | null>(null);
|
||||
const [deleteProjectFiles, setDeleteProjectFiles] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
if (confirmingDelete?.kind !== "project") {
|
||||
setDeleteProjectFiles(false);
|
||||
}
|
||||
}, [confirmingDelete]);
|
||||
|
||||
async function commitDelete() {
|
||||
const target = confirmingDelete;
|
||||
if (!target) return;
|
||||
|
|
|
|||
|
|
@ -3,27 +3,35 @@
|
|||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Progress } from "@/components/ui/progress";
|
||||
import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store";
|
||||
import { useMonitorOverlayStore } from "@/features/settings";
|
||||
import { useSystemInfo } from "@/hooks/use-system";
|
||||
import { useT } from "@/i18n";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react";
|
||||
import { motion } from "motion/react";
|
||||
import { useRef } from "react";
|
||||
import { AnimatePresence, motion, useDragControls } from "motion/react";
|
||||
import { type PointerEvent, useMemo, useState } from "react";
|
||||
|
||||
function clampPercent(value: number): number {
|
||||
return Math.max(0, Math.min(100, value));
|
||||
}
|
||||
|
||||
function usageIndicatorClass(percent: number): string {
|
||||
if (percent >= 90) return "bg-destructive";
|
||||
if (percent >= 70) return "bg-amber-500";
|
||||
if (percent >= 90) {
|
||||
return "bg-destructive";
|
||||
}
|
||||
if (percent >= 70) {
|
||||
return "bg-amber-500";
|
||||
}
|
||||
return "bg-primary";
|
||||
}
|
||||
|
||||
function usageTextClass(percent: number): string {
|
||||
if (percent >= 90) return "text-destructive";
|
||||
if (percent >= 70) return "text-amber-600 dark:text-amber-400";
|
||||
if (percent >= 90) {
|
||||
return "text-destructive";
|
||||
}
|
||||
if (percent >= 70) {
|
||||
return "text-amber-600 dark:text-amber-400";
|
||||
}
|
||||
return "text-primary";
|
||||
}
|
||||
|
||||
|
|
@ -39,9 +47,18 @@ export function FloatingMonitor() {
|
|||
const { isOpen, setIsOpen } = useMonitorOverlayStore();
|
||||
const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 });
|
||||
|
||||
const constraintsRef = useRef<HTMLDivElement>(null);
|
||||
const [constraintsElement, setConstraintsElement] =
|
||||
useState<HTMLDivElement | null>(null);
|
||||
const constraintsRef = useMemo(
|
||||
() => ({ current: constraintsElement }),
|
||||
[constraintsElement],
|
||||
);
|
||||
const dragControls = useDragControls();
|
||||
|
||||
if (!isOpen) return null;
|
||||
function startDrag(event: PointerEvent<HTMLDivElement>) {
|
||||
event.preventDefault();
|
||||
dragControls.start(event);
|
||||
}
|
||||
|
||||
const ramTotal = systemInfo.memory?.total_gb ?? 0;
|
||||
const ramAvailable = systemInfo.memory?.available_gb ?? 0;
|
||||
|
|
@ -64,99 +81,109 @@ export function FloatingMonitor() {
|
|||
const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0;
|
||||
|
||||
return (
|
||||
<div
|
||||
ref={constraintsRef}
|
||||
className="fixed inset-0 z-50 pointer-events-none"
|
||||
>
|
||||
<motion.div
|
||||
layout={true}
|
||||
drag={true}
|
||||
dragConstraints={constraintsRef}
|
||||
dragElastic={0.1}
|
||||
dragMomentum={false}
|
||||
initial={{ opacity: 0, scale: 0.9 }}
|
||||
animate={{ opacity: 1, scale: 1 }}
|
||||
exit={{ opacity: 0, scale: 0.9 }}
|
||||
className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-2rem)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none"
|
||||
>
|
||||
<div className="mb-2 flex items-center justify-between gap-2 border-b border-border/60 pb-2">
|
||||
<div className="flex min-w-0 flex-1 items-center gap-1.5 truncate text-xs font-semibold text-foreground">
|
||||
<CpuIcon className="size-3.5 shrink-0 text-primary" />
|
||||
<span className="truncate">
|
||||
{t("settings.resources.liveMonitor.title")}
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex items-center gap-1 shrink-0">
|
||||
<div className="cursor-grab rounded-md px-1 text-muted-foreground/60 transition-colors hover:bg-muted/60 hover:text-muted-foreground active:cursor-grabbing">
|
||||
<GripVerticalIcon className="size-3.5" />
|
||||
</div>
|
||||
|
||||
<Button
|
||||
size="icon-xs"
|
||||
variant="ghost"
|
||||
className="text-muted-foreground hover:text-foreground"
|
||||
onClick={() => setIsOpen(false)}
|
||||
title={t("common.close")}
|
||||
aria-label={t("common.close")}
|
||||
>
|
||||
<XIcon className="size-3" />
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<motion.div
|
||||
initial={{ opacity: 0, height: 0 }}
|
||||
animate={{ opacity: 1, height: "auto" }}
|
||||
exit={{ opacity: 0, height: 0 }}
|
||||
className="space-y-3 overflow-hidden"
|
||||
<AnimatePresence>
|
||||
{isOpen && (
|
||||
<div
|
||||
ref={setConstraintsElement}
|
||||
className="fixed inset-0 z-50 pointer-events-none"
|
||||
>
|
||||
<div className="space-y-1">
|
||||
<div className="flex justify-between text-[11px] font-medium font-mono">
|
||||
<span>{t("settings.resources.liveMonitor.ram")}</span>
|
||||
<span className={cn("tabular-nums", usageTextClass(ramPercent))}>
|
||||
{Math.round(ramPercent)}%
|
||||
</span>
|
||||
</div>
|
||||
<div className="text-xs text-muted-foreground font-mono tabular-nums">
|
||||
{formatGiB(ramUsed)} / {formatGiB(ramTotal)}
|
||||
</div>
|
||||
<Progress
|
||||
value={ramPercent}
|
||||
className="mt-1 h-1.5 rounded-full bg-muted"
|
||||
indicatorClassName={usageIndicatorClass(ramPercent)}
|
||||
/>
|
||||
</div>
|
||||
|
||||
{hasGpu && (
|
||||
<div className="space-y-1">
|
||||
<div className="flex justify-between text-[11px] font-medium font-mono">
|
||||
<span className="truncate flex-1 pr-2">
|
||||
{t("settings.resources.liveMonitor.vram")}{" "}
|
||||
{devices.length > 1
|
||||
? `(${devices.length} GPUs)`
|
||||
: `(${devices[0].name ?? "GPU"})`}
|
||||
<motion.div
|
||||
drag={true}
|
||||
dragControls={dragControls}
|
||||
dragListener={false}
|
||||
dragConstraints={constraintsRef}
|
||||
dragElastic={0}
|
||||
dragMomentum={false}
|
||||
initial={{ opacity: 0, scale: 0.9 }}
|
||||
animate={{ opacity: 1, scale: 1 }}
|
||||
exit={{ opacity: 0, scale: 0.9 }}
|
||||
className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-2rem)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none"
|
||||
>
|
||||
<div className="mb-2 flex items-center justify-between gap-2 border-b border-border/60 pb-2">
|
||||
<div className="flex min-w-0 flex-1 items-center gap-1.5 truncate text-xs font-semibold text-foreground">
|
||||
<CpuIcon className="size-3.5 shrink-0 text-primary" />
|
||||
<span className="truncate">
|
||||
{t("settings.resources.liveMonitor.title")}
|
||||
</span>
|
||||
<span
|
||||
className={cn(
|
||||
"shrink-0 tabular-nums",
|
||||
usageTextClass(vramPercent),
|
||||
)}
|
||||
</div>
|
||||
<div className="flex items-center gap-1 shrink-0">
|
||||
<div
|
||||
onPointerDown={startDrag}
|
||||
className="touch-none cursor-grab rounded-md px-1 text-muted-foreground/60 transition-colors hover:bg-muted/60 hover:text-muted-foreground active:cursor-grabbing"
|
||||
>
|
||||
{Math.round(vramPercent)}%
|
||||
</span>
|
||||
<GripVerticalIcon className="size-3.5" />
|
||||
</div>
|
||||
|
||||
<Button
|
||||
size="icon-xs"
|
||||
variant="ghost"
|
||||
className="text-muted-foreground hover:text-foreground"
|
||||
onClick={() => setIsOpen(false)}
|
||||
title={t("common.close")}
|
||||
aria-label={t("common.close")}
|
||||
>
|
||||
<XIcon className="size-3" />
|
||||
</Button>
|
||||
</div>
|
||||
<div className="text-xs text-muted-foreground font-mono tabular-nums">
|
||||
{formatGiB(vramUsed)} / {formatGiB(vramTotal)}
|
||||
</div>
|
||||
<Progress
|
||||
value={vramPercent}
|
||||
className="mt-1 h-1.5 rounded-full bg-muted"
|
||||
indicatorClassName={usageIndicatorClass(vramPercent)}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</motion.div>
|
||||
</motion.div>
|
||||
</div>
|
||||
|
||||
<motion.div
|
||||
initial={{ opacity: 0, height: 0 }}
|
||||
animate={{ opacity: 1, height: "auto" }}
|
||||
exit={{ opacity: 0, height: 0 }}
|
||||
className="space-y-3 overflow-hidden"
|
||||
>
|
||||
<div className="space-y-1">
|
||||
<div className="flex justify-between text-[11px] font-medium font-mono">
|
||||
<span>{t("settings.resources.liveMonitor.ram")}</span>
|
||||
<span
|
||||
className={cn("tabular-nums", usageTextClass(ramPercent))}
|
||||
>
|
||||
{Math.round(ramPercent)}%
|
||||
</span>
|
||||
</div>
|
||||
<div className="text-xs text-muted-foreground font-mono tabular-nums">
|
||||
{formatGiB(ramUsed)} / {formatGiB(ramTotal)}
|
||||
</div>
|
||||
<Progress
|
||||
value={ramPercent}
|
||||
className="mt-1 h-1.5 rounded-full bg-muted"
|
||||
indicatorClassName={usageIndicatorClass(ramPercent)}
|
||||
/>
|
||||
</div>
|
||||
|
||||
{hasGpu && (
|
||||
<div className="space-y-1">
|
||||
<div className="flex justify-between text-[11px] font-medium font-mono">
|
||||
<span className="truncate flex-1 pr-2">
|
||||
{t("settings.resources.liveMonitor.vram")}{" "}
|
||||
{devices.length > 1
|
||||
? `(${devices.length} GPUs)`
|
||||
: `(${devices[0].name ?? "GPU"})`}
|
||||
</span>
|
||||
<span
|
||||
className={cn(
|
||||
"shrink-0 tabular-nums",
|
||||
usageTextClass(vramPercent),
|
||||
)}
|
||||
>
|
||||
{Math.round(vramPercent)}%
|
||||
</span>
|
||||
</div>
|
||||
<div className="text-xs text-muted-foreground font-mono tabular-nums">
|
||||
{formatGiB(vramUsed)} / {formatGiB(vramTotal)}
|
||||
</div>
|
||||
<Progress
|
||||
value={vramPercent}
|
||||
className="mt-1 h-1.5 rounded-full bg-muted"
|
||||
indicatorClassName={usageIndicatorClass(vramPercent)}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</motion.div>
|
||||
</motion.div>
|
||||
</div>
|
||||
)}
|
||||
</AnimatePresence>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { getAuthToken } from "@/features/auth";
|
||||
import { resolveInitialConfig } from "@/features/model-picker";
|
||||
import { projectHasSources } from "@/features/rag/api/rag-api";
|
||||
import { apiUrl } from "@/lib/api-base";
|
||||
import { parseParamCountB } from "@/lib/model-size";
|
||||
|
|
@ -63,11 +64,17 @@ import {
|
|||
listStoredChatThreads,
|
||||
updateStoredChatThread,
|
||||
} from "../utils/chat-history-storage";
|
||||
import {
|
||||
readLastLocalModelLoad,
|
||||
recordLastLocalModelLoad,
|
||||
type LastLocalModelKind,
|
||||
} from "../utils/last-local-model-load";
|
||||
import { getImageInputUnavailableReason } from "../utils/image-input-support";
|
||||
import {
|
||||
hasClosedThinkTag,
|
||||
parseAssistantContent,
|
||||
} from "../utils/parse-assistant-content";
|
||||
import { resolveLoadMaxSeqLength } from "../presets/preset-policy";
|
||||
import {
|
||||
generateAudio,
|
||||
listCachedGguf,
|
||||
|
|
@ -1309,6 +1316,30 @@ const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi;
|
|||
const GGUF_KNOWN_QUANT_RE =
|
||||
/(UD-)?(MXFP[0-9]+(?:_[A-Z0-9]+)*|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?|TQ[0-9]+_[0-9]+|Q[0-9]+_K_[A-Z]+|Q[0-9]+_[0-9]+|Q[0-9]+_K|BF16|F16|F32)/i;
|
||||
|
||||
type AutoLoadCandidate = {
|
||||
id: string;
|
||||
kind: LastLocalModelKind;
|
||||
ggufVariant: string | null;
|
||||
maxSeqLength: number;
|
||||
successLabel: string;
|
||||
};
|
||||
|
||||
function autoLoadCandidateKey(
|
||||
kind: LastLocalModelKind,
|
||||
id: string,
|
||||
ggufVariant?: string | null,
|
||||
): string {
|
||||
return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`;
|
||||
}
|
||||
|
||||
function findCachedRepo<T extends { repo_id: string }>(
|
||||
repos: T[],
|
||||
id: string,
|
||||
): T | undefined {
|
||||
const normalized = id.toLowerCase();
|
||||
return repos.find((repo) => repo.repo_id.toLowerCase() === normalized);
|
||||
}
|
||||
|
||||
function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean {
|
||||
const normalized = filename.replace(/\\/g, "/").toLowerCase();
|
||||
const separatorIndex = normalized.lastIndexOf("/");
|
||||
|
|
@ -1357,14 +1388,18 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
const hfToken = store.hfToken || null;
|
||||
const trustRemoteCode = store.params.trustRemoteCode ?? false;
|
||||
const specSettings = resolveSpeculativeSettingsForLoad();
|
||||
const lastLoaded = readLastLocalModelLoad();
|
||||
const toastId = toast("Loading a model…", {
|
||||
description: "Auto-selecting the smallest downloaded model.",
|
||||
description: lastLoaded
|
||||
? "Loading last used model."
|
||||
: "Auto-selecting the smallest downloaded model.",
|
||||
duration: 5000,
|
||||
closeButton: true,
|
||||
});
|
||||
let blockedByTrustRemoteCode = false;
|
||||
let hadNonTrustFailure = false;
|
||||
let loadAttempts = 0;
|
||||
const skippedAutoLoadCandidates = new Set<string>();
|
||||
|
||||
async function canAutoLoad(payload: {
|
||||
model_path: string;
|
||||
|
|
@ -1389,12 +1424,225 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
async function loadAutoLoadCandidate(
|
||||
candidate: AutoLoadCandidate,
|
||||
): Promise<boolean> {
|
||||
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) {
|
||||
return false;
|
||||
}
|
||||
const currentStore = useChatRuntimeStore.getState();
|
||||
const { config } = resolveInitialConfig(candidate.id, candidate.ggufVariant);
|
||||
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
|
||||
modelId: candidate.id,
|
||||
ggufVariant: candidate.ggufVariant,
|
||||
isGguf: candidate.kind === "gguf",
|
||||
customContextLength: config.customContextLength,
|
||||
ggufContextLength: null,
|
||||
currentCheckpoint: currentStore.params.checkpoint,
|
||||
activeGgufVariant: currentStore.activeGgufVariant,
|
||||
maxSeqLength: candidate.maxSeqLength,
|
||||
presetSource: currentStore.activePresetSource,
|
||||
});
|
||||
const effectiveSpeculativeType =
|
||||
config.speculativeType ?? specSettings.speculativeType;
|
||||
const effectiveSpecDraftNMax =
|
||||
config.specDraftNMax ?? specSettings.specDraftNMax;
|
||||
const effectiveChatTemplateOverride = config.chatTemplateOverride?.trim()
|
||||
? config.chatTemplateOverride
|
||||
: null;
|
||||
if (
|
||||
!(await canAutoLoad({
|
||||
model_path: candidate.id,
|
||||
max_seq_length: effectiveMaxSeqLength,
|
||||
is_lora: false,
|
||||
gguf_variant: candidate.ggufVariant,
|
||||
}))
|
||||
) {
|
||||
skippedAutoLoadCandidates.add(
|
||||
autoLoadCandidateKey(candidate.kind, candidate.id, candidate.ggufVariant),
|
||||
);
|
||||
return false;
|
||||
}
|
||||
loadAttempts += 1;
|
||||
const loadResp = await loadModel({
|
||||
model_path: candidate.id,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: effectiveMaxSeqLength,
|
||||
load_in_4bit: true,
|
||||
is_lora: false,
|
||||
gguf_variant: candidate.ggufVariant,
|
||||
trust_remote_code: trustRemoteCode,
|
||||
chat_template_override: effectiveChatTemplateOverride,
|
||||
cache_type_kv: config.kvCacheDtype,
|
||||
speculative_type: effectiveSpeculativeType,
|
||||
spec_draft_n_max: effectiveSpecDraftNMax,
|
||||
tensor_parallel: config.tensorParallel,
|
||||
});
|
||||
saveSpeculativeType(effectiveSpeculativeType);
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined);
|
||||
const store = useChatRuntimeStore.getState();
|
||||
store.setModelRequiresTrustRemoteCode(
|
||||
loadResp.requires_trust_remote_code ?? false,
|
||||
);
|
||||
store.setParams({
|
||||
...store.params,
|
||||
maxTokens:
|
||||
candidate.kind === "gguf"
|
||||
? loadResp.context_length ?? 131072
|
||||
: effectiveMaxSeqLength,
|
||||
});
|
||||
const autoModel: ChatModelSummary = {
|
||||
id: candidate.id,
|
||||
name: loadResp.display_name ?? candidate.id,
|
||||
isVision: loadResp.is_vision ?? false,
|
||||
isLora: loadResp.is_lora ?? false,
|
||||
isGguf: loadResp.is_gguf ?? candidate.kind === "gguf",
|
||||
isAudio: loadResp.is_audio ?? false,
|
||||
audioType: loadResp.audio_type ?? null,
|
||||
hasAudioInput: loadResp.has_audio_input ?? false,
|
||||
};
|
||||
if (!store.models.some((m) => m.id === candidate.id)) {
|
||||
store.setModels([...store.models, autoModel]);
|
||||
}
|
||||
if (candidate.kind === "gguf") {
|
||||
useChatRuntimeStore.setState({
|
||||
ggufContextLength: loadResp.context_length ?? 131072,
|
||||
ggufMaxContextLength:
|
||||
loadResp.max_context_length ?? loadResp.context_length ?? 131072,
|
||||
ggufNativeContextLength: loadResp.native_context_length ?? null,
|
||||
supportsReasoning: loadResp.supports_reasoning ?? false,
|
||||
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
|
||||
reasoningEnabled: loadResp.supports_reasoning ?? false,
|
||||
...reasoningCapsFromLoad(loadResp),
|
||||
supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
|
||||
supportsTools: loadResp.supports_tools ?? false,
|
||||
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
|
||||
kvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
tensorParallel: loadResp.tensor_parallel ?? false,
|
||||
loadedTensorParallel: loadResp.tensor_parallel ?? false,
|
||||
defaultChatTemplate: loadResp.chat_template ?? null,
|
||||
chatTemplateOverride: effectiveChatTemplateOverride,
|
||||
loadedChatTemplateOverride: effectiveChatTemplateOverride,
|
||||
customContextLength: null,
|
||||
loadedIsMultimodal: isMultimodalResponse(loadResp),
|
||||
loadedIsDiffusion: loadResp.is_diffusion ?? false,
|
||||
...resolveLoadedSpeculativeSettings(loadResp),
|
||||
});
|
||||
} else {
|
||||
useChatRuntimeStore.setState({
|
||||
supportsReasoning: loadResp.supports_reasoning ?? false,
|
||||
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
|
||||
reasoningEnabled: loadResp.supports_reasoning ?? false,
|
||||
...reasoningCapsFromLoad(loadResp),
|
||||
supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
|
||||
supportsTools: loadResp.supports_tools ?? false,
|
||||
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
|
||||
kvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
tensorParallel: loadResp.tensor_parallel ?? false,
|
||||
loadedTensorParallel: loadResp.tensor_parallel ?? false,
|
||||
defaultChatTemplate: loadResp.chat_template ?? null,
|
||||
chatTemplateOverride: effectiveChatTemplateOverride,
|
||||
loadedChatTemplateOverride: effectiveChatTemplateOverride,
|
||||
customContextLength: null,
|
||||
...resolveLoadedSpeculativeSettings(loadResp),
|
||||
loadedIsMultimodal: isMultimodalResponse(loadResp),
|
||||
loadedIsDiffusion: loadResp.is_diffusion ?? false,
|
||||
});
|
||||
}
|
||||
if (!(loadResp.is_lora ?? false)) {
|
||||
recordLastLocalModelLoad({
|
||||
id: candidate.id,
|
||||
kind: candidate.kind,
|
||||
ggufVariant: candidate.ggufVariant,
|
||||
});
|
||||
}
|
||||
toast.success(candidate.successLabel, { id: toastId });
|
||||
return true;
|
||||
}
|
||||
try {
|
||||
const [ggufRepos, modelRepos] = await Promise.all([
|
||||
listCachedGguf().catch(() => []),
|
||||
listCachedModels().catch(() => []),
|
||||
]);
|
||||
|
||||
if (lastLoaded) {
|
||||
if (lastLoaded.kind === "gguf") {
|
||||
const repo = findCachedRepo(ggufRepos, lastLoaded.id);
|
||||
if (repo && lastLoaded.ggufVariant) {
|
||||
try {
|
||||
const variants = await listGgufVariants(repo.repo_id);
|
||||
const variant = variants.variants.find(
|
||||
(entry) =>
|
||||
entry.downloaded &&
|
||||
entry.quant?.toLowerCase() ===
|
||||
lastLoaded.ggufVariant?.toLowerCase() &&
|
||||
isAutoLoadableGgufVariant(entry),
|
||||
);
|
||||
if (variant) {
|
||||
toast("Loading last used model…", {
|
||||
id: toastId,
|
||||
description: `${repo.repo_id} (${variant.quant})`,
|
||||
duration: 5000,
|
||||
});
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
kind: "gguf",
|
||||
ggufVariant: variant.quant,
|
||||
maxSeqLength: 0,
|
||||
successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
|
||||
})
|
||||
) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
skippedAutoLoadCandidates.add(
|
||||
autoLoadCandidateKey("gguf", repo.repo_id, lastLoaded.ggufVariant),
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
const repo = findCachedRepo(modelRepos, lastLoaded.id);
|
||||
if (repo) {
|
||||
try {
|
||||
toast("Loading last used model…", {
|
||||
id: toastId,
|
||||
description: repo.repo_id,
|
||||
duration: 5000,
|
||||
});
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
kind: "model",
|
||||
ggufVariant: null,
|
||||
maxSeqLength: store.params.maxSeqLength,
|
||||
successLabel: `Loaded ${repo.repo_id}`,
|
||||
})
|
||||
) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
skippedAutoLoadCandidates.add(
|
||||
autoLoadCandidateKey("model", repo.repo_id),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
toast("Loading a model…", {
|
||||
id: toastId,
|
||||
description: "Auto-selecting the smallest downloaded model.",
|
||||
duration: 5000,
|
||||
});
|
||||
}
|
||||
|
||||
// GGUF first: smallest-total-size repo, then its smallest variant.
|
||||
if (ggufRepos.length > 0) {
|
||||
const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes);
|
||||
|
|
@ -1408,82 +1656,23 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
if (downloaded.length > 0) {
|
||||
const variant = downloaded[0];
|
||||
if (
|
||||
!(await canAutoLoad({
|
||||
model_path: repo.repo_id,
|
||||
max_seq_length: 0,
|
||||
is_lora: false,
|
||||
gguf_variant: variant.quant,
|
||||
}))
|
||||
skippedAutoLoadCandidates.has(
|
||||
autoLoadCandidateKey("gguf", repo.repo_id, variant.quant),
|
||||
)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
loadAttempts += 1;
|
||||
const loadResp = await loadModel({
|
||||
model_path: repo.repo_id,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: 0,
|
||||
load_in_4bit: true,
|
||||
is_lora: false,
|
||||
gguf_variant: variant.quant,
|
||||
trust_remote_code: trustRemoteCode,
|
||||
speculative_type: specSettings.speculativeType,
|
||||
spec_draft_n_max: specSettings.specDraftNMax,
|
||||
});
|
||||
saveSpeculativeType(specSettings.speculativeType);
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.setCheckpoint(repo.repo_id, variant.quant);
|
||||
const store = useChatRuntimeStore.getState();
|
||||
store.setModelRequiresTrustRemoteCode(
|
||||
loadResp.requires_trust_remote_code ?? false,
|
||||
);
|
||||
store.setParams({
|
||||
...store.params,
|
||||
maxTokens: loadResp.context_length ?? 131072,
|
||||
});
|
||||
// Add to store so the selector shows the name.
|
||||
const autoModel: ChatModelSummary = {
|
||||
id: repo.repo_id,
|
||||
name: loadResp.display_name ?? repo.repo_id,
|
||||
isVision: loadResp.is_vision ?? false,
|
||||
isLora: loadResp.is_lora ?? false,
|
||||
isGguf: loadResp.is_gguf ?? false,
|
||||
isAudio: loadResp.is_audio ?? false,
|
||||
audioType: loadResp.audio_type ?? null,
|
||||
hasAudioInput: loadResp.has_audio_input ?? false,
|
||||
};
|
||||
const existingModels = store.models;
|
||||
if (!existingModels.some((m) => m.id === repo.repo_id)) {
|
||||
store.setModels([...existingModels, autoModel]);
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
kind: "gguf",
|
||||
ggufVariant: variant.quant,
|
||||
maxSeqLength: 0,
|
||||
successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
|
||||
})
|
||||
) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
useChatRuntimeStore.setState({
|
||||
ggufContextLength: loadResp.context_length ?? 131072,
|
||||
ggufMaxContextLength:
|
||||
loadResp.max_context_length ??
|
||||
loadResp.context_length ??
|
||||
131072,
|
||||
supportsReasoning: loadResp.supports_reasoning ?? false,
|
||||
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
|
||||
reasoningEnabled: loadResp.supports_reasoning ?? false,
|
||||
...reasoningCapsFromLoad(loadResp),
|
||||
supportsPreserveThinking:
|
||||
loadResp.supports_preserve_thinking ?? false,
|
||||
supportsTools: loadResp.supports_tools ?? false,
|
||||
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
|
||||
kvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
|
||||
tensorParallel: loadResp.tensor_parallel ?? false,
|
||||
loadedTensorParallel: loadResp.tensor_parallel ?? false,
|
||||
defaultChatTemplate: loadResp.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
loadedChatTemplateOverride: null,
|
||||
loadedIsMultimodal: isMultimodalResponse(loadResp),
|
||||
...resolveLoadedSpeculativeSettings(loadResp),
|
||||
});
|
||||
toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, {
|
||||
id: toastId,
|
||||
});
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
|
|
@ -1501,64 +1690,23 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
|
||||
try {
|
||||
if (
|
||||
!(await canAutoLoad({
|
||||
model_path: repo.repo_id,
|
||||
max_seq_length: 4096,
|
||||
is_lora: false,
|
||||
gguf_variant: null,
|
||||
}))
|
||||
skippedAutoLoadCandidates.has(
|
||||
autoLoadCandidateKey("model", repo.repo_id),
|
||||
)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
loadAttempts += 1;
|
||||
const sfLoadResp = await loadModel({
|
||||
model_path: repo.repo_id,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: 4096,
|
||||
load_in_4bit: true,
|
||||
is_lora: false,
|
||||
gguf_variant: null,
|
||||
trust_remote_code: trustRemoteCode,
|
||||
speculative_type: specSettings.speculativeType,
|
||||
spec_draft_n_max: specSettings.specDraftNMax,
|
||||
});
|
||||
saveSpeculativeType(specSettings.speculativeType);
|
||||
useChatRuntimeStore.getState().setCheckpoint(repo.repo_id);
|
||||
const store = useChatRuntimeStore.getState();
|
||||
store.setModelRequiresTrustRemoteCode(
|
||||
sfLoadResp.requires_trust_remote_code ?? false,
|
||||
);
|
||||
store.setParams({ ...store.params, maxTokens: 4096 });
|
||||
useChatRuntimeStore.setState({
|
||||
supportsReasoning: sfLoadResp.supports_reasoning ?? false,
|
||||
reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false,
|
||||
reasoningEnabled: sfLoadResp.supports_reasoning ?? false,
|
||||
...reasoningCapsFromLoad(sfLoadResp),
|
||||
supportsPreserveThinking:
|
||||
sfLoadResp.supports_preserve_thinking ?? false,
|
||||
supportsTools: sfLoadResp.supports_tools ?? false,
|
||||
// Parity with the GGUF branch above.
|
||||
...resolveToolsEnabledOnLoad(sfLoadResp.supports_tools ?? false),
|
||||
defaultChatTemplate: sfLoadResp.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
loadedChatTemplateOverride: null,
|
||||
...resolveLoadedSpeculativeSettings(sfLoadResp),
|
||||
});
|
||||
const sfModel: ChatModelSummary = {
|
||||
id: repo.repo_id,
|
||||
name: sfLoadResp.display_name ?? repo.repo_id,
|
||||
isVision: sfLoadResp.is_vision ?? false,
|
||||
isLora: sfLoadResp.is_lora ?? false,
|
||||
isGguf: sfLoadResp.is_gguf ?? false,
|
||||
};
|
||||
if (!store.models.some((m) => m.id === repo.repo_id)) {
|
||||
store.setModels([...store.models, sfModel]);
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
kind: "model",
|
||||
ggufVariant: null,
|
||||
maxSeqLength: 4096,
|
||||
successLabel: `Loaded ${repo.repo_id}`,
|
||||
})
|
||||
) {
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
}
|
||||
useChatRuntimeStore.setState({
|
||||
loadedIsMultimodal: isMultimodalResponse(sfLoadResp),
|
||||
});
|
||||
toast.success(`Loaded ${repo.repo_id}`, { id: toastId });
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
} catch {
|
||||
hadNonTrustFailure = true;
|
||||
continue;
|
||||
|
|
@ -1650,6 +1798,11 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
loadedIsMultimodal: isMultimodalResponse(loadResp),
|
||||
...resolveLoadedSpeculativeSettings(loadResp),
|
||||
});
|
||||
recordLastLocalModelLoad({
|
||||
id: "unsloth/Qwen3.5-4B-MTP-GGUF",
|
||||
kind: "gguf",
|
||||
ggufVariant: "UD-Q4_K_XL",
|
||||
});
|
||||
toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId });
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
} catch {
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ import {
|
|||
mergeBackendRecommendedInference,
|
||||
resolveLoadMaxSeqLength,
|
||||
} from "../presets/preset-policy";
|
||||
import { recordLastLocalModelLoad } from "../utils/last-local-model-load";
|
||||
import {
|
||||
isMultimodalResponse,
|
||||
} from "../types/api";
|
||||
|
|
@ -804,6 +805,23 @@ export function useChatModelRuntime() {
|
|||
}
|
||||
}
|
||||
await refresh({ signal: abortCtrl.signal });
|
||||
if (
|
||||
!isLora &&
|
||||
!(loadResponse.is_lora ?? false) &&
|
||||
!nativePathToken &&
|
||||
!isLocalModelPath(modelId) &&
|
||||
!isExternalModelId(modelId)
|
||||
) {
|
||||
if (loadResponse.is_gguf || isGguf || ggufVariant) {
|
||||
recordLastLocalModelLoad({
|
||||
id: modelId,
|
||||
kind: "gguf",
|
||||
ggufVariant: ggufVariant ?? null,
|
||||
});
|
||||
} else {
|
||||
recordLastLocalModelLoad({ id: modelId, kind: "model" });
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
// Skip rollback if user cancelled -- model is already being unloaded.
|
||||
if (abortCtrl.signal.aborted) throw error;
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ export { ChatSearchDialog } from "./components/chat-search-dialog";
|
|||
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
|
||||
export type { ProjectRecord } from "./types";
|
||||
export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
|
||||
export { listStoredChatThreads } from "./utils/chat-history-storage";
|
||||
export { ArtifactCard } from "./artifacts/artifact-card";
|
||||
export {
|
||||
useChatArtifactsStore,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,86 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
export type LastLocalModelKind = "gguf" | "model";
|
||||
|
||||
export type LastLocalModelLoad = {
|
||||
id: string;
|
||||
kind: LastLocalModelKind;
|
||||
ggufVariant: string | null;
|
||||
loadedAt: number;
|
||||
};
|
||||
|
||||
const STORAGE_KEY = "unsloth.last-local-model-load.v1";
|
||||
|
||||
function storage(): Storage | null {
|
||||
try {
|
||||
return typeof localStorage === "undefined" ? null : localStorage;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function isLastLocalModelKind(value: unknown): value is LastLocalModelKind {
|
||||
return value === "gguf" || value === "model";
|
||||
}
|
||||
|
||||
export function readLastLocalModelLoad(): LastLocalModelLoad | null {
|
||||
try {
|
||||
const raw = storage()?.getItem(STORAGE_KEY);
|
||||
if (!raw) {
|
||||
return null;
|
||||
}
|
||||
const parsed = JSON.parse(raw) as Partial<LastLocalModelLoad>;
|
||||
if (
|
||||
typeof parsed.id !== "string" ||
|
||||
!parsed.id.trim() ||
|
||||
!isLastLocalModelKind(parsed.kind) ||
|
||||
typeof parsed.loadedAt !== "number"
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
if (
|
||||
parsed.kind === "gguf" &&
|
||||
(typeof parsed.ggufVariant !== "string" || !parsed.ggufVariant.trim())
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
return {
|
||||
id: parsed.id,
|
||||
kind: parsed.kind,
|
||||
ggufVariant:
|
||||
typeof parsed.ggufVariant === "string" ? parsed.ggufVariant : null,
|
||||
loadedAt: parsed.loadedAt,
|
||||
};
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
export function recordLastLocalModelLoad(input: {
|
||||
id: string;
|
||||
kind: LastLocalModelKind;
|
||||
ggufVariant?: string | null;
|
||||
}): void {
|
||||
const id = input.id.trim();
|
||||
if (!id) {
|
||||
return;
|
||||
}
|
||||
const ggufVariant = input.ggufVariant?.trim() || null;
|
||||
if (input.kind === "gguf" && !ggufVariant) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
storage()?.setItem(
|
||||
STORAGE_KEY,
|
||||
JSON.stringify({
|
||||
id,
|
||||
kind: input.kind,
|
||||
ggufVariant: input.kind === "gguf" ? ggufVariant : null,
|
||||
loadedAt: Date.now(),
|
||||
} satisfies LastLocalModelLoad),
|
||||
);
|
||||
} catch {
|
||||
// Ignore disabled storage / quota errors; auto-load falls back to size order.
|
||||
}
|
||||
}
|
||||
|
|
@ -494,3 +494,13 @@ export async function removeUnstructuredFile(
|
|||
throw new Error("Failed to remove file");
|
||||
}
|
||||
}
|
||||
|
||||
export async function removeUnstructuredBlock(blockId: string): Promise<void> {
|
||||
const res = await authFetch(
|
||||
`${DATA_DESIGNER_API_BASE}/seed/unstructured-block/${encodeURIComponent(blockId)}`,
|
||||
{ method: "DELETE" },
|
||||
);
|
||||
if (!res.ok && res.status !== 404) {
|
||||
throw new Error("Failed to remove uploaded files");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -53,6 +53,11 @@ import {
|
|||
inspectSeedDataset,
|
||||
inspectSeedUpload,
|
||||
} from "../../api";
|
||||
import { useRecipeStudioStore } from "../../stores/recipe-studio";
|
||||
import {
|
||||
makeUnstructuredUploadUid,
|
||||
resolveUnstructuredUploadBlockId,
|
||||
} from "../../utils/config-factories";
|
||||
import { resolveImagePreview } from "../../utils/image-preview";
|
||||
import type {
|
||||
GithubItemType,
|
||||
|
|
@ -597,6 +602,41 @@ export function SeedDialog({
|
|||
const mode = config.seed_source_type ?? "hf";
|
||||
const previewEmpty = getPreviewEmptyStateCopy(mode);
|
||||
|
||||
const queueUploadCleanup = useRecipeStudioStore(
|
||||
(state) => state.queueUploadCleanup,
|
||||
);
|
||||
|
||||
// config.id collides across recipes (ids reset to n1 on import); use a
|
||||
// stable per-block uid instead. Generate one synchronously so the first
|
||||
// rendered drop zone cannot upload under a legacy node id.
|
||||
const uploadUid = config.unstructured_upload_uid?.trim() ?? "";
|
||||
const unstructuredFileCount = config.unstructured_file_ids?.length ?? 0;
|
||||
const generatedUploadUidRef = useRef<string | null>(null);
|
||||
if (
|
||||
mode === "unstructured" &&
|
||||
!uploadUid &&
|
||||
unstructuredFileCount === 0 &&
|
||||
generatedUploadUidRef.current === null
|
||||
) {
|
||||
generatedUploadUidRef.current = makeUnstructuredUploadUid();
|
||||
}
|
||||
const uploadBlockId = resolveUnstructuredUploadBlockId({
|
||||
configId: config.id,
|
||||
uploadUid,
|
||||
generatedUploadUid: generatedUploadUidRef.current,
|
||||
unstructuredFileCount,
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
if (mode !== "unstructured") return;
|
||||
if (uploadUid) return;
|
||||
if (unstructuredFileCount > 0) return;
|
||||
const nextUid =
|
||||
generatedUploadUidRef.current ?? makeUnstructuredUploadUid();
|
||||
generatedUploadUidRef.current = nextUid;
|
||||
onUpdate({ unstructured_upload_uid: nextUid });
|
||||
}, [mode, uploadUid, unstructuredFileCount, onUpdate]);
|
||||
|
||||
const prevModeRef = useRef(mode);
|
||||
useEffect(() => {
|
||||
const prevMode = prevModeRef.current;
|
||||
|
|
@ -720,6 +760,11 @@ export function SeedDialog({
|
|||
subset: config.hf_subset?.trim() || undefined,
|
||||
preview_size: 10,
|
||||
});
|
||||
// Queue the block's upload directory for deletion after the next
|
||||
// save; only uid-namespaced directories qualify (single owner).
|
||||
if (uploadUid && unstructuredFileCount > 0) {
|
||||
queueUploadCleanup(uploadUid);
|
||||
}
|
||||
onUpdate({
|
||||
hf_path: response.resolved_path,
|
||||
seed_columns: response.columns,
|
||||
|
|
@ -730,6 +775,7 @@ export function SeedDialog({
|
|||
hf_split: response.split ?? "",
|
||||
hf_subset: response.subset ?? "",
|
||||
local_file_name: "",
|
||||
unstructured_upload_uid: "",
|
||||
unstructured_file_ids: [],
|
||||
unstructured_file_names: [],
|
||||
unstructured_file_sizes: [],
|
||||
|
|
@ -754,6 +800,11 @@ export function SeedDialog({
|
|||
content_base64: payload,
|
||||
preview_size: 10,
|
||||
});
|
||||
// Queue the block's upload directory for deletion after the next
|
||||
// save; only uid-namespaced directories qualify (single owner).
|
||||
if (uploadUid && unstructuredFileCount > 0) {
|
||||
queueUploadCleanup(uploadUid);
|
||||
}
|
||||
onUpdate({
|
||||
hf_path: response.resolved_path,
|
||||
seed_columns: response.columns,
|
||||
|
|
@ -765,6 +816,7 @@ export function SeedDialog({
|
|||
hf_subset: "",
|
||||
hf_split: "",
|
||||
local_file_name: localFile.name,
|
||||
unstructured_upload_uid: "",
|
||||
unstructured_file_ids: [],
|
||||
unstructured_file_names: [],
|
||||
unstructured_file_sizes: [],
|
||||
|
|
@ -789,7 +841,7 @@ export function SeedDialog({
|
|||
|
||||
const { chunkSize, chunkOverlap } = resolveChunking(config);
|
||||
const response = await inspectSeedUpload({
|
||||
block_id: config.id,
|
||||
block_id: uploadBlockId,
|
||||
file_ids: fileIds,
|
||||
file_names: fileNames,
|
||||
preview_size: 10,
|
||||
|
|
@ -827,7 +879,18 @@ export function SeedDialog({
|
|||
setIsInspecting(false);
|
||||
}
|
||||
},
|
||||
[config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles],
|
||||
[
|
||||
config,
|
||||
getCurrentLoadKey,
|
||||
localFile,
|
||||
mode,
|
||||
onUpdate,
|
||||
queueUploadCleanup,
|
||||
unstructuredFiles,
|
||||
unstructuredFileCount,
|
||||
uploadBlockId,
|
||||
uploadUid,
|
||||
],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
|
|
@ -997,7 +1060,7 @@ export function SeedDialog({
|
|||
|
||||
{mode === "unstructured" && (
|
||||
<UnstructuredDropZone
|
||||
blockId={config.id}
|
||||
blockId={uploadBlockId}
|
||||
files={unstructuredFiles}
|
||||
onFilesChange={handleUnstructuredFilesChange}
|
||||
disabled={isInspecting}
|
||||
|
|
|
|||
|
|
@ -54,11 +54,17 @@ export function UnstructuredDropZone({
|
|||
}: UnstructuredDropZoneProps) {
|
||||
const inputRef = useRef<HTMLInputElement>(null);
|
||||
const filesRef = useRef(files);
|
||||
const blockIdRef = useRef(blockId);
|
||||
const mountedRef = useRef(true);
|
||||
const [isDragOver, setIsDragOver] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
filesRef.current = files;
|
||||
}, [files]);
|
||||
blockIdRef.current = blockId;
|
||||
}, [files, blockId]);
|
||||
useEffect(() => () => {
|
||||
mountedRef.current = false;
|
||||
}, []);
|
||||
|
||||
const totalSize = files.reduce((sum, f) => sum + f.size, 0);
|
||||
|
||||
|
|
@ -134,15 +140,32 @@ export function UnstructuredDropZone({
|
|||
if (entry.status === "uploading" && entry.abortController) {
|
||||
entry.abortController.abort();
|
||||
}
|
||||
if (
|
||||
const needsServerRemove =
|
||||
entry.id &&
|
||||
entry.status === "ok" &&
|
||||
!deletedIdsRef.current.has(entry.id)
|
||||
) {
|
||||
deletedIdsRef.current.add(entry.id);
|
||||
void removeUnstructuredFile(blockId, entry.id).catch(() => {});
|
||||
}
|
||||
!deletedIdsRef.current.has(entry.id);
|
||||
onFilesChange((prev) => prev.filter((_, i) => i !== index));
|
||||
if (!needsServerRemove) return;
|
||||
deletedIdsRef.current.add(entry.id);
|
||||
removeUnstructuredFile(blockId, entry.id).catch(() => {
|
||||
// Skip if the drop zone unmounted or its block changed: the id no
|
||||
// longer belongs here and restoring would leak it into another block.
|
||||
if (!mountedRef.current || blockIdRef.current !== blockId) return;
|
||||
// Still exists server-side (counts toward quota); restore it at its
|
||||
// original position.
|
||||
deletedIdsRef.current.delete(entry.id);
|
||||
onFilesChange((prev) => {
|
||||
const next = [...prev];
|
||||
next.splice(Math.min(index, next.length), 0, {
|
||||
id: entry.id,
|
||||
name: entry.name,
|
||||
size: entry.size,
|
||||
status: "ok",
|
||||
error: "Remove failed — try again",
|
||||
});
|
||||
return next;
|
||||
});
|
||||
});
|
||||
},
|
||||
[blockId, onFilesChange],
|
||||
);
|
||||
|
|
|
|||
|
|
@ -4,11 +4,13 @@
|
|||
import { useCallback, useEffect, useMemo, useState } from "react";
|
||||
import { toastError, toastSuccess } from "@/shared/toast";
|
||||
import { normalizeNonEmptyName } from "@/utils";
|
||||
import { removeUnstructuredBlock } from "../api";
|
||||
import {
|
||||
buildSignature,
|
||||
copyTextToClipboard,
|
||||
formatSavedLabel,
|
||||
} from "../executions/execution-helpers";
|
||||
import { useRecipeStudioStore } from "../stores/recipe-studio";
|
||||
import { importRecipePayload, type RecipeSnapshot } from "../utils/import";
|
||||
import type { RecipePayloadResult } from "../utils/payload/types";
|
||||
|
||||
|
|
@ -72,7 +74,10 @@ function stripApiKeys(value: unknown): unknown {
|
|||
!Array.isArray(output.env)
|
||||
) {
|
||||
output.env = Object.fromEntries(
|
||||
Object.keys(output.env as Record<string, unknown>).map((envKey) => [envKey, ""]),
|
||||
Object.keys(output.env as Record<string, unknown>).map((envKey) => [
|
||||
envKey,
|
||||
"",
|
||||
]),
|
||||
);
|
||||
}
|
||||
return output;
|
||||
|
|
@ -82,10 +87,7 @@ function inferHfRepoIdFromPath(pathValue: unknown): string {
|
|||
if (typeof pathValue !== "string") {
|
||||
return "";
|
||||
}
|
||||
const parts = pathValue
|
||||
.trim()
|
||||
.split("/")
|
||||
.filter(Boolean);
|
||||
const parts = pathValue.trim().split("/").filter(Boolean);
|
||||
if (parts.length >= 3 && parts[0] === "datasets") {
|
||||
return `${parts[1]}/${parts[2]}`;
|
||||
}
|
||||
|
|
@ -126,8 +128,7 @@ function sanitizeSeedForShare(payload: unknown): unknown {
|
|||
typeof ui?.seed_source_type === "string" ? ui.seed_source_type : null;
|
||||
const sourceType =
|
||||
typeof source?.seed_type === "string" ? source.seed_type : null;
|
||||
const shouldResetHfState =
|
||||
sourceType === "hf" || uiSourceType === "hf";
|
||||
const shouldResetHfState = sourceType === "hf" || uiSourceType === "hf";
|
||||
const shouldResetLocalState =
|
||||
sourceType === "local" ||
|
||||
sourceType === "unstructured" ||
|
||||
|
|
@ -144,6 +145,7 @@ function sanitizeSeedForShare(payload: unknown): unknown {
|
|||
ui.seed_drop_columns = [];
|
||||
ui.seed_preview_rows = [];
|
||||
ui.local_file_name = "";
|
||||
ui.unstructured_upload_uid = "";
|
||||
ui.unstructured_file_ids = [];
|
||||
ui.unstructured_file_names = [];
|
||||
ui.unstructured_file_sizes = [];
|
||||
|
|
@ -165,6 +167,7 @@ function sanitizeSeedForShare(payload: unknown): unknown {
|
|||
ui.seed_drop_columns = [];
|
||||
ui.seed_preview_rows = [];
|
||||
ui.local_file_name = "";
|
||||
ui.unstructured_upload_uid = "";
|
||||
ui.unstructured_file_ids = [];
|
||||
ui.unstructured_file_names = [];
|
||||
ui.unstructured_file_sizes = [];
|
||||
|
|
@ -174,6 +177,43 @@ function sanitizeSeedForShare(payload: unknown): unknown {
|
|||
return root;
|
||||
}
|
||||
|
||||
// Delete queued upload directories once a save stops referencing them, so a
|
||||
// reload before autosave can never leave the saved recipe pointing at
|
||||
// already-deleted files. Skips any uid the just-saved payload still uses.
|
||||
function drainQueuedUploadCleanups(
|
||||
savedPayload: RecipePayloadResult["payload"],
|
||||
): void {
|
||||
const pending = useRecipeStudioStore.getState().pendingUploadCleanups;
|
||||
if (pending.length === 0) {
|
||||
return;
|
||||
}
|
||||
const ui =
|
||||
savedPayload && typeof savedPayload === "object"
|
||||
? (savedPayload as { ui?: Record<string, unknown> }).ui
|
||||
: undefined;
|
||||
const savedUid =
|
||||
ui && typeof ui.unstructured_upload_uid === "string"
|
||||
? ui.unstructured_upload_uid
|
||||
: "";
|
||||
const ready = pending.filter((uid) => uid !== savedUid);
|
||||
if (ready.length === 0) {
|
||||
return;
|
||||
}
|
||||
for (const uid of ready) {
|
||||
void removeUnstructuredBlock(uid)
|
||||
.then(() => {
|
||||
useRecipeStudioStore.setState((state) => ({
|
||||
pendingUploadCleanups: state.pendingUploadCleanups.filter(
|
||||
(pendingUid) => pendingUid !== uid,
|
||||
),
|
||||
}));
|
||||
})
|
||||
.catch((error) => {
|
||||
console.warn("Failed to clean up uploaded documents:", error);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
export function useRecipePersistence({
|
||||
recipeId,
|
||||
initialRecipeName,
|
||||
|
|
@ -202,8 +242,10 @@ export function useRecipePersistence({
|
|||
() => buildSignature(normalizedWorkflowName, currentPayload),
|
||||
[currentPayload, normalizedWorkflowName],
|
||||
);
|
||||
const isDirty = savedSignature.length > 0 && currentSignature !== savedSignature;
|
||||
const saveTone: SaveTone = !isDirty && Boolean(lastSavedAt) ? "success" : "error";
|
||||
const isDirty =
|
||||
savedSignature.length > 0 && currentSignature !== savedSignature;
|
||||
const saveTone: SaveTone =
|
||||
!isDirty && Boolean(lastSavedAt) ? "success" : "error";
|
||||
const savedAtLabel = formatSavedLabel(lastSavedAt);
|
||||
|
||||
useEffect(() => {
|
||||
|
|
@ -214,7 +256,9 @@ export function useRecipePersistence({
|
|||
setLastSavedAt(initialSavedAt);
|
||||
setCopied(false);
|
||||
|
||||
const parsed = importRecipePayload(JSON.stringify(initialPayload));
|
||||
const parsed = importRecipePayload(JSON.stringify(initialPayload), {
|
||||
preserveUnstructuredUploads: true,
|
||||
});
|
||||
if (parsed.snapshot) {
|
||||
loadRecipe(parsed.snapshot);
|
||||
} else {
|
||||
|
|
@ -252,6 +296,7 @@ export function useRecipePersistence({
|
|||
});
|
||||
setLastSavedAt(result.updatedAt);
|
||||
setSavedSignature(buildSignature(nextName, currentPayload));
|
||||
drainQueuedUploadCleanups(currentPayload);
|
||||
} catch (error) {
|
||||
console.error("Save recipe failed:", error);
|
||||
toastError("Save failed", "Could not save recipe.");
|
||||
|
|
@ -270,11 +315,28 @@ export function useRecipePersistence({
|
|||
return () => window.clearTimeout(timeoutId);
|
||||
}, [isDirty, persistRecipe, saveLoading]);
|
||||
|
||||
// Drain queued cleanups even when autosave is skipped: a net-zero edit (add
|
||||
// then remove an unstructured seed before the 800ms debounce) keeps isDirty
|
||||
// false, so the autosave effect never drains and the queued uid leaks its
|
||||
// upload dir. Not-dirty means currentPayload equals the saved recipe, and
|
||||
// drain skips the uid it still references, so only dirs no saved recipe
|
||||
// points at are deleted (keeps the save-first invariant).
|
||||
useEffect(() => {
|
||||
if (!initialRecipeReady || isDirty || saveLoading) {
|
||||
return;
|
||||
}
|
||||
drainQueuedUploadCleanups(currentPayload);
|
||||
}, [currentPayload, initialRecipeReady, isDirty, saveLoading]);
|
||||
|
||||
const copyRecipe = useCallback(async (): Promise<void> => {
|
||||
setCopied(false);
|
||||
try {
|
||||
const safePayload = sanitizeSeedForShare(stripApiKeys(payloadResult.payload));
|
||||
const ok = await copyTextToClipboard(JSON.stringify(safePayload, null, 2));
|
||||
const safePayload = sanitizeSeedForShare(
|
||||
stripApiKeys(payloadResult.payload),
|
||||
);
|
||||
const ok = await copyTextToClipboard(
|
||||
JSON.stringify(safePayload, null, 2),
|
||||
);
|
||||
if (!ok) {
|
||||
throw new Error("Clipboard not available.");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ import {
|
|||
} from "../utils/handles";
|
||||
import type { RecipeSnapshot } from "../utils/import";
|
||||
import { getLayoutedElements } from "../utils/layout";
|
||||
import { makeUnstructuredUploadUid } from "../utils/config-factories";
|
||||
import {
|
||||
centerModelInfraNodes,
|
||||
optimizeModelInfraEdgeHandles,
|
||||
|
|
@ -76,6 +77,12 @@ type RecipeStudioState = {
|
|||
nextId: number;
|
||||
nextY: number;
|
||||
fitViewTick: number;
|
||||
// Upload-uid directories whose owning block dropped them; server-side
|
||||
// deletion is deferred until a save no longer references them, so a
|
||||
// reload before autosave cannot leave a saved recipe pointing at
|
||||
// deleted files.
|
||||
pendingUploadCleanups: string[];
|
||||
queueUploadCleanup: (uid: string) => void;
|
||||
setSheetOpen: (open: boolean) => void;
|
||||
setSheetView: (view: SheetView) => void;
|
||||
setProcessors: (processors: RecipeProcessorConfig[]) => void;
|
||||
|
|
@ -137,6 +144,7 @@ const INITIAL_STATE = {
|
|||
nextId: 3,
|
||||
nextY: 280,
|
||||
fitViewTick: 0,
|
||||
pendingUploadCleanups: [],
|
||||
} satisfies Pick<
|
||||
RecipeStudioState,
|
||||
| "nodes"
|
||||
|
|
@ -154,6 +162,7 @@ const INITIAL_STATE = {
|
|||
| "nextId"
|
||||
| "nextY"
|
||||
| "fitViewTick"
|
||||
| "pendingUploadCleanups"
|
||||
>;
|
||||
|
||||
function buildAddedNodeState(
|
||||
|
|
@ -269,6 +278,20 @@ function isModelSemanticEdge(
|
|||
);
|
||||
}
|
||||
|
||||
// Upload uid of a seed block whose server-side directory becomes orphaned
|
||||
// when the block drops it. Only uid directories qualify (single owner);
|
||||
// legacy node-id directories can be shared by other recipes.
|
||||
function seedUploadCleanupUid(config: NodeConfig | undefined): string | null {
|
||||
if (!config || config.kind !== "seed") {
|
||||
return null;
|
||||
}
|
||||
const uid = config.unstructured_upload_uid?.trim();
|
||||
if (!uid || !config.unstructured_file_ids?.length) {
|
||||
return null;
|
||||
}
|
||||
return uid;
|
||||
}
|
||||
|
||||
export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
|
||||
...INITIAL_STATE,
|
||||
setSheetOpen: (open) => set({ sheetOpen: open }),
|
||||
|
|
@ -278,6 +301,12 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
|
|||
setDialogOpen: (open) => set({ dialogOpen: open }),
|
||||
setExecutionLocked: (locked) => set({ executionLocked: locked }),
|
||||
resetRecipe: () => set(INITIAL_STATE),
|
||||
queueUploadCleanup: (uid) =>
|
||||
set((state) =>
|
||||
state.pendingUploadCleanups.includes(uid)
|
||||
? state
|
||||
: { pendingUploadCleanups: [...state.pendingUploadCleanups, uid] },
|
||||
),
|
||||
selectConfig: (id) => set({ activeConfigId: id, dialogOpen: false }),
|
||||
openConfig: (id) => set({ activeConfigId: id, dialogOpen: true }),
|
||||
setLayoutDirection: (direction) =>
|
||||
|
|
@ -383,7 +412,18 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
|
|||
}
|
||||
return buildAddedNodeState(state, "sampler", type, position, openDialog);
|
||||
}),
|
||||
addSeedNode: (type, position, openDialog = true) =>
|
||||
addSeedNode: (type, position, openDialog = true) => {
|
||||
const current = get();
|
||||
if (!current.executionLocked) {
|
||||
// The reset below clears the block's upload uid and file list; queue
|
||||
// its server-side directory for deletion after the next save.
|
||||
const uid = seedUploadCleanupUid(
|
||||
Object.values(current.configs).find((config) => config.kind === "seed"),
|
||||
);
|
||||
if (uid) {
|
||||
current.queueUploadCleanup(uid);
|
||||
}
|
||||
}
|
||||
set((state) => {
|
||||
if (state.executionLocked) {
|
||||
return state;
|
||||
|
|
@ -413,6 +453,8 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
|
|||
hf_token: "",
|
||||
hf_endpoint: "https://huggingface.co",
|
||||
local_file_name: "",
|
||||
unstructured_upload_uid:
|
||||
nextSourceType === "unstructured" ? makeUnstructuredUploadUid() : "",
|
||||
unstructured_file_ids: [],
|
||||
unstructured_file_names: [],
|
||||
unstructured_file_sizes: [],
|
||||
|
|
@ -446,7 +488,8 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
|
|||
activeConfigId: existing.id,
|
||||
dialogOpen: openDialog,
|
||||
};
|
||||
}),
|
||||
});
|
||||
},
|
||||
addLlmNode: (type, position, openDialog = true) =>
|
||||
set((state) => {
|
||||
if (state.executionLocked) {
|
||||
|
|
@ -699,6 +742,9 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
|
|||
dialogOpen: false,
|
||||
sheetView: "root",
|
||||
fitViewTick: state.fitViewTick + 1,
|
||||
// Queued cleanups belong to the previous recipe; draining them after
|
||||
// a save of this one could delete files its saved payload still uses.
|
||||
pendingUploadCleanups: [],
|
||||
})),
|
||||
setAuxNodePosition: (id, position) =>
|
||||
set((state) => {
|
||||
|
|
@ -786,6 +832,17 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
|
|||
set(applyUpdate);
|
||||
},
|
||||
onNodesChange: (changes) => {
|
||||
const current = get();
|
||||
if (!current.executionLocked) {
|
||||
for (const change of changes) {
|
||||
if (change.type === "remove") {
|
||||
const uid = seedUploadCleanupUid(current.configs[change.id]);
|
||||
if (uid) {
|
||||
current.queueUploadCleanup(uid);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
const applyNodesChange = (state: RecipeStudioState) => {
|
||||
if (state.executionLocked) {
|
||||
return state;
|
||||
|
|
|
|||
|
|
@ -340,6 +340,8 @@ export type SeedConfig = {
|
|||
hf_token?: string;
|
||||
hf_endpoint?: string;
|
||||
local_file_name?: string;
|
||||
// ui-only: stable per-block id for uploads, since node ids collide across imports
|
||||
unstructured_upload_uid?: string;
|
||||
unstructured_file_ids?: string[];
|
||||
unstructured_file_names?: string[];
|
||||
unstructured_file_sizes?: number[];
|
||||
|
|
|
|||
|
|
@ -20,6 +20,46 @@ import type {
|
|||
} from "../types";
|
||||
import { nextName } from "./naming";
|
||||
|
||||
export function makeUnstructuredUploadUid(): string {
|
||||
if (typeof globalThis.crypto?.randomUUID === "function") {
|
||||
return globalThis.crypto.randomUUID().replace(/-/g, "").toLowerCase();
|
||||
}
|
||||
if (typeof globalThis.crypto?.getRandomValues === "function") {
|
||||
const bytes = new Uint8Array(16);
|
||||
globalThis.crypto.getRandomValues(bytes);
|
||||
return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join(
|
||||
"",
|
||||
);
|
||||
}
|
||||
let uid = "";
|
||||
while (uid.length < 32) {
|
||||
uid += Math.floor(Math.random() * 0x100000000)
|
||||
.toString(16)
|
||||
.padStart(8, "0");
|
||||
}
|
||||
return uid.slice(0, 32);
|
||||
}
|
||||
|
||||
export function resolveUnstructuredUploadBlockId({
|
||||
configId,
|
||||
uploadUid,
|
||||
generatedUploadUid,
|
||||
unstructuredFileCount,
|
||||
}: {
|
||||
configId: string;
|
||||
uploadUid: string;
|
||||
generatedUploadUid: string | null;
|
||||
unstructuredFileCount: number;
|
||||
}): string {
|
||||
if (uploadUid) {
|
||||
return uploadUid;
|
||||
}
|
||||
if (generatedUploadUid) {
|
||||
return generatedUploadUid;
|
||||
}
|
||||
return unstructuredFileCount > 0 ? configId : "";
|
||||
}
|
||||
|
||||
export function makeSamplerConfig(
|
||||
id: string,
|
||||
samplerType: SamplerType,
|
||||
|
|
@ -368,6 +408,9 @@ export function makeSeedConfig(
|
|||
hf_token: "",
|
||||
hf_endpoint: "https://huggingface.co",
|
||||
local_file_name: "",
|
||||
...(seedSourceType === "unstructured"
|
||||
? { unstructured_upload_uid: makeUnstructuredUploadUid() }
|
||||
: {}),
|
||||
unstructured_file_ids: [],
|
||||
unstructured_file_names: [],
|
||||
unstructured_file_sizes: [],
|
||||
|
|
|
|||
|
|
@ -16,11 +16,7 @@ import type {
|
|||
} from "../../types";
|
||||
import { buildEdges } from "./edges";
|
||||
import { isRecord, parseJson, readString } from "./helpers";
|
||||
import {
|
||||
parseColumn,
|
||||
parseModelConfig,
|
||||
parseModelProvider,
|
||||
} from "./parsers";
|
||||
import { parseColumn, parseModelConfig, parseModelProvider } from "./parsers";
|
||||
import { parseSeedConfig } from "./parsers/seed-config-parser";
|
||||
import { buildNodes, parseUi } from "./ui";
|
||||
import type { ImportResult } from "./types";
|
||||
|
|
@ -43,6 +39,7 @@ type UiInput = {
|
|||
seed_drop_columns?: unknown;
|
||||
seed_preview_rows?: unknown;
|
||||
local_file_name?: unknown;
|
||||
unstructured_upload_uid?: unknown;
|
||||
unstructured_file_ids?: unknown;
|
||||
unstructured_file_names?: unknown;
|
||||
unstructured_file_sizes?: unknown;
|
||||
|
|
@ -51,6 +48,10 @@ type UiInput = {
|
|||
advanced_open_by_node?: unknown;
|
||||
};
|
||||
|
||||
type ImportRecipePayloadOptions = {
|
||||
preserveUnstructuredUploads?: boolean;
|
||||
};
|
||||
|
||||
type UiMarkdownNoteNode = {
|
||||
name: string;
|
||||
markdown: string;
|
||||
|
|
@ -90,7 +91,7 @@ function parseProcessors(input: unknown): RecipeProcessorConfig[] {
|
|||
? templateRaw
|
||||
: isRecord(templateRaw)
|
||||
? JSON.stringify(templateRaw, null, 2)
|
||||
: "{\n \"text\": \"{{ column_name }}\"\n}";
|
||||
: '{\n "text": "{{ column_name }}"\n}';
|
||||
processors.push({
|
||||
id: `p${index + 1}`,
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
|
|
@ -135,9 +136,7 @@ function parseSeedDropColumns(input: unknown): string[] {
|
|||
return Array.from(values);
|
||||
}
|
||||
|
||||
function parseMcpProviders(
|
||||
input: unknown,
|
||||
): Map<string, LlmMcpProviderConfig> {
|
||||
function parseMcpProviders(input: unknown): Map<string, LlmMcpProviderConfig> {
|
||||
const providers = new Map<string, LlmMcpProviderConfig>();
|
||||
if (!Array.isArray(input)) {
|
||||
return providers;
|
||||
|
|
@ -156,13 +155,12 @@ function parseMcpProviders(
|
|||
const args = Array.isArray(item.args)
|
||||
? item.args.map((value) => String(value))
|
||||
: [];
|
||||
const envPairs =
|
||||
isRecord(item.env)
|
||||
? Object.entries(item.env).map(([key, value]) => ({
|
||||
key: String(key),
|
||||
value: String(value),
|
||||
}))
|
||||
: [];
|
||||
const envPairs = isRecord(item.env)
|
||||
? Object.entries(item.env).map(([key, value]) => ({
|
||||
key: String(key),
|
||||
value: String(value),
|
||||
}))
|
||||
: [];
|
||||
providers.set(name, {
|
||||
id: `mcp-${index + 1}`,
|
||||
name,
|
||||
|
|
@ -209,7 +207,8 @@ function parseToolConfigs(input: unknown): Map<string, LlmToolConfig> {
|
|||
allow_tools: allowTools,
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
max_tool_call_turns:
|
||||
item.max_tool_call_turns === null || item.max_tool_call_turns === undefined
|
||||
item.max_tool_call_turns === null ||
|
||||
item.max_tool_call_turns === undefined
|
||||
? "5"
|
||||
: String(item.max_tool_call_turns),
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
|
|
@ -257,7 +256,9 @@ function parseUiMarkdownNoteNodes(input: unknown): UiMarkdownNoteNode[] {
|
|||
return noteNodes;
|
||||
}
|
||||
|
||||
function parseUiToolProfileNodes(input: unknown): Map<string, Record<string, string[]>> {
|
||||
function parseUiToolProfileNodes(
|
||||
input: unknown,
|
||||
): Map<string, Record<string, string[]>> {
|
||||
const toolProfiles = new Map<string, Record<string, string[]>>();
|
||||
if (!Array.isArray(input)) {
|
||||
return toolProfiles;
|
||||
|
|
@ -312,9 +313,15 @@ function parseAdvancedOpenByNode(input: unknown): Record<string, boolean> {
|
|||
return out;
|
||||
}
|
||||
|
||||
type AdvancedOpenConfig = LlmConfig | SamplerConfig | SeedConfig | ValidatorConfig;
|
||||
type AdvancedOpenConfig =
|
||||
| LlmConfig
|
||||
| SamplerConfig
|
||||
| SeedConfig
|
||||
| ValidatorConfig;
|
||||
|
||||
function isAdvancedOpenConfig(config: NodeConfig): config is AdvancedOpenConfig {
|
||||
function isAdvancedOpenConfig(
|
||||
config: NodeConfig,
|
||||
): config is AdvancedOpenConfig {
|
||||
return (
|
||||
config.kind === "llm" ||
|
||||
config.kind === "sampler" ||
|
||||
|
|
@ -350,7 +357,8 @@ function buildToolProfileConfig(
|
|||
.map((providerName) => mcpProvidersByName.get(providerName))
|
||||
.flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])),
|
||||
// biome-ignore lint/style/useNamingConvention: ui schema
|
||||
fetched_tools_by_provider: fetchedToolsByProfileName.get(canonical.tool_alias) ?? {},
|
||||
fetched_tools_by_provider:
|
||||
fetchedToolsByProfileName.get(canonical.tool_alias) ?? {},
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
allow_tools: [...(canonical.allow_tools ?? [])],
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
|
|
@ -360,7 +368,10 @@ function buildToolProfileConfig(
|
|||
};
|
||||
}
|
||||
|
||||
export function importRecipePayload(input: string): ImportResult {
|
||||
export function importRecipePayload(
|
||||
input: string,
|
||||
options: ImportRecipePayloadOptions = {},
|
||||
): ImportResult {
|
||||
const parsed = parseJson(input);
|
||||
if (!parsed.data || !isRecord(parsed.data)) {
|
||||
return {
|
||||
|
|
@ -369,9 +380,9 @@ export function importRecipePayload(input: string): ImportResult {
|
|||
};
|
||||
}
|
||||
|
||||
const recipe = (isRecord(parsed.data.recipe)
|
||||
? parsed.data.recipe
|
||||
: parsed.data) as RecipeInput;
|
||||
const recipe = (
|
||||
isRecord(parsed.data.recipe) ? parsed.data.recipe : parsed.data
|
||||
) as RecipeInput;
|
||||
const ui = isRecord(parsed.data.ui) ? (parsed.data.ui as UiInput) : null;
|
||||
|
||||
if (!Array.isArray(recipe.columns)) {
|
||||
|
|
@ -410,21 +421,36 @@ export function importRecipePayload(input: string): ImportResult {
|
|||
.map((row) => ({ ...row }))
|
||||
: undefined;
|
||||
const uiLocalFileName = readString(ui?.local_file_name) ?? undefined;
|
||||
// Preserve file IDs/names from saved recipes (cleared at share time by sanitizeSeedForShare)
|
||||
const uiUnstructuredFileIds: string[] = Array.isArray(ui?.unstructured_file_ids)
|
||||
? (ui.unstructured_file_ids as string[]).filter((v): v is string => typeof v === "string")
|
||||
: [];
|
||||
const uiUnstructuredFileNames: string[] = Array.isArray(ui?.unstructured_file_names)
|
||||
? (ui.unstructured_file_names as string[]).filter((v): v is string => typeof v === "string")
|
||||
: [];
|
||||
const uiUnstructuredFileSizes: number[] = Array.isArray(ui?.unstructured_file_sizes)
|
||||
? (ui.unstructured_file_sizes as number[]).filter((v): v is number => typeof v === "number")
|
||||
: [];
|
||||
const preserveUnstructuredUploads =
|
||||
options.preserveUnstructuredUploads === true;
|
||||
const uiUnstructuredUploadUid = preserveUnstructuredUploads
|
||||
? (readString(ui?.unstructured_upload_uid) ?? undefined)
|
||||
: undefined;
|
||||
const uiUnstructuredFileIds: string[] =
|
||||
preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_ids)
|
||||
? (ui.unstructured_file_ids as string[]).filter(
|
||||
(v): v is string => typeof v === "string",
|
||||
)
|
||||
: [];
|
||||
const uiUnstructuredFileNames: string[] =
|
||||
preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_names)
|
||||
? (ui.unstructured_file_names as string[]).filter(
|
||||
(v): v is string => typeof v === "string",
|
||||
)
|
||||
: [];
|
||||
const uiUnstructuredFileSizes: number[] =
|
||||
preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_sizes)
|
||||
? (ui.unstructured_file_sizes as number[]).filter(
|
||||
(v): v is number => typeof v === "number",
|
||||
)
|
||||
: [];
|
||||
const uiUnstructuredChunkSize = readStringNumber(ui?.unstructured_chunk_size);
|
||||
const uiUnstructuredChunkOverlap = readStringNumber(
|
||||
ui?.unstructured_chunk_overlap,
|
||||
);
|
||||
const uiAdvancedOpenByNode = parseAdvancedOpenByNode(ui?.advanced_open_by_node);
|
||||
const uiAdvancedOpenByNode = parseAdvancedOpenByNode(
|
||||
ui?.advanced_open_by_node,
|
||||
);
|
||||
const uiMarkdownNotes = parseUiMarkdownNoteNodes(ui?.nodes);
|
||||
const uiToolProfilesByName = parseUiToolProfileNodes(ui?.nodes);
|
||||
|
||||
|
|
@ -459,11 +485,13 @@ export function importRecipePayload(input: string): ImportResult {
|
|||
: payloadSeedDropColumns,
|
||||
seed_preview_rows: uiSeedPreviewRows,
|
||||
local_file_name: uiLocalFileName,
|
||||
unstructuredUploadUid: uiUnstructuredUploadUid,
|
||||
unstructuredFileIds: uiUnstructuredFileIds,
|
||||
unstructuredFileNames: uiUnstructuredFileNames,
|
||||
unstructuredFileSizes: uiUnstructuredFileSizes,
|
||||
unstructured_chunk_size: uiUnstructuredChunkSize,
|
||||
unstructured_chunk_overlap: uiUnstructuredChunkOverlap,
|
||||
preserveUnstructuredUploads,
|
||||
});
|
||||
if (seedConfig) {
|
||||
applyAdvancedOpen(seedConfig, uiAdvancedOpenByNode);
|
||||
|
|
@ -567,12 +595,7 @@ export function importRecipePayload(input: string): ImportResult {
|
|||
const { layouts, auxNodes, edges: uiEdges, layoutDirection } = parseUi(ui);
|
||||
const resolvedLayoutDirection = layoutDirection ?? "LR";
|
||||
const nodes = buildNodes(configs, layouts);
|
||||
const edges = buildEdges(
|
||||
configs,
|
||||
nameToId,
|
||||
uiEdges,
|
||||
resolvedLayoutDirection,
|
||||
);
|
||||
const edges = buildEdges(configs, nameToId, uiEdges, resolvedLayoutDirection);
|
||||
const auxNodePositions = Object.fromEntries(
|
||||
auxNodes.flatMap((item) => {
|
||||
const llmId = nameToId.get(item.llm);
|
||||
|
|
@ -583,10 +606,7 @@ export function importRecipePayload(input: string): ImportResult {
|
|||
}),
|
||||
);
|
||||
|
||||
const maxY = nodes.reduce(
|
||||
(acc, node) => Math.max(acc, node.position.y),
|
||||
0,
|
||||
);
|
||||
const maxY = nodes.reduce((acc, node) => Math.max(acc, node.position.y), 0);
|
||||
|
||||
return {
|
||||
errors: [],
|
||||
|
|
|
|||
|
|
@ -197,17 +197,26 @@ export function parseSeedConfig(
|
|||
seed_drop_columns?: string[];
|
||||
seed_preview_rows?: Record<string, unknown>[];
|
||||
local_file_name?: string;
|
||||
unstructuredUploadUid?: string;
|
||||
unstructuredFileIds?: string[];
|
||||
unstructuredFileNames?: string[];
|
||||
unstructuredFileSizes?: number[];
|
||||
unstructured_chunk_size?: string;
|
||||
unstructured_chunk_overlap?: string;
|
||||
preserveUnstructuredUploads?: boolean;
|
||||
},
|
||||
): SeedConfig | null {
|
||||
if (!seedConfigRaw) {
|
||||
return null;
|
||||
}
|
||||
const parsed = parseSeedSettings(seedConfigRaw);
|
||||
const parsed = { ...parseSeedSettings(seedConfigRaw) };
|
||||
if (
|
||||
parsed.seed_source_type === "unstructured" &&
|
||||
options?.preserveUnstructuredUploads !== true
|
||||
) {
|
||||
parsed.hf_path = "";
|
||||
parsed.resolved_paths = [];
|
||||
}
|
||||
let sourceType: SeedSourceType = "hf";
|
||||
if (parsed.seed_source_type === "hf") {
|
||||
sourceType = "hf";
|
||||
|
|
@ -230,6 +239,9 @@ export function parseSeedConfig(
|
|||
...(options?.local_file_name !== undefined
|
||||
? { local_file_name: options.local_file_name }
|
||||
: {}),
|
||||
...(options?.unstructuredUploadUid
|
||||
? { unstructured_upload_uid: options.unstructuredUploadUid }
|
||||
: {}),
|
||||
...(options?.unstructuredFileIds !== undefined
|
||||
? { unstructured_file_ids: options.unstructuredFileIds }
|
||||
: {}),
|
||||
|
|
|
|||
|
|
@ -440,6 +440,9 @@ export function buildRecipePayload(
|
|||
unstructured_file_names: firstSeed.unstructured_file_names,
|
||||
unstructured_file_sizes: firstSeed.unstructured_file_sizes,
|
||||
}),
|
||||
...(firstSeed?.unstructured_upload_uid?.trim() && {
|
||||
unstructured_upload_uid: firstSeed.unstructured_upload_uid,
|
||||
}),
|
||||
...(firstSeed &&
|
||||
firstSeed.unstructured_chunk_size !== undefined && {
|
||||
unstructured_chunk_size: firstSeed.unstructured_chunk_size,
|
||||
|
|
|
|||
|
|
@ -71,6 +71,8 @@ export type RecipePayload = {
|
|||
seed_preview_rows?: Record<string, unknown>[];
|
||||
local_file_name?: string;
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
unstructured_upload_uid?: string;
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
unstructured_file_ids?: string[];
|
||||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
unstructured_file_names?: string[];
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ export {
|
|||
savePersonalization,
|
||||
} from "./api/personalization";
|
||||
export { setTheme, useTheme } from "./stores/theme-store";
|
||||
export { useMonitorOverlayStore } from "./stores/monitor-overlay-store";
|
||||
export type {
|
||||
Personalization,
|
||||
PersonalizationAppearance,
|
||||
|
|
|
|||
|
|
@ -409,7 +409,6 @@ export const en = {
|
|||
description: "Access Unsloth via the OpenAI-compatible API.",
|
||||
readDocs: "Read the API docs",
|
||||
noAccess: "No API access yet.",
|
||||
newBadge: "New",
|
||||
accessTokens: "Access tokens",
|
||||
loadError: "Couldn't load API access.",
|
||||
createError: "Couldn't create access token.",
|
||||
|
|
|
|||
|
|
@ -296,7 +296,6 @@ export const ja = {
|
|||
description: "OpenAI互換 API を介して Unsloth にアクセスします。",
|
||||
readDocs: "API ドキュメントを読む",
|
||||
noAccess: "まだ API アクセス権がありません。",
|
||||
newBadge: "新規",
|
||||
accessTokens: "アクセストークン",
|
||||
loadError: "API アクセス権を読み込めませんでした。",
|
||||
createError: "アクセストークンを作成できませんでした。",
|
||||
|
|
|
|||
|
|
@ -363,7 +363,6 @@ export const ptBR = {
|
|||
"Acesse o Unsloth por meio da API compatível com OpenAI.",
|
||||
readDocs: "Leia a documentação da API",
|
||||
noAccess: "Nenhum acesso à API ainda.",
|
||||
newBadge: "Novo",
|
||||
accessTokens: "Tokens de acesso",
|
||||
loadError: "Não foi possível carregar o acesso à API.",
|
||||
createError: "Não foi possível criar o token de acesso.",
|
||||
|
|
|
|||
|
|
@ -267,7 +267,6 @@ export const zhCN = {
|
|||
description: "通过兼容 OpenAI 的 API 以编程方式访问 Unsloth。",
|
||||
readDocs: "阅读 API 文档",
|
||||
noAccess: "还没有 API 访问权限。",
|
||||
newBadge: "新",
|
||||
accessTokens: "访问 token",
|
||||
loadError: "无法加载 API 访问权限。",
|
||||
createError: "无法创建访问 token。",
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import argparse
|
|||
import atexit
|
||||
import errno
|
||||
import fnmatch
|
||||
import glob
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
|
|
@ -265,6 +266,7 @@ class HostInfo:
|
|||
has_physical_nvidia: bool
|
||||
has_usable_nvidia: bool
|
||||
has_rocm: bool = False
|
||||
has_intel_gpu: bool = False
|
||||
rocm_gfx_target: str | None = None
|
||||
# (major, minor) from platform.mac_ver(); None off macOS or if unparseable.
|
||||
# Skips a macos prebuilt whose minimum-OS exceeds this host.
|
||||
|
|
@ -1284,162 +1286,6 @@ def synthetic_checksums_for_release(
|
|||
)
|
||||
|
||||
|
||||
def parse_direct_linux_release_bundle(
|
||||
repo: str, release: dict[str, Any]
|
||||
) -> PublishedReleaseBundle | None:
|
||||
release_tag = release.get("tag_name")
|
||||
if not isinstance(release_tag, str) or not release_tag:
|
||||
return None
|
||||
|
||||
assets = release_asset_map(release)
|
||||
artifacts: list[PublishedLlamaArtifact] = []
|
||||
inferred_labels: list[str] = []
|
||||
|
||||
linux_asset_re = re.compile(
|
||||
r"^app-(?P<label>.+)-(?P<target>linux-x64(?:-cpu)?|linux-x64-cuda\d+-(?:older|newer|portable))\.tar\.gz$"
|
||||
)
|
||||
for asset_name in sorted(assets):
|
||||
match = linux_asset_re.fullmatch(asset_name)
|
||||
if not match:
|
||||
continue
|
||||
inferred_labels.append(match.group("label"))
|
||||
target = match.group("target")
|
||||
if target in {"linux-x64", "linux-x64-cpu"}:
|
||||
artifacts.append(
|
||||
PublishedLlamaArtifact(
|
||||
asset_name = asset_name,
|
||||
install_kind = "linux-cpu",
|
||||
runtime_line = None,
|
||||
coverage_class = None,
|
||||
supported_sms = [],
|
||||
min_sm = None,
|
||||
max_sm = None,
|
||||
bundle_profile = None,
|
||||
rank = 1000,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
bundle_profile = target.removeprefix("linux-x64-")
|
||||
profile = _resolve_linux_bundle_profile(bundle_profile)
|
||||
if profile is None:
|
||||
continue
|
||||
artifacts.append(
|
||||
PublishedLlamaArtifact(
|
||||
asset_name = asset_name,
|
||||
install_kind = "linux-cuda",
|
||||
runtime_line = str(profile["runtime_line"]),
|
||||
coverage_class = str(profile["coverage_class"]),
|
||||
supported_sms = [str(value) for value in profile["supported_sms"]],
|
||||
min_sm = int(profile["min_sm"]),
|
||||
max_sm = int(profile["max_sm"]),
|
||||
bundle_profile = bundle_profile,
|
||||
rank = int(profile["rank"]),
|
||||
)
|
||||
)
|
||||
|
||||
if not artifacts:
|
||||
return None
|
||||
|
||||
upstream_tag = (
|
||||
release_tag
|
||||
if is_release_tag_like(release_tag)
|
||||
else inferred_labels[0]
|
||||
if len(set(inferred_labels)) == 1 and inferred_labels
|
||||
else release_tag
|
||||
)
|
||||
selection_log = [
|
||||
f"published_release: repo={repo}",
|
||||
f"published_release: tag={release_tag}",
|
||||
f"published_release: upstream_tag={upstream_tag}",
|
||||
"published_release: direct_asset_scan=linux",
|
||||
]
|
||||
return PublishedReleaseBundle(
|
||||
repo = repo,
|
||||
release_tag = release_tag,
|
||||
upstream_tag = upstream_tag,
|
||||
assets = assets,
|
||||
manifest_asset_name = DEFAULT_PUBLISHED_MANIFEST_ASSET,
|
||||
artifacts = artifacts,
|
||||
selection_log = selection_log,
|
||||
)
|
||||
|
||||
|
||||
def direct_linux_release_plan(
|
||||
release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
|
||||
) -> InstallReleasePlan | None:
|
||||
bundle = parse_direct_linux_release_bundle(repo, release)
|
||||
if bundle is None:
|
||||
return None
|
||||
if not direct_release_matches_request(
|
||||
release_tag = bundle.release_tag,
|
||||
llama_tag = bundle.upstream_tag,
|
||||
requested_tag = requested_tag,
|
||||
):
|
||||
return None
|
||||
|
||||
attempts: list[AssetChoice] = []
|
||||
if host.has_usable_nvidia:
|
||||
# Prefer the cudart major Studio loads at runtime (torch's bundled
|
||||
# libcudart), not the newest on disk. Otherwise a stray cuda13
|
||||
# runtime outranks the torch cuda12 the binary links against.
|
||||
torch_preference = detect_torch_cuda_runtime_preference(host)
|
||||
selection = linux_cuda_choice_from_release(
|
||||
host,
|
||||
bundle,
|
||||
preferred_runtime_line = torch_preference.runtime_line,
|
||||
selection_preamble = torch_preference.selection_log,
|
||||
)
|
||||
if selection is not None:
|
||||
attempts.extend(selection.attempts)
|
||||
elif not host.has_rocm:
|
||||
# A ROCm-only host gets no CPU asset: leaving attempts empty lets the
|
||||
# raise below trigger a HIP source build instead of shipping a CPU
|
||||
# binary on a GPU host (this ggml-org path has no per-gfx ROCm asset).
|
||||
cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu")
|
||||
if cpu_choice is not None:
|
||||
attempts.append(cpu_choice)
|
||||
# NVIDIA hosts whose CUDA selection produced nothing fall through to the
|
||||
# raise below (mirroring the ROCm policy above): the caller then walks
|
||||
# back to an older release that still ships a usable CUDA line instead of
|
||||
# silently installing a CPU binary on a GPU host. Today's walk-back only
|
||||
# works because partial releases ship no CPU bundle; this keeps it working
|
||||
# if a future partial release does.
|
||||
if not attempts:
|
||||
raise PrebuiltFallback("no compatible Linux prebuilt asset was found")
|
||||
approved_checksums = synthetic_checksums_for_release(
|
||||
repo,
|
||||
bundle.release_tag,
|
||||
bundle.upstream_tag,
|
||||
)
|
||||
resolved_upstream_tag = bundle.upstream_tag
|
||||
if DEFAULT_PUBLISHED_SHA256_ASSET in bundle.assets and not is_release_tag_like(
|
||||
bundle.upstream_tag
|
||||
):
|
||||
approved_checksums = load_approved_release_checksums(repo, bundle.release_tag)
|
||||
# Require exact source provenance for branch/pull/commit releases.
|
||||
# Mirrors validated_checksums_for_bundle so incomplete metadata fails
|
||||
# closed instead of degrading to the legacy branch-as-tag source
|
||||
# hydration path this PR eliminates.
|
||||
if (
|
||||
not approved_checksums.source_commit
|
||||
or exact_source_archive_hash(approved_checksums) is None
|
||||
or source_clone_url_from_checksums(approved_checksums) is None
|
||||
):
|
||||
raise PrebuiltFallback(
|
||||
f"approved checksum asset {DEFAULT_PUBLISHED_SHA256_ASSET} for "
|
||||
f"{repo}@{bundle.release_tag} did not contain exact source provenance"
|
||||
)
|
||||
attempts = apply_approved_hashes(attempts, approved_checksums)
|
||||
return InstallReleasePlan(
|
||||
requested_tag = requested_tag,
|
||||
llama_tag = resolved_upstream_tag,
|
||||
release_tag = bundle.release_tag,
|
||||
attempts = attempts,
|
||||
approved_checksums = approved_checksums,
|
||||
)
|
||||
|
||||
|
||||
def direct_upstream_release_plan(
|
||||
release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
|
||||
) -> InstallReleasePlan | None:
|
||||
|
|
@ -1482,6 +1328,24 @@ def direct_upstream_release_plan(
|
|||
install_kind = "windows-hip",
|
||||
)
|
||||
)
|
||||
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. Gate
|
||||
# on no PHYSICAL NVIDIA (not just no usable one): a host that hid NVIDIA
|
||||
# via CUDA_VISIBLE_DEVICES must not reach Vulkan, which ignores that mask
|
||||
# and could enumerate the reserved card. Falls through to CPU below.
|
||||
elif host.has_intel_gpu and not host.has_physical_nvidia:
|
||||
vulkan_asset = f"llama-{release_tag}-bin-win-vulkan-x64.zip"
|
||||
vulkan_url = assets.get(vulkan_asset)
|
||||
if vulkan_url:
|
||||
attempts.append(
|
||||
AssetChoice(
|
||||
repo = repo,
|
||||
tag = release_tag,
|
||||
name = vulkan_asset,
|
||||
url = vulkan_url,
|
||||
source_label = "upstream",
|
||||
install_kind = "windows-vulkan",
|
||||
)
|
||||
)
|
||||
cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip"
|
||||
cpu_url = assets.get(cpu_asset)
|
||||
if cpu_url:
|
||||
|
|
@ -1545,6 +1409,23 @@ def direct_upstream_release_plan(
|
|||
# ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm
|
||||
# asset, so they fall through to the empty-attempts raise (HIP source
|
||||
# build) rather than silently getting a CPU binary on a GPU host.
|
||||
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. The
|
||||
# elif already excludes usable NVIDIA and ROCm; also require no PHYSICAL
|
||||
# NVIDIA so a CUDA-hidden card isn't reached through Vulkan (CPU below).
|
||||
if host.has_intel_gpu and not host.has_physical_nvidia:
|
||||
vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-x64.tar.gz"
|
||||
vulkan_url = assets.get(vulkan_asset)
|
||||
if vulkan_url:
|
||||
attempts.append(
|
||||
AssetChoice(
|
||||
repo = repo,
|
||||
tag = release_tag,
|
||||
name = vulkan_asset,
|
||||
url = vulkan_url,
|
||||
source_label = "upstream",
|
||||
install_kind = "linux-vulkan",
|
||||
)
|
||||
)
|
||||
asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz"
|
||||
asset_url = assets.get(asset_name)
|
||||
if asset_url:
|
||||
|
|
@ -1564,6 +1445,23 @@ def direct_upstream_release_plan(
|
|||
# selector returned 0 attempts and the installer fell back to a
|
||||
# source build on every Linux ARM64 host (DGX Spark, Ampere
|
||||
# Altra, GitHub-hosted ubuntu-24.04-arm runners, etc.).
|
||||
# Intel (or other non-NVIDIA/non-AMD) GPU: prefer the Vulkan prebuilt,
|
||||
# mirroring the x86_64 branch. Upstream ships bin-ubuntu-vulkan-arm64.
|
||||
# No physical NVIDIA: don't reach a CUDA-hidden card through Vulkan.
|
||||
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
|
||||
vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-arm64.tar.gz"
|
||||
vulkan_url = assets.get(vulkan_asset)
|
||||
if vulkan_url:
|
||||
attempts.append(
|
||||
AssetChoice(
|
||||
repo = repo,
|
||||
tag = release_tag,
|
||||
name = vulkan_asset,
|
||||
url = vulkan_url,
|
||||
source_label = "upstream",
|
||||
install_kind = "linux-vulkan",
|
||||
)
|
||||
)
|
||||
asset_name = f"llama-{release_tag}-bin-ubuntu-arm64.tar.gz"
|
||||
asset_url = assets.get(asset_name)
|
||||
if asset_url:
|
||||
|
|
@ -2864,6 +2762,64 @@ def _pick_rocm_gfx_target(out: str) -> str | None:
|
|||
return _tokens[0]
|
||||
|
||||
|
||||
# Display-adapter device class: one NNNN subkey per installed display driver
|
||||
# config, each carrying the driver's DriverDesc and PCI MatchingDeviceId.
|
||||
_WINDOWS_DISPLAY_CLASS_KEY = (
|
||||
r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}"
|
||||
)
|
||||
|
||||
|
||||
def windows_intel_gpu_in_registry() -> bool:
|
||||
"""Whether the Windows registry lists an Intel display adapter.
|
||||
|
||||
In-process Windows counterpart of the Linux DRM vendor-id check (0x8086),
|
||||
with weaker semantics: the class key lists installed display-driver
|
||||
configs, which can outlive removed hardware, where sysfs lists present
|
||||
devices. A stale Intel entry at worst routes to the upstream Vulkan
|
||||
prebuilt instead of the fork CPU bundle: inference still works (the
|
||||
Vulkan build runs on CPU when no Vulkan device exists), at the cost of
|
||||
fork-only extras such as the DiffusionGemma visual server. detect_host's
|
||||
PowerShell + WMI probe can silently miss a real Intel GPU: a cold
|
||||
powershell.exe start plus the first CIM query routinely exceeds the 15s
|
||||
budget on hosts with slow AV scanning or a degraded WMI repository, and
|
||||
the probe swallows the timeout (#4452, Arc A770 routed to the CPU
|
||||
prebuilt). Reading the display-adapter class key needs no subprocess and
|
||||
answers in microseconds. Matches the PCI vendor id in MatchingDeviceId
|
||||
(ven_8086) or an Intel DriverDesc.
|
||||
"""
|
||||
try:
|
||||
import winreg
|
||||
except ImportError:
|
||||
return False
|
||||
try:
|
||||
with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, _WINDOWS_DISPLAY_CLASS_KEY) as class_key:
|
||||
for index in range(winreg.QueryInfoKey(class_key)[0]):
|
||||
try:
|
||||
name = winreg.EnumKey(class_key, index)
|
||||
if not name.isdigit():
|
||||
# "Properties" is ACL-restricted and not an adapter.
|
||||
continue
|
||||
with winreg.OpenKey(class_key, name) as adapter_key:
|
||||
for value_name, needle in (
|
||||
("MatchingDeviceId", "ven_8086"),
|
||||
("DriverDesc", "intel"),
|
||||
):
|
||||
try:
|
||||
value, _ = winreg.QueryValueEx(adapter_key, value_name)
|
||||
except OSError:
|
||||
continue
|
||||
if needle in str(value).lower():
|
||||
return True
|
||||
except OSError:
|
||||
continue
|
||||
except Exception:
|
||||
# Advisory probe: any unexpected failure must degrade to the CIM
|
||||
# fallback, never crash the installer (mirrors detect_host's own
|
||||
# swallow around the CIM probe).
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def detect_host() -> HostInfo:
|
||||
system = platform.system()
|
||||
machine = platform.machine().lower()
|
||||
|
|
@ -3075,6 +3031,46 @@ def detect_host() -> HostInfo:
|
|||
# Note: amdhip64.dll presence alone is NOT treated as GPU evidence
|
||||
# since the HIP SDK can be installed without an AMD GPU.
|
||||
|
||||
# Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs
|
||||
# vendor id (0x8086); Windows reads the display-adapter registry class,
|
||||
# then falls back to the WMI video controller list. Only probed with no
|
||||
# usable NVIDIA and no ROCm (matching the Vulkan branches), keeping the
|
||||
# probe (notably the Windows powershell call) off that path.
|
||||
has_intel_gpu = False
|
||||
if not has_usable_nvidia and not has_rocm:
|
||||
if is_linux:
|
||||
for _vendor_file in glob.glob("/sys/class/drm/card*/device/vendor"):
|
||||
try:
|
||||
with open(_vendor_file) as _vf:
|
||||
if _vf.read().strip().lower() == "0x8086":
|
||||
has_intel_gpu = True
|
||||
break
|
||||
except OSError:
|
||||
continue
|
||||
elif is_windows:
|
||||
# Registry first (in-process; see windows_intel_gpu_in_registry).
|
||||
# The CIM query stays as the fallback when the registry shows no
|
||||
# Intel adapter.
|
||||
has_intel_gpu = windows_intel_gpu_in_registry()
|
||||
if not has_intel_gpu:
|
||||
_ps = shutil.which("powershell") or shutil.which("pwsh")
|
||||
if _ps:
|
||||
try:
|
||||
_result = run_capture(
|
||||
[
|
||||
_ps,
|
||||
"-NoProfile",
|
||||
"-Command",
|
||||
"Get-CimInstance Win32_VideoController | "
|
||||
"Select-Object -ExpandProperty Name",
|
||||
],
|
||||
timeout = 15,
|
||||
)
|
||||
if _result.returncode == 0 and "intel" in _result.stdout.lower():
|
||||
has_intel_gpu = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return HostInfo(
|
||||
system = system,
|
||||
machine = machine,
|
||||
|
|
@ -3090,6 +3086,7 @@ def detect_host() -> HostInfo:
|
|||
has_physical_nvidia = has_physical_nvidia,
|
||||
has_usable_nvidia = has_usable_nvidia,
|
||||
has_rocm = has_rocm,
|
||||
has_intel_gpu = has_intel_gpu,
|
||||
rocm_gfx_target = rocm_gfx_target,
|
||||
macos_version = macos_version,
|
||||
)
|
||||
|
|
@ -3126,6 +3123,7 @@ def _apply_host_overrides(
|
|||
has_physical_nvidia = False,
|
||||
has_rocm = False,
|
||||
rocm_gfx_target = None,
|
||||
has_intel_gpu = False,
|
||||
)
|
||||
gfx = _normalize_forwarded_gfx(override_rocm_gfx)
|
||||
if gfx:
|
||||
|
|
@ -3866,6 +3864,23 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
|
|||
"falling back to source build with HIP support"
|
||||
)
|
||||
|
||||
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. No
|
||||
# physical NVIDIA (not just no usable one): a CUDA-hidden card must not
|
||||
# be reached through Vulkan, which ignores CUDA_VISIBLE_DEVICES.
|
||||
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
|
||||
vulkan_name = f"llama-{llama_tag}-bin-ubuntu-vulkan-x64.tar.gz"
|
||||
if vulkan_name in upstream_assets:
|
||||
log(f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}")
|
||||
return AssetChoice(
|
||||
repo = UPSTREAM_REPO,
|
||||
tag = llama_tag,
|
||||
name = vulkan_name,
|
||||
url = upstream_assets[vulkan_name],
|
||||
source_label = "upstream",
|
||||
install_kind = "linux-vulkan",
|
||||
)
|
||||
log("Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU")
|
||||
|
||||
upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz"
|
||||
if upstream_name not in upstream_assets:
|
||||
raise PrebuiltFallback("upstream Linux CPU asset was not found")
|
||||
|
|
@ -3908,6 +3923,24 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
|
|||
)
|
||||
log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU")
|
||||
|
||||
# Intel (or other non-NVIDIA/non-AMD) GPU on Windows: use Vulkan. No
|
||||
# physical NVIDIA so a CUDA-hidden card isn't reached through Vulkan.
|
||||
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
|
||||
vulkan_name = f"llama-{llama_tag}-bin-win-vulkan-x64.zip"
|
||||
if vulkan_name in upstream_assets:
|
||||
log(
|
||||
f"Intel GPU detected on Windows -- using upstream Vulkan prebuilt {vulkan_name}"
|
||||
)
|
||||
return AssetChoice(
|
||||
repo = UPSTREAM_REPO,
|
||||
tag = llama_tag,
|
||||
name = vulkan_name,
|
||||
url = upstream_assets[vulkan_name],
|
||||
source_label = "upstream",
|
||||
install_kind = "windows-vulkan",
|
||||
)
|
||||
log("Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU")
|
||||
|
||||
upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip"
|
||||
if upstream_name not in upstream_assets:
|
||||
raise PrebuiltFallback("upstream Windows CPU asset was not found")
|
||||
|
|
@ -4503,6 +4536,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
|
|||
"linux-arm64-cuda",
|
||||
"linux-rocm",
|
||||
"linux-arm64",
|
||||
"linux-vulkan",
|
||||
}:
|
||||
return ["llama-server", "llama-quantize", "llama-diffusion-gemma-visual-server", "lib*.so*"]
|
||||
if choice.install_kind in {"macos-arm64", "macos-x64"}:
|
||||
|
|
@ -4516,6 +4550,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
|
|||
"windows-cpu",
|
||||
"windows-cuda",
|
||||
"windows-hip",
|
||||
"windows-vulkan",
|
||||
"windows-rocm",
|
||||
"windows-arm64",
|
||||
}:
|
||||
|
|
@ -5731,8 +5766,10 @@ def validate_server(
|
|||
"linux-cuda",
|
||||
"linux-arm64-cuda",
|
||||
"linux-rocm",
|
||||
"linux-vulkan",
|
||||
"windows-cuda",
|
||||
"windows-hip",
|
||||
"windows-vulkan",
|
||||
"windows-rocm",
|
||||
"macos-arm64",
|
||||
}
|
||||
|
|
@ -6354,6 +6391,20 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
|
|||
["libmtmd.so*"],
|
||||
["libggml-hip.so*"],
|
||||
]
|
||||
if choice.install_kind == "linux-vulkan":
|
||||
return [
|
||||
["libllama-common.so*"],
|
||||
["libllama.so*"],
|
||||
["libggml.so*"],
|
||||
["libggml-base.so*"],
|
||||
# Match the sibling globs (linux-cuda/-rocm): x64 bundles ship
|
||||
# arch-suffixed libggml-cpu-<variant>.so, arm64 may ship a bare
|
||||
# libggml-cpu.so; the '-' form missed the latter and re-flagged
|
||||
# the install unhealthy on every check.
|
||||
["libggml-cpu*.so*"],
|
||||
["libmtmd.so*"],
|
||||
["libggml-vulkan.so*"],
|
||||
]
|
||||
if choice.install_kind in {"windows-cpu", "windows-arm64"}:
|
||||
return [["llama.dll"]]
|
||||
if choice.install_kind == "windows-cuda":
|
||||
|
|
@ -6373,6 +6424,8 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
|
|||
return groups
|
||||
if choice.install_kind in {"windows-hip", "windows-rocm"}:
|
||||
return [["llama.dll"], ["*hip*.dll"]]
|
||||
if choice.install_kind == "windows-vulkan":
|
||||
return [["llama.dll"], ["ggml-vulkan.dll"]]
|
||||
return []
|
||||
|
||||
|
||||
|
|
@ -6654,6 +6707,89 @@ def validate_prebuilt_attempts(
|
|||
raise PrebuiltFallback("no prebuilt bundle passed validation")
|
||||
|
||||
|
||||
def force_vulkan_requested() -> bool:
|
||||
"""Whether UNSLOTH_FORCE_VULKAN opts this host into the Vulkan llama.cpp
|
||||
prebuilt instead of its detected CUDA/ROCm backend (e.g. so an AMD user can
|
||||
run the Vulkan build for inference). Scoped to the llama.cpp backend; the
|
||||
torch/training stack installs separately and still sees the real GPU.
|
||||
"""
|
||||
return os.environ.get("UNSLOTH_FORCE_VULKAN", "").strip().lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
)
|
||||
|
||||
|
||||
def _vulkan_only_host(host: HostInfo) -> HostInfo:
|
||||
"""Rewrite ``host`` so the asset selectors take their Vulkan branch.
|
||||
|
||||
That branch fires on ``has_intel_gpu and not nvidia and not rocm``, so clear
|
||||
the CUDA/ROCm flags and raise the integrated-GPU flag. The synthetic flag
|
||||
never leaves install planning -- it only routes the llama.cpp prebuilt
|
||||
choice, not the torch/training stack.
|
||||
"""
|
||||
return dataclasses_replace(
|
||||
host,
|
||||
has_usable_nvidia = False,
|
||||
has_physical_nvidia = False,
|
||||
has_rocm = False,
|
||||
has_intel_gpu = True,
|
||||
)
|
||||
|
||||
|
||||
def _route_to_vulkan_prebuilt(
|
||||
host: HostInfo, published_repo: str, published_release_tag: str, *, force_cpu: bool
|
||||
) -> tuple[HostInfo, str, str]:
|
||||
"""Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
|
||||
|
||||
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes
|
||||
from UPSTREAM_REPO. Two triggers route here, both suppressed under
|
||||
--cpu-fallback (the explicit "give me CPU" last resort wins):
|
||||
* UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend;
|
||||
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose
|
||||
of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
|
||||
Applied by BOTH the install path and the --resolve-prebuilt probe so the
|
||||
"is a prebuilt available" answer matches what actually gets installed.
|
||||
|
||||
Returns the (possibly rewritten) host, repo, and release tag.
|
||||
"""
|
||||
forced = force_vulkan_requested()
|
||||
# Gate auto-routing on no PHYSICAL NVIDIA, not merely no usable one: a mixed
|
||||
# NVIDIA+Intel host that hides NVIDIA with CUDA_VISIBLE_DEVICES=""/-1 keeps
|
||||
# has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores
|
||||
# CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the
|
||||
# reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides.
|
||||
auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
|
||||
if force_cpu or not (forced or auto_intel):
|
||||
return host, published_repo, published_release_tag
|
||||
if host.is_macos:
|
||||
if forced:
|
||||
log(
|
||||
"UNSLOTH_FORCE_VULKAN is set but ignored on macOS "
|
||||
"(Metal is used; there is no Vulkan prebuilt)"
|
||||
)
|
||||
return host, published_repo, published_release_tag
|
||||
if forced:
|
||||
log(
|
||||
"UNSLOTH_FORCE_VULKAN is set; installing the upstream Vulkan "
|
||||
"llama.cpp prebuilt instead of the detected GPU backend"
|
||||
)
|
||||
# Forcing may override a detected NVIDIA/ROCm host, so normalize it to
|
||||
# Vulkan-only; an auto-detected Intel host already is.
|
||||
host = _vulkan_only_host(host)
|
||||
else:
|
||||
log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt")
|
||||
# Swapping the fork for upstream invalidates a fork release pin: the two use
|
||||
# different tag namespaces (fork b9596-mix-<sha> vs upstream b9596), so a
|
||||
# pinned fork tag would make the upstream resolver query a nonexistent
|
||||
# release and fall back to source. Drop it and let the upstream resolver
|
||||
# pick by the requested llama tag. A pin already on an explicit upstream repo
|
||||
# (repo unchanged here) is preserved.
|
||||
if published_repo != UPSTREAM_REPO:
|
||||
published_release_tag = ""
|
||||
return host, UPSTREAM_REPO, published_release_tag
|
||||
|
||||
|
||||
def diffusion_visual_server_backfill_needed(
|
||||
install_dir: Path, host: HostInfo, choice: AssetChoice
|
||||
) -> bool:
|
||||
|
|
@ -6696,6 +6832,9 @@ def install_prebuilt(
|
|||
override_rocm_gfx = override_rocm_gfx,
|
||||
force_cpu = force_cpu,
|
||||
)
|
||||
host, published_repo, published_release_tag = _route_to_vulkan_prebuilt(
|
||||
host, published_repo, published_release_tag, force_cpu = force_cpu
|
||||
)
|
||||
choice: AssetChoice | None = None
|
||||
try:
|
||||
with install_lock(install_lock_path(install_dir)):
|
||||
|
|
@ -6708,7 +6847,9 @@ def install_prebuilt(
|
|||
f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install"
|
||||
)
|
||||
# Single resolver: every fork host selects from the release manifest;
|
||||
# an explicit ggml-org override selects by asset filename instead.
|
||||
# an explicit ggml-org override selects by asset filename instead. A
|
||||
# forced-Vulkan host already has published_repo pointed at
|
||||
# UPSTREAM_REPO above, so the resolver takes the Vulkan asset branch.
|
||||
requested_tag, release_plans = resolve_simple_install_release_plans(
|
||||
llama_tag,
|
||||
host,
|
||||
|
|
@ -6994,10 +7135,14 @@ def main() -> int:
|
|||
override_rocm_gfx = args.rocm_gfx,
|
||||
force_cpu = args.cpu_fallback,
|
||||
)
|
||||
repo = args.published_repo
|
||||
# Same Vulkan routing the install path applies, so the probe's answer
|
||||
# matches what would install (an Intel/forced-Vulkan host -> upstream).
|
||||
host, repo, release_tag = _route_to_vulkan_prebuilt(
|
||||
host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback
|
||||
)
|
||||
try:
|
||||
_requested, plans = resolve_simple_install_release_plans(
|
||||
args.resolve_prebuilt, host, repo, args.published_release_tag or ""
|
||||
args.resolve_prebuilt, host, repo, release_tag
|
||||
)
|
||||
choice = plans[0].attempts[0] if plans and plans[0].attempts else None
|
||||
if choice is None:
|
||||
|
|
|
|||
|
|
@ -2621,7 +2621,18 @@ function Fast-Install {
|
|||
param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
|
||||
if ($UseUv) {
|
||||
$VenvPy = (Get-Command python).Source
|
||||
$result = & uv pip install --python $VenvPy @Args_ 2>&1
|
||||
# An explicit --index-url must win. Inherited uv index env vars otherwise
|
||||
# override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop
|
||||
# them only for index-pinned installs; mirrors still apply elsewhere.
|
||||
$saved = @{}
|
||||
if (@($Args_) -contains '--index-url') {
|
||||
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
|
||||
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
|
||||
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
|
||||
}
|
||||
}
|
||||
try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 }
|
||||
finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } }
|
||||
if ($LASTEXITCODE -eq 0) { return }
|
||||
}
|
||||
& python -m pip install @Args_ 2>&1
|
||||
|
|
|
|||
122
tests/python/test_fast_sentence_transformer_embedding_parity.py
Normal file
122
tests/python/test_fast_sentence_transformer_embedding_parity.py
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team.
|
||||
"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text
|
||||
like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a
|
||||
"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building
|
||||
via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings;
|
||||
`_create_transformer_module` uses `Transformer.load(...)` instead.
|
||||
|
||||
Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST
|
||||
is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity,
|
||||
opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_transformer_load_signature_supports_unsloth_kwargs():
|
||||
"""Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs
|
||||
the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back
|
||||
to Transformer(...) there, so mirror that gate and skip."""
|
||||
models = pytest.importorskip("sentence_transformers.models")
|
||||
load = getattr(models.Transformer, "load", None)
|
||||
assert callable(load), (
|
||||
"sentence_transformers Transformer.load is missing; the #6881 fix in "
|
||||
"unsloth.models.sentence_transformer._create_transformer_module depends on it."
|
||||
)
|
||||
params = inspect.signature(load).parameters
|
||||
accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
|
||||
# Mirror _create_transformer_module's hub_capable gate.
|
||||
hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision"))
|
||||
if not hub_capable:
|
||||
pytest.skip(
|
||||
"legacy Transformer.load(input_path); production path falls back to Transformer(...)"
|
||||
)
|
||||
unsupported = [
|
||||
k
|
||||
for k in ("token", "cache_folder", "revision", "trust_remote_code")
|
||||
if not (accepts_var_kw or k in params)
|
||||
]
|
||||
assert not unsupported, (
|
||||
f"installed sentence_transformers Transformer.load no longer accepts {unsupported} "
|
||||
f"and has no **kwargs; update _create_transformer_module (#6881) before it silently "
|
||||
f"falls back to Transformer(...)."
|
||||
)
|
||||
|
||||
|
||||
def _probe_texts():
|
||||
return [
|
||||
"roasted chickpeas in 20 kg bags",
|
||||
"The capital of France is Paris.",
|
||||
"A fast brown fox jumps over the lazy dog.",
|
||||
"recette de tarte aux pommes traditionnelle",
|
||||
]
|
||||
|
||||
|
||||
def test_fast_sentence_transformer_matches_stock_st():
|
||||
"""End-to-end: FastSentenceTransformer embeddings and tokenization must match a
|
||||
stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and
|
||||
GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners."""
|
||||
model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL")
|
||||
if not model_id:
|
||||
pytest.skip(
|
||||
"set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model "
|
||||
"(HF id or local path) to run the #6881 parity test"
|
||||
)
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner")
|
||||
np = pytest.importorskip("numpy")
|
||||
pytest.importorskip("sentence_transformers")
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
device = "cuda"
|
||||
# Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native
|
||||
# embedders such as EmbeddingGemma (Gemma3), which would mask real parity.
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
texts = _probe_texts()
|
||||
max_seq_length = 256
|
||||
|
||||
# Control FIRST, before importing unsloth, so its global import patches never
|
||||
# touch the stock reference (mirrors the issue's "restart runtime" repro).
|
||||
ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype})
|
||||
ctrl.max_seq_length = max_seq_length
|
||||
ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist()
|
||||
ctrl_emb = np.asarray(
|
||||
ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
|
||||
)
|
||||
|
||||
import unsloth # noqa: F401
|
||||
from unsloth import FastSentenceTransformer
|
||||
|
||||
fast = FastSentenceTransformer.from_pretrained(
|
||||
model_id,
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = dtype,
|
||||
load_in_4bit = False,
|
||||
load_in_16bit = True,
|
||||
)
|
||||
fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist()
|
||||
fast_emb = np.asarray(
|
||||
fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
|
||||
)
|
||||
|
||||
# Identical tokenization = no chat-template wrapping slipped in (the #6881 defect).
|
||||
assert fast_ids == ctrl_ids, (
|
||||
f"tokenization diverged (chat-template wrapping regressed?):\n"
|
||||
f" stock: {ctrl_ids}\n fast: {fast_ids}"
|
||||
)
|
||||
|
||||
cos = (ctrl_emb * fast_emb).sum(1) / (
|
||||
np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1)
|
||||
)
|
||||
assert float(cos.min()) > 0.99, (
|
||||
f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 "
|
||||
f"(per-text {[round(float(c), 5) for c in cos]})"
|
||||
)
|
||||
46
tests/python/test_remove_special_tokens_no_bos.py
Normal file
46
tests/python/test_remove_special_tokens_no_bos.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_remove_special_tokens():
|
||||
# Extract remove_special_tokens without importing unsloth (importing unsloth
|
||||
# needs unsloth_zoo / a GPU). The function is pure Python and uses no imports,
|
||||
# so it execs cleanly in an empty namespace.
|
||||
source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py"
|
||||
tree = ast.parse(source.read_text(encoding = "utf-8"))
|
||||
funcs = [
|
||||
node
|
||||
for node in tree.body
|
||||
if isinstance(node, ast.FunctionDef) and node.name == "remove_special_tokens"
|
||||
]
|
||||
namespace = {}
|
||||
module = ast.Module(body = funcs, type_ignores = [])
|
||||
ast.fix_missing_locations(module)
|
||||
exec(compile(module, str(source), "exec"), namespace)
|
||||
return namespace["remove_special_tokens"]
|
||||
|
||||
|
||||
class _StubTokenizer:
|
||||
def __init__(self, bos_token):
|
||||
self.bos_token = bos_token
|
||||
|
||||
|
||||
def test_no_bos_tokenizer_does_not_crash():
|
||||
# Tokenizers such as Qwen2 / Qwen2.5, GPT-2, Falcon and GPT-NeoX have no BOS
|
||||
# token, so tokenizer.bos_token is None. remove_special_tokens must leave the
|
||||
# prompt untouched instead of raising
|
||||
# "TypeError: startswith first arg must be str or a tuple of str, not NoneType".
|
||||
remove_special_tokens = _load_remove_special_tokens()
|
||||
assert remove_special_tokens(_StubTokenizer(None), "Hello world") == "Hello world"
|
||||
|
||||
|
||||
def test_double_bos_is_stripped():
|
||||
# A tokenizer with a BOS token still has a single leading BOS removed.
|
||||
remove_special_tokens = _load_remove_special_tokens()
|
||||
assert remove_special_tokens(_StubTokenizer("<s>"), "<s>Hello world") == "Hello world"
|
||||
|
||||
|
||||
def test_prompt_without_leading_bos_unchanged():
|
||||
# A BOS-bearing tokenizer leaves a prompt that does not start with BOS alone.
|
||||
remove_special_tokens = _load_remove_special_tokens()
|
||||
assert remove_special_tokens(_StubTokenizer("<s>"), "Hello world") == "Hello world"
|
||||
97
tests/python/test_to_sharegpt_optional_none.py
Normal file
97
tests/python/test_to_sharegpt_optional_none.py
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
import ast
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_formatter_builders():
|
||||
# Extract _parse_combined_prompt and _create_formatter without importing
|
||||
# unsloth (importing unsloth needs unsloth_zoo / a GPU). Both are pure
|
||||
# Python and only use the `re` module.
|
||||
source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py"
|
||||
tree = ast.parse(source.read_text(encoding = "utf-8"))
|
||||
wanted = {"_parse_combined_prompt", "_create_formatter"}
|
||||
funcs = [
|
||||
node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted
|
||||
]
|
||||
namespace = {"re": re}
|
||||
module = ast.Module(body = funcs, type_ignores = [])
|
||||
ast.fix_missing_locations(module)
|
||||
exec(compile(module, str(source), "exec"), namespace)
|
||||
return namespace["_parse_combined_prompt"], namespace["_create_formatter"]
|
||||
|
||||
|
||||
class _StubDataset:
|
||||
def __init__(self, column_names):
|
||||
self.column_names = column_names
|
||||
|
||||
|
||||
def _render(merged_prompt, columns, batch):
|
||||
parse, create = _load_formatter_builders()
|
||||
possible_columns, final_optional_prompts = parse(merged_prompt, _StubDataset(columns))
|
||||
processor = create(possible_columns, final_optional_prompts, "text")
|
||||
return processor(batch)["text"]
|
||||
|
||||
|
||||
def test_optional_block_missing_second_column_does_not_render_none():
|
||||
# A [[...]] block may reference several columns; only the first gates the
|
||||
# block. A later column that is None must not render as the literal "None".
|
||||
merged_prompt = "Location: [[{city}, {country}]] end"
|
||||
out = _render(
|
||||
merged_prompt,
|
||||
["city", "country"],
|
||||
{"city": ["Paris"], "country": [None]},
|
||||
)
|
||||
assert out[0] == "Location: Paris, end"
|
||||
assert "None" not in out[0]
|
||||
|
||||
|
||||
def test_optional_block_all_columns_present_unchanged():
|
||||
merged_prompt = "Location: [[{city}, {country}]] end"
|
||||
out = _render(
|
||||
merged_prompt,
|
||||
["city", "country"],
|
||||
{"city": ["Paris"], "country": ["France"]},
|
||||
)
|
||||
assert out[0] == "Location: Paris, France end"
|
||||
|
||||
|
||||
def test_optional_block_gating_column_empty_is_dropped():
|
||||
# When the gating (first) column is empty the whole block is omitted; this
|
||||
# behaviour is unchanged by the None coercion.
|
||||
merged_prompt = "Location: [[{city}, {country}]] end"
|
||||
out = _render(
|
||||
merged_prompt,
|
||||
["city", "country"],
|
||||
{"city": [""], "country": ["France"]},
|
||||
)
|
||||
assert out[0] == "Location: end"
|
||||
|
||||
|
||||
def test_single_column_optional_block_gated_out_on_none():
|
||||
# Single-column blocks were already gated correctly (the sole column is the
|
||||
# gate); confirm they stay unaffected.
|
||||
merged_prompt = "Name: [[{name}]]!"
|
||||
out = _render(merged_prompt, ["name"], {"name": [None, "Bob"]})
|
||||
assert out == ["Name: !", "Name: Bob!"]
|
||||
|
||||
|
||||
def test_required_column_none_does_not_render_none():
|
||||
# A required (non-[[...]]) column that is None must not render as the
|
||||
# literal "None" either; coercion happens at the row source, so both the
|
||||
# required and optional branches are covered.
|
||||
merged_prompt = "Location: {city}, {country} end"
|
||||
out = _render(
|
||||
merged_prompt,
|
||||
["city", "country"],
|
||||
{"city": ["Paris"], "country": [None]},
|
||||
)
|
||||
assert out[0] == "Location: Paris, end"
|
||||
assert "None" not in out[0]
|
||||
|
||||
|
||||
def test_optional_block_falsy_but_present_gating_value_still_renders():
|
||||
# The gate keeps a block whenever the first column is not "". A falsy but
|
||||
# real value (0) must not be treated as absent, so the block still renders.
|
||||
merged_prompt = "Count: [[{n}]]!"
|
||||
out = _render(merged_prompt, ["n"], {"n": [0]})
|
||||
assert out[0] == "Count: 0!"
|
||||
|
|
@ -14,6 +14,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/
|
|||
_REPO_ROOT = _TESTS_DIR.parent # unsloth/
|
||||
_INSTALL_SH = _REPO_ROOT / "install.sh"
|
||||
_INSTALL_PS1 = _REPO_ROOT / "install.ps1"
|
||||
_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1"
|
||||
_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
|
||||
|
||||
|
||||
|
|
@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged:
|
|||
assert '"torch>=2.4,<2.11.0"' in self._ps1
|
||||
|
||||
|
||||
class TestInstallPs1UvDefaultIndex:
|
||||
"""Installer-managed torch indexes must override inherited uv defaults."""
|
||||
|
||||
_ps1 = _read(_INSTALL_PS1)
|
||||
|
||||
def test_torch_installs_use_default_index(self):
|
||||
assert "--default-index $TorchIndexUrl" in self._ps1
|
||||
assert "--default-index $ROCmIndexUrl" in self._ps1
|
||||
|
||||
def test_torch_installs_do_not_use_deprecated_index_url(self):
|
||||
assert "--index-url $TorchIndexUrl" not in self._ps1
|
||||
assert "--index-url $ROCmIndexUrl" not in self._ps1
|
||||
|
||||
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
|
||||
# Extra-index vars outrank --default-index, so pinned installs must clear them.
|
||||
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
|
||||
assert var in self._ps1
|
||||
assert 'Remove-Item "Env:$n"' in self._ps1
|
||||
|
||||
|
||||
class TestSetupPs1FastInstallIndex:
|
||||
"""setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning."""
|
||||
|
||||
_ps1 = _read(_SETUP_PS1)
|
||||
|
||||
def test_fast_install_clears_all_uv_index_env_vars(self):
|
||||
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
|
||||
assert var in self._ps1
|
||||
# Must truly remove the vars (child sees no value), not set them empty.
|
||||
assert 'Remove-Item "Env:$n"' in self._ps1
|
||||
|
||||
|
||||
class TestInstallShUvDefaultIndex:
|
||||
"""Linux/Mac installer torch indexes must override inherited uv defaults."""
|
||||
|
||||
_sh = _read(_INSTALL_SH)
|
||||
|
||||
def test_torch_installs_use_default_index(self):
|
||||
assert '--default-index "$TORCH_INDEX_URL"' in self._sh
|
||||
|
||||
def test_torch_installs_do_not_use_deprecated_index_url(self):
|
||||
assert '--index-url "$TORCH_INDEX_URL"' not in self._sh
|
||||
|
||||
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
|
||||
# --default-index installs run with all uv index env vars unset via `env -u`.
|
||||
assert (
|
||||
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
|
||||
)
|
||||
|
||||
|
||||
# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
|
||||
class TestTorchConstraintShell:
|
||||
"""Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""
|
||||
|
|
|
|||
|
|
@ -312,7 +312,7 @@ if [ "$SKIP_TORCH" = true ]; then
|
|||
else
|
||||
echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
|
||||
uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
|
||||
--index-url "$TORCH_INDEX_URL"
|
||||
--default-index "$TORCH_INDEX_URL"
|
||||
fi
|
||||
TORCH_EOF
|
||||
|
||||
|
|
|
|||
|
|
@ -2684,68 +2684,6 @@ class TestBlackwellCuda124Exclusion:
|
|||
assert kept == [cpu]
|
||||
|
||||
|
||||
# N.1c3. direct_linux_release_plan -- no silent CPU on NVIDIA hosts
|
||||
|
||||
|
||||
class TestDirectLinuxNvidiaCpuGate:
|
||||
"""A linux-cpu-only release on an NVIDIA host must raise (caller walks back to a usable CUDA line), not silently CPU-install. CPU-only hosts keep the CPU bundle."""
|
||||
|
||||
def _bundle_cpu_only(self):
|
||||
return make_release(
|
||||
[
|
||||
make_artifact(
|
||||
"llama-b8508-bin-ubuntu-x64.tar.gz",
|
||||
install_kind = "linux-cpu",
|
||||
runtime_line = None,
|
||||
coverage_class = None,
|
||||
supported_sms = [],
|
||||
min_sm = None,
|
||||
max_sm = None,
|
||||
bundle_profile = None,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def _patch(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
INSTALL_LLAMA_PREBUILT,
|
||||
"parse_direct_linux_release_bundle",
|
||||
lambda repo, release: self._bundle_cpu_only(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
INSTALL_LLAMA_PREBUILT,
|
||||
"detect_torch_cuda_runtime_preference",
|
||||
lambda host: CudaRuntimePreference(runtime_line = None, selection_log = []),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
INSTALL_LLAMA_PREBUILT,
|
||||
"detected_linux_runtime_lines",
|
||||
lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
|
||||
)
|
||||
|
||||
def test_nvidia_host_without_cuda_line_raises_for_walkback(self, monkeypatch):
|
||||
self._patch(monkeypatch)
|
||||
host = make_host(driver_cuda_version = (13, 1), compute_caps = ["100"])
|
||||
with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt"):
|
||||
INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
|
||||
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
|
||||
)
|
||||
|
||||
def test_cpu_host_still_gets_cpu_bundle(self, monkeypatch):
|
||||
self._patch(monkeypatch)
|
||||
host = make_host(
|
||||
nvidia_smi = None,
|
||||
driver_cuda_version = None,
|
||||
compute_caps = [],
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
)
|
||||
plan = INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
|
||||
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
|
||||
)
|
||||
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
|
||||
|
||||
|
||||
class TestLinuxPublishedAttemptsNvidiaCpuGate:
|
||||
"""Live fork-manifest path: an NVIDIA host whose CUDA selection finds nothing gets an empty attempt list (source-builds with CUDA), not the manifest CPU bundle. CPU-only hosts still get the CPU bundle."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1264,22 +1264,124 @@ with sync_playwright() as p:
|
|||
# placeholder, and /api/health goes unreachable shortly after.
|
||||
# ─────────────────────────────────────────────────────
|
||||
step("Shutdown via account menu")
|
||||
# Re-login with NEW2 for a valid /api/shutdown token (CLI rotation
|
||||
# invalidated the old one). The stale token can make the SPA auth guard
|
||||
# abort this goto with ERR_ABORTED, or redirect to the same /login URL
|
||||
# ("interrupted by another navigation"); resolve on domcontentloaded and
|
||||
# tolerate either -- the pw-field wait below confirms we are on /login.
|
||||
_tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
|
||||
# Start fresh after the CLI rotation invalidates this browser session.
|
||||
# Stay in the SAME context: macOS Chromium runs --single-process, where
|
||||
# closing the last context kills the browser and a second context cannot
|
||||
# be created. Open the new page before closing the old one; the context
|
||||
# init script covers the new page.
|
||||
try:
|
||||
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
|
||||
ctx.clear_cookies()
|
||||
except Exception as exc:
|
||||
if not any(t in str(exc) for t in _tolerated_nav):
|
||||
raise
|
||||
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
|
||||
pw_field = page.locator("#password")
|
||||
pw_field.wait_for(state = "visible", timeout = 60_000)
|
||||
pw_field.fill(NEW2)
|
||||
page.locator('button[type="submit"]').click()
|
||||
info(f"WARN clearing stale session cookies failed: {exc!r}")
|
||||
# Auth tokens live in localStorage, and /login's guest guard redirects on
|
||||
# their mere presence, so drop them before navigating.
|
||||
try:
|
||||
page.evaluate(
|
||||
"['unsloth_auth_token', 'unsloth_auth_refresh_token']"
|
||||
".forEach((key) => localStorage.removeItem(key))"
|
||||
)
|
||||
except Exception as exc:
|
||||
info(f"WARN clearing stale auth tokens failed: {exc!r}")
|
||||
_fresh_page = ctx.new_page()
|
||||
_fresh_page.set_default_timeout(60_000)
|
||||
_fresh_page.on("pageerror", lambda e: page_errors.append(str(e)))
|
||||
_fresh_page.on("console", _on_console)
|
||||
try:
|
||||
page.close()
|
||||
except Exception:
|
||||
pass
|
||||
page = _fresh_page
|
||||
|
||||
# Re-login with NEW2 for a valid /api/shutdown token. Route changes can
|
||||
# still abort or interrupt this navigation, so the field wait below is the
|
||||
# final confirmation that we reached /login.
|
||||
_tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
|
||||
# A slow CI runner can make this re-login navigation time out even with the
|
||||
# server healthy, so retry the whole goto/wait/fill/submit sequence (mirrors
|
||||
# the change-password retry above). wait_for_health is a diagnostic pre-gate.
|
||||
wait_for_health(BASE, timeout = 30.0, info = info)
|
||||
relogin_err: Exception | None = None
|
||||
for _relogin_attempt in range(3):
|
||||
try:
|
||||
try:
|
||||
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
|
||||
except Exception as exc:
|
||||
if not any(t in str(exc) for t in _tolerated_nav):
|
||||
raise
|
||||
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
|
||||
pw_field = page.locator("#password")
|
||||
pw_field.wait_for(state = "visible", timeout = 60_000)
|
||||
pw_field.fill(NEW2)
|
||||
# Wait on the login POST so a transient 4xx/5xx is caught and retried
|
||||
# here, not swallowed until the out-of-loop composer wait.
|
||||
status, _ = click_and_wait_for_response(
|
||||
page,
|
||||
url_substr = "/api/auth/login",
|
||||
method = "POST",
|
||||
do_click = lambda: page.locator('button[type="submit"]').click(),
|
||||
timeout_ms = 30_000,
|
||||
info = lambda m: print(f"[ui] {m}", flush = True),
|
||||
)
|
||||
if status is not None and status >= 400:
|
||||
raise AssertionError(
|
||||
f"login POST returned {status}; see console_errors={console_errors[:1]!r}"
|
||||
)
|
||||
relogin_err = None
|
||||
break
|
||||
except Exception as e:
|
||||
relogin_err = e
|
||||
try:
|
||||
cur_url = page.url
|
||||
except Exception:
|
||||
cur_url = "<page closed>"
|
||||
print(
|
||||
f"[ui] re-login attempt {_relogin_attempt + 1} failed: "
|
||||
f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; "
|
||||
f"page_errors={len(page_errors)} console_errors={len(console_errors)}",
|
||||
flush = True,
|
||||
)
|
||||
if console_errors:
|
||||
print(
|
||||
f"[ui] first console.error: {console_errors[0][:200]!r}",
|
||||
flush = True,
|
||||
)
|
||||
if page_errors:
|
||||
print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True)
|
||||
try:
|
||||
shoot(f"18-relogin-attempt-{_relogin_attempt + 1}-fail")
|
||||
except Exception:
|
||||
pass
|
||||
if _relogin_attempt < 2:
|
||||
# ERR_NO_BUFFER_SPACE needs the OS to recover socket
|
||||
# buffers; back off 5s then 15s before retrying.
|
||||
if "ERR_NO_BUFFER_SPACE" in str(e):
|
||||
backoff_s = 5 if _relogin_attempt == 0 else 15
|
||||
print(
|
||||
f"[ui] ENOBUFS detected; sleeping {backoff_s}s "
|
||||
f"before retry to let OS recover socket buffers...",
|
||||
flush = True,
|
||||
)
|
||||
time.sleep(backoff_s)
|
||||
# Replace the page if it died; otherwise next iteration's
|
||||
# page.goto() handles the reload.
|
||||
old_page = page
|
||||
page = recover_or_replace_page(
|
||||
page,
|
||||
ctx,
|
||||
default_timeout_ms = 60_000,
|
||||
info = lambda m: print(f"[ui] recovery: {m}", flush = True),
|
||||
)
|
||||
# A freshly created replacement page loses the pageerror/console
|
||||
# listeners; re-attach so error tracking survives recovery.
|
||||
if page is not old_page:
|
||||
page.on("pageerror", lambda e: page_errors.append(str(e)))
|
||||
page.on("console", _on_console)
|
||||
if relogin_err is not None:
|
||||
raise relogin_err
|
||||
# Composer mount confirms the rotated session is authenticated. Kept OUTSIDE the
|
||||
# retry: the loop breaks right after submit, so we never re-goto /login once login
|
||||
# has set tokens -- that would hit the guest guard, redirect to /chat, and make a
|
||||
# merely-slow composer look like a broken login.
|
||||
composer = page.locator('textarea[aria-label="Message input"]')
|
||||
composer.wait_for(state = "visible", timeout = 60_000)
|
||||
shoot("18-relogin-with-NEW2")
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from pathlib import Path
|
|||
|
||||
|
||||
SOURCE_PATH = Path(__file__).resolve().parents[2] / "studio" / "backend" / "routes" / "inference.py"
|
||||
SRC = SOURCE_PATH.read_text()
|
||||
SRC = SOURCE_PATH.read_text(encoding = "utf-8")
|
||||
_TREE = ast.parse(SRC)
|
||||
|
||||
|
||||
|
|
@ -166,16 +166,20 @@ def test_chat_completions_streams_avoid_starlette_task_group():
|
|||
|
||||
|
||||
def test_openai_passthrough_stream_avoids_starlette_task_group():
|
||||
top = _async_function("_openai_passthrough_stream")
|
||||
functions = [
|
||||
_async_function("_openai_passthrough_stream"),
|
||||
_async_function("_openai_passthrough_stream_admitted"),
|
||||
]
|
||||
legacy_calls = []
|
||||
same_task_calls = 0
|
||||
for sub in ast.walk(top):
|
||||
if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)):
|
||||
continue
|
||||
if sub.func.id == "StreamingResponse":
|
||||
legacy_calls.append(sub.lineno)
|
||||
if sub.func.id == "_SameTaskStreamingResponse":
|
||||
same_task_calls += 1
|
||||
for fn in functions:
|
||||
for sub in ast.walk(fn):
|
||||
if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)):
|
||||
continue
|
||||
if sub.func.id == "StreamingResponse":
|
||||
legacy_calls.append(sub.lineno)
|
||||
if sub.func.id == "_SameTaskStreamingResponse":
|
||||
same_task_calls += 1
|
||||
assert not legacy_calls, (
|
||||
"OpenAI passthrough streams must use _SameTaskStreamingResponse, "
|
||||
"not Starlette's legacy task-group StreamingResponse. Lines: "
|
||||
|
|
@ -197,7 +201,7 @@ def test_direct_llama_server_streams_install_disconnect_watcher():
|
|||
"openai_completions",
|
||||
"_responses_stream",
|
||||
"_anthropic_passthrough_stream",
|
||||
"_openai_passthrough_stream",
|
||||
"_openai_passthrough_stream_admitted",
|
||||
}
|
||||
missing = [
|
||||
name
|
||||
|
|
|
|||
47
tests/test_bad_mappings_redirect.py
Normal file
47
tests/test_bad_mappings_redirect.py
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
"""Regression test for BAD_MAPPINGS redirecting oversized dynamic quants.
|
||||
|
||||
get_model_name previously applied BAD_MAPPINGS only to the resolver's output,
|
||||
but several listed names (the `-unsloth-bnb-4bit` dynamic quants, plus any name
|
||||
the resolver doesn't map) come back as None, so their BAD_MAPPINGS entries were
|
||||
dead and the oversized model loaded. Asserting over every entry catches all of
|
||||
them. The mapper table and the resolver have no heavy imports of their own,
|
||||
so we exec the import-free mapper module and ast-extract the resolver functions
|
||||
rather than importing unsloth (which needs a GPU).
|
||||
"""
|
||||
|
||||
import ast
|
||||
import os
|
||||
|
||||
_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models")
|
||||
|
||||
|
||||
def _load_get_model_name():
|
||||
mapper_ns = {}
|
||||
with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f:
|
||||
exec(compile(f.read(), "mapper.py", "exec"), mapper_ns)
|
||||
|
||||
with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f:
|
||||
tree = ast.parse(f.read())
|
||||
|
||||
namespace = dict(mapper_ns)
|
||||
namespace["SUPPORTS_FOURBIT"] = True
|
||||
namespace["_env_says_offline"] = lambda: True
|
||||
namespace["_get_new_mapper"] = lambda: ({}, {}, {})
|
||||
|
||||
wanted = {"__get_model_name", "_resolve_with_mappers", "get_model_name"}
|
||||
for node in tree.body:
|
||||
if isinstance(node, ast.Assign) and any(
|
||||
getattr(target, "id", None) == "BAD_MAPPINGS" for target in node.targets
|
||||
):
|
||||
exec(compile(ast.Module([node], []), "<bad_mappings>", "exec"), namespace)
|
||||
elif isinstance(node, ast.FunctionDef) and node.name in wanted:
|
||||
exec(compile(ast.Module([node], []), node.name, "exec"), namespace)
|
||||
|
||||
return namespace["get_model_name"], namespace["BAD_MAPPINGS"]
|
||||
|
||||
|
||||
def test_bad_mappings_redirect_every_listed_name():
|
||||
get_model_name, bad_mappings = _load_get_model_name()
|
||||
assert bad_mappings, "BAD_MAPPINGS should not be empty"
|
||||
for name, expected in bad_mappings.items():
|
||||
assert get_model_name(name, load_in_4bit = True) == expected, name
|
||||
63
tests/test_fast_gemv_dispatch.py
Normal file
63
tests/test_fast_gemv_dispatch.py
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is
|
||||
already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the
|
||||
bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up
|
||||
# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection
|
||||
# errors instead of producing a skip. Any other import error still surfaces as a failure.
|
||||
pytest.importorskip("bitsandbytes")
|
||||
|
||||
import unsloth # noqa: F401 (sets UNSLOTH_IS_PRESENT before transformers)
|
||||
from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES
|
||||
|
||||
_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None
|
||||
|
||||
|
||||
def _proj(weight, weight_scale = None):
|
||||
proj = SimpleNamespace(weight = weight, bias = None, merged = False)
|
||||
if weight_scale is not None:
|
||||
proj.weight_scale = weight_scale
|
||||
return proj
|
||||
|
||||
|
||||
def test_bf16_weight_scale_not_used_as_quant_state():
|
||||
"""A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None."""
|
||||
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2))
|
||||
W, W_quant = get_lora_parameters_bias(proj)[:2]
|
||||
assert W_quant is None
|
||||
|
||||
|
||||
def test_fp8_weight_keeps_scale():
|
||||
"""An actual fp8 weight still resolves its weight_scale as the quant state."""
|
||||
if _FP8 is None:
|
||||
pytest.skip("no float8 dtype in this torch build")
|
||||
scale = torch.rand(2, 2)
|
||||
proj = _proj(torch.randn(4, 4).to(_FP8), scale)
|
||||
W, W_quant = get_lora_parameters_bias(proj)[:2]
|
||||
assert W_quant is scale
|
||||
|
||||
|
||||
def test_plain_bf16_has_no_quant_state():
|
||||
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16))
|
||||
W, W_quant = get_lora_parameters_bias(proj)[:2]
|
||||
assert W_quant is None
|
||||
249
tests/test_fp8_device_context.py
Normal file
249
tests/test_fp8_device_context.py
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
FP8_SOURCE = REPO_ROOT / "unsloth" / "kernels" / "fp8.py"
|
||||
|
||||
|
||||
class _FakeDeviceModule:
|
||||
def __init__(self, device_count: int) -> None:
|
||||
self._device_count = device_count
|
||||
self.device_calls = []
|
||||
|
||||
def device_count(self) -> int:
|
||||
return self._device_count
|
||||
|
||||
def device(self, device):
|
||||
self.device_calls.append(device)
|
||||
return ("device-context", device)
|
||||
|
||||
|
||||
class _FakeTorch:
|
||||
Tensor = object
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cuda_device_count: int,
|
||||
xpu_device_count: int = 0,
|
||||
) -> None:
|
||||
self.cuda = _FakeDeviceModule(cuda_device_count)
|
||||
self.xpu = _FakeDeviceModule(xpu_device_count)
|
||||
|
||||
|
||||
class _LaunchVisitor(ast.NodeVisitor):
|
||||
def __init__(self) -> None:
|
||||
self.guarded_launches: set[str] = set()
|
||||
self.unguarded_launches: set[str] = set()
|
||||
self._inside_fp8_device_context = 0
|
||||
|
||||
def visit_With(self, node: ast.With) -> None:
|
||||
enters_context = any(
|
||||
isinstance(item.context_expr, ast.Call)
|
||||
and isinstance(item.context_expr.func, ast.Name)
|
||||
and item.context_expr.func.id == "_fp8_triton_device_context"
|
||||
for item in node.items
|
||||
)
|
||||
if enters_context:
|
||||
self._inside_fp8_device_context += 1
|
||||
for statement in node.body:
|
||||
self.visit(statement)
|
||||
if enters_context:
|
||||
self._inside_fp8_device_context -= 1
|
||||
|
||||
def visit_Call(self, node: ast.Call) -> None:
|
||||
launch_name = self._triton_launch_name(node)
|
||||
if launch_name is not None:
|
||||
if self._inside_fp8_device_context:
|
||||
self.guarded_launches.add(launch_name)
|
||||
else:
|
||||
self.unguarded_launches.add(launch_name)
|
||||
self.generic_visit(node)
|
||||
|
||||
@staticmethod
|
||||
def _triton_launch_name(node: ast.Call) -> str | None:
|
||||
if isinstance(node.func, ast.Name) and node.func.id == "triton_quantize_fp8_block":
|
||||
return node.func.id
|
||||
if not isinstance(node.func, ast.Subscript):
|
||||
return None
|
||||
if not isinstance(node.func.value, ast.Name):
|
||||
return None
|
||||
return node.func.value.id
|
||||
|
||||
|
||||
def _load_device_context_helper(fake_torch: _FakeTorch):
|
||||
source = FP8_SOURCE.read_text()
|
||||
tree = ast.parse(source)
|
||||
for node in tree.body:
|
||||
if isinstance(node, ast.FunctionDef) and node.name == "_fp8_triton_device_context":
|
||||
namespace = {"torch": fake_torch, "nullcontext": nullcontext}
|
||||
exec(ast.get_source_segment(source, node), namespace)
|
||||
return namespace["_fp8_triton_device_context"]
|
||||
raise AssertionError("_fp8_triton_device_context was not found")
|
||||
|
||||
|
||||
def test_fp8_device_context_selects_cuda_tensor_device_on_multi_gpu() -> None:
|
||||
fake_torch = _FakeTorch(cuda_device_count = 2)
|
||||
helper = _load_device_context_helper(fake_torch)
|
||||
tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda"))
|
||||
|
||||
context = helper(tensor)
|
||||
|
||||
assert context == ("device-context", tensor.device)
|
||||
assert fake_torch.cuda.device_calls == [tensor.device]
|
||||
|
||||
|
||||
def test_fp8_device_context_is_noop_for_single_cuda_device() -> None:
|
||||
fake_torch = _FakeTorch(cuda_device_count = 1)
|
||||
helper = _load_device_context_helper(fake_torch)
|
||||
tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda"))
|
||||
|
||||
context = helper(tensor)
|
||||
|
||||
assert isinstance(context, nullcontext)
|
||||
assert fake_torch.cuda.device_calls == []
|
||||
|
||||
|
||||
def test_fp8_device_context_selects_xpu_tensor_device_on_multi_gpu() -> None:
|
||||
fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 2)
|
||||
helper = _load_device_context_helper(fake_torch)
|
||||
tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu"))
|
||||
|
||||
context = helper(tensor)
|
||||
|
||||
assert context == ("device-context", tensor.device)
|
||||
assert fake_torch.xpu.device_calls == [tensor.device]
|
||||
|
||||
|
||||
def test_fp8_device_context_is_noop_for_single_xpu_device() -> None:
|
||||
fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 1)
|
||||
helper = _load_device_context_helper(fake_torch)
|
||||
tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu"))
|
||||
|
||||
context = helper(tensor)
|
||||
|
||||
assert isinstance(context, nullcontext)
|
||||
assert fake_torch.xpu.device_calls == []
|
||||
|
||||
|
||||
def test_fp8_device_context_is_noop_for_non_cuda_tensor() -> None:
|
||||
fake_torch = _FakeTorch(cuda_device_count = 8)
|
||||
helper = _load_device_context_helper(fake_torch)
|
||||
tensor = SimpleNamespace(device = SimpleNamespace(type = "cpu"))
|
||||
|
||||
context = helper(tensor)
|
||||
|
||||
assert isinstance(context, nullcontext)
|
||||
assert fake_torch.cuda.device_calls == []
|
||||
|
||||
|
||||
def test_fp8_triton_launches_enter_tensor_device_context() -> None:
|
||||
tree = ast.parse(FP8_SOURCE.read_text())
|
||||
function_names = {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)}
|
||||
assert "_fp8_triton_device_context" in function_names
|
||||
|
||||
visitor = _LaunchVisitor()
|
||||
visitor.visit(tree)
|
||||
|
||||
expected_launches = {
|
||||
"weight_dequant_kernel",
|
||||
"act_quant_kernel",
|
||||
"_w8a8_block_fp8_matmul",
|
||||
"triton_quantize_fp8_block",
|
||||
}
|
||||
assert expected_launches <= visitor.guarded_launches
|
||||
assert not (expected_launches & visitor.unguarded_launches)
|
||||
|
||||
|
||||
def _require_two_cuda_devices():
|
||||
torch = pytest.importorskip("torch")
|
||||
pytest.importorskip("triton")
|
||||
|
||||
if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
|
||||
pytest.skip("requires at least two CUDA devices")
|
||||
return torch
|
||||
|
||||
|
||||
def test_weight_dequant_block_runs_on_tensor_device_when_current_device_differs() -> None:
|
||||
torch = _require_two_cuda_devices()
|
||||
from unsloth.kernels.fp8 import weight_dequant_block
|
||||
|
||||
previous_device = torch.cuda.current_device()
|
||||
try:
|
||||
torch.cuda.set_device(0)
|
||||
x = torch.arange(256 * 256, device = "cuda:1", dtype = torch.float32).reshape(256, 256)
|
||||
scales = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device = "cuda:1", dtype = torch.float32)
|
||||
|
||||
actual = weight_dequant_block(x, scales, block_size = 128, dtype = torch.float32)
|
||||
|
||||
expanded_scales = scales.repeat_interleave(128, dim = 0).repeat_interleave(128, dim = 1)
|
||||
expected = x * expanded_scales
|
||||
|
||||
assert actual.device == x.device
|
||||
assert torch.cuda.current_device() == 0
|
||||
torch.testing.assert_close(actual, expected)
|
||||
finally:
|
||||
torch.cuda.set_device(previous_device)
|
||||
|
||||
|
||||
def test_act_quant_runs_on_tensor_device_when_current_device_differs() -> None:
|
||||
torch = _require_two_cuda_devices()
|
||||
if not hasattr(torch, "float8_e4m3fn"):
|
||||
pytest.skip("requires torch.float8_e4m3fn")
|
||||
if torch.cuda.get_device_capability(1)[0] < 9:
|
||||
pytest.skip("requires FP8-capable CUDA hardware")
|
||||
|
||||
from unsloth.kernels.fp8 import act_quant
|
||||
|
||||
previous_device = torch.cuda.current_device()
|
||||
try:
|
||||
torch.cuda.set_device(0)
|
||||
x = torch.arange(256, device = "cuda:1", dtype = torch.float32).reshape(2, 128)
|
||||
|
||||
y, scales = act_quant(x, block_size = 128)
|
||||
|
||||
assert y.device == x.device
|
||||
assert scales.device == x.device
|
||||
assert torch.cuda.current_device() == 0
|
||||
finally:
|
||||
torch.cuda.set_device(previous_device)
|
||||
|
||||
|
||||
def test_w8a8_block_fp8_matmul_triton_runs_on_tensor_device_when_current_device_differs() -> None:
|
||||
torch = _require_two_cuda_devices()
|
||||
if not hasattr(torch, "float8_e4m3fn"):
|
||||
pytest.skip("requires torch.float8_e4m3fn")
|
||||
if torch.cuda.get_device_capability(1)[0] < 9:
|
||||
pytest.skip("requires FP8-capable CUDA hardware")
|
||||
|
||||
from unsloth.kernels.fp8 import w8a8_block_fp8_matmul_triton
|
||||
|
||||
previous_device = torch.cuda.current_device()
|
||||
try:
|
||||
torch.cuda.set_device(0)
|
||||
A = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn)
|
||||
B = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn)
|
||||
As = torch.ones((128, 1), device = "cuda:1", dtype = torch.float32)
|
||||
Bs = torch.ones((1, 1), device = "cuda:1", dtype = torch.float32)
|
||||
|
||||
actual = w8a8_block_fp8_matmul_triton(
|
||||
A,
|
||||
B,
|
||||
As,
|
||||
Bs,
|
||||
block_size = [128, 128],
|
||||
output_dtype = torch.float32,
|
||||
)
|
||||
|
||||
expected = torch.full((128, 128), 128.0, device = "cuda:1", dtype = torch.float32)
|
||||
assert actual.device == A.device
|
||||
assert torch.cuda.current_device() == 0
|
||||
torch.testing.assert_close(actual, expected)
|
||||
finally:
|
||||
torch.cuda.set_device(previous_device)
|
||||
358
tests/test_fp8_restore_dropped_scale.py
Normal file
358
tests/test_fp8_restore_dropped_scale.py
Normal file
|
|
@ -0,0 +1,358 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Restoring dropped block-fp8 `weight_scale_inv` tensors on load (#6200).
|
||||
|
||||
Some block-scale fp8 checkpoints leave a Linear (e.g. `mlp.gate_proj`) unconverted, so its raw
|
||||
quantized values land in a plain bf16 weight and its `weight_scale_inv` is dropped, producing a
|
||||
garbage un-scaled weight. `_restore_dropped_fp8_scales` dequantizes such orphaned weights in place
|
||||
using the scale from the checkpoint. Runs offline on CPU with synthetic checkpoints.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from safetensors.torch import save_file
|
||||
|
||||
# Import unsloth first to set UNSLOTH_IS_PRESENT env var.
|
||||
import unsloth
|
||||
from unsloth.models.loader_utils import _restore_dropped_fp8_scales, _FP8_DTYPES
|
||||
|
||||
|
||||
_SHARD = "model-00001-of-00001.safetensors"
|
||||
_FP8 = _FP8_DTYPES[0] if _FP8_DTYPES else None
|
||||
|
||||
|
||||
def _write_checkpoint(
|
||||
path,
|
||||
tensors,
|
||||
filename = _SHARD,
|
||||
include_index = True,
|
||||
):
|
||||
save_file(tensors, os.path.join(path, filename))
|
||||
if include_index:
|
||||
weight_map = {name: filename for name in tensors}
|
||||
with open(os.path.join(path, "model.safetensors.index.json"), "w") as f:
|
||||
json.dump({"weight_map": weight_map}, f)
|
||||
|
||||
|
||||
def _fp8_config(block = (2, 2)):
|
||||
return SimpleNamespace(
|
||||
quantization_config = {
|
||||
"quant_method": "fp8",
|
||||
"weight_block_size": list(block),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _fp8_anchor():
|
||||
"""A module carrying a real fp8 weight, so the model looks like a genuine fp8 load."""
|
||||
m = nn.Linear(2, 2, bias = False)
|
||||
m.weight = nn.Parameter(torch.randn(2, 2).to(_FP8), requires_grad = False)
|
||||
return m
|
||||
|
||||
|
||||
def _bf16_linear(out_f, in_f, raw):
|
||||
m = nn.Linear(in_f, out_f, bias = False).to(torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
m.weight.copy_(raw)
|
||||
return m
|
||||
|
||||
|
||||
def _expand(scale, block, shape):
|
||||
bs0, bs1 = block
|
||||
expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
|
||||
return expanded[: shape[0], : shape[1]]
|
||||
|
||||
|
||||
def test_restore_dequantizes_orphaned_scale():
|
||||
"""A plain bf16 weight whose scale was dropped is dequantized in place."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
torch.manual_seed(0)
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(
|
||||
d,
|
||||
{
|
||||
"layer.weight": raw.to(torch.float32),
|
||||
"layer.weight_scale_inv": scale,
|
||||
},
|
||||
)
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_skips_already_fp8_weight():
|
||||
"""A correctly converted fp8 weight is skipped, never double-scaled."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
weight = torch.randn(4, 4).to(_FP8)
|
||||
before = weight.clone()
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.layer = nn.Linear(4, 4, bias = False)
|
||||
model.layer.weight = nn.Parameter(weight, requires_grad = False)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": torch.rand(2, 2)})
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 0 and skipped == 1
|
||||
assert torch.equal(model.layer.weight.data.float(), before.float())
|
||||
|
||||
|
||||
def test_skips_offloaded_meta_weight():
|
||||
"""A disk-offloaded layer (weight on the meta device) is skipped without error or restore."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = nn.Linear(4, 4, bias = False)
|
||||
# Simulate an offloaded weight living on the meta device.
|
||||
model.layer.weight = nn.Parameter(
|
||||
torch.empty(4, 4, dtype = torch.bfloat16, device = "meta"), requires_grad = False
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(
|
||||
d,
|
||||
{
|
||||
"layer.weight": raw.to(torch.float32),
|
||||
"layer.weight_scale_inv": scale,
|
||||
},
|
||||
)
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 0
|
||||
assert model.layer.weight.device.type == "meta"
|
||||
|
||||
|
||||
def test_noop_when_fully_dequantized():
|
||||
"""If the model has no fp8 weights at all (e.g. load_in_16bit dequantize), do not rescale."""
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.layer = _bf16_linear(4, 4, raw) # no fp8 anchor -> looks dequantized
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert (restored, skipped) == (0, 0)
|
||||
assert torch.equal(model.layer.weight.data, raw) # untouched
|
||||
|
||||
|
||||
def test_non_block_divisible_shape():
|
||||
"""Block scale is expanded then sliced to a non-divisible weight shape."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(3, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(3, 4, raw) # weight shape [3, 4]
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (3, 4))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_transposed_scale_layout():
|
||||
"""A scale stored in the transposed block grid is transposed before use."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 2, dtype = torch.bfloat16) # weight [4, 2] -> grid (2, 1)
|
||||
scale_correct = torch.rand(2, 1, dtype = torch.float32) + 0.1
|
||||
scale_stored = scale_correct.t().contiguous() # stored transposed as (1, 2)
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 2, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale_stored})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale_correct, (2, 2), (4, 2))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_single_file_checkpoint_without_index():
|
||||
"""Unsharded model.safetensors (no index) is still scanned for dropped scales."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(
|
||||
d, {"layer.weight_scale_inv": scale}, filename = "model.safetensors", include_index = False
|
||||
)
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
|
||||
assert torch.equal(model.layer.weight.data, expected)
|
||||
|
||||
|
||||
def test_scalar_block_size_config():
|
||||
"""A scalar weight_block_size (not a list) is handled without error."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = SimpleNamespace(
|
||||
quantization_config = {"quant_method": "fp8", "weight_block_size": 2}
|
||||
)
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
|
||||
|
||||
def test_text_only_prefix_mapping():
|
||||
"""Checkpoint keys with a language_model prefix match the stripped text-only module names."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(2, 2, dtype = torch.bfloat16)
|
||||
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
|
||||
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.model = nn.Module()
|
||||
model.model.gate_proj = _bf16_linear(2, 2, raw) # module lacks the language_model prefix
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
# checkpoint key carries the language_model wrapper the text-only load stripped
|
||||
_write_checkpoint(d, {"model.language_model.gate_proj.weight_scale_inv": scale})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
|
||||
assert torch.equal(model.model.gate_proj.weight.data, expected)
|
||||
|
||||
|
||||
def test_skips_variant_load():
|
||||
"""A variant load (variant="fp8") is skipped to avoid applying default-checkpoint scales."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(4, 4, dtype = torch.bfloat16)
|
||||
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, raw)
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
result = _restore_dropped_fp8_scales(model, d, local_files_only = True, variant = "fp8")
|
||||
assert result == (0, 0)
|
||||
assert torch.equal(model.layer.weight.data, raw) # untouched
|
||||
|
||||
|
||||
def test_vlm_language_model_model_alias():
|
||||
"""A checkpoint key language_model.model.* matches a model.language_model.* module."""
|
||||
if _FP8 is None:
|
||||
return
|
||||
raw = torch.randn(2, 2, dtype = torch.bfloat16)
|
||||
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.model = nn.Module()
|
||||
model.model.language_model = nn.Module()
|
||||
model.model.language_model.gate_proj = _bf16_linear(
|
||||
2, 2, raw
|
||||
) # -> model.language_model.gate_proj
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"language_model.model.gate_proj.weight_scale_inv": scale})
|
||||
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
|
||||
assert restored == 1
|
||||
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
|
||||
assert torch.equal(model.model.language_model.gate_proj.weight.data, expected)
|
||||
|
||||
|
||||
def test_noop_without_scale_keys():
|
||||
if _FP8 is None:
|
||||
return
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight": torch.randn(4, 4)})
|
||||
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
|
||||
|
||||
|
||||
def test_noop_without_index_or_single_file():
|
||||
if _FP8 is None:
|
||||
return
|
||||
model = nn.Module()
|
||||
model.config = _fp8_config((2, 2))
|
||||
model.anchor = _fp8_anchor()
|
||||
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
|
||||
|
||||
|
||||
def test_noop_when_not_block_fp8():
|
||||
"""A non-fp8 (or non-block) quantization config is ignored."""
|
||||
scale = torch.rand(2, 2)
|
||||
model = nn.Module()
|
||||
model.config = SimpleNamespace(quantization_config = {"quant_method": "compressed-tensors"})
|
||||
model.layer = nn.Linear(4, 4, bias = False)
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
|
||||
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
|
||||
46
tests/test_gemma_2b_mapper_key.py
Normal file
46
tests/test_gemma_2b_mapper_key.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
"""Regression test for the duplicate ``unsloth/gemma-2b-bnb-4bit`` key in
|
||||
``unsloth/models/mapper.py``.
|
||||
|
||||
The 4bit instruction-tuned Gemma 2B entry was accidentally keyed with the base
|
||||
model's repo name, so ``__INT_TO_FLOAT_MAPPER`` held two identical
|
||||
``unsloth/gemma-2b-bnb-4bit`` keys. Python keeps only the last value for a
|
||||
duplicate literal key, so the base 4bit repo resolved to the *instruct* model,
|
||||
the base model lost its reverse (4x-faster) mapping, and
|
||||
``unsloth/gemma-2b-it-bnb-4bit`` was never registered at all.
|
||||
|
||||
``mapper.py`` has no imports, so we exec it directly and inspect the built
|
||||
mappers without importing ``unsloth`` (which requires a GPU).
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
MAPPER_PATH = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models", "mapper.py")
|
||||
|
||||
|
||||
def _load_mappers():
|
||||
with open(MAPPER_PATH) as f:
|
||||
source = f.read()
|
||||
namespace = {}
|
||||
exec(compile(source, MAPPER_PATH, "exec"), namespace)
|
||||
return namespace
|
||||
|
||||
|
||||
def test_gemma_2b_base_and_instruct_4bit_are_distinct():
|
||||
namespace = _load_mappers()
|
||||
int_to_float = namespace["INT_TO_FLOAT_MAPPER"]
|
||||
float_to_int = namespace["FLOAT_TO_INT_MAPPER"]
|
||||
|
||||
# The base 4bit repo must resolve to the base model, not the instruct one.
|
||||
assert int_to_float["unsloth/gemma-2b-bnb-4bit"] == "unsloth/gemma-2b"
|
||||
|
||||
# The instruct 4bit repo must be registered and resolve to the instruct model.
|
||||
assert "unsloth/gemma-2b-it-bnb-4bit" in int_to_float
|
||||
assert int_to_float["unsloth/gemma-2b-it-bnb-4bit"] == "unsloth/gemma-2b-it"
|
||||
|
||||
# The base model must reverse-map back to the base 4bit repo.
|
||||
assert float_to_int["unsloth/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit"
|
||||
assert float_to_int["google/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit"
|
||||
|
||||
# The instruct model must reverse-map to the instruct 4bit repo.
|
||||
assert float_to_int["unsloth/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit"
|
||||
assert float_to_int["google/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit"
|
||||
|
|
@ -257,6 +257,39 @@ def test_recompute_helper_scales_on_cpu():
|
|||
), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
|
||||
|
||||
|
||||
def test_extended_rope_scaling_keeps_llama3_and_carries_theta():
|
||||
# Long-context extension keeps native llama3, but falls back to linear for every other
|
||||
# type (the patched attention constructor only rebuilds linear/llama3/longrope), and the
|
||||
# linear dict carries rope_theta so transformers v5 does not fall back to base 10000.
|
||||
from types import SimpleNamespace
|
||||
|
||||
from unsloth.models.llama import _extended_rope_scaling
|
||||
|
||||
# llama3 model: keep native scaling, do not synthesize linear.
|
||||
scaling, native = _extended_rope_scaling(_make_config(LLAMA3_ROPE_SCALING), 2.0)
|
||||
assert (
|
||||
scaling is None and native == "llama3"
|
||||
), "must keep native llama3 scaling instead of overwriting it with linear."
|
||||
|
||||
# yarn is not rebuildable by the patcher -> keep the safe linear fallback, not native.
|
||||
yarn = SimpleNamespace(rope_scaling = {"rope_type": "yarn", "factor": 2.0}, rope_theta = 500000.0)
|
||||
scaling, _ = _extended_rope_scaling(yarn, 2.0)
|
||||
assert scaling == {
|
||||
"type": "linear",
|
||||
"factor": 2.0,
|
||||
"rope_theta": 500000.0,
|
||||
}, f"yarn must fall back to linear (patcher cannot rebuild it), got {scaling}."
|
||||
|
||||
# plain RoPE with theta only under v5 rope_parameters: linear must carry rope_theta.
|
||||
v5 = SimpleNamespace(rope_parameters = {"rope_type": "default", "rope_theta": 1000000.0})
|
||||
scaling, _ = _extended_rope_scaling(v5, 2.0)
|
||||
assert scaling == {
|
||||
"type": "linear",
|
||||
"factor": 2.0,
|
||||
"rope_theta": 1000000.0,
|
||||
}, f"linear override dropped rope_theta on v5 (got {scaling}); base would fall back to 10000."
|
||||
|
||||
|
||||
def test_extended_rotary_reads_config_factor():
|
||||
# LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8
|
||||
# (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405).
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ ST_TAGS = [
|
|||
"v5.2.3",
|
||||
"v5.3.0",
|
||||
"v5.4.1",
|
||||
"v5.5.1",
|
||||
"v5.6.0",
|
||||
"master",
|
||||
]
|
||||
|
||||
|
|
@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str):
|
|||
)
|
||||
|
||||
|
||||
# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881).
|
||||
@pytest.mark.parametrize("tag", ST_TAGS)
|
||||
def test_st_transformer_load_accepts_unsloth_kwargs(tag: str):
|
||||
"""unsloth builds saved ST models via Transformer.load(...) so the saved
|
||||
modality_config is honored (#6881). If .load stops accepting the hub kwargs it
|
||||
passes (and has no **kwargs), update the fix before it silently regresses. Not
|
||||
locating .load is a SKIP (may be inherited); the live test guards the install."""
|
||||
candidates = [
|
||||
"sentence_transformers/models/Transformer.py",
|
||||
"sentence_transformers/models/transformer.py",
|
||||
"sentence_transformers/base/modules/transformer.py",
|
||||
"sentence_transformers/base/modules/module.py",
|
||||
]
|
||||
for p in candidates:
|
||||
src = fetch_text("UKPLab/sentence-transformers", tag, p)
|
||||
if src is None or not has_def(src, "load", "func"):
|
||||
continue
|
||||
m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S)
|
||||
if m is None:
|
||||
continue
|
||||
sig = m.group(1)
|
||||
accepts_var_kw = "**" in sig
|
||||
missing = [
|
||||
kw
|
||||
for kw in ("token", "cache_folder", "revision", "trust_remote_code")
|
||||
if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig))
|
||||
]
|
||||
assert not missing, (
|
||||
f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no "
|
||||
f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module "
|
||||
f"(#6881) before it silently falls back to Transformer(...)."
|
||||
)
|
||||
return
|
||||
pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)")
|
||||
|
||||
|
||||
# sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls.
|
||||
@pytest.mark.parametrize("tag", ST_TAGS)
|
||||
def test_st_util_helpers(tag: str):
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules]
|
|||
# Fix some issues before importing other packages
|
||||
from .import_fixes import (
|
||||
fix_message_factory_issue,
|
||||
fix_torch_check_is_size,
|
||||
check_fbgemm_gpu_version,
|
||||
disable_broken_causal_conv1d,
|
||||
disable_broken_vllm,
|
||||
|
|
@ -72,6 +73,7 @@ fix_bitsandbytes_rocm_arch_detection()
|
|||
disable_broken_causal_conv1d()
|
||||
disable_broken_vllm()
|
||||
fix_message_factory_issue()
|
||||
fix_torch_check_is_size()
|
||||
check_fbgemm_gpu_version()
|
||||
torchvision_compatibility_check()
|
||||
fix_diffusers_warnings()
|
||||
|
|
@ -81,6 +83,7 @@ del fix_bitsandbytes_rocm_arch_detection
|
|||
del disable_broken_causal_conv1d
|
||||
del disable_broken_vllm
|
||||
del fix_message_factory_issue
|
||||
del fix_torch_check_is_size
|
||||
del check_fbgemm_gpu_version
|
||||
del torchvision_compatibility_check
|
||||
del fix_diffusers_warnings
|
||||
|
|
@ -173,6 +176,7 @@ from .import_fixes import (
|
|||
fix_vllm_guided_decoding_params,
|
||||
fix_vllm_pdl_blackwell,
|
||||
fix_triton_compiled_kernel_missing_attrs,
|
||||
fix_dynamo_config_thread_visibility,
|
||||
patch_trunc_normal_precision_issue,
|
||||
ignore_logger_messages,
|
||||
patch_ipykernel_hf_xet,
|
||||
|
|
@ -203,6 +207,10 @@ fix_vllm_guided_decoding_params()
|
|||
fix_trl_vllm_ascend()
|
||||
fix_vllm_pdl_blackwell()
|
||||
fix_triton_compiled_kernel_missing_attrs()
|
||||
# Must run before unsloth_zoo's patch_torch_compile and the gpt-oss temporary
|
||||
# patches raise the dynamo recompile limits, so those settings reach the
|
||||
# autograd worker threads on torch >= 2.12.
|
||||
fix_dynamo_config_thread_visibility()
|
||||
patch_trunc_normal_precision_issue()
|
||||
ignore_logger_messages()
|
||||
patch_ipykernel_hf_xet()
|
||||
|
|
@ -233,6 +241,7 @@ del fix_vllm_guided_decoding_params
|
|||
del fix_trl_vllm_ascend
|
||||
del fix_vllm_pdl_blackwell
|
||||
del fix_triton_compiled_kernel_missing_attrs
|
||||
del fix_dynamo_config_thread_visibility
|
||||
del patch_trunc_normal_precision_issue
|
||||
del ignore_logger_messages
|
||||
del patch_ipykernel_hf_xet
|
||||
|
|
|
|||
|
|
@ -2103,8 +2103,9 @@ def get_chat_template(
|
|||
|
||||
def remove_special_tokens(tokenizer, prompt):
|
||||
# Removes double BOS token
|
||||
if prompt.startswith(tokenizer.bos_token):
|
||||
prompt = prompt[len(tokenizer.bos_token):]
|
||||
bos_token = getattr(tokenizer, "bos_token", None)
|
||||
if bos_token is not None and prompt.startswith(bos_token):
|
||||
prompt = prompt[len(bos_token):]
|
||||
return prompt
|
||||
|
||||
|
||||
|
|
@ -2185,7 +2186,16 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name
|
|||
|
||||
texts = []
|
||||
for row_idx in range(n_rows):
|
||||
row_values = {column: examples[column][row_idx] for column in columns}
|
||||
# Coerce missing (None) columns to "" so they do not render as the
|
||||
# literal string "None" in the emitted text. In a [[...]] block only
|
||||
# the first column gates the block, so a later column can still be
|
||||
# None here; required columns can be None too. Coercing at the source
|
||||
# covers both; since None is now "", the gate below only needs to
|
||||
# test for "" (an empty first column still drops the block).
|
||||
row_values = {
|
||||
column: ("" if (value := examples[column][row_idx]) is None else value)
|
||||
for column in columns
|
||||
}
|
||||
formatter_values = {}
|
||||
|
||||
for formatter_template in formatter_templates:
|
||||
|
|
@ -2196,7 +2206,7 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name
|
|||
continue
|
||||
|
||||
_, optional_name, prompt, needed_columns = formatter_template
|
||||
if row_values[needed_columns[0]] not in (None, ""):
|
||||
if row_values[needed_columns[0]] != "":
|
||||
prompt_values = {column: row_values[column] for column in needed_columns}
|
||||
formatter_values[optional_name] = prompt.format(**prompt_values)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -172,6 +172,10 @@ if not UNSLOTH_ENABLE_LOGGING:
|
|||
# Deprecation warnings from torchao
|
||||
warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated")
|
||||
warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated")
|
||||
# torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant)
|
||||
warnings.filterwarnings(
|
||||
"ignore", message = r"_check_is_size will be removed", category = FutureWarning
|
||||
)
|
||||
|
||||
# TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752)
|
||||
warnings.filterwarnings(
|
||||
|
|
@ -253,6 +257,30 @@ if not UNSLOTH_ENABLE_LOGGING:
|
|||
)
|
||||
|
||||
|
||||
def fix_torch_check_is_size():
|
||||
"""Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit
|
||||
dequant calls it). The FutureWarning is silenced in suppress_cuda_printf."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if hasattr(torch, "_check_is_size"):
|
||||
return
|
||||
|
||||
def _check_is_size(
|
||||
i,
|
||||
message = None,
|
||||
*,
|
||||
max = None,
|
||||
):
|
||||
torch._check(i >= 0, message)
|
||||
if max is not None:
|
||||
torch._check(i <= max, message)
|
||||
|
||||
torch._check_is_size = _check_is_size
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
# Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype'
|
||||
# MUST do this at the start primarily due to tensorflow causing issues
|
||||
def fix_message_factory_issue():
|
||||
|
|
@ -1064,6 +1092,135 @@ def fix_triton_compiled_kernel_missing_attrs():
|
|||
)
|
||||
|
||||
|
||||
def fix_dynamo_config_thread_visibility():
|
||||
"""torch 2.12 made torch._dynamo/_inductor config overrides thread-local
|
||||
(ContextVars), so `config.recompile_limit = 1024` set on the main thread is
|
||||
invisible to the autograd worker threads that run backward. Gradient
|
||||
checkpointing recompiles fullgraph gpt-oss kernels there against the default
|
||||
limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config
|
||||
assignments into the process-global entry default (torch <= 2.11 semantics).
|
||||
config.patch(...) and config.load_config(...) also assign via __setattr__ but
|
||||
are thread-local by design, so skip mirroring while inside one (tracked per
|
||||
thread). No-op below torch 2.12 and on any torch without this internal layout.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if Version(torch.__version__) < Version("2.12.0"):
|
||||
return
|
||||
import torch._dynamo.config as _dynamo_config
|
||||
from torch.utils._config_module import ConfigModule
|
||||
from contextvars import ContextVar
|
||||
except Exception:
|
||||
return
|
||||
|
||||
try:
|
||||
probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None)
|
||||
if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar):
|
||||
# Overrides are not context-local on this torch; nothing to fix.
|
||||
return
|
||||
original_setattr = ConfigModule.__setattr__
|
||||
if getattr(original_setattr, "__unsloth_patched__", False):
|
||||
return
|
||||
except Exception:
|
||||
return
|
||||
|
||||
mirrored_modules = ("torch._dynamo.config", "torch._inductor.config")
|
||||
|
||||
# config.patch(...) and config.load_config(...) also assign via __setattr__, but
|
||||
# their writes are thread-local by design; a per-thread depth counter marks them
|
||||
# so they are not mirrored into the process-global default.
|
||||
import threading
|
||||
|
||||
_scoped_depth = threading.local()
|
||||
|
||||
def _in_scoped_write():
|
||||
return getattr(_scoped_depth, "n", 0) > 0
|
||||
|
||||
def _bump(delta):
|
||||
_scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta
|
||||
|
||||
original_patch = ConfigModule.patch
|
||||
if not getattr(original_patch, "__unsloth_patched__", False):
|
||||
|
||||
@functools.wraps(original_patch)
|
||||
def _patched_patch(self, *args, **kwargs):
|
||||
ctx = original_patch(self, *args, **kwargs)
|
||||
try:
|
||||
cls = type(ctx) # patch() builds a fresh ConfigPatch class each call
|
||||
if not getattr(cls, "__unsloth_patch_wrapped__", False):
|
||||
_enter0, _exit0 = cls.__enter__, cls.__exit__
|
||||
|
||||
def _enter(s, _e = _enter0):
|
||||
_bump(1)
|
||||
try:
|
||||
return _e(s)
|
||||
finally:
|
||||
_bump(-1)
|
||||
|
||||
def _exit(
|
||||
s,
|
||||
*a,
|
||||
_x = _exit0,
|
||||
):
|
||||
_bump(1)
|
||||
try:
|
||||
return _x(s, *a)
|
||||
finally:
|
||||
_bump(-1)
|
||||
|
||||
cls.__enter__, cls.__exit__ = _enter, _exit
|
||||
cls.__unsloth_patch_wrapped__ = True
|
||||
except Exception:
|
||||
pass
|
||||
return ctx
|
||||
|
||||
_patched_patch.__unsloth_patched__ = True
|
||||
ConfigModule.patch = _patched_patch
|
||||
|
||||
# load_config restores a saved config by calling setattr per key (thread-local).
|
||||
original_load_config = getattr(ConfigModule, "load_config", None)
|
||||
if callable(original_load_config) and not getattr(
|
||||
original_load_config, "__unsloth_patched__", False
|
||||
):
|
||||
|
||||
@functools.wraps(original_load_config)
|
||||
def _patched_load_config(self, *args, **kwargs):
|
||||
_bump(1)
|
||||
try:
|
||||
return original_load_config(self, *args, **kwargs)
|
||||
finally:
|
||||
_bump(-1)
|
||||
|
||||
_patched_load_config.__unsloth_patched__ = True
|
||||
ConfigModule.load_config = _patched_load_config
|
||||
|
||||
@functools.wraps(original_setattr)
|
||||
def _patched_setattr(self, name, value):
|
||||
original_setattr(self, name, value)
|
||||
if _in_scoped_write():
|
||||
return # transient patch / load_config write: keep it thread-local
|
||||
# Aliases (cache_size_limit -> recompile_limit) re-enter with the real name.
|
||||
if self.__dict__.get("__name__", None) in mirrored_modules:
|
||||
try:
|
||||
entry = self.__dict__["_config"].get(name, None)
|
||||
if entry is not None and entry.alias is None:
|
||||
entry.default = value
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
_patched_setattr.__unsloth_patched__ = True
|
||||
ConfigModule.__setattr__ = _patched_setattr
|
||||
|
||||
# No replay of existing overrides: unsloth installs this before it sets any
|
||||
# dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying
|
||||
# would also bake a still-active config.patch override into the global default.
|
||||
logger.info(
|
||||
"Unsloth: Patched torch config modules so dynamo/inductor settings "
|
||||
"(e.g. recompile_limit) apply across threads on torch >= 2.12."
|
||||
)
|
||||
|
||||
|
||||
def patch_trunc_normal_precision_issue():
|
||||
"""
|
||||
Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.
|
||||
|
|
@ -1323,8 +1480,7 @@ def fix_vllm_pdl_blackwell():
|
|||
|
||||
if patched:
|
||||
logger.info(
|
||||
f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - "
|
||||
f"patched: {', '.join(patched)}"
|
||||
f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}"
|
||||
)
|
||||
else:
|
||||
# Just set the env var - vLLM might be an older version without supports_pdl
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import os
|
||||
from contextlib import nullcontext
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import triton
|
||||
|
|
@ -24,6 +25,15 @@ from unsloth_zoo.temporary_patches.common import torch_compile
|
|||
|
||||
torch_matmul = torch.matmul
|
||||
|
||||
|
||||
def _fp8_triton_device_context(tensor: torch.Tensor):
|
||||
if tensor.device.type == "cuda" and torch.cuda.device_count() > 1:
|
||||
return torch.cuda.device(tensor.device)
|
||||
if tensor.device.type == "xpu" and hasattr(torch, "xpu") and torch.xpu.device_count() > 1:
|
||||
return torch.xpu.device(tensor.device)
|
||||
return nullcontext()
|
||||
|
||||
|
||||
try:
|
||||
from transformers.integrations.finegrained_fp8 import FP8Linear
|
||||
except:
|
||||
|
|
@ -95,7 +105,8 @@ def weight_dequant_block(
|
|||
triton.cdiv(M, meta["BLOCK_SIZE"]),
|
||||
triton.cdiv(N, meta["BLOCK_SIZE"]),
|
||||
)
|
||||
weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size)
|
||||
with _fp8_triton_device_context(x):
|
||||
weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size)
|
||||
return y
|
||||
|
||||
|
||||
|
|
@ -149,7 +160,8 @@ def act_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, tor
|
|||
def grid(meta):
|
||||
return (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
|
||||
|
||||
act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size)
|
||||
with _fp8_triton_device_context(x):
|
||||
act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size)
|
||||
return y, s
|
||||
|
||||
|
||||
|
|
@ -274,32 +286,33 @@ def w8a8_block_fp8_matmul_triton(
|
|||
def grid(META):
|
||||
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
|
||||
|
||||
_w8a8_block_fp8_matmul[grid](
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
As,
|
||||
Bs,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_n,
|
||||
block_k,
|
||||
A.stride(-2),
|
||||
A.stride(-1),
|
||||
B.stride(1),
|
||||
B.stride(0),
|
||||
C.stride(-2),
|
||||
C.stride(-1),
|
||||
As.stride(-2),
|
||||
As.stride(-1),
|
||||
Bs.stride(1),
|
||||
Bs.stride(0),
|
||||
BLOCK_SIZE_M = BLOCK_SIZE_M,
|
||||
BLOCK_SIZE_N = BLOCK_SIZE_N,
|
||||
BLOCK_SIZE_K = BLOCK_SIZE_K,
|
||||
GROUP_SIZE_M = 8,
|
||||
)
|
||||
with _fp8_triton_device_context(A):
|
||||
_w8a8_block_fp8_matmul[grid](
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
As,
|
||||
Bs,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
block_n,
|
||||
block_k,
|
||||
A.stride(-2),
|
||||
A.stride(-1),
|
||||
B.stride(1),
|
||||
B.stride(0),
|
||||
C.stride(-2),
|
||||
C.stride(-1),
|
||||
As.stride(-2),
|
||||
As.stride(-1),
|
||||
Bs.stride(1),
|
||||
Bs.stride(0),
|
||||
BLOCK_SIZE_M = BLOCK_SIZE_M,
|
||||
BLOCK_SIZE_N = BLOCK_SIZE_N,
|
||||
BLOCK_SIZE_K = BLOCK_SIZE_K,
|
||||
GROUP_SIZE_M = 8,
|
||||
)
|
||||
return C
|
||||
|
||||
|
||||
|
|
@ -311,13 +324,14 @@ def torchao_block_matmul(
|
|||
block_size: tuple[int, int],
|
||||
output_dtype: torch.dtype = torch.bfloat16,
|
||||
):
|
||||
out = torchao_blockwise_gemm(
|
||||
act_q.contiguous(),
|
||||
act_scale.contiguous(),
|
||||
weight_q.contiguous(),
|
||||
weight_scale.contiguous(),
|
||||
block_size = block_size[1],
|
||||
)
|
||||
with _fp8_triton_device_context(act_q):
|
||||
out = torchao_blockwise_gemm(
|
||||
act_q.contiguous(),
|
||||
act_scale.contiguous(),
|
||||
weight_q.contiguous(),
|
||||
weight_scale.contiguous(),
|
||||
block_size = block_size[1],
|
||||
)
|
||||
return out.to(output_dtype)
|
||||
|
||||
|
||||
|
|
@ -540,7 +554,8 @@ class FP8_fbgemm_block_linear(torch.autograd.Function):
|
|||
f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {bs_n, bs_k}"
|
||||
)
|
||||
|
||||
xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None)
|
||||
with _fp8_triton_device_context(X):
|
||||
xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None)
|
||||
# TODO: WARNING - diverges from baseline for high X values, producing
|
||||
# gibberish / high starting loss. Do not use until resolved; kept for a
|
||||
# future headstart.
|
||||
|
|
|
|||
|
|
@ -282,6 +282,21 @@ def QUANT_STATE(W):
|
|||
return getattr(W, "quant_state", None)
|
||||
|
||||
|
||||
# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a
|
||||
# quant state when the weight itself is still fp8. compressed-tensors layers expose an
|
||||
# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around;
|
||||
# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv /
|
||||
# fast_dequantize path, which then reads a missing `absmax` and crashes.
|
||||
_FP8_WEIGHT_DTYPES = tuple(
|
||||
dtype
|
||||
for dtype in (
|
||||
getattr(torch, "float8_e4m3fn", None),
|
||||
getattr(torch, "float8_e5m2", None),
|
||||
)
|
||||
if dtype is not None
|
||||
)
|
||||
|
||||
|
||||
def get_lora_parameters(proj):
|
||||
"""Return (weight, weight quant_state, lora A, lora B, lora scale).
|
||||
With QAT enabled, also fake-quantizes the base layer and lora weights.
|
||||
|
|
@ -298,9 +313,11 @@ def get_lora_parameters(proj):
|
|||
if weight_fake_quantizer is not None:
|
||||
W = weight_fake_quantizer(W)
|
||||
|
||||
# Get quant state for 4bit or FP8
|
||||
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
|
||||
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
|
||||
# must not carry a scale as its quant state or fast_gemv will crash on it.
|
||||
W_quant = getattr(W, "quant_state", None)
|
||||
if W_quant is None:
|
||||
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
|
||||
W_quant = getattr(base_layer, "weight_scale_inv", None)
|
||||
if W_quant is None:
|
||||
W_quant = getattr(base_layer, "weight_scale", None)
|
||||
|
|
@ -349,9 +366,11 @@ def get_lora_parameters_bias(proj):
|
|||
) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
|
||||
W = base_layer.weight
|
||||
|
||||
# Get quant state for 4bit or FP8
|
||||
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
|
||||
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
|
||||
# must not carry a scale as its quant state or fast_gemv will crash on it.
|
||||
W_quant = getattr(W, "quant_state", None)
|
||||
if W_quant is None:
|
||||
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
|
||||
W_quant = getattr(base_layer, "weight_scale_inv", None)
|
||||
if W_quant is None:
|
||||
W_quant = getattr(base_layer, "weight_scale", None)
|
||||
|
|
|
|||
|
|
@ -28,7 +28,11 @@ from ._utils import (
|
|||
is_bfloat16_supported,
|
||||
get_quant_type,
|
||||
)
|
||||
from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
|
||||
from .loader_utils import (
|
||||
_exclude_rope_inv_freq_from_ddp,
|
||||
_get_fp8_mode_and_check_settings,
|
||||
_restore_dropped_fp8_scales,
|
||||
)
|
||||
from ..utils.packing import (
|
||||
get_packed_info_from_kwargs,
|
||||
mask_packed_sequence_boundaries,
|
||||
|
|
@ -1651,6 +1655,26 @@ def _rope_scaling_as_dict(rope_scaling):
|
|||
return {}
|
||||
|
||||
|
||||
def _extended_rope_scaling(config, factor):
|
||||
"""RoPE scaling to extend a model past its native window. Keeps native llama3 as-is
|
||||
(linear extension is far worse for long context); everything else gets linear. Returns
|
||||
(scaling_or_None, type): None keeps llama3. The linear dict carries rope_theta so
|
||||
transformers v5 (which stores it under rope_parameters) keeps the real base, not 10000.
|
||||
Only llama3 is preserved because patch_llama_rope_scaling can only rebuild linear/llama3/
|
||||
longrope and its longrope branch needs a top-level original_max_position_embeddings."""
|
||||
existing = _rope_scaling_as_dict(
|
||||
getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {}
|
||||
)
|
||||
existing_type = existing.get("rope_type") or existing.get("type")
|
||||
if existing_type == "llama3":
|
||||
return None, existing_type
|
||||
return {
|
||||
"type": "linear",
|
||||
"factor": factor,
|
||||
"rope_theta": _get_rope_theta(config),
|
||||
}, existing_type
|
||||
|
||||
|
||||
def _llama3_inv_freq_from_config(
|
||||
config,
|
||||
rope_scaling,
|
||||
|
|
@ -2518,34 +2542,33 @@ class FastLlamaModel:
|
|||
max_seq_length = model_max_seq_length
|
||||
|
||||
if (rope_scaling is None) and (max_seq_length > model_max_seq_length):
|
||||
rope_scaling = max_seq_length / model_max_seq_length
|
||||
factor = max_seq_length / model_max_seq_length
|
||||
|
||||
if fast_inference:
|
||||
raise NotImplementedError(
|
||||
"Unsloth: Fast inference does not yet work with RoPE Scaling."
|
||||
)
|
||||
|
||||
logger.warning_once(
|
||||
f"Unsloth: {model_name} can only handle sequence lengths of at most "
|
||||
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
|
||||
f"{round(rope_scaling, 3)}, it can be magically be extended to "
|
||||
f"{max_seq_length}!"
|
||||
)
|
||||
|
||||
# Warn RoPE scaling isn't allowed
|
||||
if not has_rope_scaling:
|
||||
raise RuntimeError(
|
||||
f"However, {model_name} doesn't support RoPE Scaling!\n"
|
||||
"Please file a feature request at https://github.com/unslothai/unsloth."
|
||||
linear_scaling, native_type = _extended_rope_scaling(model_config, factor)
|
||||
if linear_scaling is not None:
|
||||
logger.warning_once(
|
||||
f"Unsloth: {model_name} can only handle sequence lengths of at most "
|
||||
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
|
||||
f"{round(factor, 3)}, it can be magically be extended to "
|
||||
f"{max_seq_length}!"
|
||||
)
|
||||
if not has_rope_scaling:
|
||||
raise RuntimeError(
|
||||
f"However, {model_name} doesn't support RoPE Scaling!\n"
|
||||
"Please file a feature request at https://github.com/unslothai/unsloth."
|
||||
)
|
||||
kwargs["rope_scaling"] = linear_scaling
|
||||
else:
|
||||
# Native llama3 scaling already handles long context; just widen the window.
|
||||
logger.warning_once(
|
||||
f"Unsloth: extending {model_name} to {max_seq_length} using its native "
|
||||
f"{native_type} RoPE scaling."
|
||||
)
|
||||
|
||||
rope_scaling = {
|
||||
"type": "linear",
|
||||
"factor": rope_scaling,
|
||||
}
|
||||
|
||||
# Add to kwargs
|
||||
kwargs["rope_scaling"] = rope_scaling
|
||||
|
||||
from .loader_utils import (
|
||||
check_and_disable_bitsandbytes_loading,
|
||||
|
|
@ -2659,6 +2682,18 @@ class FastLlamaModel:
|
|||
offload_embedding = False,
|
||||
fast_inference = fast_inference,
|
||||
)
|
||||
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
|
||||
_restore_dropped_fp8_scales(
|
||||
model,
|
||||
model_name,
|
||||
local_files_only = kwargs.get("local_files_only", False),
|
||||
token = token,
|
||||
# Weights load from the default branch (revision not forwarded), so read scales from there too.
|
||||
revision = None,
|
||||
subfolder = kwargs.get("subfolder"),
|
||||
cache_dir = kwargs.get("cache_dir"),
|
||||
variant = kwargs.get("variant"),
|
||||
)
|
||||
elif not fast_inference:
|
||||
if user_config is not None:
|
||||
# Transformers 5.x @strict model init rejects extra kwargs next
|
||||
|
|
@ -2697,6 +2732,18 @@ class FastLlamaModel:
|
|||
offload_embedding = False,
|
||||
fast_inference = False,
|
||||
)
|
||||
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
|
||||
_restore_dropped_fp8_scales(
|
||||
model,
|
||||
model_name,
|
||||
local_files_only = kwargs.get("local_files_only", False),
|
||||
token = token,
|
||||
# Weights load from the default branch (revision not forwarded), so read scales from there too.
|
||||
revision = None,
|
||||
subfolder = kwargs.get("subfolder"),
|
||||
cache_dir = kwargs.get("cache_dir"),
|
||||
variant = kwargs.get("variant"),
|
||||
)
|
||||
model.fast_generate = make_fast_generate_wrapper(model.generate)
|
||||
model.fast_generate_batches = None
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -239,6 +239,11 @@ def get_model_name(
|
|||
and new_model_name.lower() in BAD_MAPPINGS
|
||||
):
|
||||
new_model_name = BAD_MAPPINGS[new_model_name.lower()]
|
||||
elif new_model_name is None and model_name.lower() in BAD_MAPPINGS:
|
||||
# Some bad names (e.g. the `-unsloth-bnb-4bit` dynamic quants) are keys
|
||||
# of the mappers, not values, so the resolver returns None for them and
|
||||
# the remap above is skipped; remap the input name directly instead.
|
||||
new_model_name = BAD_MAPPINGS[model_name.lower()]
|
||||
|
||||
if (
|
||||
new_model_name is None
|
||||
|
|
@ -362,6 +367,287 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
|
|||
pass
|
||||
|
||||
|
||||
_FP8_DTYPES = tuple(
|
||||
dtype
|
||||
for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
|
||||
if dtype is not None
|
||||
)
|
||||
|
||||
|
||||
def _fp8_block_size_from_config(model):
|
||||
"""Return the [block_out, block_in] block size of an fp8 checkpoint, or None if not block-fp8."""
|
||||
config = getattr(model, "config", None)
|
||||
quant = getattr(config, "quantization_config", None)
|
||||
if quant is None:
|
||||
return None
|
||||
if hasattr(quant, "to_dict"):
|
||||
quant = quant.to_dict()
|
||||
if not isinstance(quant, dict):
|
||||
return None
|
||||
if quant.get("quant_method") != "fp8":
|
||||
return None
|
||||
block = quant.get("weight_block_size")
|
||||
if not block:
|
||||
return None
|
||||
if isinstance(block, (int, float)):
|
||||
block = [block, block]
|
||||
elif isinstance(block, (list, tuple)):
|
||||
if len(block) == 1:
|
||||
block = [block[0], block[0]]
|
||||
elif len(block) < 2:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
return [int(block[0]), int(block[1])]
|
||||
|
||||
|
||||
def _load_fp8_weight_map(
|
||||
model_name,
|
||||
local_files_only,
|
||||
token,
|
||||
revision = None,
|
||||
subfolder = None,
|
||||
cache_dir = None,
|
||||
):
|
||||
"""Return the checkpoint's tensor->file map, using the same snapshot the load used.
|
||||
|
||||
Prefers the sharded `model.safetensors.index.json`; falls back to a single `model.safetensors`
|
||||
(every tensor maps to that one file) so unsharded checkpoints are covered too.
|
||||
"""
|
||||
|
||||
def _local_path(filename):
|
||||
return (
|
||||
os.path.join(model_name, subfolder, filename)
|
||||
if subfolder
|
||||
else os.path.join(model_name, filename)
|
||||
)
|
||||
|
||||
def _remote_path(filename):
|
||||
from huggingface_hub import hf_hub_download
|
||||
return hf_hub_download(
|
||||
model_name,
|
||||
filename,
|
||||
revision = revision,
|
||||
subfolder = subfolder,
|
||||
cache_dir = cache_dir,
|
||||
local_files_only = local_files_only,
|
||||
token = token,
|
||||
)
|
||||
|
||||
index_file = "model.safetensors.index.json"
|
||||
single_file = "model.safetensors"
|
||||
is_local = os.path.isdir(model_name)
|
||||
|
||||
# Sharded checkpoint.
|
||||
if is_local and os.path.exists(_local_path(index_file)):
|
||||
index_path = _local_path(index_file)
|
||||
elif not is_local:
|
||||
try:
|
||||
index_path = _remote_path(index_file)
|
||||
except Exception:
|
||||
index_path = None
|
||||
else:
|
||||
index_path = None
|
||||
if index_path is not None:
|
||||
import json
|
||||
with open(index_path, "r") as f:
|
||||
return json.load(f).get("weight_map", None)
|
||||
|
||||
# Unsharded single file: map every tensor to it.
|
||||
try:
|
||||
if is_local and os.path.exists(_local_path(single_file)):
|
||||
single_path = _local_path(single_file)
|
||||
elif not is_local:
|
||||
single_path = _remote_path(single_file)
|
||||
else:
|
||||
return None
|
||||
from safetensors import safe_open
|
||||
with safe_open(single_path, framework = "pt") as f:
|
||||
return {key: single_file for key in f.keys()}
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_fp8_shard(
|
||||
model_name,
|
||||
shard,
|
||||
local_files_only,
|
||||
token,
|
||||
revision = None,
|
||||
subfolder = None,
|
||||
cache_dir = None,
|
||||
):
|
||||
"""Resolve a checkpoint shard filename to a local path (repo id or local dir)."""
|
||||
if os.path.isdir(model_name):
|
||||
return (
|
||||
os.path.join(model_name, subfolder, shard)
|
||||
if subfolder
|
||||
else os.path.join(model_name, shard)
|
||||
)
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
return hf_hub_download(
|
||||
model_name,
|
||||
shard,
|
||||
revision = revision,
|
||||
subfolder = subfolder,
|
||||
cache_dir = cache_dir,
|
||||
local_files_only = local_files_only,
|
||||
token = token,
|
||||
)
|
||||
|
||||
|
||||
def _match_fp8_module(module_by_name, base):
|
||||
"""Resolve a checkpoint module name to a live module, allowing for VLM key remappings.
|
||||
|
||||
VLM loads can name the text tower differently from the checkpoint keys: `text_only=True`
|
||||
strips the `language_model.` wrapper (so `model.language_model.layers.*` -> `model.layers.*`),
|
||||
and full VLM loads may expose `model.language_model.*` while the checkpoint stores
|
||||
`language_model.model.*`. Try the raw key first, then a few safe remappings.
|
||||
"""
|
||||
if base in module_by_name:
|
||||
return module_by_name[base]
|
||||
candidates = []
|
||||
if "language_model." in base:
|
||||
candidates.append(base.replace("language_model.", "", 1)) # text-only: drop wrapper
|
||||
if "language_model.model." in base:
|
||||
candidates.append(base.replace("language_model.model.", "model.language_model.", 1))
|
||||
if base.startswith("language_model."):
|
||||
candidates.append("model." + base) # add model. prefix
|
||||
for candidate in candidates:
|
||||
if candidate in module_by_name:
|
||||
return module_by_name[candidate]
|
||||
return None
|
||||
|
||||
|
||||
def _restore_dropped_fp8_scales(
|
||||
model,
|
||||
model_name,
|
||||
*,
|
||||
local_files_only = False,
|
||||
token = None,
|
||||
revision = None,
|
||||
subfolder = None,
|
||||
cache_dir = None,
|
||||
variant = None,
|
||||
):
|
||||
"""Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load.
|
||||
|
||||
On some block-scale fp8 checkpoints (e.g. Qwen3.6-27B-FP8, issue #6200) transformers fails to
|
||||
convert a Linear (such as `mlp.gate_proj`) to an fp8 module, loading the raw quantized values
|
||||
into a plain bf16 weight and discarding its `weight_scale_inv` as an unexpected key. The weight
|
||||
is then used un-scaled, producing a garbage model. For every checkpoint scale whose live weight
|
||||
is not fp8, dequantize the orphaned weight in place. Modules that were converted correctly keep
|
||||
an fp8 weight and are skipped, so a healthy checkpoint is a no-op. Returns (restored, skipped).
|
||||
"""
|
||||
try:
|
||||
block = _fp8_block_size_from_config(model)
|
||||
if block is None or not _FP8_DTYPES:
|
||||
return (0, 0)
|
||||
# A variant load reads variant-named files; skip to avoid applying default scales to them.
|
||||
if variant:
|
||||
return (0, 0)
|
||||
# No fp8 params means the checkpoint was dequantized on purpose (e.g. load_in_16bit);
|
||||
# re-applying a scale would corrupt those already-correct 16bit weights, so do nothing.
|
||||
if not any(p.dtype in _FP8_DTYPES for p in model.parameters()):
|
||||
return (0, 0)
|
||||
weight_map = _load_fp8_weight_map(
|
||||
model_name, local_files_only, token, revision, subfolder, cache_dir
|
||||
)
|
||||
if not weight_map:
|
||||
return (0, 0)
|
||||
|
||||
scale_keys = {k: v for k, v in weight_map.items() if k.endswith(".weight_scale_inv")}
|
||||
if not scale_keys:
|
||||
return (0, 0)
|
||||
|
||||
module_by_name = dict(model.named_modules())
|
||||
bs0, bs1 = block
|
||||
restored = 0
|
||||
skipped = 0
|
||||
failed = 0
|
||||
offloaded = 0
|
||||
shard_cache = {}
|
||||
for scale_key, shard in scale_keys.items():
|
||||
base = scale_key[: -len(".weight_scale_inv")]
|
||||
module = _match_fp8_module(module_by_name, base)
|
||||
if module is None:
|
||||
continue
|
||||
weight = getattr(module, "weight", None)
|
||||
if not isinstance(weight, torch.Tensor) or weight.ndim != 2:
|
||||
continue
|
||||
if weight.device.type == "meta":
|
||||
# Disk-offloaded layer: weight lives on meta until forward, so it cannot be
|
||||
# scaled in place here. Count and warn rather than silently leave it unscaled.
|
||||
offloaded += 1
|
||||
continue
|
||||
if weight.dtype in _FP8_DTYPES:
|
||||
# Correctly converted fp8 module: the fp8 path already handles the scale.
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
# Errors after this point are per-tensor: warn and continue, never abort or hide them.
|
||||
try:
|
||||
if shard not in shard_cache:
|
||||
from safetensors import safe_open
|
||||
shard_path = _resolve_fp8_shard(
|
||||
model_name,
|
||||
shard,
|
||||
local_files_only,
|
||||
token,
|
||||
revision,
|
||||
subfolder,
|
||||
cache_dir,
|
||||
)
|
||||
shard_cache[shard] = safe_open(shard_path, framework = "pt")
|
||||
scale = shard_cache[shard].get_tensor(scale_key).to(torch.float32)
|
||||
|
||||
out_features, in_features = weight.shape
|
||||
out_blocks = (out_features + bs0 - 1) // bs0
|
||||
in_blocks = (in_features + bs1 - 1) // bs1
|
||||
if tuple(scale.shape) == (out_blocks, in_blocks):
|
||||
pass
|
||||
elif tuple(scale.shape) == (in_blocks, out_blocks) and out_blocks != in_blocks:
|
||||
# Transposed block layout: same handling as the fp8 forward path.
|
||||
scale = scale.t().contiguous()
|
||||
else:
|
||||
# Shape does not match the block grid: skip rather than apply a wrong scale.
|
||||
continue
|
||||
scale = scale.to(weight.device)
|
||||
with torch.no_grad():
|
||||
if out_features % bs0 == 0 and in_features % bs1 == 0:
|
||||
# Memory-frugal path: multiply block views in place against the broadcast
|
||||
# fp32 scale, avoiding a full expanded scale and fp32 copy that could OOM.
|
||||
# The in-place multiply promotes to fp32, matching the fallback exactly.
|
||||
module.weight.data.view(out_blocks, bs0, in_blocks, bs1).mul_(
|
||||
scale[:, None, :, None]
|
||||
)
|
||||
else:
|
||||
scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(
|
||||
bs1, dim = 1
|
||||
)[:out_features, :in_features]
|
||||
module.weight.data = (weight.to(torch.float32) * scale_expanded).to(
|
||||
weight.dtype
|
||||
)
|
||||
restored += 1
|
||||
except Exception:
|
||||
failed += 1
|
||||
continue
|
||||
|
||||
if restored > 0:
|
||||
print(f"Unsloth: Restored {restored} dropped FP8 weight_scale_inv tensor(s) on load")
|
||||
if failed > 0:
|
||||
print(f"Unsloth: {failed} dropped FP8 weight_scale_inv tensor(s) could not be restored")
|
||||
if offloaded > 0:
|
||||
print(
|
||||
f"Unsloth: {offloaded} dropped FP8 weight_scale_inv tensor(s) skipped because the "
|
||||
"layer is disk-offloaded; load without disk offload so the scales can be restored"
|
||||
)
|
||||
return (restored, skipped)
|
||||
except Exception:
|
||||
return (0, 0)
|
||||
|
||||
|
||||
def check_and_disable_bitsandbytes_loading(
|
||||
model_config,
|
||||
load_in_4bit = True,
|
||||
|
|
|
|||
|
|
@ -134,7 +134,7 @@ __INT_TO_FLOAT_MAPPER = \
|
|||
"unsloth/gemma-7b-it",
|
||||
"google/gemma-7b-it",
|
||||
),
|
||||
"unsloth/gemma-2b-bnb-4bit" : (
|
||||
"unsloth/gemma-2b-it-bnb-4bit" : (
|
||||
"unsloth/gemma-2b-it",
|
||||
"google/gemma-2b-it",
|
||||
),
|
||||
|
|
|
|||
|
|
@ -990,7 +990,17 @@ class FastSentenceTransformer(FastModel):
|
|||
return None
|
||||
|
||||
@staticmethod
|
||||
def _create_transformer_module(model_name, model, tokenizer, max_seq_length, trust_remote_code):
|
||||
def _create_transformer_module(
|
||||
model_name,
|
||||
model,
|
||||
tokenizer,
|
||||
max_seq_length,
|
||||
trust_remote_code,
|
||||
token = None,
|
||||
cache_dir = None,
|
||||
revision = None,
|
||||
module_subfolder = "",
|
||||
):
|
||||
"""Helper to create and configure a Transformer module."""
|
||||
from sentence_transformers.models import Transformer
|
||||
|
||||
|
|
@ -1077,7 +1087,45 @@ class FastSentenceTransformer(FastModel):
|
|||
elif "tokenizer_args" in transformer_init_params:
|
||||
transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy()
|
||||
|
||||
transformer_module = Transformer(model_name, **transformer_kwargs)
|
||||
# Build via Transformer.load so the saved modality_config is honored: plain
|
||||
# Transformer(...) makes ST 5.x infer a "message" modality for chat-template
|
||||
# models (e.g. Qwen3-Embedding), chat-wrapping inputs and degrading embeddings
|
||||
# (#6881). Only use .load when it resolves a Hub id (accepts the kwargs or
|
||||
# **kwargs); legacy ST 3.x/4.x load(input_path) is local-only with no modality
|
||||
# bug, so fall back to the constructor.
|
||||
transformer_module = None
|
||||
transformer_load = getattr(Transformer, "load", None)
|
||||
has_modules_json = (
|
||||
FastSentenceTransformer._module_path(
|
||||
model_name, token, cache_dir = cache_dir, revision = revision
|
||||
)
|
||||
is not None
|
||||
)
|
||||
if callable(transformer_load) and has_modules_json:
|
||||
load_params = inspect.signature(transformer_load).parameters
|
||||
accepts_var_kw = any(
|
||||
p.kind is inspect.Parameter.VAR_KEYWORD for p in load_params.values()
|
||||
)
|
||||
hub_capable = accepts_var_kw or any(
|
||||
key in load_params for key in ("token", "cache_folder", "revision")
|
||||
)
|
||||
if hub_capable:
|
||||
load_kwargs = {
|
||||
"token": token,
|
||||
"cache_folder": cache_dir,
|
||||
"revision": revision,
|
||||
"trust_remote_code": trust_remote_code,
|
||||
**transformer_kwargs,
|
||||
}
|
||||
# Resolve config/tokenizer from the module's saved subfolder
|
||||
# (modules.json "path"), like stock ST; "" (root) is a no-op.
|
||||
if module_subfolder:
|
||||
load_kwargs["subfolder"] = module_subfolder
|
||||
if not accepts_var_kw:
|
||||
load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params}
|
||||
transformer_module = Transformer.load(model_name, **load_kwargs)
|
||||
if transformer_module is None:
|
||||
transformer_module = Transformer(model_name, **transformer_kwargs)
|
||||
finally:
|
||||
# Restore original Auto* loading immediately
|
||||
AutoModel.from_pretrained = original_model_from_pretrained
|
||||
|
|
@ -1191,6 +1239,10 @@ class FastSentenceTransformer(FastModel):
|
|||
tokenizer,
|
||||
max_seq_length,
|
||||
trust_remote_code,
|
||||
token,
|
||||
cache_dir,
|
||||
revision,
|
||||
module_subfolder = module_config.get("path") or "",
|
||||
)
|
||||
modules[name] = transformer_module
|
||||
else:
|
||||
|
|
@ -1226,7 +1278,14 @@ class FastSentenceTransformer(FastModel):
|
|||
)
|
||||
|
||||
transformer_module = FastSentenceTransformer._create_transformer_module(
|
||||
model_name, model, tokenizer, max_seq_length, trust_remote_code
|
||||
model_name,
|
||||
model,
|
||||
tokenizer,
|
||||
max_seq_length,
|
||||
trust_remote_code,
|
||||
token,
|
||||
cache_dir,
|
||||
revision,
|
||||
)
|
||||
modules["0"] = transformer_module
|
||||
|
||||
|
|
|
|||
|
|
@ -41,7 +41,11 @@ from ._utils import (
|
|||
set_task_config_attr,
|
||||
)
|
||||
from ._utils import *
|
||||
from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
|
||||
from .loader_utils import (
|
||||
_exclude_rope_inv_freq_from_ddp,
|
||||
_get_fp8_mode_and_check_settings,
|
||||
_restore_dropped_fp8_scales,
|
||||
)
|
||||
from ..save import patch_saving_functions
|
||||
from ..models.loader_utils import is_distributed
|
||||
from unsloth_zoo.gradient_checkpointing import (
|
||||
|
|
@ -1192,6 +1196,17 @@ class FastBaseModel:
|
|||
offload_embedding = offload_embedding,
|
||||
fast_inference = fast_inference,
|
||||
)
|
||||
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
|
||||
_restore_dropped_fp8_scales(
|
||||
model,
|
||||
model_name,
|
||||
local_files_only = local_files_only,
|
||||
token = token,
|
||||
revision = kwargs.get("revision"),
|
||||
subfolder = kwargs.get("subfolder"),
|
||||
cache_dir = kwargs.get("cache_dir"),
|
||||
variant = kwargs.get("variant"),
|
||||
)
|
||||
if hasattr(model, "generate"):
|
||||
model.fast_generate = make_fast_generate_wrapper(model.generate)
|
||||
model.fast_generate_batches = error_out_no_vllm
|
||||
|
|
|
|||
|
|
@ -133,6 +133,21 @@ _YOLO_OPTION = typer.Option(
|
|||
"flag/config. Any of the three spellings works for any agent."
|
||||
),
|
||||
)
|
||||
_PERSIST_OPTION = typer.Option(
|
||||
False,
|
||||
"--persist/--no-persist",
|
||||
help = (
|
||||
"Keep this agent's Unsloth-managed session dir so you can resume it later. "
|
||||
"codex/openclaw/hermes/pi have their whole home relocated into an Unsloth dir "
|
||||
"that is a throwaway temp dir (wiped on exit) by default; with --persist it "
|
||||
"lives under the Unsloth agents dir and survives, so their own resume can reopen "
|
||||
"it. claude and opencode keep sessions in your own stores (~/.claude, "
|
||||
"~/.local/share/opencode), so they already resume regardless. To reopen a "
|
||||
"session, pass the agent's own resume command through, e.g. "
|
||||
"`unsloth start codex --persist resume` or `claude --resume <id>`; those flow to "
|
||||
"the agent unchanged."
|
||||
),
|
||||
)
|
||||
|
||||
# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no
|
||||
# such flag (config only) and are handled in their config writers, so they are absent.
|
||||
|
|
@ -989,6 +1004,12 @@ def _refresh_windows_path() -> None:
|
|||
os.environ["PATH"] = os.pathsep.join(entries)
|
||||
|
||||
|
||||
def _install_source(install_hint: str) -> Optional[str]:
|
||||
"""The first http(s) URL an install hint fetches, or None (e.g. an npm install)."""
|
||||
match = re.search(r"https?://[^\s'\")]+", install_hint)
|
||||
return match.group(0) if match else None
|
||||
|
||||
|
||||
def _install_agent(name: str, install_hint: str) -> Optional[str]:
|
||||
# Missing agent under --launch: offer to run its documented install command, then
|
||||
# re-resolve it on PATH. Consent-based (we never auto-run a remote install script
|
||||
|
|
@ -997,7 +1018,18 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]:
|
|||
if not sys.stdin.isatty():
|
||||
return None
|
||||
typer.echo(f"`{name}` is not installed.")
|
||||
if not typer.confirm(f"Install it now with `{install_hint}`?", default = False):
|
||||
# Make the supply-chain risk explicit before the prompt: these are the vendors'
|
||||
# own installers (curl | bash, irm | iex, npm), run with the user's privileges,
|
||||
# and nothing checks a signature or hash on the fetched content. Naming the source
|
||||
# turns a blind "yes" into informed consent.
|
||||
source = _install_source(install_hint)
|
||||
warning = (
|
||||
f"This will download and RUN a script from {source} with your privileges"
|
||||
if source
|
||||
else f"This will RUN `{install_hint}` with your privileges"
|
||||
)
|
||||
typer.secho(f"{warning}; there is no signature or hash check.", fg = "yellow", err = True)
|
||||
if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False):
|
||||
return None
|
||||
# Run each hint through the shell it is written for: PowerShell (irm | iex, or npm)
|
||||
# on Windows, /bin/sh (curl | bash, or npm) everywhere else.
|
||||
|
|
@ -1116,15 +1148,20 @@ def _agents_config_root() -> Path:
|
|||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _session_config(agent: str, launch: bool):
|
||||
def _session_config(
|
||||
agent: str,
|
||||
launch: bool,
|
||||
persist: bool = False,
|
||||
):
|
||||
"""Yield a private directory for an agent's session config (never the user's own).
|
||||
|
||||
launch: an ephemeral temp dir removed after the agent process exits, so nothing
|
||||
persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later
|
||||
on this machine), reused across runs. Either way the user's real ~/.<agent>
|
||||
config is left untouched.
|
||||
launch (default): an ephemeral temp dir removed after the agent process exits, so
|
||||
nothing persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run
|
||||
later on this machine), reused across runs. persist (from --persist): use that same
|
||||
stable dir even for a launch, so the agent's session survives the exit and can be
|
||||
resumed next time. Either way the user's real ~/.<agent> config is left untouched.
|
||||
"""
|
||||
if launch:
|
||||
if launch and not persist:
|
||||
path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-"))
|
||||
try:
|
||||
yield path
|
||||
|
|
@ -1436,6 +1473,7 @@ def claude(
|
|||
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
|
||||
serve: bool = _SERVE_OPTION,
|
||||
yolo: bool = _YOLO_OPTION,
|
||||
persist: bool = _PERSIST_OPTION,
|
||||
):
|
||||
"""Point Claude Code at the running Studio server and start it."""
|
||||
base, key, entry = _connect(
|
||||
|
|
@ -1480,6 +1518,9 @@ def claude(
|
|||
# --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions.
|
||||
# IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a
|
||||
# sandbox is detected, and we don't want to falsely claim one on the user's host.
|
||||
# claude keeps its history in ~/.claude/projects, which --settings/env never
|
||||
# relocate, so a session already survives exit; resume it with `claude --continue`
|
||||
# or `--resume <id>` passed through.
|
||||
command = [
|
||||
"claude",
|
||||
"--model",
|
||||
|
|
@ -1516,6 +1557,7 @@ def codex(
|
|||
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
|
||||
serve: bool = _SERVE_OPTION,
|
||||
yolo: bool = _YOLO_OPTION,
|
||||
persist: bool = _PERSIST_OPTION,
|
||||
):
|
||||
"""Point OpenAI Codex at the running Studio server and start it."""
|
||||
base, key, entry = _connect(
|
||||
|
|
@ -1541,7 +1583,7 @@ def codex(
|
|||
*_yolo_command_flags("codex", yolo),
|
||||
*ctx.args,
|
||||
]
|
||||
with _session_config("codex", launch) as home:
|
||||
with _session_config("codex", launch, persist = persist) as home:
|
||||
write_codex_config(base, entry, home)
|
||||
env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)}
|
||||
_run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex")
|
||||
|
|
@ -1559,6 +1601,7 @@ def openclaw(
|
|||
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
|
||||
serve: bool = _SERVE_OPTION,
|
||||
yolo: bool = _YOLO_OPTION,
|
||||
persist: bool = _PERSIST_OPTION,
|
||||
):
|
||||
"""Point OpenClaw at the running Studio server and start it."""
|
||||
base, key, entry = _connect(
|
||||
|
|
@ -1584,7 +1627,7 @@ def openclaw(
|
|||
if os.name == "nt"
|
||||
else "curl -fsSL https://openclaw.ai/install.sh | bash"
|
||||
)
|
||||
with _session_config("openclaw", launch) as cfg:
|
||||
with _session_config("openclaw", launch, persist = persist) as cfg:
|
||||
config_path = cfg / "openclaw.json"
|
||||
# key lives in the config, not the env; --yolo writes the exec policy here too.
|
||||
write_openclaw_config(base, key, entry, config_path, yolo = yolo)
|
||||
|
|
@ -1605,6 +1648,7 @@ def opencode(
|
|||
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
|
||||
serve: bool = _SERVE_OPTION,
|
||||
yolo: bool = _YOLO_OPTION,
|
||||
persist: bool = _PERSIST_OPTION,
|
||||
):
|
||||
"""Point OpenCode at the running Studio server and start it."""
|
||||
base, key, entry = _connect(
|
||||
|
|
@ -1628,7 +1672,9 @@ def opencode(
|
|||
command = ["opencode", "--model", opencode_model]
|
||||
else:
|
||||
command = ["opencode"]
|
||||
with _session_config("opencode", launch) as cfg:
|
||||
# opencode keeps sessions in ~/.local/share/opencode (never relocated), so resume
|
||||
# already survives exit; reopen the last one by passing `opencode --continue` through.
|
||||
with _session_config("opencode", launch, persist = persist) as cfg:
|
||||
config_path = cfg / "opencode.json"
|
||||
# OPENCODE_CONFIG is an overlay (loaded between the user's global and project
|
||||
# configs), so this adds the Unsloth provider/model for the session without
|
||||
|
|
@ -1680,6 +1726,7 @@ def hermes(
|
|||
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
|
||||
serve: bool = _SERVE_OPTION,
|
||||
yolo: bool = _YOLO_OPTION,
|
||||
persist: bool = _PERSIST_OPTION,
|
||||
):
|
||||
"""Point Hermes (Nous Research) at the running Studio server and start it."""
|
||||
base, key, entry = _connect(
|
||||
|
|
@ -1691,7 +1738,7 @@ def hermes(
|
|||
)
|
||||
command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args]
|
||||
install_hint = _hermes_install_hint()
|
||||
with _session_config("hermes", launch) as home:
|
||||
with _session_config("hermes", launch, persist = persist) as home:
|
||||
# HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state)
|
||||
# like CODEX_HOME, so the user's ~/.hermes is left untouched for the session.
|
||||
write_hermes_config(base, entry, home / "config.yaml")
|
||||
|
|
@ -1711,6 +1758,7 @@ def pi(
|
|||
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
|
||||
serve: bool = _SERVE_OPTION,
|
||||
yolo: bool = _YOLO_OPTION,
|
||||
persist: bool = _PERSIST_OPTION,
|
||||
):
|
||||
"""Point Pi (coding agent) at the running Studio server and start it."""
|
||||
base, key, entry = _connect(
|
||||
|
|
@ -1735,7 +1783,7 @@ def pi(
|
|||
# --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs
|
||||
# no install scripts), so accepting the prompt skips dependency lifecycle scripts.
|
||||
install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent"
|
||||
with _session_config("pi", launch) as home:
|
||||
with _session_config("pi", launch, persist = persist) as home:
|
||||
# Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers
|
||||
# it over $HOME/.pi/agent), so pin it at the session dir: an inherited
|
||||
# PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real
|
||||
|
|
|
|||
|
|
@ -128,6 +128,32 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch):
|
|||
assert ran == [["powershell", "-NoProfile", "-Command", install_hint]]
|
||||
|
||||
|
||||
def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys):
|
||||
# Before the confirm, a remote installer must name the URL it fetches so the
|
||||
# user consents to a specific source rather than blindly accepting.
|
||||
monkeypatch.setattr(start.os, "name", "nt")
|
||||
monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
|
||||
monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) # decline: nothing runs
|
||||
hint = "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup"
|
||||
assert start._install_agent("hermes", hint) is None
|
||||
err = capsys.readouterr().err
|
||||
assert "https://hermes-agent.nousresearch.com/install.ps1" in err
|
||||
assert "download and RUN" in err
|
||||
assert "signature or hash" in err
|
||||
|
||||
|
||||
def test_install_agent_warns_for_package_installer(monkeypatch, capsys):
|
||||
# An npm-style installer has no URL to fetch, but still runs with the user's
|
||||
# privileges, so the warning names the command instead.
|
||||
monkeypatch.setattr(start.os, "name", "posix")
|
||||
monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
|
||||
monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False)
|
||||
assert start._install_agent("codex", "npm install -g @openai/codex") is None
|
||||
err = capsys.readouterr().err
|
||||
assert "npm install -g @openai/codex" in err
|
||||
assert "with your privileges" in err
|
||||
|
||||
|
||||
def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch):
|
||||
monkeypatch.setattr(start.os, "name", "nt")
|
||||
|
||||
|
|
@ -2522,3 +2548,145 @@ def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path
|
|||
with start._session_config("codex", launch = False) as home2:
|
||||
assert home2 == home
|
||||
assert (home2 / "sessions" / "live.sqlite").read_text() == "state"
|
||||
|
||||
|
||||
# ── --persist: persist the agent session so it can be resumed ────────────────
|
||||
def test_session_config_persist_uses_stable_dir_and_survives(monkeypatch, tmp_path):
|
||||
# --persist routes a launch to the stable Unsloth agents dir (the one --no-launch
|
||||
# already uses) instead of a throwaway temp dir, and never wipes it on exit.
|
||||
monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents")
|
||||
with start._session_config("codex", launch = True, persist = True) as home:
|
||||
assert home == tmp_path / "agents" / "codex"
|
||||
(home / "marker").write_text("kept")
|
||||
assert home.exists()
|
||||
assert (home / "marker").read_text() == "kept"
|
||||
|
||||
|
||||
def test_session_config_default_launch_is_ephemeral():
|
||||
# Default launch (no --persist) still uses a throwaway temp dir wiped on exit.
|
||||
with start._session_config("codex", launch = True) as home:
|
||||
assert home.exists()
|
||||
assert "unsloth-codex-" in home.name
|
||||
assert not home.exists()
|
||||
|
||||
|
||||
# The temp-dir agents: --persist points each one's home/state env at the stable dir;
|
||||
# without it, at an ephemeral temp path. opencode is handled separately (only its
|
||||
# config overlay is relocated; its session data was never in the temp dir).
|
||||
_RESUME_ENV_VAR = {
|
||||
"codex": "CODEX_HOME",
|
||||
"openclaw": "OPENCLAW_STATE_DIR",
|
||||
"hermes": "HERMES_HOME",
|
||||
"pi": "HOME",
|
||||
}
|
||||
|
||||
|
||||
def _capture_launch(monkeypatch, argv):
|
||||
captured = {}
|
||||
|
||||
def run(
|
||||
command,
|
||||
env = None,
|
||||
**kwargs,
|
||||
):
|
||||
captured["command"] = command
|
||||
captured["env"] = env
|
||||
return SimpleNamespace(returncode = 0)
|
||||
|
||||
monkeypatch.setattr(start.subprocess, "run", run)
|
||||
result = CliRunner().invoke(start.start_app, argv)
|
||||
assert result.exit_code == 0, result.output
|
||||
return captured
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
|
||||
def test_resume_persists_agent_home_to_stable_dir(agent, fake_studio, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
|
||||
captured = _capture_launch(monkeypatch, [agent, "--persist"])
|
||||
stable = tmp_path / "agents" / agent
|
||||
assert captured["env"][_RESUME_ENV_VAR[agent]] == str(stable)
|
||||
# The stable dir survives the agent exit, so the session can be resumed.
|
||||
assert stable.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
|
||||
def test_default_launch_home_is_ephemeral(agent, fake_studio, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
|
||||
captured = _capture_launch(monkeypatch, [agent])
|
||||
home = captured["env"][_RESUME_ENV_VAR[agent]]
|
||||
assert f"unsloth-{agent}-" in home
|
||||
assert str(tmp_path / "agents") not in home
|
||||
|
||||
|
||||
def test_resume_opencode_config_in_stable_dir(fake_studio, tmp_path, monkeypatch):
|
||||
# opencode's session data lives in ~/.local/share/opencode (never relocated), so
|
||||
# resume already survives exit; --persist also stabilizes its config overlay dir.
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
|
||||
captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
|
||||
stable = tmp_path / "agents" / "opencode"
|
||||
assert captured["env"]["OPENCODE_CONFIG"] == str(stable / "opencode.json")
|
||||
assert stable.exists()
|
||||
|
||||
|
||||
def test_persist_bare_codex_launch_has_no_resume_token(fake_studio, monkeypatch):
|
||||
# A bare `--persist` only persists the session dir; it must NOT auto-append a native
|
||||
# resume token, or the very first launch (no session yet) would send codex down its
|
||||
# no-session error path. The user resumes explicitly: `unsloth start codex --persist resume`.
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
|
||||
captured = _capture_launch(monkeypatch, ["codex", "--persist"])
|
||||
assert "resume" not in captured["command"]
|
||||
# command[0] is the resolved executable path; assert the argv after it.
|
||||
assert captured["command"][1:] == ["--oss", "--profile", start._CODEX_PROFILE]
|
||||
|
||||
|
||||
def test_persist_bare_opencode_launch_has_no_resume_token(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
|
||||
captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
|
||||
assert "--continue" not in captured["command"]
|
||||
assert captured["command"][1:] == ["--model", f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"]
|
||||
|
||||
|
||||
def test_persist_bare_claude_launch_has_no_resume_token(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
|
||||
monkeypatch.setattr(start, "_claude_flags", lambda: [])
|
||||
captured = _capture_launch(monkeypatch, ["claude", "--persist"])
|
||||
assert "--continue" not in captured["command"]
|
||||
assert captured["command"][1:] == ["--model", MODEL["id"]]
|
||||
|
||||
|
||||
def test_resume_with_passthrough_does_not_auto_append(fake_studio, monkeypatch):
|
||||
# When the caller drives their own subcommand, --persist only persists the dir; it
|
||||
# must not inject a resume token that would collide with the user's command.
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
|
||||
captured = _capture_launch(monkeypatch, ["codex", "--persist", "exec", "hello"])
|
||||
assert "resume" not in captured["command"]
|
||||
assert captured["command"][-2:] == ["exec", "hello"]
|
||||
|
||||
|
||||
def test_default_launch_has_no_resume_token(fake_studio, monkeypatch):
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
|
||||
captured = _capture_launch(monkeypatch, ["codex"])
|
||||
assert "resume" not in captured["command"]
|
||||
|
||||
|
||||
def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch):
|
||||
# openclaw/hermes persist their session dir but have no non-interactive resume
|
||||
# selector, so --persist must not append a token; their own picker resumes.
|
||||
for agent in ("openclaw", "hermes"):
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}")
|
||||
captured = _capture_launch(monkeypatch, [agent, "--persist"])
|
||||
assert "resume" not in captured["command"]
|
||||
assert "--continue" not in captured["command"]
|
||||
|
||||
|
||||
def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch):
|
||||
# The persistence flag is --persist, NOT --resume, so an agent's own
|
||||
# `--resume <id>` (e.g. `unsloth start claude --resume <guid>`) still flows
|
||||
# through to the agent verbatim and is not swallowed as a Studio option.
|
||||
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
|
||||
monkeypatch.setattr(start, "_claude_flags", lambda: [])
|
||||
captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"])
|
||||
assert captured["command"][-2:] == ["--resume", "some-session-guid"]
|
||||
# Studio never auto-appends its own resume token when the user drives resume.
|
||||
assert captured["command"].count("--resume") == 1
|
||||
assert "--continue" not in captured["command"]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue