Merge main into feat/studio-parallel-slots-ui for PR #7447
Resolve the conflicts against the parallel-chats work now on main (#7455). - studio/backend/run.py: drop this branch's duplicate _PARALLEL_* block. main already defines those above run_server() with _PARALLEL_DEFAULT_PLAIN = 4, and keeping the second definition shadowed it back to 1, reverting the new default for direct backend launches. - studio/backend/run.py and unsloth_cli/commands/studio.py: keep the per-load override note in the --parallel help, drop the now stale "unsloth studio run defaults to 4" contrast now that both defaults are 4. - studio/backend/core/inference/llama_cpp.py: take main's wording for the --kv-unified clamp comment. - tests/studio/test_model_picker_contracts.py: keep both tests, which each side appended at the same spot. Also guard the per-load slot resolution. _load_model_impl and validate_model read fastapi_request.app.state directly, which raises for a direct caller whose request carries no app; it is reached via getattr now, matching the fallback already used by _openai_llama_admission_capacity. Without that, three of main's load tests fail with AttributeError once the branches meet. test_active_generations.py is 79 passed, matching main.
This commit is contained in:
Daniel Han
commit
73136370cd
178 changed files with 22670 additions and 1836 deletions
6
.github/workflows/consolidated-tests-ci.yml
vendored
6
.github/workflows/consolidated-tests-ci.yml
vendored
|
|
@ -7,7 +7,7 @@
|
|||
#
|
||||
# Why a separate workflow:
|
||||
# - studio-backend-ci.yml's "Repo tests (CPU)" job already auto-discovers
|
||||
# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 16
|
||||
# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 17
|
||||
# Bucket-A tests below live inside those --ignore dirs (CPU-runnable but
|
||||
# historically excluded with their GPU siblings); pulling them out into
|
||||
# a sibling job keeps the existing 760-passed baseline stable while we
|
||||
|
|
@ -274,6 +274,7 @@ jobs:
|
|||
tests/saving/test_export_dispatch.py \
|
||||
tests/saving/test_imatrix_export.py \
|
||||
tests/saving/test_gguf_single_pass_export.py \
|
||||
tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \
|
||||
tests/utils/test_attention_masks.py \
|
||||
tests/utils/test_trunc_normal_patch.py \
|
||||
tests/python/test_fast_language_model_text_only.py
|
||||
|
|
@ -365,6 +366,7 @@ jobs:
|
|||
tests/saving/test_export_dispatch.py \
|
||||
tests/saving/test_imatrix_export.py \
|
||||
tests/saving/test_gguf_single_pass_export.py \
|
||||
tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \
|
||||
tests/utils/test_attention_masks.py \
|
||||
tests/utils/test_trunc_normal_patch.py \
|
||||
tests/python/test_fast_language_model_text_only.py \
|
||||
|
|
@ -2129,7 +2131,7 @@ jobs:
|
|||
pip show unsloth_zoo
|
||||
echo "::endgroup::"
|
||||
echo "Consolidated job done. Coverage:"
|
||||
echo " - 16 unsloth Bucket-A tests under tests/saving/ + tests/utils/"
|
||||
echo " - 17 unsloth Bucket-A tests under tests/saving/ + tests/utils/"
|
||||
echo " - unsloth_zoo @ ${UNSLOTH_ZOO_REF} pytest tests/ (5 GPU cases deselected)"
|
||||
echo " - unsloth_zoo.compiler.test_apply_fused_lm_head"
|
||||
|
||||
|
|
|
|||
40
.github/workflows/studio-update-smoke.yml
vendored
40
.github/workflows/studio-update-smoke.yml
vendored
|
|
@ -146,6 +146,46 @@ jobs:
|
|||
kill "$PID" 2>/dev/null || true
|
||||
echo "post-update Unsloth /api/health OK"
|
||||
|
||||
- name: A complete install reports itself complete
|
||||
run: |
|
||||
set -o pipefail
|
||||
unsloth studio verify-install
|
||||
unsloth studio desktop-capabilities --json | tee /tmp/caps.json
|
||||
jq -e '.studio_install_ok == true' /tmp/caps.json
|
||||
jq -e '.desktop_manageability_version >= 2' /tmp/caps.json
|
||||
|
||||
- name: An incomplete install must not report itself ready
|
||||
# An installer killed part-way leaves a working CLI but no studio.txt
|
||||
# deps, which the old preflight called ManagedReady. The manifest is
|
||||
# written last, so removing it reproduces that state.
|
||||
run: |
|
||||
set -o pipefail
|
||||
# install.sh's default root, resolved explicitly: `python` on PATH
|
||||
# here is setup-python's, not the managed venv.
|
||||
MANIFEST="$HOME/.unsloth/studio/unsloth_studio/unsloth_install_manifest.json"
|
||||
test -f "$MANIFEST" || { echo "::error::installer never wrote $MANIFEST"; exit 1; }
|
||||
rm -f "$MANIFEST"
|
||||
unsloth studio desktop-capabilities --json | tee /tmp/caps_bad.json
|
||||
jq -e '.studio_install_ok == false' /tmp/caps_bad.json
|
||||
if unsloth studio verify-install; then
|
||||
echo "::error::verify-install passed on an install with no manifest"
|
||||
exit 1
|
||||
fi
|
||||
echo "incomplete install correctly reported not-ready"
|
||||
|
||||
- name: Update repairs an incomplete install
|
||||
# `--local` bypasses setup.sh's PyPI version compare, so this asserts
|
||||
# the repair OUTCOME. The non-local fast path the desktop Repair button
|
||||
# uses is covered by tests/studio/install/test_setup_fast_path_guard.py.
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
set -o pipefail
|
||||
unsloth studio update --local 2>&1 | tee logs/update_repair.log
|
||||
unsloth studio verify-install
|
||||
unsloth studio desktop-capabilities --json | jq -e '.studio_install_ok == true'
|
||||
echo "update repaired the incomplete install"
|
||||
|
||||
- name: Uninstall and verify clean
|
||||
# Round-trip the installer through scripts/uninstall.sh: confirms the
|
||||
# uninstaller actually finds and removes everything install.sh +
|
||||
|
|
|
|||
25
.github/workflows/wheel-smoke.yml
vendored
25
.github/workflows/wheel-smoke.yml
vendored
|
|
@ -127,6 +127,31 @@ jobs:
|
|||
cd /tmp
|
||||
/tmp/v/bin/python -c "from studio.backend.main import app; print('Unsloth backend OK:', app.title)"
|
||||
|
||||
- name: CLI without the Studio stack guides instead of tracebacking
|
||||
# The smoke above installs studio.txt first, so it cannot catch a wheel
|
||||
# that ships studio/ without declaring what it imports (#4701, #5260,
|
||||
# #7147). Drop only structlog to reuse that venv without a re-download.
|
||||
run: |
|
||||
set -eu
|
||||
/tmp/v/bin/pip uninstall -y structlog >/dev/null
|
||||
cd /tmp
|
||||
status=0
|
||||
for args in "export ./nope ./out" "list-checkpoints"; do
|
||||
echo "--- unsloth $args"
|
||||
out=$(/tmp/v/bin/unsloth $args 2>&1 || true)
|
||||
printf '%s\n' "$out"
|
||||
case "$out" in
|
||||
*Traceback*)
|
||||
echo "FAIL: raw traceback instead of guidance"; status=1 ;;
|
||||
esac
|
||||
case "$out" in
|
||||
*'unsloth studio update'*) ;;
|
||||
*) echo "FAIL: no remediation in the message"; status=1 ;;
|
||||
esac
|
||||
done
|
||||
/tmp/v/bin/pip install -q structlog >/dev/null
|
||||
exit "$status"
|
||||
|
||||
- name: Upload wheel on failure
|
||||
if: failure()
|
||||
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
|
||||
|
|
|
|||
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -238,4 +238,5 @@ package-lock.json
|
|||
!studio/package-lock.json
|
||||
llama.cpp/
|
||||
# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo.
|
||||
/~/
|
||||
~/
|
||||
/temp/
|
||||
|
|
|
|||
49
install.ps1
49
install.ps1
|
|
@ -28,6 +28,14 @@ function Install-UnslothStudio {
|
|||
}
|
||||
}
|
||||
|
||||
function Clear-TauriInstallError {
|
||||
param([string]$Message)
|
||||
if ($TauriMode) {
|
||||
Write-TauriLog "ERROR_CLEAR" $Message
|
||||
[Console]::Error.WriteLine("[TAURI:ERROR_CLEAR] $Message")
|
||||
}
|
||||
}
|
||||
|
||||
function Format-TauriDiagBool {
|
||||
param([bool]$Value)
|
||||
if ($Value) { return "true" }
|
||||
|
|
@ -86,7 +94,7 @@ function Install-UnslothStudio {
|
|||
[int]$Code = 1
|
||||
)
|
||||
if ($Code -eq 0) { $Code = 1 }
|
||||
Write-TauriLog "ERROR" $Message
|
||||
Write-TauriLog "ERROR_DEFAULT" $Message
|
||||
if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) {
|
||||
Restore-StudioVenvRollback
|
||||
}
|
||||
|
|
@ -485,7 +493,8 @@ function Install-UnslothStudio {
|
|||
# Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
|
||||
function Invoke-InstallCommand {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][ScriptBlock]$Command
|
||||
[Parameter(Mandatory = $true)][ScriptBlock]$Command,
|
||||
[string]$Label = "install command"
|
||||
)
|
||||
# Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
|
||||
# for --default-index, clear the uv index env vars (restore in finally) and set
|
||||
|
|
@ -504,6 +513,7 @@ function Install-UnslothStudio {
|
|||
try {
|
||||
# Reset to avoid stale values from prior native commands.
|
||||
$global:LASTEXITCODE = 0
|
||||
Write-TauriLog "OUTPUT_CLEAR" $Label
|
||||
if ($script:UnslothVerbose) {
|
||||
# Merge stderr into stdout so progress/warning output stays visible
|
||||
# without flipping $? on successful native commands (PS 5.1 treats
|
||||
|
|
@ -518,7 +528,13 @@ function Install-UnslothStudio {
|
|||
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
|
||||
}
|
||||
}
|
||||
return [int]$LASTEXITCODE
|
||||
$exitCode = [int]$LASTEXITCODE
|
||||
if ($exitCode -eq 0) {
|
||||
Clear-TauriInstallError "$Label recovered"
|
||||
} else {
|
||||
Write-TauriLog "ERROR_OUTPUT" "$Label failed (exit code $exitCode)"
|
||||
}
|
||||
return $exitCode
|
||||
} finally {
|
||||
$ErrorActionPreference = $prevEap
|
||||
if ($savedUvIndex) {
|
||||
|
|
@ -549,7 +565,7 @@ function Install-UnslothStudio {
|
|||
}
|
||||
$attempt = 1
|
||||
while ($true) {
|
||||
$code = Invoke-InstallCommand $Command
|
||||
$code = Invoke-InstallCommand -Command $Command -Label $Label
|
||||
if ($code -eq 0) { return 0 }
|
||||
if ($attempt -ge $maxAttempts) { return $code }
|
||||
substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow"
|
||||
|
|
@ -1603,7 +1619,7 @@ exit 0
|
|||
if (-not (Test-Path -LiteralPath $VenvPython)) {
|
||||
step "venv" "creating Python $($DetectedPython.Version) virtual environment"
|
||||
substep "$VenvDir"
|
||||
$venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" }
|
||||
$venvExit = Invoke-InstallCommand -Label "create virtual environment" { uv venv $VenvDir --python "$($DetectedPython.Path)" }
|
||||
if ($venvExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to create virtual environment (exit code $venvExit)" $venvExit)
|
||||
|
|
@ -2375,7 +2391,7 @@ exit 0
|
|||
}
|
||||
if ($StudioLocalInstall) {
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
$overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
|
||||
|
|
@ -2464,7 +2480,7 @@ exit 0
|
|||
|
||||
if ($StudioLocalInstall) {
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
$overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
|
||||
|
|
@ -2487,7 +2503,7 @@ exit 0
|
|||
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
|
||||
}
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
$overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
|
||||
|
|
@ -2535,7 +2551,7 @@ exit 0
|
|||
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
|
||||
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
|
||||
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
|
||||
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
|
||||
$torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch (ROCm)" { 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)
|
||||
|
|
@ -2544,7 +2560,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>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
|
||||
$torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch ($expectedTorchTag)" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --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)
|
||||
|
|
@ -2645,6 +2661,9 @@ exit 0
|
|||
# an inherited value would put llama.cpp in the wrong place.
|
||||
$previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME
|
||||
$hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome)
|
||||
$previousTauriMode = $env:UNSLOTH_TAURI_MODE
|
||||
$hadPreviousTauriMode = ($null -ne $previousTauriMode)
|
||||
$env:UNSLOTH_TAURI_MODE = if ($TauriMode) { "1" } else { "0" }
|
||||
if ($StudioRedirectMode -eq 'env') {
|
||||
$env:UNSLOTH_STUDIO_HOME = $StudioHome
|
||||
} else {
|
||||
|
|
@ -2674,14 +2693,22 @@ exit 0
|
|||
} else {
|
||||
Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue
|
||||
}
|
||||
if ($hadPreviousTauriMode) {
|
||||
$env:UNSLOTH_TAURI_MODE = $previousTauriMode
|
||||
} else {
|
||||
Remove-Item Env:UNSLOTH_TAURI_MODE -ErrorAction SilentlyContinue
|
||||
}
|
||||
Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue
|
||||
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
|
||||
Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue
|
||||
}
|
||||
if ($setupExit -ne 0) {
|
||||
Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
|
||||
if (-not $TauriMode) {
|
||||
Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
|
||||
}
|
||||
return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit)
|
||||
}
|
||||
Clear-TauriInstallError "studio setup completed"
|
||||
|
||||
# ── Expose `unsloth` via a shim dir containing only unsloth.exe ──
|
||||
# We do NOT add the venv Scripts dir to PATH (it also holds python.exe
|
||||
|
|
|
|||
198
install.sh
198
install.sh
|
|
@ -207,18 +207,37 @@ run_install_cmd() {
|
|||
# command's exit code across the pipe without relying on pipefail
|
||||
# (this script runs under plain sh).
|
||||
_rcf=$(mktemp)
|
||||
{ "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output
|
||||
tauri_stream_log stdout "OUTPUT_CLEAR" "$_label"
|
||||
{
|
||||
if "$@" 2>&1; then
|
||||
_cmd_rc=0
|
||||
else
|
||||
_cmd_rc=$?
|
||||
fi
|
||||
printf '%s' "$_cmd_rc" > "$_rcf"
|
||||
} | _redact_install_output
|
||||
_rc=$(cat "$_rcf" 2>/dev/null || echo 1)
|
||||
rm -f "$_rcf"
|
||||
[ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0
|
||||
_rc=${_rc:-1}
|
||||
if [ "$_rc" -eq 0 ] 2>/dev/null; then
|
||||
tauri_clear_install_error "$_label recovered"
|
||||
return 0
|
||||
fi
|
||||
tauri_stream_log stdout "ERROR_OUTPUT" "$_label failed (exit code $_rc)"
|
||||
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
|
||||
return "$_rc"
|
||||
fi
|
||||
_log=$(mktemp)
|
||||
"$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; }
|
||||
tauri_stream_log stderr "OUTPUT_CLEAR" "$_label"
|
||||
"$@" >"$_log" 2>&1 && {
|
||||
rm -f "$_log"
|
||||
tauri_clear_install_error "$_label recovered"
|
||||
return 0
|
||||
}
|
||||
_rc=$?
|
||||
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
|
||||
_redact_install_output "$_log" >&2
|
||||
tauri_stream_log stderr "ERROR_OUTPUT" "$_label failed (exit code $_rc)"
|
||||
rm -f "$_log"
|
||||
return $_rc
|
||||
}
|
||||
|
|
@ -257,6 +276,51 @@ run_install_cmd_retry() {
|
|||
done
|
||||
}
|
||||
|
||||
# True when the runtime target is gfx906 (MI50/Radeon VII): the prebuilt AMD
|
||||
# bitsandbytes wheel carries no gfx906 kernels, and force-reinstalling it would
|
||||
# clobber a user's source-built bnb (the only 4-bit path on this arch) on every
|
||||
# `studio update`. So skip the auto-install and leave whatever bnb is present.
|
||||
# _gfx906_target is set during torch-index resolution; also honor an explicit
|
||||
# UNSLOTH_ROCM_GFX_ARCH so a pinned-index install still skips. The override is
|
||||
# normalized (gfx906:sramecc-:xnack- -> gfx906) so a copied HIP gcnArchName counts.
|
||||
_is_gfx906_bnb_skip() {
|
||||
[ "${_gfx906_target:-false}" = true ] && return 0
|
||||
_bnb_gfx_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]')
|
||||
_bnb_gfx_env=${_bnb_gfx_env%%:*}
|
||||
[ "$_bnb_gfx_env" = "gfx906" ] && return 0
|
||||
# A pinned index (UNSLOTH_TORCH_INDEX_URL/_FAMILY) skips the reroute block that
|
||||
# sets _gfx906_target, so a real gfx906 host with a pinned rocm6.3 index and no
|
||||
# UNSLOTH_ROCM_GFX_ARCH would otherwise clobber a source-built bnb. Probe here
|
||||
# in that gap; skip only when gfx906 is the SOLE distinct arch (mixed hosts
|
||||
# opt in via the env var, mirroring the reroute block's de-dup rule).
|
||||
if [ -z "$_bnb_gfx_env" ] && [ "${_torch_index_pinned:-false}" = true ]; then
|
||||
_bnb_gfx_probe=$(_probe_amd_gfx_arch | awk 'NF && !seen[$0]++')
|
||||
[ "$_bnb_gfx_probe" = "gfx906" ] && return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
# `pip install unsloth` resolves its unconditional bitsandbytes dep to a generic
|
||||
# CUDA wheel (no gfx906 kernels) once we skip the prebuilt one. Snapshot bnb before
|
||||
# the unsloth install, then drop a freshly pulled wheel afterwards while leaving a
|
||||
# pre-existing source build in place.
|
||||
_gfx906_bnb_installed() {
|
||||
"$_VENV_PY" -c "import importlib.util as u, sys; sys.exit(0 if u.find_spec('bitsandbytes') else 1)" >/dev/null 2>&1
|
||||
}
|
||||
_gfx906_bnb_snapshot() {
|
||||
_gfx906_bnb_absent_before=false
|
||||
_is_gfx906_bnb_skip || return 0
|
||||
_gfx906_bnb_installed || _gfx906_bnb_absent_before=true
|
||||
}
|
||||
_gfx906_bnb_prune() {
|
||||
_is_gfx906_bnb_skip || return 0
|
||||
[ "${_gfx906_bnb_absent_before:-false}" = true ] || return 0
|
||||
_gfx906_bnb_installed || return 0
|
||||
substep "gfx906: removing generic bitsandbytes pulled in as a dependency (no gfx906 kernels; build from source for 4-bit QLoRA)" "$C_WARN"
|
||||
uv pip uninstall --python "$_VENV_PY" bitsandbytes >/dev/null 2>&1 \
|
||||
|| "$_VENV_PY" -m pip uninstall -y bitsandbytes >/dev/null 2>&1 || true
|
||||
}
|
||||
|
||||
# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main
|
||||
# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2
|
||||
# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the
|
||||
|
|
@ -338,6 +402,34 @@ tauri_log() {
|
|||
fi
|
||||
}
|
||||
|
||||
tauri_stream_log() {
|
||||
_tsl_stream="$1"
|
||||
_tsl_tag="$2"
|
||||
shift 2
|
||||
if [ "$TAURI_MODE" = true ]; then
|
||||
if [ "$_tsl_stream" = stderr ]; then
|
||||
printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*" >&2
|
||||
else
|
||||
printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*"
|
||||
fi
|
||||
fi
|
||||
}
|
||||
|
||||
rollback_substep() {
|
||||
if [ "$TAURI_MODE" = true ]; then
|
||||
tauri_log "PROGRESS" "$1"
|
||||
else
|
||||
substep "$@"
|
||||
fi
|
||||
}
|
||||
|
||||
tauri_clear_install_error() {
|
||||
if [ "$TAURI_MODE" = true ]; then
|
||||
tauri_log "ERROR_CLEAR" "$1"
|
||||
printf '[TAURI:ERROR_CLEAR] %s\n' "$1" >&2
|
||||
fi
|
||||
}
|
||||
|
||||
tauri_diag_marker() {
|
||||
_diag_gpu_branch="${1:-unknown}"
|
||||
_diag_torch_index_family="${2:-none}"
|
||||
|
|
@ -498,10 +590,10 @@ _restore_studio_venv_replacement() {
|
|||
_VENV_ROLLBACK_ACTIVE=false
|
||||
return 0
|
||||
}
|
||||
substep "restoring previous environment after failed install..." "$C_WARN"
|
||||
rollback_substep "restoring previous environment after failed install..." "$C_WARN"
|
||||
rm -rf "$_VENV_ROLLBACK_TARGET"
|
||||
if mv "$_VENV_ROLLBACK_DIR" "$_VENV_ROLLBACK_TARGET"; then
|
||||
substep "restored previous environment"
|
||||
rollback_substep "restored previous environment"
|
||||
_VENV_ROLLBACK_ACTIVE=false
|
||||
_VENV_ROLLBACK_DIR=""
|
||||
else
|
||||
|
|
@ -3296,10 +3388,20 @@ case "$_torch_index_leaf" in
|
|||
if (n > 0) print vals[idx]
|
||||
}')
|
||||
fi
|
||||
# An explicit UNSLOTH_ROCM_GFX_ARCH=gfx906 pins the runtime target to the
|
||||
# MI50 / Radeon VII path and must win over Strix probe-order detection on a
|
||||
# mixed Strix + MI50 host, so the Strix reroute is suppressed when it is set.
|
||||
# Normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) and
|
||||
# trim whitespace (mirrors the Python .strip()) so the feature-flag suffix or
|
||||
# a stray newline does not defeat the exact gfx906 comparisons below.
|
||||
_gfx906_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]')
|
||||
_gfx906_env=${_gfx906_env%%:*}
|
||||
_strix_gfx=""
|
||||
case "$_runtime_gfx" in
|
||||
gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;;
|
||||
esac
|
||||
if [ "$_gfx906_env" != "gfx906" ]; then
|
||||
case "$_runtime_gfx" in
|
||||
gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;;
|
||||
esac
|
||||
fi
|
||||
# Skip rocm7.13+ generic indexes: they already ship the fixes, so the
|
||||
# arch build (rocm7.13) would be a downgrade rather than a rescue.
|
||||
if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then
|
||||
|
|
@ -3327,6 +3429,57 @@ case "$_torch_index_leaf" in
|
|||
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
|
||||
_amd_gpu_radeon=false
|
||||
fi
|
||||
# ── MI50 / Radeon VII (gfx906, Vega 20): legacy community-supported path ──
|
||||
# Newer rocm wheel families bundle ROCm libraries whose Tensile kernels
|
||||
# dropped gfx906 (rocBLAS "TensileLibrary.dat ... not read for gfx906",
|
||||
# ROCm/TheRock#1844), so a rocm6.4+/7.x index installs a torch that fails
|
||||
# at the first BLAS call. The rocm6.3 index is the last one whose wheels
|
||||
# run on gfx906 (torch 2.7.0 verified on MI50 32GB; up to 2.9 in community
|
||||
# use). Reroute any newer picked index; leave rocm6.0-6.3 alone.
|
||||
#
|
||||
# Target resolution: an explicit UNSLOTH_ROCM_GFX_ARCH wins (lets a host
|
||||
# whose rocminfo/amd-smi emit no gfx token still opt in; _gfx906_env was
|
||||
# lowercased above, before the Strix block it suppresses). Otherwise only
|
||||
# treat gfx906 as the target when it is the SOLE distinct arch present:
|
||||
# _gfx_all is de-duplicated by visible index, which loses per-device
|
||||
# ordinals on a mixed host, so a non-gfx906 selection must never be
|
||||
# downgraded to rocm6.3 -- such hosts set UNSLOTH_ROCM_GFX_ARCH to opt in.
|
||||
_gfx906_target=false
|
||||
if [ -n "$_gfx906_env" ]; then
|
||||
[ "$_gfx906_env" = "gfx906" ] && _gfx906_target=true
|
||||
elif [ -n "$_gfx_all" ]; then
|
||||
_gfx906_uniq=$(printf '%s\n' "$_gfx_all" | awk 'NF && !seen[$0]++')
|
||||
[ "$_gfx906_uniq" = "gfx906" ] && _gfx906_target=true
|
||||
fi
|
||||
# gfx906 always trains from the PyTorch rocm6.3 wheels, never the Radeon repo
|
||||
# (repo.radeon.com wheels carry no gfx906 BLAS kernels). Clear the Radeon
|
||||
# marketing-name flag as soon as gfx906 is the target -- even when the host
|
||||
# already picks rocm6.0-6.3 and the reroute below is a no-op -- so a Radeon VII
|
||||
# does not divert to the radeon branch on those versions.
|
||||
if [ "$_gfx906_target" = true ]; then
|
||||
_amd_gpu_radeon=false
|
||||
fi
|
||||
if [ "$_gfx906_target" = true ] && ! _rocm_leaf_below "$_torch_index_leaf" 6 4; then
|
||||
echo "" >&2
|
||||
echo " [WARN] gfx906 (MI50 / Radeon VII / Vega 20) detected -- routing torch to the" >&2
|
||||
echo " [WARN] rocm6.3 index: it is the last wheel family that runs on gfx906 (newer" >&2
|
||||
echo " [WARN] rocm wheels ship without gfx906 BLAS kernels and fail at first use)." >&2
|
||||
echo " [WARN] gfx906 is a community-maintained legacy path: 16-bit LoRA and full" >&2
|
||||
echo " [WARN] finetuning work out of the box; bitsandbytes 4-bit QLoRA requires a" >&2
|
||||
echo " [WARN] source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd)." >&2
|
||||
echo "" >&2
|
||||
_amd_gfx906_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
|
||||
while [ "${_amd_gfx906_base%/}" != "$_amd_gfx906_base" ]; do
|
||||
_amd_gfx906_base="${_amd_gfx906_base%/}"
|
||||
done
|
||||
TORCH_INDEX_URL="${_amd_gfx906_base}/rocm6.3"
|
||||
# Reset to the default (<2.11) window: a rocm7.2 pick raised the floor
|
||||
# to 2.11 above, which the rocm6.3 index (torch <= 2.9.x) cannot satisfy.
|
||||
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
|
||||
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
|
||||
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
|
||||
# (_amd_gpu_radeon already cleared above for every gfx906 target.)
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
fi # _torch_index_pinned guard (Radeon + Strix reroute)
|
||||
|
|
@ -3553,6 +3706,7 @@ for _p in ('torch', 'torchvision', 'torchaudio'):
|
|||
if [ "$_MIGRATED" = true ]; then
|
||||
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
|
||||
# existing torch/CUDA unless the ROCm repair below fires.
|
||||
_gfx906_bnb_snapshot
|
||||
substep "upgrading unsloth in migrated environment..."
|
||||
if [ "$SKIP_TORCH" = true ]; then
|
||||
# No-torch: install unsloth + unsloth-zoo with --no-deps (current
|
||||
|
|
@ -3594,13 +3748,18 @@ if [ "$_MIGRATED" = true ]; then
|
|||
# existing ROCm installs gain the AMD bitsandbytes build without a
|
||||
# fresh reinstall.
|
||||
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
if _is_gfx906_bnb_skip; then
|
||||
substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
|
||||
else
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
fi
|
||||
# Repair ROCm torch if overwritten during migrated install
|
||||
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
|
||||
if [ -z "$_has_hip" ]; then
|
||||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
_install_torch_default_index --force-reinstall
|
||||
fi
|
||||
_gfx906_bnb_prune
|
||||
fi
|
||||
elif [ -n "$TORCH_INDEX_URL" ]; then
|
||||
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
|
||||
|
|
@ -3791,8 +3950,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
# host stays in GGUF-only mode rather than pulling in bitsandbytes,
|
||||
# which is only useful once torch is present for training.
|
||||
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
if _is_gfx906_bnb_skip; then
|
||||
substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
|
||||
else
|
||||
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
|
||||
fi
|
||||
fi
|
||||
_gfx906_bnb_snapshot
|
||||
# Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed
|
||||
tauri_log "STEP" "Installing Unsloth"
|
||||
substep "installing unsloth (this may take a few minutes)..."
|
||||
|
|
@ -3843,6 +4007,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
substep "repairing ROCm torch (overwritten by dependency resolution)..."
|
||||
_install_torch_default_index --force-reinstall
|
||||
fi
|
||||
_gfx906_bnb_prune
|
||||
fi
|
||||
else
|
||||
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
|
||||
|
|
@ -3937,6 +4102,7 @@ if [ -n "$VENV_ABS_BIN" ]; then
|
|||
fi
|
||||
|
||||
if ! command -v bash >/dev/null 2>&1; then
|
||||
tauri_log "ERROR" "bash is required to run studio setup"
|
||||
step "setup" "bash is required to run studio setup" "$C_ERR"
|
||||
substep "Please install bash and re-run install.sh"
|
||||
exit 1
|
||||
|
|
@ -3975,6 +4141,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
|
|||
STUDIO_LOCAL_REPO="$_REPO_ROOT" \
|
||||
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
|
||||
UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \
|
||||
UNSLOTH_TAURI_MODE="$TAURI_MODE" \
|
||||
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
|
||||
else
|
||||
# Explicitly reset STUDIO_LOCAL_INSTALL / STUDIO_LOCAL_REPO so a stale
|
||||
|
|
@ -3990,9 +4157,14 @@ else
|
|||
STUDIO_LOCAL_REPO= \
|
||||
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
|
||||
UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \
|
||||
UNSLOTH_TAURI_MODE="$TAURI_MODE" \
|
||||
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
|
||||
fi
|
||||
|
||||
if [ "$_SETUP_EXIT" -eq 0 ]; then
|
||||
tauri_clear_install_error "studio setup completed"
|
||||
fi
|
||||
|
||||
# ── Make 'unsloth' available via $_LOCAL_BIN (resolved earlier) ──
|
||||
# Env-mode: $_LOCAL_BIN is $STUDIO_HOME/bin; skip shell-rc PATH append so we
|
||||
# don't pollute the user's profile with a workspace-scoped path.
|
||||
|
|
@ -4048,7 +4220,11 @@ fi
|
|||
# PATH and shortcuts are already set up so the user can fix and retry.
|
||||
if [ "$_SETUP_EXIT" -ne 0 ]; then
|
||||
echo ""
|
||||
step "error" "studio setup failed (exit code $_SETUP_EXIT)" "$C_ERR"
|
||||
if [ "$TAURI_MODE" = true ]; then
|
||||
tauri_log "ERROR_DEFAULT" "studio setup failed (exit code $_SETUP_EXIT)"
|
||||
else
|
||||
step "error" "studio setup failed (exit code $_SETUP_EXIT)" "$C_ERR"
|
||||
fi
|
||||
echo ""
|
||||
exit "$_SETUP_EXIT"
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -30,6 +30,12 @@ dependencies = [
|
|||
"pydantic",
|
||||
"pyyaml",
|
||||
"nest-asyncio",
|
||||
# Every CLI command imports studio.backend.*, which reaches structlog at
|
||||
# module level. The rest of the server stack lives in the studio extra.
|
||||
"structlog>=24.1.0",
|
||||
# unsloth_cli/__init__.py reaches click via commands/start.py, so every
|
||||
# command needs it. typer supplied it until 0.27 dropped the dependency.
|
||||
"click>=8.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
|
|
@ -68,6 +74,33 @@ include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"]
|
|||
exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"]
|
||||
|
||||
[project.optional-dependencies]
|
||||
# Studio's server stack, mirroring studio/backend/requirements/studio.txt.
|
||||
# test_studio_extra_matches_requirements.py catches drift.
|
||||
studio = [
|
||||
"typer",
|
||||
"fastapi",
|
||||
"uvicorn",
|
||||
"pydantic",
|
||||
"packaging",
|
||||
"matplotlib==3.10.9",
|
||||
"pandas",
|
||||
"nest_asyncio",
|
||||
"datasets==4.3.0",
|
||||
"pyjwt",
|
||||
"huggingface-hub==0.36.2",
|
||||
"structlog>=24.1.0",
|
||||
"diceware",
|
||||
"ddgs",
|
||||
"cryptography>=42.0.0",
|
||||
"boto3>=1.34.0",
|
||||
"httpx>=0.27.0",
|
||||
"fastmcp>=3.0.2",
|
||||
"sqlite-vec==0.1.9",
|
||||
"pymupdf==1.27.2.3",
|
||||
"pymupdf4llm==0.3.4",
|
||||
"python-docx==1.2.0",
|
||||
]
|
||||
|
||||
triton = [
|
||||
"triton>=3.0.0 ; ('linux' in sys_platform)",
|
||||
"triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
|
||||
|
|
|
|||
|
|
@ -164,6 +164,21 @@ async def get_current_subject_allow_password_change(
|
|||
)
|
||||
|
||||
|
||||
# The literal the examples ship with; pasted unedited more often than a revoked key.
|
||||
API_KEY_PLACEHOLDER = f"{API_KEY_PREFIX}YOUR_KEY"
|
||||
|
||||
|
||||
def _invalid_api_key_detail(token: str) -> str:
|
||||
"""Why the key failed. Only the example placeholder is called out; every real
|
||||
key gets one indistinguishable message, so this leaks no key existence."""
|
||||
if token == API_KEY_PLACEHOLDER:
|
||||
return (
|
||||
"This is the placeholder key from the example. Create an API key in "
|
||||
f"Unsloth Studio under Settings > API and use it in place of {API_KEY_PLACEHOLDER}."
|
||||
)
|
||||
return "Invalid or expired API key"
|
||||
|
||||
|
||||
async def _get_current_subject(
|
||||
credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool
|
||||
) -> str:
|
||||
|
|
@ -176,7 +191,7 @@ async def _get_current_subject(
|
|||
if username is None:
|
||||
raise HTTPException(
|
||||
status_code = status.HTTP_401_UNAUTHORIZED,
|
||||
detail = "Invalid or expired API key",
|
||||
detail = _invalid_api_key_detail(token),
|
||||
)
|
||||
return username
|
||||
|
||||
|
|
|
|||
|
|
@ -6,12 +6,14 @@
|
|||
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.
|
||||
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>\\t<name>`` 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. ``name`` is ggml's device description (the marketing name, e.g.
|
||||
"AMD Radeon RX 9070 XT"); empty when the registry lookup fails.
|
||||
|
||||
Uses only the standard library so it stays runnable as a bare script.
|
||||
"""
|
||||
|
|
@ -24,15 +26,30 @@ import sys
|
|||
_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.
|
||||
def _igpu_flags_and_names(base, lib, count: int) -> tuple[list[bool], list[str]]:
|
||||
"""Per-device integrated-GPU flags and descriptions 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.
|
||||
i``), so reg index == device ordinal. Returns all-False / empty-name on any
|
||||
failure so the reader never over-caps a discrete card and the memory
|
||||
readings still get through.
|
||||
"""
|
||||
flags = [False] * count
|
||||
names = [""] * count
|
||||
|
||||
# The name lookup is bound OUTSIDE the type-detection try: a ggml-base
|
||||
# without ggml_backend_dev_description (older/custom build) must degrade to
|
||||
# unnamed devices, not abort before the iGPU flags are read (which would
|
||||
# count an iGPU's shared RAM as VRAM).
|
||||
describe = None
|
||||
try:
|
||||
base.ggml_backend_dev_description.restype = ctypes.c_char_p
|
||||
base.ggml_backend_dev_description.argtypes = [ctypes.c_void_p]
|
||||
describe = base.ggml_backend_dev_description
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
|
||||
lib.ggml_backend_vk_reg.argtypes = []
|
||||
|
|
@ -45,17 +62,31 @@ def _igpu_flags(base, lib, count: int) -> list[bool]:
|
|||
|
||||
reg = lib.ggml_backend_vk_reg()
|
||||
if not reg:
|
||||
return flags
|
||||
return flags, names
|
||||
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
|
||||
if describe is not None:
|
||||
try:
|
||||
desc = describe(dev)
|
||||
if desc:
|
||||
# Tabs/newlines would corrupt the line protocol;
|
||||
# spaces are safe.
|
||||
names[i] = (
|
||||
desc.decode("utf-8", errors = "replace")
|
||||
.replace("\t", " ")
|
||||
.replace("\n", " ")
|
||||
.strip()
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
# Best-effort: any failure degrades to "discrete" so the memory
|
||||
# readings still get through instead of crashing the probe.
|
||||
# Best-effort: any failure degrades to "discrete"/"unnamed" so the
|
||||
# memory readings still get through instead of crashing the probe.
|
||||
pass
|
||||
return flags
|
||||
return flags, names
|
||||
|
||||
|
||||
def main() -> int:
|
||||
|
|
@ -63,6 +94,14 @@ def main() -> int:
|
|||
return 0
|
||||
bindir = sys.argv[1]
|
||||
|
||||
# Device names can be non-ASCII (localized drivers); the platform-default
|
||||
# stdout encoding (e.g. cp1252) would raise on them and lose the whole
|
||||
# inventory. The reader decodes UTF-8 with the same error mode.
|
||||
try:
|
||||
sys.stdout.reconfigure(encoding = "utf-8", errors = "replace")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 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.
|
||||
|
|
@ -96,12 +135,12 @@ def main() -> int:
|
|||
]
|
||||
|
||||
count = lib.ggml_backend_vk_get_device_count()
|
||||
igpu = _igpu_flags(base, lib, count)
|
||||
igpu, names = _igpu_flags_and_names(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))
|
||||
rows.append("%d\t%d\t%d\t%d\t%s" % (i, free.value, int(igpu[i]), total.value, names[i]))
|
||||
sys.stdout.write("\n".join(rows))
|
||||
return 0
|
||||
|
||||
|
|
|
|||
|
|
@ -172,6 +172,136 @@ def anthropic_messages_to_openai(
|
|||
return result
|
||||
|
||||
|
||||
_ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS = {
|
||||
"bash": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {"type": "string"},
|
||||
"restart": {"type": "boolean"},
|
||||
},
|
||||
"anyOf": [
|
||||
{"required": ["command"]},
|
||||
{"properties": {"restart": {"const": True}}, "required": ["restart"]},
|
||||
],
|
||||
},
|
||||
"text_editor": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"enum": ["view", "str_replace", "create", "insert"],
|
||||
},
|
||||
"path": {"type": "string"},
|
||||
"view_range": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"old_str": {"type": "string"},
|
||||
"new_str": {"type": "string"},
|
||||
"file_text": {"type": "string"},
|
||||
"insert_line": {"type": "integer"},
|
||||
"insert_text": {"type": "string"},
|
||||
},
|
||||
"required": ["command", "path"],
|
||||
},
|
||||
"computer": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"action": {"type": "string"},
|
||||
"coordinate": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"text": {"type": "string"},
|
||||
"duration": {"type": "number"},
|
||||
"scroll_direction": {"type": "string"},
|
||||
"scroll_amount": {"type": "integer"},
|
||||
"start_coordinate": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"key": {"type": "string"},
|
||||
},
|
||||
"required": ["action"],
|
||||
"additionalProperties": True,
|
||||
},
|
||||
"memory": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"enum": ["view", "create", "str_replace", "insert", "delete", "rename"],
|
||||
},
|
||||
"path": {"type": "string"},
|
||||
"view_range": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"minItems": 2,
|
||||
"maxItems": 2,
|
||||
},
|
||||
"file_text": {"type": "string"},
|
||||
"old_str": {"type": "string"},
|
||||
"new_str": {"type": "string"},
|
||||
"insert_line": {"type": "integer"},
|
||||
"insert_text": {"type": "string"},
|
||||
"old_path": {"type": "string"},
|
||||
"new_path": {"type": "string"},
|
||||
},
|
||||
"required": ["command"],
|
||||
},
|
||||
}
|
||||
|
||||
_ANTHROPIC_SCHEMA_CLIENT_TOOL_DESCRIPTIONS = {
|
||||
"bash": "Run a command in the caller-owned persistent bash session, or restart it.",
|
||||
"text_editor": "View, create, or edit files in the caller-owned filesystem.",
|
||||
"computer": "Interact with the caller-owned computer using an action and its parameters.",
|
||||
"memory": "Store and retrieve files in the caller-owned persistent memory directory.",
|
||||
}
|
||||
|
||||
|
||||
def anthropic_schema_client_tool_kind(tool) -> Optional[str]:
|
||||
"""Return the kind of a schema-less Anthropic client tool, if recognized."""
|
||||
td = tool if isinstance(tool, dict) else tool.model_dump()
|
||||
if td.get("input_schema") is not None:
|
||||
return None
|
||||
type_ = td.get("type")
|
||||
if not isinstance(type_, str):
|
||||
return None
|
||||
kind, separator, version = type_.rpartition("_")
|
||||
if (
|
||||
separator
|
||||
and kind in _ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS
|
||||
and len(version) == 8
|
||||
and version.isdigit()
|
||||
):
|
||||
return kind
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_schema_client_tool_parameters(td: dict, kind: str) -> dict:
|
||||
parameters = _ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS[kind]
|
||||
if kind != "text_editor":
|
||||
return parameters
|
||||
|
||||
version = td["type"].rpartition("_")[2]
|
||||
commands = list(parameters["properties"]["command"]["enum"])
|
||||
if version < "20250429":
|
||||
commands.append("undo_edit")
|
||||
return {
|
||||
**parameters,
|
||||
"properties": {
|
||||
**parameters["properties"],
|
||||
"command": {**parameters["properties"]["command"], "enum": commands},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def anthropic_tools_to_openai(tools: list) -> list[dict]:
|
||||
"""Convert Anthropic client tools to OpenAI function-tool format."""
|
||||
result = []
|
||||
|
|
@ -179,6 +309,9 @@ def anthropic_tools_to_openai(tools: list) -> list[dict]:
|
|||
td = t if isinstance(t, dict) else t.model_dump()
|
||||
name = td.get("name")
|
||||
input_schema = td.get("input_schema")
|
||||
schema_client_kind = anthropic_schema_client_tool_kind(td)
|
||||
if schema_client_kind is not None:
|
||||
input_schema = _anthropic_schema_client_tool_parameters(td, schema_client_kind)
|
||||
if not name or input_schema is None:
|
||||
continue
|
||||
result.append(
|
||||
|
|
@ -186,7 +319,8 @@ def anthropic_tools_to_openai(tools: list) -> list[dict]:
|
|||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"description": td.get("description", ""),
|
||||
"description": td.get("description")
|
||||
or _ANTHROPIC_SCHEMA_CLIENT_TOOL_DESCRIPTIONS.get(schema_client_kind, ""),
|
||||
"parameters": input_schema,
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
|
|
@ -18,6 +19,14 @@ _MAX_PROMPT_CHARS = 12000
|
|||
_MAX_REPLY_CHARS = 12000
|
||||
_PREVIEW_CHARS = 360
|
||||
|
||||
# Opt-in startup kill switch for Studio's in-memory API monitor.
|
||||
_DISABLE_ENV = "UNSLOTH_STUDIO_DISABLE_API_MONITOR"
|
||||
_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
|
||||
|
||||
|
||||
def _api_monitor_disabled() -> bool:
|
||||
return os.environ.get(_DISABLE_ENV, "").strip().lower() in _TRUE_VALUES
|
||||
|
||||
|
||||
def _trim(text: Optional[str], limit: int) -> str:
|
||||
if not text:
|
||||
|
|
@ -52,6 +61,13 @@ class ApiMonitorEntry:
|
|||
total_tokens: Optional[int] = None
|
||||
total_tokens_authoritative: bool = False
|
||||
error: Optional[str] = None
|
||||
# "request" (HTTP call) or "lifecycle" (model load/unload: event/reason, not a prompt; shared).
|
||||
kind: str = "request"
|
||||
event: Optional[str] = None
|
||||
reason: Optional[str] = None
|
||||
shared: bool = False
|
||||
# 0-100 for a running download row; None when not applicable.
|
||||
progress: Optional[float] = None
|
||||
|
||||
def snapshot(self, *, include_details: bool = True) -> dict[str, Any]:
|
||||
duration_ms = None
|
||||
|
|
@ -85,6 +101,10 @@ class ApiMonitorEntry:
|
|||
"completion_tokens": self.completion_tokens,
|
||||
"total_tokens": self.total_tokens,
|
||||
"error": self.error,
|
||||
"kind": self.kind,
|
||||
"event": self.event,
|
||||
"reason": self.reason,
|
||||
"progress": self.progress,
|
||||
}
|
||||
if include_details:
|
||||
payload["prompt"] = self.prompt
|
||||
|
|
@ -93,10 +113,16 @@ class ApiMonitorEntry:
|
|||
|
||||
|
||||
class ApiMonitor:
|
||||
def __init__(self, max_entries: int = _MAX_ENTRIES):
|
||||
def __init__(
|
||||
self,
|
||||
max_entries: int = _MAX_ENTRIES,
|
||||
*,
|
||||
enabled: bool = True,
|
||||
):
|
||||
self._entries: deque[ApiMonitorEntry] = deque()
|
||||
self._max_entries = max(0, max_entries)
|
||||
self._lock = threading.Lock()
|
||||
self._enabled = enabled
|
||||
|
||||
def start(
|
||||
self,
|
||||
|
|
@ -108,6 +134,8 @@ class ApiMonitor:
|
|||
context_length: Optional[int] = None,
|
||||
subject: Optional[str] = None,
|
||||
) -> str:
|
||||
if not self._enabled:
|
||||
return ""
|
||||
now = time.time()
|
||||
entry = ApiMonitorEntry(
|
||||
id = f"apireq_{uuid.uuid4().hex[:12]}",
|
||||
|
|
@ -127,6 +155,75 @@ class ApiMonitor:
|
|||
self._trim_terminal_locked()
|
||||
return entry.id
|
||||
|
||||
def record_lifecycle(
|
||||
self,
|
||||
*,
|
||||
event: str,
|
||||
model: str,
|
||||
reason: Optional[str] = None,
|
||||
running: bool = False,
|
||||
) -> str:
|
||||
"""Record a model load/unload alongside the request traffic that caused it.
|
||||
|
||||
``running=True`` opens the row for the caller to close with :meth:`finish` /
|
||||
:meth:`fail`; an unload is terminal on arrival. Rows are shared (visible to
|
||||
every subject) and share the request retention budget.
|
||||
"""
|
||||
if not self._enabled:
|
||||
return ""
|
||||
now = time.time()
|
||||
entry = ApiMonitorEntry(
|
||||
id = f"apievt_{uuid.uuid4().hex[:12]}",
|
||||
endpoint = f"model.{event}",
|
||||
method = "",
|
||||
model = model or "default",
|
||||
prompt = "",
|
||||
status = "running" if running else "completed",
|
||||
started_at = now,
|
||||
updated_at = now,
|
||||
started_monotonic = time.monotonic(),
|
||||
finished_at = None if running else now,
|
||||
finished_monotonic = None if running else time.monotonic(),
|
||||
kind = "lifecycle",
|
||||
event = event,
|
||||
reason = reason,
|
||||
shared = True,
|
||||
)
|
||||
with self._lock:
|
||||
self._entries.appendleft(entry)
|
||||
self._trim_terminal_locked()
|
||||
return entry.id
|
||||
|
||||
def relabel(self, entry_id: Optional[str], model: str) -> None:
|
||||
"""Rename an open lifecycle row once the load resolves its real id: up front
|
||||
the caller only has the load path, which may be an HF snapshot dir."""
|
||||
if not entry_id or not model:
|
||||
return
|
||||
with self._lock:
|
||||
entry = self._find_locked(entry_id)
|
||||
if entry is not None:
|
||||
entry.model = model
|
||||
entry.updated_at = time.time()
|
||||
|
||||
def set_progress(self, entry_id: Optional[str], progress: Optional[float]) -> None:
|
||||
"""Update an open download row's percentage (clamped to 0-100)."""
|
||||
if not entry_id or progress is None:
|
||||
return
|
||||
with self._lock:
|
||||
entry = self._find_locked(entry_id)
|
||||
if entry is not None and entry.status == "running":
|
||||
entry.progress = min(100.0, max(0.0, float(progress)))
|
||||
entry.updated_at = time.time()
|
||||
|
||||
def discard(self, entry_id: Optional[str]) -> None:
|
||||
"""Drop a row that turned out not to be an event (an already-satisfied load)."""
|
||||
if not entry_id:
|
||||
return
|
||||
with self._lock:
|
||||
entry = self._find_locked(entry_id)
|
||||
if entry is not None:
|
||||
self._entries.remove(entry)
|
||||
|
||||
def append_reply(self, entry_id: Optional[str], text: str) -> None:
|
||||
if not entry_id or not text:
|
||||
return
|
||||
|
|
@ -212,6 +309,18 @@ class ApiMonitor:
|
|||
self._entries.appendleft(entry)
|
||||
self._trim_terminal_locked()
|
||||
|
||||
def fail_open(self, entry_id: Optional[str], error: str) -> None:
|
||||
"""Fail only a still-open row: unlike :meth:`fail`, a catch-all in a
|
||||
``finally`` cannot stamp an error onto a request that already succeeded."""
|
||||
if not entry_id:
|
||||
return
|
||||
with self._lock:
|
||||
entry = self._find_locked(entry_id)
|
||||
if entry is None or entry.finished_at is not None:
|
||||
return
|
||||
# Same lock as the check, so a finish() cannot land in between.
|
||||
self._fail_locked(entry, error)
|
||||
|
||||
def fail(self, entry_id: Optional[str], error: str) -> None:
|
||||
if not entry_id:
|
||||
return
|
||||
|
|
@ -224,15 +333,18 @@ class ApiMonitor:
|
|||
if error:
|
||||
entry.error = _trim(error, 1000)
|
||||
return
|
||||
now = time.time()
|
||||
entry.status = "error"
|
||||
entry.error = _trim(error, 1000)
|
||||
entry.updated_at = now
|
||||
entry.finished_at = now
|
||||
entry.finished_monotonic = time.monotonic()
|
||||
self._entries.remove(entry)
|
||||
self._entries.appendleft(entry)
|
||||
self._trim_terminal_locked()
|
||||
self._fail_locked(entry, error)
|
||||
|
||||
def _fail_locked(self, entry: ApiMonitorEntry, error: str) -> None:
|
||||
now = time.time()
|
||||
entry.status = "error"
|
||||
entry.error = _trim(error, 1000)
|
||||
entry.updated_at = now
|
||||
entry.finished_at = now
|
||||
entry.finished_monotonic = time.monotonic()
|
||||
self._entries.remove(entry)
|
||||
self._entries.appendleft(entry)
|
||||
self._trim_terminal_locked()
|
||||
|
||||
def snapshot(
|
||||
self,
|
||||
|
|
@ -244,7 +356,7 @@ class ApiMonitor:
|
|||
return [
|
||||
entry.snapshot(include_details = include_details)
|
||||
for entry in self._entries
|
||||
if subject is None or entry.subject == subject
|
||||
if self._visible(entry, subject)
|
||||
]
|
||||
|
||||
def get(
|
||||
|
|
@ -257,22 +369,29 @@ class ApiMonitor:
|
|||
entry = self._find_locked(entry_id)
|
||||
if entry is None:
|
||||
return None
|
||||
if subject is not None and entry.subject != subject:
|
||||
if not self._visible(entry, subject):
|
||||
return None
|
||||
return entry.snapshot(include_details = True)
|
||||
|
||||
def active_count(self, *, subject: Optional[str] = None) -> int:
|
||||
# Lifecycle rows show as "running" while loading but are not in-flight API requests.
|
||||
with self._lock:
|
||||
return sum(
|
||||
1
|
||||
for entry in self._entries
|
||||
if entry.status == "running" and (subject is None or entry.subject == subject)
|
||||
if entry.status == "running"
|
||||
and entry.kind != "lifecycle"
|
||||
and (subject is None or entry.subject == subject)
|
||||
)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._entries.clear()
|
||||
|
||||
@staticmethod
|
||||
def _visible(entry: ApiMonitorEntry, subject: Optional[str]) -> bool:
|
||||
return subject is None or entry.subject == subject or entry.shared
|
||||
|
||||
def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]:
|
||||
for entry in self._entries:
|
||||
if entry.id == entry_id:
|
||||
|
|
@ -292,4 +411,4 @@ class ApiMonitor:
|
|||
self._entries = kept
|
||||
|
||||
|
||||
api_monitor = ApiMonitor()
|
||||
api_monitor = ApiMonitor(enabled = not _api_monitor_disabled())
|
||||
|
|
|
|||
|
|
@ -326,6 +326,58 @@ def _normalize_tool_call_arguments(messages: list) -> list:
|
|||
return out if mutated else messages
|
||||
|
||||
|
||||
def _take_tool_result(pending: list, call_id) -> Optional[dict]:
|
||||
if call_id:
|
||||
for i, result in enumerate(pending):
|
||||
if result.get("tool_call_id") == call_id:
|
||||
return pending.pop(i)
|
||||
for i, result in enumerate(pending):
|
||||
if not result.get("tool_call_id"):
|
||||
return pending.pop(i)
|
||||
return None
|
||||
|
||||
|
||||
def _split_parallel_tool_calls(messages: list) -> list:
|
||||
"""Llama 3.x templates render one call per message, so split parallel calls
|
||||
into consecutive single-call messages, each followed by its own result."""
|
||||
if not any(isinstance(m, dict) and len(m.get("tool_calls") or ()) > 1 for m in messages):
|
||||
return messages
|
||||
|
||||
out: list = []
|
||||
i = 0
|
||||
total = len(messages)
|
||||
while i < total:
|
||||
msg = messages[i]
|
||||
calls = msg.get("tool_calls") if isinstance(msg, dict) else None
|
||||
if not calls or len(calls) <= 1:
|
||||
out.append(msg)
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Tool results right after this message answer its calls.
|
||||
j = i + 1
|
||||
pending: list = []
|
||||
while (
|
||||
j < total
|
||||
and isinstance(messages[j], dict)
|
||||
and messages[j].get("role") in ("tool", "ipython")
|
||||
):
|
||||
pending.append(messages[j])
|
||||
j += 1
|
||||
|
||||
for idx, call in enumerate(calls):
|
||||
piece = {**msg, "tool_calls": [call]}
|
||||
if idx:
|
||||
piece["content"] = ""
|
||||
out.append(piece)
|
||||
result = _take_tool_result(pending, call.get("id") if isinstance(call, dict) else None)
|
||||
if result is not None:
|
||||
out.append(result)
|
||||
out.extend(pending)
|
||||
i = j
|
||||
return out
|
||||
|
||||
|
||||
def apply_chat_template_for_generation(
|
||||
tokenizer,
|
||||
messages: list,
|
||||
|
|
@ -378,13 +430,21 @@ def apply_chat_template_for_generation(
|
|||
try:
|
||||
return _render(messages)
|
||||
except Exception:
|
||||
# Strict tool templates reject the JSON-string ``arguments`` form via
|
||||
# TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
|
||||
# Original messages render first, so working templates stay byte-identical.
|
||||
# Retry with repairs applied cumulatively. Originals render first, so
|
||||
# working templates stay byte-identical.
|
||||
candidates: list = []
|
||||
normalized = _normalize_tool_call_arguments(messages)
|
||||
if normalized is messages:
|
||||
raise
|
||||
return _render(normalized)
|
||||
if normalized is not messages:
|
||||
candidates.append(normalized)
|
||||
split = _split_parallel_tool_calls(normalized)
|
||||
if split is not normalized:
|
||||
candidates.append(split)
|
||||
for candidate in candidates:
|
||||
try:
|
||||
return _render(candidate)
|
||||
except Exception:
|
||||
continue
|
||||
raise
|
||||
|
||||
|
||||
def render_native_template(
|
||||
|
|
|
|||
|
|
@ -2281,8 +2281,13 @@ class InferenceBackend:
|
|||
except Exception as e:
|
||||
logger.warning(f"Could not fully reset model state for {model_name}: {e}")
|
||||
|
||||
def reset_generation_state(self):
|
||||
"""Reset any cached generation state to prevent hanging after errors"""
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
"""Reset any cached generation state to prevent hanging after errors
|
||||
|
||||
``caller_cancel_event`` is accepted for signature parity with the
|
||||
orchestrator, which uses it to drop a reset from a request that never
|
||||
started. Nothing here cancels a live generation, so it is unused.
|
||||
"""
|
||||
try:
|
||||
# Clear cached state for ALL loaded models
|
||||
for model_name in self.models.keys():
|
||||
|
|
|
|||
|
|
@ -13,37 +13,85 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
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"
|
||||
# dataclass(slots = True) halves per-instance overhead. Measured as perf-neutral
|
||||
# here, not a speed win: it costs a little on construction and gains it back on
|
||||
# access. It is 3.10+ and this package declares >=3.9, so gate it rather than
|
||||
# dropping it outright. Empty on 3.9 means a plain dataclass.
|
||||
_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {}
|
||||
|
||||
|
||||
ADMISSION_CONTROL_ENV = "UNSLOTH_LLAMA_ADMISSION_CONTROL"
|
||||
ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_TIMEOUT"
|
||||
ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_LLAMA_ADMISSION_KEEPALIVE_INTERVAL"
|
||||
ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_LLAMA_ADMISSION_MAX_QUEUE"
|
||||
ADMISSION_QUEUE_PER_SLOT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_PER_SLOT"
|
||||
|
||||
# The UNSLOTH_OPENAI_COMPAT_* spellings predate this queue being shared with the
|
||||
# Anthropic /v1/messages route (same llama-server slots). Still honored; the
|
||||
# neutral name above wins when both are set.
|
||||
_LEGACY_ENV = {
|
||||
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
|
||||
# None: a queued request waits for its slot indefinitely rather than timing out.
|
||||
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None
|
||||
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0
|
||||
DEFAULT_ADMISSION_MAX_QUEUE = 64
|
||||
# None: no absolute cap, the wait line is sized from the pool instead.
|
||||
DEFAULT_ADMISSION_MAX_QUEUE = None
|
||||
# Wait line = 16 x the serving slots, so it tracks --parallel (4 slots -> 64
|
||||
# waiters, 8 -> 128). Purely a memory guard; waiting itself is never timed out.
|
||||
DEFAULT_ADMISSION_QUEUE_PER_SLOT = 16
|
||||
# Floor for the scaled line, so a 1-slot backend (plain `unsloth studio`, or any
|
||||
# load downshifted to fit VRAM) keeps the depth it had before scaling existed
|
||||
# rather than dropping to 16 and rejecting callers that used to queue.
|
||||
DEFAULT_ADMISSION_MIN_QUEUE = 64
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
@dataclass(frozen = True, **_SLOTS)
|
||||
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
|
||||
queue_per_slot: Optional[int] = DEFAULT_ADMISSION_QUEUE_PER_SLOT
|
||||
# Unconditional floor on the scaled line. The env path clears it when the
|
||||
# operator sets QUEUE_PER_SLOT, so only the default multiplier is floored.
|
||||
min_queue: Optional[int] = DEFAULT_ADMISSION_MIN_QUEUE
|
||||
|
||||
def queue_limit(self, capacity: int) -> Optional[int]:
|
||||
"""How many callers may line up for a pool of ``capacity`` slots.
|
||||
|
||||
An explicit ``max_queue`` wins; otherwise the line scales with the slots
|
||||
so it follows ``--parallel``. The default multiplier is floored, so a
|
||||
1-slot backend does not end up shallower than it was before scaling. None
|
||||
(or any non-positive setting) means an unbounded line.
|
||||
"""
|
||||
if self.max_queue is not None:
|
||||
return self.max_queue if self.max_queue > 0 else None
|
||||
if not self.queue_per_slot or self.queue_per_slot <= 0:
|
||||
return None
|
||||
scaled = self.queue_per_slot * max(1, capacity)
|
||||
return max(self.min_queue, scaled) if self.min_queue else scaled
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
@dataclass(frozen = True, **_SLOTS)
|
||||
class LlamaAdmissionSnapshot:
|
||||
key: str
|
||||
capacity: int
|
||||
active: int
|
||||
queued: int
|
||||
free: int = 0
|
||||
|
||||
|
||||
class LlamaAdmissionError(Exception):
|
||||
|
|
@ -69,8 +117,17 @@ class LlamaAdmissionCancelled(LlamaAdmissionError):
|
|||
pass
|
||||
|
||||
|
||||
def _bool_env(name: str, default: bool) -> bool:
|
||||
def _raw_env(name: str) -> Optional[str]:
|
||||
"""Value for a canonical name, falling back to its legacy spelling."""
|
||||
value = os.environ.get(name)
|
||||
if value is None or not value.strip():
|
||||
legacy = _LEGACY_ENV.get(name)
|
||||
value = os.environ.get(legacy) if legacy else None
|
||||
return value
|
||||
|
||||
|
||||
def _bool_env(name: str, default: bool) -> bool:
|
||||
value = _raw_env(name)
|
||||
if value is None or not value.strip():
|
||||
return default
|
||||
value = value.strip().lower()
|
||||
|
|
@ -82,7 +139,7 @@ def _bool_env(name: str, default: bool) -> bool:
|
|||
|
||||
|
||||
def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]:
|
||||
value = os.environ.get(name)
|
||||
value = _raw_env(name)
|
||||
if value is None or not value.strip():
|
||||
return default
|
||||
try:
|
||||
|
|
@ -93,7 +150,7 @@ def _optional_positive_float_env(name: str, default: Optional[float]) -> Optiona
|
|||
|
||||
|
||||
def _positive_float_env(name: str, default: float) -> float:
|
||||
value = os.environ.get(name)
|
||||
value = _raw_env(name)
|
||||
if value is None or not value.strip():
|
||||
return default
|
||||
try:
|
||||
|
|
@ -103,19 +160,38 @@ def _positive_float_env(name: str, default: float) -> float:
|
|||
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
|
||||
def _queue_limits_from_env() -> tuple[Optional[int], Optional[int], Optional[int]]:
|
||||
"""(max_queue, queue_per_slot, min_queue) from the environment.
|
||||
|
||||
An absolute MAX_QUEUE wins outright; MAX_QUEUE=0 asks for an unbounded line.
|
||||
Unset leaves the per-slot multiplier in charge (itself 0 for unbounded). The
|
||||
floor applies only to the default multiplier: setting QUEUE_PER_SLOT means
|
||||
the operator wants that exact depth, however shallow.
|
||||
"""
|
||||
# Explicit means it parsed, not just that something was set: a typo falls back
|
||||
# to the default multiplier, so it has to keep the default's floor too.
|
||||
raw_per_slot = _raw_env(ADMISSION_QUEUE_PER_SLOT_ENV)
|
||||
try:
|
||||
parsed = int(value.strip())
|
||||
per_slot = int((raw_per_slot or "").strip())
|
||||
except ValueError:
|
||||
return default
|
||||
return parsed if parsed > 0 else None
|
||||
per_slot, min_queue = DEFAULT_ADMISSION_QUEUE_PER_SLOT, DEFAULT_ADMISSION_MIN_QUEUE
|
||||
else:
|
||||
per_slot, min_queue = (per_slot if per_slot > 0 else None), None
|
||||
raw = _raw_env(ADMISSION_MAX_QUEUE_ENV)
|
||||
if raw is None or not raw.strip():
|
||||
return None, per_slot, min_queue
|
||||
try:
|
||||
parsed = int(raw.strip())
|
||||
except ValueError:
|
||||
return None, per_slot, min_queue
|
||||
return (parsed, None, None) if parsed > 0 else (None, None, None)
|
||||
|
||||
|
||||
def llama_admission_config_from_env() -> LlamaAdmissionConfig:
|
||||
max_queue, queue_per_slot, min_queue = _queue_limits_from_env()
|
||||
return LlamaAdmissionConfig(
|
||||
queue_per_slot = queue_per_slot,
|
||||
min_queue = min_queue,
|
||||
enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED),
|
||||
queue_timeout_s = _optional_positive_float_env(
|
||||
ADMISSION_QUEUE_TIMEOUT_ENV,
|
||||
|
|
@ -125,14 +201,11 @@ def llama_admission_config_from_env() -> LlamaAdmissionConfig:
|
|||
ADMISSION_KEEPALIVE_INTERVAL_ENV,
|
||||
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
|
||||
),
|
||||
max_queue = _optional_positive_int_env(
|
||||
ADMISSION_MAX_QUEUE_ENV,
|
||||
DEFAULT_ADMISSION_MAX_QUEUE,
|
||||
),
|
||||
max_queue = max_queue,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
@dataclass(**_SLOTS)
|
||||
class _Waiter:
|
||||
loop: asyncio.AbstractEventLoop
|
||||
future: asyncio.Future
|
||||
|
|
@ -141,20 +214,100 @@ class _Waiter:
|
|||
|
||||
|
||||
class LlamaAdmissionLease:
|
||||
def __init__(self, queue: Optional["LlamaAdmissionQueue"]):
|
||||
__slots__ = ("_queue", "_slot", "_released", "_release_lock", "_parked")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
queue: Optional["LlamaAdmissionQueue"],
|
||||
slot: Optional[int] = None,
|
||||
):
|
||||
self._queue = queue
|
||||
self._slot = slot
|
||||
self._released = False
|
||||
self._release_lock = threading.Lock()
|
||||
self._parked = False
|
||||
|
||||
@property
|
||||
def slot(self) -> Optional[int]:
|
||||
"""Pool slot this lease holds, or None when admission is disabled."""
|
||||
return self._slot
|
||||
|
||||
def park(self) -> None:
|
||||
"""Hand the slot back while this holder waits on something off the GPU.
|
||||
|
||||
A run stopped on a tool approval prompt is not decoding, so holding its
|
||||
slot would let unanswered prompts fill the pool while llama-server idles.
|
||||
The lease itself stays valid: releasing it after a park is still correct.
|
||||
"""
|
||||
queue = self._queue
|
||||
slot = None
|
||||
with self._release_lock:
|
||||
if queue is None or self._released or self._parked:
|
||||
return
|
||||
self._parked = True
|
||||
slot, self._slot = self._slot, None
|
||||
queue.park(slot)
|
||||
|
||||
def unpark(self) -> None:
|
||||
"""Drop the parked state without reclaiming a slot.
|
||||
|
||||
For a holder that is tearing down: it will not decode again. Resuming
|
||||
holders must use ``unpark_async``, which waits for a slot instead of
|
||||
going back to llama-server past the admission limit.
|
||||
"""
|
||||
with self._release_lock:
|
||||
if not self._parked:
|
||||
return
|
||||
self._parked = False
|
||||
if self._queue is not None:
|
||||
self._queue.unpark()
|
||||
|
||||
async def unpark_async(
|
||||
self,
|
||||
*,
|
||||
cancel_event = None,
|
||||
poll_s: float = 0.02,
|
||||
) -> None:
|
||||
"""Take a slot back, waiting until the pool has room.
|
||||
|
||||
``park`` gave the slot to a waiter, so by the time the user answers the
|
||||
prompt someone else may be decoding in it. Resuming regardless put two
|
||||
holders on a one-slot server. Gives up if the caller is cancelled, since
|
||||
the holder is then leaving anyway and must not be stuck here.
|
||||
"""
|
||||
queue = self._queue
|
||||
if queue is None or not self._parked:
|
||||
return
|
||||
slot = await queue.acquire_parked_slot(cancel_event = cancel_event, poll_s = poll_s)
|
||||
stranded = None
|
||||
with self._release_lock:
|
||||
# release() may have run during the wait; it clears the flag and does
|
||||
# the unpark itself, so only the caller that clears it here repeats one.
|
||||
parked, self._parked = self._parked, False
|
||||
if self._released:
|
||||
# Released while waiting: this lease will never hand the slot
|
||||
# back, so return it here rather than strand it for good.
|
||||
stranded = slot
|
||||
else:
|
||||
self._slot = slot
|
||||
if parked:
|
||||
queue.unpark()
|
||||
if stranded is not None:
|
||||
queue.release(stranded)
|
||||
|
||||
def release(self) -> None:
|
||||
queue = None
|
||||
parked = False
|
||||
with self._release_lock:
|
||||
if self._released:
|
||||
return
|
||||
self._released = True
|
||||
queue = self._queue
|
||||
parked, self._parked = self._parked, False
|
||||
if queue is not None:
|
||||
queue.release()
|
||||
if parked:
|
||||
queue.unpark()
|
||||
queue.release(self._slot)
|
||||
|
||||
async def __aenter__(self) -> "LlamaAdmissionLease":
|
||||
return self
|
||||
|
|
@ -164,6 +317,8 @@ class LlamaAdmissionLease:
|
|||
|
||||
|
||||
class LlamaAdmissionReservation:
|
||||
__slots__ = ("_queue", "_lease", "_waiter", "snapshot")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -195,6 +350,13 @@ class LlamaAdmissionReservation:
|
|||
return self._lease
|
||||
|
||||
async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]:
|
||||
"""Wait up to ``timeout_s`` for a slot.
|
||||
|
||||
A timeout leaves this reservation queued so the caller can poll again.
|
||||
Any exit that abandons the wait for good must call ``cancel()``, or the
|
||||
slot granted later is delivered to a future nobody reads and is never
|
||||
released.
|
||||
"""
|
||||
lease = self.lease_nowait()
|
||||
if lease is not None:
|
||||
return lease
|
||||
|
|
@ -229,12 +391,74 @@ class LlamaAdmissionReservation:
|
|||
|
||||
|
||||
class LlamaAdmissionQueue:
|
||||
"""A fixed pool of generation slots for one llama-server, plus a FIFO wait line.
|
||||
|
||||
The pool mirrors llama-server's own ``--parallel`` slots: ``capacity`` slot ids
|
||||
are each either free or held by exactly one caller. A caller that finds every
|
||||
slot busy waits in arrival order and is handed the next slot to free, so no
|
||||
caller is starved. This bounds only the callers that reserve: chat completions
|
||||
and messages do, while /v1/completions, Studio's own chat endpoint and RAG
|
||||
captioning all reach llama-server directly, so it is not a global cap.
|
||||
Waiting is unbounded in time by default (``queue_timeout_s``
|
||||
None); the wait line itself is bounded, and only how many may line up before
|
||||
new arrivals are rejected. By default that is ``16 x slots`` floored at 64,
|
||||
not unlimited: an unbounded line takes ``max_queue`` or ``queue_per_slot``
|
||||
set to 0. See ``LlamaAdmissionConfig.queue_limit``.
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"key",
|
||||
"_lock",
|
||||
"_capacity",
|
||||
"_free",
|
||||
"_in_use",
|
||||
"_held",
|
||||
"_waiters",
|
||||
"_parked",
|
||||
"_unpark_tickets",
|
||||
"_unpark_seq",
|
||||
)
|
||||
|
||||
def __init__(self, key: str):
|
||||
self.key = key
|
||||
self._lock = threading.Lock()
|
||||
self._active = 0
|
||||
self._capacity = 1
|
||||
self._free: list[int] = [0]
|
||||
# Held slots as a bitmask: one int instead of a set, so the pool costs the
|
||||
# same whether it is idle or saturated. _held is its popcount, kept as a
|
||||
# counter because int.bit_count() is 3.10+ and this package targets 3.9.
|
||||
self._in_use = 0
|
||||
self._held = 0
|
||||
self._waiters: Deque[_Waiter] = deque()
|
||||
# Holders parked on a tool approval prompt. They hold no slot, so this only
|
||||
# keeps the queue off the idle-eviction list while they are away.
|
||||
self._parked = 0
|
||||
# FIFO tickets for holders resuming from a park (see acquire_parked_slot). A
|
||||
# bare count deadlocked: every approved holder blocked every other one.
|
||||
self._unpark_tickets: Deque[int] = deque()
|
||||
self._unpark_seq = 0
|
||||
|
||||
def _resize_pool_locked(self, capacity: int) -> None:
|
||||
# Slots past a shrunk capacity retire when their holder releases them.
|
||||
if capacity == self._capacity:
|
||||
return
|
||||
self._capacity = capacity
|
||||
self._free = [slot for slot in range(capacity) if not self._in_use >> slot & 1]
|
||||
|
||||
def _can_admit_locked(self, reserved: int) -> bool:
|
||||
# Slots still held above a shrunk capacity keep occupying the backend, so
|
||||
# count every held slot against the ceiling, not just the ids below it.
|
||||
# ``reserved`` holds slots back for approved holders waiting to resume;
|
||||
# without it a stream of new arrivals took the next slot, forever.
|
||||
return bool(self._free) and (self._held + reserved) < self._capacity
|
||||
|
||||
def _take_slot_locked(self, reserved: int) -> Optional[int]:
|
||||
if not self._can_admit_locked(reserved):
|
||||
return None
|
||||
slot = self._free.pop()
|
||||
self._in_use |= 1 << slot
|
||||
self._held += 1
|
||||
return slot
|
||||
|
||||
def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation:
|
||||
capacity = max(1, int(capacity or 1))
|
||||
|
|
@ -242,22 +466,25 @@ class LlamaAdmissionQueue:
|
|||
return LlamaAdmissionReservation(
|
||||
queue = None,
|
||||
lease = LlamaAdmissionLease(None),
|
||||
snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0),
|
||||
snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0, capacity),
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
with self._lock:
|
||||
self._capacity = capacity
|
||||
self._prune_waiters_locked()
|
||||
self._resize_pool_locked(capacity)
|
||||
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:
|
||||
if not self._waiters:
|
||||
slot = self._take_slot_locked(len(self._unpark_tickets))
|
||||
if slot is not None:
|
||||
# No snapshot here: callers read it through snapshot_now(),
|
||||
# which re-reads the queue, so building one per admitted
|
||||
# request would be pure allocation on the hot path.
|
||||
return LlamaAdmissionReservation(
|
||||
queue = self,
|
||||
lease = LlamaAdmissionLease(self, slot),
|
||||
)
|
||||
limit = config.queue_limit(self._capacity)
|
||||
if limit is not None and self._live_waiters_locked() >= limit:
|
||||
raise LlamaAdmissionQueueFull(
|
||||
"llama-server generation queue is full",
|
||||
snapshot = self._snapshot_locked(),
|
||||
|
|
@ -270,15 +497,74 @@ class LlamaAdmissionQueue:
|
|||
return LlamaAdmissionReservation(
|
||||
queue = self,
|
||||
waiter = waiter,
|
||||
snapshot = self._snapshot_locked(),
|
||||
)
|
||||
|
||||
def release(self) -> None:
|
||||
def _release_slot_locked(self, slot: Optional[int]) -> None:
|
||||
# A slot id at or past a shrunk capacity retires instead of returning.
|
||||
if slot is None or not self._in_use >> slot & 1:
|
||||
return
|
||||
self._in_use &= ~(1 << slot)
|
||||
self._held -= 1
|
||||
if slot < self._capacity:
|
||||
self._free.append(slot)
|
||||
|
||||
def release(self, slot: Optional[int]) -> None:
|
||||
with self._lock:
|
||||
if self._active > 0:
|
||||
self._active -= 1
|
||||
self._release_slot_locked(slot)
|
||||
self._grant_waiters_locked()
|
||||
|
||||
def park(self, slot: Optional[int]) -> None:
|
||||
"""Return a parked holder's slot to the pool. See ``LlamaAdmissionLease.park``."""
|
||||
with self._lock:
|
||||
self._parked += 1
|
||||
self._release_slot_locked(slot)
|
||||
self._grant_waiters_locked()
|
||||
|
||||
def unpark(self) -> None:
|
||||
with self._lock:
|
||||
if self._parked > 0:
|
||||
self._parked -= 1
|
||||
|
||||
async def acquire_parked_slot(
|
||||
self,
|
||||
*,
|
||||
cancel_event = None,
|
||||
poll_s: float = 0.02,
|
||||
) -> Optional[int]:
|
||||
"""Wait for a slot for a holder resuming from a park, None if cancelled.
|
||||
|
||||
Ordered by ticket rather than counted, so approvals resume in the order
|
||||
they came back: counting them made every approved holder block every
|
||||
other one, and with nothing decoding that never resolved.
|
||||
"""
|
||||
with self._lock:
|
||||
self._unpark_seq += 1
|
||||
ticket = self._unpark_seq
|
||||
self._unpark_tickets.append(ticket)
|
||||
try:
|
||||
while True:
|
||||
with self._lock:
|
||||
ahead = 0
|
||||
for queued in self._unpark_tickets:
|
||||
if queued == ticket:
|
||||
break
|
||||
ahead += 1
|
||||
# Only the approvals ahead of this one hold slots back from it.
|
||||
slot = self._take_slot_locked(ahead)
|
||||
if slot is not None:
|
||||
return slot
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return None
|
||||
await asyncio.sleep(poll_s)
|
||||
finally:
|
||||
with self._lock:
|
||||
try:
|
||||
self._unpark_tickets.remove(ticket)
|
||||
except ValueError:
|
||||
pass
|
||||
# This ticket was holding a slot back from the wait line.
|
||||
self._grant_waiters_locked()
|
||||
|
||||
def cancel(self, waiter: _Waiter) -> None:
|
||||
lease_to_release = None
|
||||
with self._lock:
|
||||
|
|
@ -291,7 +577,13 @@ class LlamaAdmissionQueue:
|
|||
lease_to_release = waiter.granted_lease
|
||||
waiter.granted_lease = None
|
||||
if not waiter.future.done():
|
||||
waiter.loop.call_soon_threadsafe(waiter.future.cancel)
|
||||
try:
|
||||
waiter.loop.call_soon_threadsafe(waiter.future.cancel)
|
||||
except RuntimeError:
|
||||
# Loop gone. Routes call cancel() from finally blocks, so
|
||||
# raising here would both mask their exception and skip the
|
||||
# release below, stranding the slot for the process lifetime.
|
||||
pass
|
||||
if lease_to_release is not None:
|
||||
lease_to_release.release()
|
||||
|
||||
|
|
@ -303,20 +595,32 @@ class LlamaAdmissionQueue:
|
|||
def is_idle(self) -> bool:
|
||||
with self._lock:
|
||||
self._prune_waiters_locked()
|
||||
return self._active == 0 and not self._waiters
|
||||
# A parked holder owns no slot but is coming back to this queue, so
|
||||
# evicting it here would resume it against a fresh 1-slot pool.
|
||||
return self._in_use == 0 and not self._waiters and not self._parked
|
||||
|
||||
def _grant_waiters_locked(self) -> None:
|
||||
self._prune_waiters_locked()
|
||||
while self._waiters and self._active < self._capacity:
|
||||
# Dead waiters are skipped as they are popped, so no prune is needed here.
|
||||
while self._waiters and self._can_admit_locked(len(self._unpark_tickets)):
|
||||
waiter = self._waiters.popleft()
|
||||
if waiter.cancelled or waiter.future.done():
|
||||
continue
|
||||
self._active += 1
|
||||
lease = LlamaAdmissionLease(self)
|
||||
slot = self._take_slot_locked(len(self._unpark_tickets))
|
||||
lease = LlamaAdmissionLease(self, slot)
|
||||
waiter.granted_lease = lease
|
||||
waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
|
||||
try:
|
||||
waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
|
||||
except RuntimeError:
|
||||
# Waiter's loop is gone. Reclaim the slot; leaving the bit set
|
||||
# would strand it, since _free is rebuilt from the bitmask.
|
||||
waiter.granted_lease = None
|
||||
self._release_slot_locked(slot)
|
||||
|
||||
def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None:
|
||||
# Runs on the waiter's own loop thread, which is also the only thread that
|
||||
# cancels that reservation, so waiter state is safe to touch unlocked here.
|
||||
# release() may be called from any thread, but only reaches this via
|
||||
# call_soon_threadsafe. Cancelling off-loop would need this under _lock.
|
||||
if waiter.cancelled or waiter.future.done():
|
||||
waiter.granted_lease = None
|
||||
if not waiter.future.done():
|
||||
|
|
@ -331,16 +635,32 @@ class LlamaAdmissionQueue:
|
|||
lease.release()
|
||||
|
||||
def _prune_waiters_locked(self) -> None:
|
||||
# Rebuilding the deque on every reserve/release dominated the hot path, so
|
||||
# only pay it when a waiter actually died out of band (an externally
|
||||
# cancelled future); cancel() already drops its own waiter eagerly.
|
||||
for waiter in self._waiters:
|
||||
if waiter.cancelled or waiter.future.done():
|
||||
break
|
||||
else:
|
||||
return
|
||||
self._waiters = deque(
|
||||
waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done()
|
||||
)
|
||||
|
||||
def _live_waiters_locked(self) -> int:
|
||||
self._prune_waiters_locked()
|
||||
return len(self._waiters)
|
||||
|
||||
def _snapshot_locked(self) -> LlamaAdmissionSnapshot:
|
||||
return LlamaAdmissionSnapshot(
|
||||
key = self.key,
|
||||
capacity = self._capacity,
|
||||
active = self._active,
|
||||
active = self._held,
|
||||
queued = len(self._waiters),
|
||||
# What another caller could actually take, so the admission log never
|
||||
# shows free slots next to queued requests: after a shrink, ids below
|
||||
# the new capacity can be free while holdovers still fill the ceiling.
|
||||
free = min(len(self._free), max(0, self._capacity - self._held)),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -98,6 +98,7 @@ from core.inference.tool_call_parser import (
|
|||
from core.inference.tool_loop_controller import (
|
||||
ToolLoopController,
|
||||
append_deferred_nudges,
|
||||
awaiting_approval_status,
|
||||
tool_event_provenance,
|
||||
)
|
||||
from state.tool_approvals import (
|
||||
|
|
@ -3135,8 +3136,9 @@ class LlamaCppBackend:
|
|||
prefer_rocr masks at the ROCr/HSA layer instead (clearing HIP). A HIP mask
|
||||
filters only AFTER the HSA runtime enumerates every agent, and that
|
||||
enumeration segfaults at startup on a GPU the build has no kernels for
|
||||
(e.g. a gfx1103 iGPU under a gfx110X prebuilt), before llama-server logs a
|
||||
line. ROCR drops the device at the driver layer, consuming physical ids.
|
||||
(e.g. a gfx1036 iGPU under a gfx103X prebuilt: that bundle maps only
|
||||
gfx1030/1031/1032/1034), before llama-server logs a line. ROCR drops the
|
||||
device at the driver layer, consuming physical ids.
|
||||
The CPU-only sentinel ("-1") has no portable ROCR spelling, so it keeps
|
||||
the HIP mask. Windows keeps the HIP mask too: ROCR_VISIBLE_DEVICES is a
|
||||
Linux ROCr variable (Windows HIP has no ROCr layer), so the ROCR pin
|
||||
|
|
@ -3520,18 +3522,17 @@ class LlamaCppBackend:
|
|||
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.
|
||||
def _run_vulkan_probe(binary: Optional[str] = None) -> list[dict]:
|
||||
"""Run ``_vulkan_probe.py`` and parse its per-device lines.
|
||||
|
||||
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.
|
||||
Returns raw (uncapped) rows sorted by index:
|
||||
``{"index", "free_mib", "total_mib", "is_igpu", "name"}``. 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.
|
||||
``name`` is ggml's device description; "" from an older 4-column probe.
|
||||
[] when no Vulkan build or device is reachable.
|
||||
"""
|
||||
binary = binary or LlamaCppBackend._find_llama_server_binary()
|
||||
if not binary:
|
||||
|
|
@ -3556,10 +3557,13 @@ class LlamaCppBackend:
|
|||
)
|
||||
probe_script = Path(__file__).with_name("_vulkan_probe.py")
|
||||
try:
|
||||
# UTF-8 to match the probe's stdout reconfigure: device names can be
|
||||
# non-ASCII, and the platform-default decode (cp1252) could throw.
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(probe_script), str(binary_dir)],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
encoding = "utf-8",
|
||||
errors = "replace",
|
||||
timeout = 15,
|
||||
env = env,
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
|
|
@ -3573,21 +3577,56 @@ class LlamaCppBackend:
|
|||
logger.debug(f"vulkan GPU probe failed: {e}")
|
||||
return []
|
||||
|
||||
gpus: list[tuple[int, int, int]] = []
|
||||
rows: list[dict] = []
|
||||
for line in result.stdout.strip().splitlines():
|
||||
parts = line.split("\t")
|
||||
if len(parts) != 4:
|
||||
# 4 columns from an older probe (no name); 5 with the name column.
|
||||
if len(parts) not in (4, 5):
|
||||
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)
|
||||
rows.append(
|
||||
{
|
||||
"index": int(parts[0]),
|
||||
"free_mib": int(parts[1]) // (1024 * 1024),
|
||||
"is_igpu": parts[2] == "1",
|
||||
"total_mib": int(parts[3]) // (1024 * 1024),
|
||||
"name": parts[4].strip() if len(parts) == 5 else "",
|
||||
}
|
||||
)
|
||||
except ValueError:
|
||||
continue
|
||||
rows.sort(key = lambda r: r["index"])
|
||||
return rows
|
||||
|
||||
@staticmethod
|
||||
def vulkan_device_inventory(binary: Optional[str] = None) -> list[dict]:
|
||||
"""UI-facing Vulkan device list: the devices llama-server will actually
|
||||
use, with real totals (an iGPU keeps its shared-RAM total here -- the
|
||||
caller labels it, unlike the fit which zeroes it). Same rows as
|
||||
``_run_vulkan_probe``; names fall back to ``Vulkan<i>``.
|
||||
"""
|
||||
rows = LlamaCppBackend._run_vulkan_probe(binary)
|
||||
for row in rows:
|
||||
if not row["name"]:
|
||||
row["name"] = f"Vulkan{row['index']}"
|
||||
return rows
|
||||
|
||||
@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.
|
||||
|
||||
Fit-oriented view of ``_run_vulkan_probe``: returns (device_index,
|
||||
free_mib, total_mib) sorted by index. 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.
|
||||
"""
|
||||
gpus: list[tuple[int, int, int]] = []
|
||||
for row in LlamaCppBackend._run_vulkan_probe(binary):
|
||||
idx, free_mib, is_igpu = row["index"], row["free_mib"], row["is_igpu"]
|
||||
# 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 row["total_mib"]
|
||||
capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu)
|
||||
if capped < free_mib:
|
||||
logger.info(
|
||||
|
|
@ -3596,7 +3635,6 @@ class LlamaCppBackend:
|
|||
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: "
|
||||
|
|
@ -6534,10 +6572,10 @@ class LlamaCppBackend:
|
|||
binary = self._find_llama_server_binary()
|
||||
is_vulkan_backend = self._is_vulkan_backend(binary)
|
||||
|
||||
# Without --kv-unified an explicit --parallel N splits -c into windows of -c/N, so a
|
||||
# build lacking the flag would shrink every context window for a feature it cannot
|
||||
# serve: use one slot. After the requested count is captured (the echo still reports
|
||||
# it), before the KV estimates (the fit matches what launches).
|
||||
# Without --kv-unified an explicit --parallel N splits -c into windows of -c/N, so on a
|
||||
# build lacking the flag the default of 4 would quarter every context window for a
|
||||
# feature it cannot serve: fall back to one slot. Ahead of the KV estimates so the
|
||||
# fit matches what launches.
|
||||
if (
|
||||
n_parallel > 1
|
||||
and binary
|
||||
|
|
@ -6674,12 +6712,23 @@ class LlamaCppBackend:
|
|||
# Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server;
|
||||
# serve them with the diffusion runner (same OpenAI-compat interface).
|
||||
if self._is_diffusion:
|
||||
# Final defense: route and pre-teardown preflights reject before Phase 1.
|
||||
if is_vulkan_backend and gpu_ids:
|
||||
raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR)
|
||||
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
|
||||
# prior load (this path skips the command builder that clears it).
|
||||
self._layer_preserves_tensor_intent = False
|
||||
# On a Vulkan build gpu_ids are ggml Vulkan ordinals, but the diffusion
|
||||
# runner selects its device by CUDA physical index (_diffusion_gpu_arg
|
||||
# forwards gpu_ids[0] as a CUDA/DG_GPU token) with no mapping to them.
|
||||
# The route rejects a CONFIRMED-diffusion pick up front; an uncached GGUF
|
||||
# only classified as diffusion post-download still reaches here with a
|
||||
# pin, so drop it and serve on the default device (like an unpinned load).
|
||||
if gpu_ids and is_vulkan_backend:
|
||||
logger.warning(
|
||||
"Ignoring gpu_ids %s for diffusion GGUF on a Vulkan build: "
|
||||
"the diffusion runner cannot map ggml Vulkan ordinals; "
|
||||
"serving on the default device.",
|
||||
gpu_ids,
|
||||
)
|
||||
gpu_ids = None
|
||||
with self._lock:
|
||||
if self._cancel_event.is_set():
|
||||
logger.info("Load cancelled before diffusion server start")
|
||||
|
|
@ -8275,7 +8324,7 @@ class LlamaCppBackend:
|
|||
env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
# Mask on AMD at the ROCr/HSA layer: HIP-only masking still
|
||||
# enumerates every agent first, which segfaults on a deselected
|
||||
# unsupported GPU (e.g. gfx1103 iGPU under a gfx110X prebuilt).
|
||||
# unsupported GPU (e.g. gfx1036 iGPU under a gfx103X prebuilt).
|
||||
self._emit_child_gpu_visibility(
|
||||
env, ",".join(str(i) for i in gpu_indices), prefer_rocr = True
|
||||
)
|
||||
|
|
@ -11101,6 +11150,7 @@ class LlamaCppBackend:
|
|||
from core.inference.tools import (
|
||||
build_rag_autoinject,
|
||||
execute_tool,
|
||||
has_text_only_provisional_card,
|
||||
is_always_safe_tool,
|
||||
is_high_risk_tool_call,
|
||||
)
|
||||
|
|
@ -11527,6 +11577,9 @@ class LlamaCppBackend:
|
|||
permission_mode == "auto"
|
||||
and is_always_safe_tool(current_name)
|
||||
)
|
||||
# A text-preview card still streams while gated;
|
||||
# hiding it blanks the chat.
|
||||
and not has_text_only_provisional_card(current_name)
|
||||
)
|
||||
# Keep small-argument tools on the normal path.
|
||||
_args_len = len(
|
||||
|
|
@ -11628,20 +11681,27 @@ class LlamaCppBackend:
|
|||
# TEXT call to a provisional card. Gated on an enabled-name
|
||||
# sniff + size floor so prose/small calls spawn no pane; id
|
||||
# matches the first call so the final tool_start reconciles.
|
||||
if (
|
||||
not has_structured_tc
|
||||
and not _confirm_gated_iteration
|
||||
and _text_args_call_start >= 0
|
||||
):
|
||||
if not has_structured_tc and _text_args_call_start >= 0:
|
||||
if not _text_args_id:
|
||||
_call_text = content_accum[_text_args_call_start:]
|
||||
_sniffed = _sniff_text_tool_name(
|
||||
_call_text, _enabled_tool_names
|
||||
)
|
||||
if _sniffed and (
|
||||
_sniffed == "render_html"
|
||||
or len(_call_text)
|
||||
>= _PROVISIONAL_ARGS_MIN_CHARS
|
||||
# Structured-path rule: gated calls
|
||||
# stream only from a text-preview card.
|
||||
if (
|
||||
_sniffed
|
||||
and not (
|
||||
_confirm_gated_iteration
|
||||
and not has_text_only_provisional_card(
|
||||
_sniffed
|
||||
)
|
||||
)
|
||||
and (
|
||||
_sniffed == "render_html"
|
||||
or len(_call_text)
|
||||
>= _PROVISIONAL_ARGS_MIN_CHARS
|
||||
)
|
||||
):
|
||||
_text_args_id = "call_0"
|
||||
_text_args_name = _sniffed
|
||||
|
|
@ -12230,18 +12290,31 @@ class LlamaCppBackend:
|
|||
start_event["awaiting_confirmation"] = needs_confirm
|
||||
|
||||
try:
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
# Gated calls are not running yet; a "Running ..." badge
|
||||
# counting up while it waits on a human reads as a hang.
|
||||
yield {
|
||||
"type": "status",
|
||||
"text": (
|
||||
awaiting_approval_status(decision.tool_name)
|
||||
if needs_confirm
|
||||
else decision.status_text
|
||||
),
|
||||
}
|
||||
yield start_event
|
||||
|
||||
if (
|
||||
decision_slot is not None
|
||||
and wait_tool_decision(
|
||||
_decision = (
|
||||
wait_tool_decision(
|
||||
decision_slot,
|
||||
approval_id,
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
== "deny"
|
||||
):
|
||||
if decision_slot is not None
|
||||
else None
|
||||
)
|
||||
if _decision is not None and _decision != "deny":
|
||||
# Approved: now it really is running.
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
if _decision == "deny":
|
||||
decision_slot = None
|
||||
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
|
||||
yield {
|
||||
|
|
@ -12809,10 +12882,15 @@ class LlamaCppBackend:
|
|||
min_p: float = 0.0,
|
||||
max_new_tokens: int = 2048,
|
||||
repetition_penalty: float = 1.1,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
) -> tuple:
|
||||
"""
|
||||
Generate TTS audio via llama-server /completion + codec decode.
|
||||
Returns (wav_bytes, sample_rate).
|
||||
|
||||
``cancel_event`` lets a Stop or a forced model swap end the request: the
|
||||
decode is one blocking POST, so a watcher closes the client out from under
|
||||
it rather than polling. Raises RuntimeError once cancelled.
|
||||
"""
|
||||
if audio_type not in self._TTS_PROMPTS:
|
||||
raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")
|
||||
|
|
@ -12834,15 +12912,47 @@ class LlamaCppBackend:
|
|||
if need_ids:
|
||||
payload["n_probs"] = 1
|
||||
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise RuntimeError("Audio generation cancelled")
|
||||
|
||||
with httpx.Client(
|
||||
timeout = httpx.Timeout(300, connect = 10),
|
||||
headers = self._auth_headers,
|
||||
trust_env = False,
|
||||
) as client:
|
||||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||||
finished = threading.Event()
|
||||
watcher: Optional[threading.Thread] = None
|
||||
if cancel_event is not None:
|
||||
|
||||
def _close_when_cancelled() -> None:
|
||||
while not finished.wait(0.05):
|
||||
if cancel_event.is_set():
|
||||
# Closing mid-request makes the blocking post raise
|
||||
# httpx.RequestError, the only way out of it.
|
||||
with contextlib.suppress(Exception):
|
||||
client.close()
|
||||
return
|
||||
|
||||
watcher = threading.Thread(target = _close_when_cancelled, daemon = True)
|
||||
watcher.start()
|
||||
try:
|
||||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||||
except httpx.RequestError:
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise RuntimeError("Audio generation cancelled") from None
|
||||
raise
|
||||
finally:
|
||||
finished.set()
|
||||
if watcher is not None:
|
||||
watcher.join(timeout = 0.5)
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}")
|
||||
|
||||
# The codec decode below is GPU work with no interruption point, so check here:
|
||||
# cancelling after this only wastes the decode it cannot stop.
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise RuntimeError("Audio generation cancelled")
|
||||
|
||||
data = resp.json()
|
||||
token_ids = (
|
||||
[p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
|
||||
|
|
|
|||
|
|
@ -345,6 +345,22 @@ def _loaded_identity(backend):
|
|||
return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised)
|
||||
|
||||
|
||||
def _note_idle_unload_event(freed) -> None:
|
||||
"""Monitor row for an idle auto-unload. Best-effort; uses the stash's
|
||||
advertised repo id so the row never shows the on-disk load path."""
|
||||
try:
|
||||
from core.inference.api_monitor import api_monitor
|
||||
from core.inference.model_ids import public_model_id
|
||||
|
||||
identifier, variant, advertised = (list(freed) + [None, None, None])[:3]
|
||||
label = public_model_id(advertised or identifier) or "model"
|
||||
if variant and ":" not in label:
|
||||
label = f"{label}:{variant}"
|
||||
api_monitor.record_lifecycle(event = "unload", model = label, reason = "idle")
|
||||
except Exception as exc:
|
||||
logger.debug("idle unload monitor event failed: %s", exc)
|
||||
|
||||
|
||||
async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
|
||||
"""Unload the loaded GGUF once idle past the configured TTL. Inert when off."""
|
||||
from utils.openai_auto_switch_settings import (
|
||||
|
|
@ -407,6 +423,8 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
|
|||
elif manifest:
|
||||
_delete_resume_files(manifest)
|
||||
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)
|
||||
# An idle unload stashes for reload and skips note_model_unloaded.
|
||||
_note_idle_unload_event(freed)
|
||||
seen_model = None
|
||||
except Exception as exc:
|
||||
logger.debug("idle_unload_loop iteration failed: %s", exc)
|
||||
|
|
|
|||
|
|
@ -128,6 +128,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
|
|||
parse_ctx_override(out)
|
||||
parse_cache_override(out)
|
||||
parse_split_mode_override(out)
|
||||
parse_gpu_layers_override(out)
|
||||
return out
|
||||
|
||||
|
||||
|
|
@ -203,9 +204,8 @@ _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
|
|||
# inherited -ngl is respected (the offload_overridden path), so this group is
|
||||
# opt-in, not default. Layer flags are shared with llama_cpp's override
|
||||
# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them).
|
||||
_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset(
|
||||
{"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}
|
||||
)
|
||||
_GPU_LAYER_FLAGS: frozenset[str] = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers"})
|
||||
_LAYER_OFFLOAD_FLAGS: frozenset[str] = _GPU_LAYER_FLAGS | frozenset({"-fit", "--fit"})
|
||||
_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"})
|
||||
_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS
|
||||
|
||||
|
|
@ -316,6 +316,26 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
|
|||
return _last_flag_value(args, _CACHE_FLAGS)
|
||||
|
||||
|
||||
def parse_gpu_layers_override(args: Optional[Iterable[str]]) -> Optional[int]:
|
||||
"""Return the last user-supplied GPU layer count from extras.
|
||||
|
||||
Manual GPU memory mode strips llama.cpp offload flags because the
|
||||
first-class load fields own them. Callers use this parser first to preserve
|
||||
an explicit ``-ngl`` / ``--gpu-layers`` / ``--n-gpu-layers`` value when
|
||||
translating the extras into those fields.
|
||||
"""
|
||||
raw_value = _last_flag_value(args, _GPU_LAYER_FLAGS)
|
||||
if raw_value is None:
|
||||
return None
|
||||
try:
|
||||
value = int(raw_value)
|
||||
except ValueError as exc:
|
||||
raise ValueError("llama-server GPU layers flag requires an integer value") from exc
|
||||
if value < -1:
|
||||
raise ValueError("llama-server GPU layers flag requires an integer value of at least -1")
|
||||
return value
|
||||
|
||||
|
||||
def parse_cache_override_per_axis(
|
||||
args: Optional[Iterable[str]],
|
||||
) -> tuple[Optional[str], Optional[str]]:
|
||||
|
|
|
|||
|
|
@ -34,6 +34,15 @@ class _LocalGgufEntry:
|
|||
_CACHE_TTL_S = 5.0
|
||||
_lock = threading.Lock()
|
||||
_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {})
|
||||
# Not _lock: that is held for the whole scan, so the request path would wait on it.
|
||||
_warm_lock = threading.Lock()
|
||||
# Repos that finished downloading but are not in the published index yet: nothing
|
||||
# else covers them until the next scan, and the request path must not call them absent.
|
||||
_just_downloaded: set[str] = set()
|
||||
_warming = False
|
||||
_last_scan_s = 0.0
|
||||
# Rescan at most a tenth of the time: on the TTL alone a slow scan would run continuously.
|
||||
_WARM_DUTY = 10.0
|
||||
|
||||
|
||||
def _is_abs_path_id(value: str) -> bool:
|
||||
|
|
@ -103,17 +112,26 @@ def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]:
|
|||
load_dir = _resolve_load_dir(p)
|
||||
variants, _ = list_local_gguf_variants(str(load_dir))
|
||||
quants = tuple(v.quant for v in variants if getattr(v, "quant", None))
|
||||
return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None
|
||||
if not quants:
|
||||
return None
|
||||
# That call orders by descending size, so the head is the biggest quant (often
|
||||
# F16). Downstream reads [0], and a bare id must mean whichever quant a plain
|
||||
# load would take: answering with the largest can evict a model and then OOM.
|
||||
from core.inference.openai_auto_download import preferred_quant
|
||||
|
||||
best = preferred_quant(quants)
|
||||
if best and quants[0] != best:
|
||||
quants = (best, *(q for q in quants if q != best))
|
||||
return _LocalGgufEntry(loader_id, str(load_dir), quants)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def info_has_local_gguf(info) -> bool:
|
||||
"""True when *info* (a LocalModelInfo) points to on-disk GGUF weights the
|
||||
auto-switch path can load. Read from the files, not ``info.model_format``: the
|
||||
HF-cache scanner leaves model_format unset for GGUF snapshots, so a
|
||||
model_format filter would drop every cached GGUF. Lets /v1/models advertise
|
||||
exactly what /v1 can serve."""
|
||||
def local_gguf_quants(info) -> Optional[tuple[str, ...]]:
|
||||
"""On-disk quant labels for *info*, or None when it is not a servable local
|
||||
GGUF. Read from the files, not ``info.model_format``: the HF-cache scanner
|
||||
leaves that unset for GGUF snapshots, so filtering on it drops every cached
|
||||
GGUF. One scan tells /v1/models what it can serve and which quant to name."""
|
||||
from pathlib import Path
|
||||
|
||||
path = getattr(info, "path", None)
|
||||
|
|
@ -123,8 +141,14 @@ def info_has_local_gguf(info) -> bool:
|
|||
if isinstance(path, str) and any(
|
||||
seg in (".studio_links", "ollama_links") for seg in Path(path).parts
|
||||
):
|
||||
return False
|
||||
return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None
|
||||
return None
|
||||
entry = _local_gguf_entry(getattr(info, "id", "") or "", info)
|
||||
return entry.variants if entry is not None else None
|
||||
|
||||
|
||||
def info_has_local_gguf(info) -> bool:
|
||||
"""True when *info* points to on-disk GGUF weights the auto-switch path can load."""
|
||||
return local_gguf_quants(info) is not None
|
||||
|
||||
|
||||
def _build_index() -> dict[str, _LocalGgufEntry]:
|
||||
|
|
@ -287,6 +311,36 @@ def _sibling_revision_entries(raw_id: str, loader_id: str):
|
|||
yield sibling.name, entry
|
||||
|
||||
|
||||
def note_downloaded(repo_id: Optional[str]) -> None:
|
||||
"""Record a repo as present ahead of the scan that will index it."""
|
||||
if not repo_id:
|
||||
return
|
||||
with _lock:
|
||||
_just_downloaded.add(repo_id.strip().lower())
|
||||
|
||||
|
||||
def recently_downloaded(repo_id: str) -> bool:
|
||||
"""Whether *repo_id* finished downloading since the last completed scan."""
|
||||
if not isinstance(repo_id, str) or not repo_id.strip():
|
||||
return False
|
||||
return repo_id.strip().lower() in _just_downloaded
|
||||
|
||||
|
||||
def invalidate_index() -> None:
|
||||
"""Mark the cached scan stale so the next resolve sees a just-finished download
|
||||
instead of waiting out the TTL.
|
||||
|
||||
Keeps the entries: the request path reads this cache without scanning, so
|
||||
emptying it would leave it with no evidence about any local model until the
|
||||
rebuild lands, and a bare request for one would be answered by whatever is
|
||||
resident. Only a completed download invalidates, and that only adds, so the
|
||||
retained entries stay true.
|
||||
"""
|
||||
global _scan
|
||||
with _lock:
|
||||
_scan = (0.0, _scan[1])
|
||||
|
||||
|
||||
def _index() -> dict[str, _LocalGgufEntry]:
|
||||
global _scan
|
||||
# Build under the lock so concurrent callers with an expired cache don't all
|
||||
|
|
@ -301,23 +355,74 @@ def _index() -> dict[str, _LocalGgufEntry]:
|
|||
# an install with many local models can itself exceed the TTL, which would
|
||||
# store the cache already expired and make every request rebuild the index.
|
||||
_scan = (time.monotonic(), fresh)
|
||||
# The scan supersedes the notes: whatever landed is in the index now.
|
||||
_just_downloaded.clear()
|
||||
return fresh
|
||||
|
||||
|
||||
def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]:
|
||||
def index_is_built() -> bool:
|
||||
"""Whether a scan has ever completed, freshness aside.
|
||||
|
||||
Lock-free on purpose: ``_lock`` is held for the whole scan, so taking it would
|
||||
park the request path on the scan it is trying to stay off. Safe because
|
||||
``_scan`` is only ever rebound, never mutated.
|
||||
"""
|
||||
return bool(_scan[0])
|
||||
|
||||
|
||||
def warm_index_soon() -> None:
|
||||
"""(Re)build the index off the request path when it is missing or past its TTL.
|
||||
|
||||
The only refresh for callers using ``allow_scan=False``. Covers a stale index,
|
||||
not just an absent one: a model downloaded through the Hub UI or dropped into a
|
||||
scan folder has no invalidation hook and would otherwise stay invisible to them
|
||||
for the life of the process. Never blocks, and never touches ``_lock``.
|
||||
"""
|
||||
global _warming
|
||||
if time.monotonic() - _scan[0] < max(_CACHE_TTL_S, _last_scan_s * _WARM_DUTY):
|
||||
return
|
||||
with _warm_lock:
|
||||
if _warming:
|
||||
return
|
||||
_warming = True
|
||||
|
||||
def _run() -> None:
|
||||
global _warming, _last_scan_s
|
||||
started = time.monotonic()
|
||||
try:
|
||||
_index()
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
_last_scan_s = time.monotonic() - started
|
||||
with _warm_lock:
|
||||
_warming = False
|
||||
|
||||
threading.Thread(target = _run, name = "local-model-index-warm", daemon = True).start()
|
||||
|
||||
|
||||
def resolve_local_gguf(
|
||||
requested: str, *, allow_scan: bool = True
|
||||
) -> Optional[tuple[str, Optional[str], str]]:
|
||||
"""Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None.
|
||||
|
||||
``load_path`` is the concrete on-disk path to hand /load (so it never fetches
|
||||
a remote), ``loader_id`` is the advertised id used as the launch-override key.
|
||||
``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first
|
||||
(so ids containing a colon still resolve); else the last ``:VARIANT`` is split
|
||||
off and resolves only when that quant is on disk.
|
||||
off and resolves only when that quant is on disk, unless it names no quant at
|
||||
all (an Ollama-style ":latest"), which means the repo.
|
||||
|
||||
``allow_scan=False`` answers from the last built index and never rebuilds, for
|
||||
the request path: the scan walks several model dirs and HF caches, takes seconds
|
||||
on a large install, and holds a lock everyone queues behind. Stale is fine there,
|
||||
since disk barely moves and a finished download calls :func:`invalidate_index`.
|
||||
"""
|
||||
if not isinstance(requested, str) or not requested.strip():
|
||||
return None
|
||||
requested = requested.strip()
|
||||
try:
|
||||
index = _index()
|
||||
index = _index() if allow_scan else _scan[1]
|
||||
entry = index.get(requested.lower())
|
||||
if entry is not None:
|
||||
variant = entry.variants[0] if entry.variants else None
|
||||
|
|
@ -333,8 +438,44 @@ def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str
|
|||
for v in entry.variants:
|
||||
if v.lower() == wanted:
|
||||
return entry.load_path, v, entry.loader_id
|
||||
return None
|
||||
from core.inference.openai_auto_download import looks_like_quant
|
||||
|
||||
if looks_like_quant(variant):
|
||||
return None
|
||||
# ":latest" or ":8b" names no file, so it means the repo; a real quant that
|
||||
# is not on disk still misses, or a swap would serve the wrong weights.
|
||||
return entry.load_path, (entry.variants[0] if entry.variants else None), entry.loader_id
|
||||
except Exception:
|
||||
# Best-effort: any resolver failure falls through to the loaded model,
|
||||
# so a malformed name can never turn a servable request into a 500.
|
||||
return None
|
||||
|
||||
|
||||
MISS_MODEL_NOT_FOUND = "model_not_found"
|
||||
MISS_VARIANT_NOT_FOUND = "variant_not_found"
|
||||
|
||||
|
||||
def describe_local_miss(requested: str) -> tuple[str, tuple[str, ...]]:
|
||||
"""Why :func:`resolve_local_gguf` missed, so an error can say "wrong quant"
|
||||
instead of "no such model".
|
||||
|
||||
``(MISS_VARIANT_NOT_FOUND, <local quants>)`` when the repo is downloaded but the
|
||||
requested ``:VARIANT`` is not, else ``(MISS_MODEL_NOT_FOUND, ())``. Fail-safe: a
|
||||
scan failure reports the generic miss rather than raising into the handler.
|
||||
"""
|
||||
if not isinstance(requested, str) or not requested.strip():
|
||||
return MISS_MODEL_NOT_FOUND, ()
|
||||
base, sep, variant = requested.strip().rpartition(":")
|
||||
from core.inference.openai_auto_download import looks_like_quant
|
||||
|
||||
# Split like the resolver or the two disagree: a tag naming no quant means the
|
||||
# repo there, so reporting a missing quant for it would name one nobody asked for.
|
||||
if not sep or not looks_like_quant(variant):
|
||||
return MISS_MODEL_NOT_FOUND, ()
|
||||
try:
|
||||
entry = _index().get(base.strip().lower())
|
||||
except Exception:
|
||||
return MISS_MODEL_NOT_FOUND, ()
|
||||
if entry is None or not entry.variants:
|
||||
return MISS_MODEL_NOT_FOUND, ()
|
||||
return MISS_VARIANT_NOT_FOUND, entry.variants
|
||||
|
|
|
|||
|
|
@ -1189,7 +1189,8 @@ class MLXInferenceBackend:
|
|||
**gen_kwargs,
|
||||
)
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
# caller_cancel_event: signature parity with the orchestrator; unused here.
|
||||
import mlx.core as mx
|
||||
import gc
|
||||
|
||||
|
|
|
|||
|
|
@ -39,10 +39,29 @@ def _looks_like_path(identifier: str) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def hf_cache_repo_id(path: Optional[str]) -> Optional[str]:
|
||||
"""``.../models--org--name/snapshots/<sha>`` -> ``org/name``, else None.
|
||||
|
||||
A model loaded from the HF cache is identified by its snapshot dir, whose
|
||||
basename is a commit hash; recover the repo id so callers don't show that.
|
||||
"""
|
||||
if not path:
|
||||
return None
|
||||
parts = str(path).replace("\\", "/").split("/")
|
||||
for index, part in enumerate(parts):
|
||||
# Only inside the real cache layout: a "models--" name alone is not a repo id.
|
||||
if part.startswith("models--") and parts[index + 1 : index + 2] == ["snapshots"]:
|
||||
return part[len("models--") :].replace("--", "/")
|
||||
return None
|
||||
|
||||
|
||||
def public_model_id(identifier: Optional[str]) -> Optional[str]:
|
||||
"""Return a clean, path-free public id for *identifier*.
|
||||
|
||||
- Local GGUF path -> the file stem with ``.gguf`` stripped, e.g.
|
||||
- HF cache path -> the repo id it came from, e.g.
|
||||
``~/.cache/huggingface/hub/models--unsloth--X-GGUF/snapshots/<sha>`` ->
|
||||
``unsloth/X-GGUF``.
|
||||
- Other local GGUF path -> the file stem with ``.gguf`` stripped, e.g.
|
||||
``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``.
|
||||
- HF repo id (``org/model``) and already-clean names -> returned unchanged.
|
||||
- ``None`` / empty -> returned unchanged.
|
||||
|
|
@ -51,6 +70,9 @@ def public_model_id(identifier: Optional[str]) -> Optional[str]:
|
|||
return identifier
|
||||
if not _looks_like_path(identifier):
|
||||
return identifier
|
||||
repo_id = hf_cache_repo_id(identifier)
|
||||
if repo_id:
|
||||
return repo_id
|
||||
name = os.path.basename(identifier.replace("\\", "/").rstrip("/"))
|
||||
if name.lower().endswith(_GGUF_SUFFIX):
|
||||
name = name[: -len(_GGUF_SUFFIX)]
|
||||
|
|
|
|||
812
studio/backend/core/inference/openai_auto_download.py
Normal file
812
studio/backend/core/inference/openai_auto_download.py
Normal file
|
|
@ -0,0 +1,812 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Opt-in: fetch a GGUF a /v1 request names but this server doesn't have.
|
||||
|
||||
Auto-switch only loads models already on disk. With
|
||||
``openai_api_auto_download_model`` on, a miss that looks like a real Hub repo is
|
||||
fetched in the background and the request is told to retry rather than held
|
||||
open: a quant is routinely tens of GB, far longer than any client (or the
|
||||
Cloudflare edge on ``--secure``) will wait, and the inference lifecycle gate must
|
||||
not be held meanwhile. The resident model keeps serving, and the retry that lands
|
||||
after the download goes through the ordinary auto-switch path.
|
||||
|
||||
Admission is deliberately narrow, since a request only needs an API key:
|
||||
- ``namespace/name`` only, and only when the Hub confirms GGUF weights. A
|
||||
namespace is not evidence of intent (LiteLLM and OpenRouter address every
|
||||
provider that way), so ``gpt-4`` and ``anthropic/claude-3.5-sonnet`` alike
|
||||
fall through to the resident model as before.
|
||||
- GGUF only, decided from the remote file list, not the repo name: GGUF runs
|
||||
under llama.cpp, which never imports repo Python.
|
||||
- ``auto_map`` is refused, so ``trust_remote_code`` is only ever granted
|
||||
deliberately in the UI, never by an API call.
|
||||
- One download at a time, so a key holder cannot fan out fetches.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Keep the Hub probe short so a slow Hub can't stall the request path.
|
||||
_MODEL_INFO_TIMEOUT_S = 8.0
|
||||
# auth_check and hf_hub_download take no timeout of their own and run while the
|
||||
# provisional slot is held, so an unresponsive Hub would pin the single flight. The
|
||||
# code probe fetches up to three configs, so it gets more room than the auth call.
|
||||
_CODE_PROBE_TIMEOUT_S = 20.0
|
||||
# Headroom left free after the download, so filling the disk can't wedge the box.
|
||||
_DISK_RESERVE_BYTES = 5 * 1024**3
|
||||
_WATCH_POLL_S = 2.0
|
||||
# A stalled watcher must not pin the single-flight slot forever.
|
||||
_MAX_WATCH_S = 24 * 60 * 60
|
||||
# Past the watch window the row is resolved, so poll only to see whether the
|
||||
# worker is still alive and still owns the slot.
|
||||
_TIMED_OUT_POLL_S = 60.0
|
||||
_RETRY_AFTER_S = 30
|
||||
# Long enough for a client honouring Retry-After to come back and be told, short
|
||||
# enough that one that never returns cannot hold the slot.
|
||||
_FAILED_HOLD_S = 3 * _RETRY_AFTER_S
|
||||
_MAX_LISTED_VARIANTS = 8
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class AutoDownloadRefusal:
|
||||
"""Why this request cannot be served yet; the route raises it in the
|
||||
surface's own error envelope."""
|
||||
|
||||
status: int
|
||||
code: str
|
||||
message: str
|
||||
retry_after: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Active:
|
||||
repo_id: str
|
||||
# None while the Hub probe is still deciding which quant to fetch.
|
||||
variant: Optional[str] = None
|
||||
expected_bytes: int = 0
|
||||
monitor_id: Optional[str] = None
|
||||
started_at: float = 0.0
|
||||
# Set when the worker failed. Held until a retry surfaces it: Retry-After is far
|
||||
# longer than the watcher poll, so the client would restart the same failing download.
|
||||
error: Optional[str] = None
|
||||
failed_at: float = 0.0
|
||||
|
||||
|
||||
_lock = threading.Lock()
|
||||
_active: Optional[_Active] = None
|
||||
|
||||
# Repos the Hub says are not servable, so a "vendor/model" miss doesn't re-probe every request.
|
||||
_NOT_SERVABLE_TTL_S = 10 * 60
|
||||
_NOT_SERVABLE_MAX = 256
|
||||
_cache_lock = threading.Lock()
|
||||
_not_servable: dict[str, float] = {}
|
||||
|
||||
|
||||
def _public_label(repo_id: str, variant: Optional[str]) -> str:
|
||||
return f"{repo_id}:{variant}" if variant else repo_id
|
||||
|
||||
|
||||
def split_model_ref(requested: str) -> tuple[str, Optional[str]]:
|
||||
"""``org/repo:QUANT`` -> ``("org/repo", "QUANT")``; no suffix -> variant None.
|
||||
|
||||
Splits on the last colon. A slash-bearing suffix is only a variant when a real
|
||||
Hub repo precedes it: "build/llama-13b" is a subdirectory GGUF key the catalog
|
||||
advertises, while "C:/models/x.gguf" leaves a drive letter that is no repo id.
|
||||
"""
|
||||
text = (requested or "").strip()
|
||||
base, sep, suffix = text.rpartition(":")
|
||||
if not sep or not base or not suffix:
|
||||
return text, None
|
||||
stripped = base.strip()
|
||||
if "/" in suffix:
|
||||
from hub.utils.paths import is_valid_repo_id
|
||||
if "/" not in stripped or not is_valid_repo_id(stripped):
|
||||
return text, None
|
||||
return stripped, suffix.strip()
|
||||
|
||||
|
||||
def is_downloadable_ref(requested: str) -> bool:
|
||||
"""Whether *requested* is shaped like a Hub repo we may fetch.
|
||||
|
||||
Requires an explicit namespace: keeps ``gpt-4`` and other foreign ids falling
|
||||
through, and stops ModelConfig.from_identifier's bare-name ``unsloth/``
|
||||
prefixing from turning an unrelated label into a real repo.
|
||||
"""
|
||||
from hub.utils.paths import is_valid_repo_id
|
||||
|
||||
repo_id, variant = split_model_ref(requested)
|
||||
if "/" not in repo_id or not is_valid_repo_id(repo_id):
|
||||
return False
|
||||
if variant is not None:
|
||||
from hub.utils.paths import is_valid_gguf_variant
|
||||
return is_valid_gguf_variant(variant)
|
||||
return True
|
||||
|
||||
|
||||
def looks_like_quant(variant: Optional[str]) -> bool:
|
||||
"""Whether a ``:suffix`` names a GGUF quant rather than a foreign tag.
|
||||
|
||||
Neither a namespace nor a colon proves a request was meant for this server
|
||||
(``vendor/model`` is LiteLLM/OpenRouter, ``name:latest`` is Ollama). A real
|
||||
quant label does.
|
||||
"""
|
||||
import re
|
||||
|
||||
from utils.models.model_config import _GGUF_KNOWN_QUANT_RE
|
||||
|
||||
if not variant:
|
||||
return False
|
||||
# _extract_quant_label can append a bpw modifier (IQ4_XS-3.53bpw); still a quant.
|
||||
label = re.sub(r"-[0-9]+(?:\.[0-9]+)?bpw$", "", variant.strip(), flags = re.IGNORECASE)
|
||||
return _GGUF_KNOWN_QUANT_RE.fullmatch(label) is not None
|
||||
|
||||
|
||||
def _hub_token(hf_token: Optional[str]):
|
||||
"""The caller's token, or an explicit False. None makes huggingface_hub fall
|
||||
back to a cached login (here the server owner's); only False is anonymous."""
|
||||
return hf_token or False
|
||||
|
||||
|
||||
def _servable_key(repo_id: str, hf_token: Optional[str]) -> str:
|
||||
"""Cache key, per credential.
|
||||
|
||||
The Hub 404s a private repo the caller cannot see, so a tokenless verdict says
|
||||
nothing about a caller who has one. Digested, so no token is held here.
|
||||
"""
|
||||
import hashlib
|
||||
|
||||
seen_as = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else "anon"
|
||||
return f"{repo_id.lower()}\n{seen_as}"
|
||||
|
||||
|
||||
def _mark_not_servable(repo_id: str, hf_token: Optional[str]) -> None:
|
||||
with _cache_lock:
|
||||
if len(_not_servable) >= _NOT_SERVABLE_MAX:
|
||||
_not_servable.clear()
|
||||
_not_servable[_servable_key(repo_id, hf_token)] = time.monotonic() + _NOT_SERVABLE_TTL_S
|
||||
|
||||
|
||||
def _is_not_servable(repo_id: str, hf_token: Optional[str]) -> bool:
|
||||
key = _servable_key(repo_id, hf_token)
|
||||
with _cache_lock:
|
||||
expires = _not_servable.get(key)
|
||||
if expires is None:
|
||||
return False
|
||||
if expires <= time.monotonic():
|
||||
del _not_servable[key]
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _gated_refusal(repo_id: str) -> AutoDownloadRefusal:
|
||||
return AutoDownloadRefusal(
|
||||
status = 403,
|
||||
code = "model_access_denied",
|
||||
message = (
|
||||
f"'{repo_id}' is gated on Hugging Face. Accept its licence, then retry with "
|
||||
"your own token in the X-Unsloth-HF-Token header: automatic download never "
|
||||
"uses this server's Hugging Face identity."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _bounded_probe(fn, *args, timeout: float, default):
|
||||
"""Run a blocking Hub probe off the loop, bounding only the wait.
|
||||
|
||||
The thread is left to finish (a blocking socket read cannot be cancelled); the
|
||||
caller takes *default*, chosen per call site so a timeout errs the safe way.
|
||||
"""
|
||||
try:
|
||||
return await asyncio.wait_for(asyncio.to_thread(fn, *args), timeout)
|
||||
except (TimeoutError, asyncio.TimeoutError):
|
||||
logger.debug("hub probe %s timed out after %ss", getattr(fn, "__name__", fn), timeout)
|
||||
return default
|
||||
|
||||
|
||||
def _auth_denied(repo_id: str, hf_token: Optional[str]) -> bool:
|
||||
"""Whether this token lacks file access to a gated repo. False when the
|
||||
check is inconclusive: the download's own auth is the real gate."""
|
||||
from hub.utils.hf_errors import hf_error_status
|
||||
|
||||
try:
|
||||
from huggingface_hub import auth_check
|
||||
auth_check(repo_id, token = _hub_token(hf_token))
|
||||
except Exception as exc:
|
||||
return hf_error_status(exc) in (401, 403)
|
||||
return False
|
||||
|
||||
|
||||
def _gguf_variants(siblings) -> dict[str, int]:
|
||||
"""Quant label -> bytes the download will actually fetch.
|
||||
|
||||
Mirrors list_gguf_variants for the selectable labels: companions (mmproj/MTP)
|
||||
and big-endian builds are not quants, and sharded quants sum across shards.
|
||||
Bytes come from the download plan, which folds companions back into every
|
||||
quant, so the disk reserve is measured against what the worker fetches.
|
||||
"""
|
||||
from hub.utils.gguf import extract_quant_label as canonical_quant_label
|
||||
from hub.utils.gguf_plan import build_gguf_variant_plans
|
||||
from utils.models.model_config import (
|
||||
_extract_quant_label,
|
||||
_is_big_endian_gguf_path,
|
||||
_is_mmproj,
|
||||
_is_mtp_drafter,
|
||||
)
|
||||
|
||||
siblings = list(siblings or [])
|
||||
plans = build_gguf_variant_plans(siblings)
|
||||
sizes: dict[str, int] = {}
|
||||
for sibling in siblings:
|
||||
name = getattr(sibling, "rfilename", "") or ""
|
||||
if not name.lower().endswith(".gguf"):
|
||||
continue
|
||||
quant = _extract_quant_label(name)
|
||||
if not looks_like_quant(quant):
|
||||
# With no recognized quant token the extractors part ways: this one takes
|
||||
# the last hyphenated segment ("7b" of llama-7b) while the plan and worker
|
||||
# key the whole stem, so advertising ours dispatches an unresolvable variant.
|
||||
quant = canonical_quant_label(name) or quant
|
||||
if _is_mmproj(name) or _is_mtp_drafter(name) or _is_big_endian_gguf_path(name, quant):
|
||||
continue
|
||||
plan = plans.get(quant.lower())
|
||||
if plan is not None:
|
||||
sizes[quant] = plan.download_size_bytes
|
||||
else:
|
||||
sizes[quant] = sizes.get(quant, 0) + int(getattr(sibling, "size", 0) or 0)
|
||||
return sizes
|
||||
|
||||
|
||||
def _remaining_bytes(repo_id: str, plan, expected_bytes: int) -> int:
|
||||
"""Bytes still to fetch: a resumed quant or a companion shared with another
|
||||
quant is already on disk, and charging for it can 507 a download that fits."""
|
||||
try:
|
||||
from hub.utils.download_registry import existing_blob_bytes
|
||||
|
||||
hashes = frozenset(
|
||||
file.sha256 for file in getattr(plan, "expected_files", ()) or () if file.sha256
|
||||
)
|
||||
if not hashes:
|
||||
return expected_bytes
|
||||
return max(0, expected_bytes - existing_blob_bytes("model", repo_id, hashes))
|
||||
except Exception:
|
||||
return expected_bytes
|
||||
|
||||
|
||||
def _enough_disk(need_bytes: int) -> tuple[bool, int]:
|
||||
"""(fits, free_bytes). Fail-open on an unreadable cache root: the download
|
||||
worker runs its own preflight, this only adds the reserve margin."""
|
||||
try:
|
||||
from hub.utils.hf_cache_state import hf_cache_root
|
||||
|
||||
root = hf_cache_root(create = True)
|
||||
if root is None:
|
||||
return True, 0
|
||||
free = shutil.disk_usage(root).free
|
||||
except Exception:
|
||||
return True, 0
|
||||
return free >= need_bytes + _DISK_RESERVE_BYTES, free
|
||||
|
||||
|
||||
def _gb(num_bytes: int) -> str:
|
||||
return f"{num_bytes / 1024**3:.1f} GB"
|
||||
|
||||
|
||||
async def _job_state(repo_id: str, variant: Optional[str]) -> tuple[str, Optional[str]]:
|
||||
from hub.services.models import downloads
|
||||
try:
|
||||
status = await downloads.get_download_status_response(repo_id, variant or "")
|
||||
return status.state, status.error
|
||||
except Exception as exc:
|
||||
# "unknown", not "idle": idle ends the watch, and a failed probe proves nothing.
|
||||
logger.debug("auto-download: status probe failed for %r: %s", repo_id, exc)
|
||||
return "unknown", None
|
||||
|
||||
|
||||
async def _progress_percent(
|
||||
repo_id: str, variant: Optional[str], expected_bytes: int, hf_token: Optional[str]
|
||||
) -> Optional[float]:
|
||||
"""0-100, or None. The hub service reports a 0-1 fraction, so scale it."""
|
||||
from hub.services.models import downloads
|
||||
try:
|
||||
payload = await downloads.get_gguf_download_progress_response(
|
||||
repo_id, variant or "", expected_bytes, hf_token
|
||||
)
|
||||
fraction = payload.get("progress")
|
||||
if not isinstance(fraction, (int, float)):
|
||||
return None
|
||||
return min(100.0, max(0.0, float(fraction) * 100.0))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _release(active: Optional[_Active]) -> None:
|
||||
"""Free the single-flight slot, but only while *active* still owns it.
|
||||
|
||||
Keying on ``repo_id`` alone let a stale operation clear a newer one: variant A
|
||||
errors, an adopting request frees the slot, a retry starts B, then A's watcher
|
||||
matches the repo and clears B, admitting a second download alongside it.
|
||||
"""
|
||||
global _active
|
||||
if active is None:
|
||||
return
|
||||
with _lock:
|
||||
if _active is active:
|
||||
_active = None
|
||||
|
||||
|
||||
async def _watch(active: _Active, hf_token: Optional[str]) -> None:
|
||||
"""Poll a dispatched job so the monitor row resolves and the resolver cache
|
||||
is dropped the moment the weights land."""
|
||||
from core.inference import api_monitor as monitor_module
|
||||
|
||||
api_monitor = monitor_module.api_monitor
|
||||
deadline = time.monotonic() + _MAX_WATCH_S
|
||||
timed_out = False
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(_TIMED_OUT_POLL_S if timed_out else _WATCH_POLL_S)
|
||||
state, error = await _job_state(active.repo_id, active.variant)
|
||||
if state in ("running", "cancelling", "unknown"):
|
||||
if timed_out:
|
||||
# A running worker still owns the slot: releasing on the clock alone
|
||||
# would admit a second multi-GB download beside it. "unknown" cannot
|
||||
# confirm it is alive, so release then, or a broken probe wedges us.
|
||||
if state == "unknown":
|
||||
return
|
||||
continue
|
||||
if time.monotonic() >= deadline:
|
||||
api_monitor.fail_open(active.monitor_id, "Download timed out")
|
||||
timed_out = True
|
||||
continue
|
||||
# Only "running" has progress; the others are still in flight, so keep the slot.
|
||||
if state == "running":
|
||||
api_monitor.set_progress(
|
||||
active.monitor_id,
|
||||
await _progress_percent(
|
||||
active.repo_id, active.variant, active.expected_bytes, hf_token
|
||||
),
|
||||
)
|
||||
continue
|
||||
if state == "cancelled":
|
||||
api_monitor.finish(active.monitor_id, status = "cancelled")
|
||||
return
|
||||
if state == "complete":
|
||||
# No invalidate here: finalize_worker_exit already dropped the cache and
|
||||
# warmed it; a second would mark that fresh scan stale and push a
|
||||
# synchronous rescan onto the client's retry.
|
||||
api_monitor.finish(active.monitor_id, status = "completed")
|
||||
elif state == "idle":
|
||||
# The job vanished without a terminal state (worker killed).
|
||||
api_monitor.fail_open(active.monitor_id, "Download did not complete")
|
||||
else:
|
||||
api_monitor.fail_open(active.monitor_id, error or f"Download {state}")
|
||||
# Keep the slot so the next retry is told it failed instead of
|
||||
# silently restarting the same download.
|
||||
active.error = error or f"Download {state}"
|
||||
active.failed_at = time.monotonic()
|
||||
return
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("auto-download: watcher failed for %r: %s", active.repo_id, exc)
|
||||
api_monitor.fail_open(active.monitor_id, "Download tracking failed")
|
||||
finally:
|
||||
if not active.failed_at:
|
||||
_release(active)
|
||||
|
||||
|
||||
def _downloading_refusal(label: str, percent: Optional[float]) -> AutoDownloadRefusal:
|
||||
progress = f" ({percent:.0f}% done)" if percent is not None else ""
|
||||
return AutoDownloadRefusal(
|
||||
status = 503,
|
||||
code = "model_downloading",
|
||||
message = (f"Downloading '{label}'{progress}. Retry shortly. Track it in Unsloth Studio."),
|
||||
retry_after = _RETRY_AFTER_S,
|
||||
)
|
||||
|
||||
|
||||
async def _is_downloadable_model(repo_id: str, hf_token: Optional[str]) -> bool:
|
||||
"""Whether the Hub has this repo with GGUF weights we could fetch.
|
||||
|
||||
Only asked while another download holds the slot, to tell a second download
|
||||
apart from an ordinary foreign label. Any failure answers False: refusing
|
||||
would strand normal traffic for the length of the download.
|
||||
"""
|
||||
if _is_not_servable(repo_id, hf_token):
|
||||
return False
|
||||
|
||||
def _probe():
|
||||
from huggingface_hub import HfApi
|
||||
return HfApi(token = _hub_token(hf_token)).model_info(repo_id, timeout = _MODEL_INFO_TIMEOUT_S)
|
||||
|
||||
try:
|
||||
info = await asyncio.to_thread(_probe)
|
||||
except Exception:
|
||||
return False
|
||||
# The same filter admission uses, not a bare .gguf test: mmproj, MTP drafters and
|
||||
# big-endian builds are companions, not quants. Answering otherwise would hold an
|
||||
# ordinary foreign label at model_download_busy for an unrelated download.
|
||||
servable = bool(_gguf_variants(getattr(info, "siblings", None)))
|
||||
if not servable:
|
||||
_mark_not_servable(repo_id, hf_token)
|
||||
return servable
|
||||
|
||||
|
||||
async def maybe_auto_download(
|
||||
requested_model: str,
|
||||
*,
|
||||
hf_token: Optional[str] = None,
|
||||
require_vision: bool = False,
|
||||
) -> Optional[AutoDownloadRefusal]:
|
||||
"""Start (or report on) a background fetch of *requested_model*.
|
||||
|
||||
Returns None when the request should carry on unchanged, or a refusal the
|
||||
caller must raise. Only called after the local resolver has already missed.
|
||||
|
||||
``require_vision`` refuses a target with no mmproj companion rather than spend
|
||||
gigabytes on weights that cannot answer the request; the local capability guard
|
||||
only ever sees an already-downloaded model.
|
||||
"""
|
||||
global _active
|
||||
|
||||
repo_id, wanted_variant = split_model_ref(requested_model)
|
||||
if not is_downloadable_ref(requested_model):
|
||||
return None
|
||||
if _is_not_servable(repo_id, hf_token) and not looks_like_quant(wanted_variant):
|
||||
return None
|
||||
|
||||
# Settle the single-flight slot before the network, so retries during a download stay cheap.
|
||||
busy: Optional[_Active] = None
|
||||
with _lock:
|
||||
current = _active
|
||||
if current is not None and current.failed_at:
|
||||
# A held failure only owns the slot until someone is told about it.
|
||||
if current.repo_id != repo_id and time.monotonic() - current.failed_at > _FAILED_HOLD_S:
|
||||
_active = current = None
|
||||
if current is not None and current.repo_id == repo_id:
|
||||
adopted = current
|
||||
elif current is not None:
|
||||
adopted = None
|
||||
busy = current
|
||||
else:
|
||||
adopted = None
|
||||
provisional = _Active(repo_id = repo_id, started_at = time.time())
|
||||
_active = provisional
|
||||
|
||||
if busy is not None:
|
||||
# Refusing before the probe blocks ordinary drop-in traffic: a namespaced label
|
||||
# that is no downloadable GGUF repo (LiteLLM/OpenRouter style) would be told to
|
||||
# wait out a multi-hour download. Only a downloadable label is a 2nd download.
|
||||
if not await _is_downloadable_model(repo_id, hf_token):
|
||||
return None
|
||||
return AutoDownloadRefusal(
|
||||
status = 503,
|
||||
code = "model_download_busy",
|
||||
message = (
|
||||
f"Already downloading '{_public_label(busy.repo_id, busy.variant)}'. "
|
||||
f"Retry '{requested_model}' once it finishes."
|
||||
),
|
||||
retry_after = _RETRY_AFTER_S,
|
||||
)
|
||||
|
||||
if adopted is not None:
|
||||
if adopted.variant is None:
|
||||
# Still probing: no job yet, and a stale whole-repo error would free the probe's slot.
|
||||
return _downloading_refusal(adopted.repo_id, None)
|
||||
state, error = await _job_state(adopted.repo_id, adopted.variant)
|
||||
if state in ("running", "cancelling", "unknown"):
|
||||
return _downloading_refusal(
|
||||
_public_label(adopted.repo_id, adopted.variant),
|
||||
await _progress_percent(
|
||||
adopted.repo_id, adopted.variant, adopted.expected_bytes, hf_token
|
||||
),
|
||||
)
|
||||
if state == "error" or adopted.error:
|
||||
error = error or adopted.error
|
||||
# Surface once, then free the slot so a retry can start over.
|
||||
_release(adopted)
|
||||
return AutoDownloadRefusal(
|
||||
status = 502,
|
||||
code = "model_download_failed",
|
||||
message = f"Downloading '{requested_model}' failed: {error or 'unknown error'}",
|
||||
)
|
||||
# complete/idle/cancelled: the watcher is about to free the slot, so retry once more.
|
||||
return _downloading_refusal(
|
||||
_public_label(adopted.repo_id, adopted.variant),
|
||||
100.0 if state == "complete" else None,
|
||||
)
|
||||
|
||||
try:
|
||||
return await _admit_and_start(
|
||||
repo_id, wanted_variant, requested_model, hf_token, provisional, require_vision
|
||||
)
|
||||
except BaseException:
|
||||
# Not `except Exception`: a cancel mid-probe would otherwise wedge the provisional slot.
|
||||
_release(provisional)
|
||||
raise
|
||||
|
||||
|
||||
async def _admit_and_start(
|
||||
repo_id: str,
|
||||
wanted_variant: Optional[str],
|
||||
requested_model: str,
|
||||
hf_token: Optional[str],
|
||||
active: _Active,
|
||||
require_vision: bool = False,
|
||||
) -> Optional[AutoDownloadRefusal]:
|
||||
from hub.utils.hf_errors import hf_error_status
|
||||
|
||||
def _probe():
|
||||
from huggingface_hub import HfApi
|
||||
return HfApi(token = _hub_token(hf_token)).model_info(
|
||||
repo_id, files_metadata = True, timeout = _MODEL_INFO_TIMEOUT_S
|
||||
)
|
||||
|
||||
try:
|
||||
info = await asyncio.to_thread(_probe)
|
||||
except Exception as exc:
|
||||
_release(active)
|
||||
status = hf_error_status(exc)
|
||||
if status == 401:
|
||||
return AutoDownloadRefusal(
|
||||
status = 401,
|
||||
code = "model_access_denied",
|
||||
message = (
|
||||
f"Hugging Face rejected the token sent for '{repo_id}'. Replace the "
|
||||
"X-Unsloth-HF-Token header with a valid token; retrying will not help."
|
||||
),
|
||||
)
|
||||
if status == 403:
|
||||
return _gated_refusal(repo_id)
|
||||
if status == 404:
|
||||
_mark_not_servable(repo_id, hf_token)
|
||||
# Unknown to the Hub reads as a foreign label; only an explicit quant makes it ours.
|
||||
if not looks_like_quant(wanted_variant):
|
||||
return None
|
||||
# A private repo reads as absent without a token; don't confirm either way.
|
||||
return AutoDownloadRefusal(
|
||||
status = 404,
|
||||
code = "model_not_found",
|
||||
message = (
|
||||
f"'{repo_id}' was not found on Hugging Face, or is not accessible. "
|
||||
"If it is private, send a token in the X-Unsloth-HF-Token header."
|
||||
),
|
||||
)
|
||||
logger.warning("auto-download: Hub lookup failed for %r: %s", repo_id, exc)
|
||||
return AutoDownloadRefusal(
|
||||
status = 503,
|
||||
code = "model_lookup_failed",
|
||||
message = f"Could not reach Hugging Face to look up '{repo_id}'. Retry shortly.",
|
||||
retry_after = _RETRY_AFTER_S,
|
||||
)
|
||||
|
||||
# Inconclusive on timeout: the download's own auth is the real gate.
|
||||
if getattr(info, "gated", False) and await _bounded_probe(
|
||||
_auth_denied, repo_id, hf_token, timeout = _MODEL_INFO_TIMEOUT_S, default = False
|
||||
):
|
||||
# Metadata for a gated repo is not file access; unchecked, the config read below lies.
|
||||
_release(active)
|
||||
return _gated_refusal(repo_id)
|
||||
|
||||
variants = _gguf_variants(getattr(info, "siblings", None))
|
||||
if not variants:
|
||||
_release(active)
|
||||
_mark_not_servable(repo_id, hf_token)
|
||||
if not looks_like_quant(wanted_variant):
|
||||
return None
|
||||
return AutoDownloadRefusal(
|
||||
status = 400,
|
||||
code = "model_not_supported",
|
||||
message = (
|
||||
f"'{repo_id}' has no GGUF weights. Automatic download serves GGUF only; "
|
||||
"load other formats from Unsloth Studio."
|
||||
),
|
||||
)
|
||||
|
||||
# trust_remote_code gate: _config_has_auto_map is tri-state, so refuse on True and on None.
|
||||
from utils.security.consent import _config_has_auto_map
|
||||
|
||||
# _hub_token, not the raw token: None lets huggingface_hub fall back to a cached
|
||||
# server login, so a caller-named repo would be probed with this server's identity.
|
||||
# Defaults to None on timeout, which refuses: unchecked is not cleared.
|
||||
has_auto_map = await _bounded_probe(
|
||||
_config_has_auto_map,
|
||||
repo_id,
|
||||
_hub_token(hf_token),
|
||||
timeout = _CODE_PROBE_TIMEOUT_S,
|
||||
default = None,
|
||||
)
|
||||
if has_auto_map is not False:
|
||||
_release(active)
|
||||
unknown = has_auto_map is None
|
||||
return AutoDownloadRefusal(
|
||||
status = 403,
|
||||
code = "remote_code_consent_required",
|
||||
message = (
|
||||
f"'{repo_id}' "
|
||||
+ (
|
||||
"could not be checked for custom code"
|
||||
if unknown
|
||||
else "ships custom code that runs on load"
|
||||
)
|
||||
+ ". Load it once in Unsloth Studio to review and approve it, then retry."
|
||||
),
|
||||
)
|
||||
|
||||
variant = _match_variant(wanted_variant, variants)
|
||||
if variant is None:
|
||||
_release(active)
|
||||
listed = sorted(variants)
|
||||
shown = ", ".join(listed[:_MAX_LISTED_VARIANTS])
|
||||
extra = len(listed) - _MAX_LISTED_VARIANTS
|
||||
return AutoDownloadRefusal(
|
||||
status = 404,
|
||||
code = "model_not_found",
|
||||
message = (
|
||||
f"'{repo_id}' has no quant '{wanted_variant}'. Available quants: "
|
||||
f"{shown}{f' and {extra} more' if extra > 0 else ''}."
|
||||
),
|
||||
)
|
||||
|
||||
expected_bytes = variants[variant]
|
||||
from hub.utils.gguf_plan import build_gguf_variant_plans
|
||||
|
||||
plan = build_gguf_variant_plans(list(getattr(info, "siblings", None) or [])).get(
|
||||
variant.lower()
|
||||
)
|
||||
if require_vision and not (plan and plan.mmproj_filenames):
|
||||
_release(active)
|
||||
return AutoDownloadRefusal(
|
||||
status = 400,
|
||||
code = "invalid_value",
|
||||
message = (
|
||||
f"'{_public_label(repo_id, variant)}' ships no mmproj companion, so it "
|
||||
"cannot answer the image or audio input in this request. It was not "
|
||||
"downloaded."
|
||||
),
|
||||
)
|
||||
|
||||
need_bytes = _remaining_bytes(repo_id, plan, expected_bytes)
|
||||
fits, free = _enough_disk(need_bytes)
|
||||
if not fits:
|
||||
_release(active)
|
||||
return AutoDownloadRefusal(
|
||||
status = 507,
|
||||
code = "insufficient_disk_space",
|
||||
message = (
|
||||
f"'{_public_label(repo_id, variant)}' needs {_gb(need_bytes)} plus "
|
||||
f"{_gb(_DISK_RESERVE_BYTES)} headroom, but only {_gb(free)} is free."
|
||||
),
|
||||
)
|
||||
|
||||
return await _dispatch(repo_id, variant, expected_bytes, requested_model, hf_token, active)
|
||||
|
||||
|
||||
def preferred_quant(labels) -> Optional[str]:
|
||||
"""The quant a plain load would pick from *labels*, or None.
|
||||
|
||||
The one ranking for "which quant did they mean": local resolution, remote
|
||||
admission and /v1/models must agree, or a bare id means a different quant
|
||||
depending on which of them answered it.
|
||||
"""
|
||||
from utils.models.model_config import _pick_best_gguf
|
||||
|
||||
# _pick_best_gguf ranks filenames and matches upper-case tokens, so feed "<LABEL>.gguf".
|
||||
synthetic: dict[str, str] = {}
|
||||
for name in labels:
|
||||
synthetic.setdefault(f"{name.upper()}.gguf", name)
|
||||
best = _pick_best_gguf(list(synthetic))
|
||||
return synthetic.get(best) if best else None
|
||||
|
||||
|
||||
def _match_variant(wanted: Optional[str], variants: dict[str, int]) -> Optional[str]:
|
||||
"""Resolve the requested quant against what the repo actually has.
|
||||
|
||||
An explicit quant matches case-insensitively and must exist: never quietly
|
||||
substitute another, unlike the loader's low-disk fallback. A bare repo id, or a
|
||||
tag that names no quant (":latest", ":8b"), uses the same preference order as a
|
||||
manual load, matching what the local resolver does with the same tag.
|
||||
"""
|
||||
if wanted:
|
||||
# Exact first, whatever shape: a repo of generically named GGUFs has real
|
||||
# variants like "llama-13b" that are valid worker keys but not quant-shaped,
|
||||
# and defaulting past one would fetch a model nobody asked for.
|
||||
lowered = {name.lower(): name for name in variants}
|
||||
exact = lowered.get(wanted.strip().lower())
|
||||
if exact is not None or looks_like_quant(wanted):
|
||||
# A quant-shaped suffix that matches nothing is a miss, never a swap.
|
||||
return exact
|
||||
return preferred_quant(variants)
|
||||
|
||||
|
||||
async def _dispatch(
|
||||
repo_id: str,
|
||||
variant: str,
|
||||
expected_bytes: int,
|
||||
requested_model: str,
|
||||
hf_token: Optional[str],
|
||||
active: _Active,
|
||||
) -> AutoDownloadRefusal:
|
||||
global _active
|
||||
|
||||
from core.inference.api_monitor import api_monitor
|
||||
from hub.schemas.downloads import DownloadModelRequest
|
||||
from hub.services.models import downloads
|
||||
|
||||
label = _public_label(repo_id, variant)
|
||||
busy = AutoDownloadRefusal(
|
||||
status = 503,
|
||||
code = "model_download_busy",
|
||||
message = f"'{repo_id}' is already being downloaded or loaded. Retry shortly.",
|
||||
retry_after = _RETRY_AFTER_S,
|
||||
)
|
||||
try:
|
||||
dispatched = await downloads.download_model_response(
|
||||
DownloadModelRequest(repo_id = repo_id, gguf_variant = variant),
|
||||
hf_token,
|
||||
allow_ambient_token = False,
|
||||
)
|
||||
except Exception as exc:
|
||||
_release(active)
|
||||
status = getattr(exc, "status_code", None)
|
||||
if status == 409:
|
||||
# A manual load or hub download already owns this repo.
|
||||
return busy
|
||||
logger.warning("auto-download: could not start %r: %s", label, exc)
|
||||
return AutoDownloadRefusal(
|
||||
status = 502,
|
||||
code = "model_download_failed",
|
||||
message = f"Could not start downloading '{requested_model}'.",
|
||||
)
|
||||
|
||||
# accepted=False means no worker launched, so report the conflict instead of taking the slot.
|
||||
if isinstance(dispatched, dict) and not dispatched.get("accepted", True):
|
||||
_release(active)
|
||||
logger.info("auto-download: dispatch refused for %s (%s)", label, dispatched.get("state"))
|
||||
return busy
|
||||
|
||||
monitor_id = api_monitor.record_lifecycle(
|
||||
event = "download", model = label, reason = "api", running = True
|
||||
)
|
||||
with _lock:
|
||||
if _active is active:
|
||||
active.variant = variant
|
||||
active.expected_bytes = expected_bytes
|
||||
active.monitor_id = monitor_id
|
||||
tracked = active
|
||||
else:
|
||||
# Released underneath us: track the job we started, but never stomp a newer owner.
|
||||
tracked = _Active(repo_id, variant, expected_bytes, monitor_id, time.time())
|
||||
if _active is None:
|
||||
_active = tracked
|
||||
|
||||
asyncio.create_task(_watch(tracked, hf_token))
|
||||
logger.info("auto-download: started %s (%s)", label, _gb(expected_bytes))
|
||||
return AutoDownloadRefusal(
|
||||
status = 503,
|
||||
code = "model_downloading",
|
||||
message = (
|
||||
f"Downloading '{label}' ({_gb(expected_bytes)}). Retry shortly. "
|
||||
"Track it in Unsloth Studio."
|
||||
),
|
||||
retry_after = _RETRY_AFTER_S,
|
||||
)
|
||||
|
||||
|
||||
def reset_for_tests() -> None:
|
||||
global _active
|
||||
with _lock:
|
||||
_active = None
|
||||
with _cache_lock:
|
||||
_not_servable.clear()
|
||||
|
|
@ -104,6 +104,14 @@ class InferenceOrchestrator:
|
|||
# so a generate queued behind the cancelled one is skipped, not run.
|
||||
self._drain_event: Any = None
|
||||
self._gen_lock = threading.Lock() # Serializes generation
|
||||
# Cancel event of the request holding _gen_lock: lets a Stop tell whether it owns the
|
||||
# running generation or is queued behind it (the worker's event is shared).
|
||||
self._active_cancel_events: list = []
|
||||
self._executing_cancel_events: list = []
|
||||
self._active_cancel_lock = threading.Lock()
|
||||
# Held across claim + _send_cmd so claim order matches the subprocess dequeue order,
|
||||
# which _owns_worker relies on.
|
||||
self._send_order_lock = threading.Lock()
|
||||
# Set during a switch so a generation winning the _gen_lock handoff bails
|
||||
# instead of starting on the outgoing model.
|
||||
self._unload_pending = False
|
||||
|
|
@ -112,6 +120,13 @@ class InferenceOrchestrator:
|
|||
# bypass _gen_lock, send commands directly, read from per-request
|
||||
# mailboxes routed by a dispatcher thread on request_id.
|
||||
self._mailboxes: dict[str, queue.Queue] = {}
|
||||
# request_id -> cancel event, so the dispatcher can move worker ownership as it routes.
|
||||
# Consumers read their mailbox whenever they get to it, so only the dispatcher sees
|
||||
# responses in the order the worker produced them.
|
||||
self._request_cancel_events: dict[str, object] = {}
|
||||
# Mailboxes for the _gen_lock generations. Kept apart from _mailboxes because that map
|
||||
# means "compare requests are in flight" to the unload and distributed paths.
|
||||
self._direct_mailboxes: dict[str, queue.Queue] = {}
|
||||
self._mailbox_lock = threading.Lock()
|
||||
self._dispatcher_thread: Optional[threading.Thread] = None
|
||||
self._dispatcher_stop = threading.Event()
|
||||
|
|
@ -321,9 +336,27 @@ class InferenceOrchestrator:
|
|||
self._resp_queue = None
|
||||
self._cancel_event = None
|
||||
self._drain_event = None
|
||||
self._reset_worker_scoped_state()
|
||||
logger.info("Inference subprocess shut down")
|
||||
return True
|
||||
|
||||
def _reset_worker_scoped_state(self) -> None:
|
||||
"""Drop bookkeeping that only means anything for the worker that just died.
|
||||
|
||||
Ownership is scoped by cancel-event identity alone, so a consumer still blocked
|
||||
on its mailbox when the process was replaced stayed recorded as the executor. A
|
||||
generation on the fresh worker then failed _owns_worker and could not be stopped.
|
||||
Mailboxes go too: nothing will ever route to them, and a stale one reads as
|
||||
compare activity to the unload path.
|
||||
"""
|
||||
with self._active_cancel_lock:
|
||||
self._active_cancel_events.clear()
|
||||
self._executing_cancel_events.clear()
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes.clear()
|
||||
self._direct_mailboxes.clear()
|
||||
self._request_cancel_events.clear()
|
||||
|
||||
def _cleanup(self):
|
||||
"""atexit handler."""
|
||||
self._shutdown_subprocess(timeout = 5.0)
|
||||
|
|
@ -463,6 +496,74 @@ class InferenceOrchestrator:
|
|||
except (EOFError, OSError, ValueError):
|
||||
return events
|
||||
|
||||
def _direct_reader(self, request_id: str):
|
||||
"""Response reader for a _gen_lock generation, safe once compare exists.
|
||||
|
||||
The dispatcher and this reader would otherwise both consume _resp_queue. A
|
||||
dispatcher started mid-stream took our responses and dropped them as
|
||||
unaddressed (truncating or hanging the chat), and this reader, already blocked
|
||||
on the queue, could take a compare request's response before that dispatcher
|
||||
saw it. Registering a mailbox fixes the first; handing foreign responses to
|
||||
their own mailbox fixes the second.
|
||||
|
||||
Returns (read_one, drain, release).
|
||||
"""
|
||||
mailbox: queue.Queue = queue.Queue()
|
||||
with self._mailbox_lock:
|
||||
self._direct_mailboxes[request_id] = mailbox
|
||||
|
||||
def read_one(timeout: float = 1.0):
|
||||
try:
|
||||
return mailbox.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
thread = self._dispatcher_thread
|
||||
if thread is not None and thread.is_alive():
|
||||
# It owns the queue now, and it routes to us.
|
||||
try:
|
||||
return mailbox.get(timeout = timeout)
|
||||
except queue.Empty:
|
||||
return None
|
||||
resp = self._read_resp(timeout = timeout)
|
||||
if resp is None:
|
||||
return None
|
||||
rid = resp.get("request_id")
|
||||
if rid and rid != request_id:
|
||||
with self._mailbox_lock:
|
||||
other = self._mailboxes.get(rid) or self._direct_mailboxes.get(rid)
|
||||
owner = self._request_cancel_events.get(rid)
|
||||
if other is not None:
|
||||
# We beat the dispatcher to this response, so make its ownership move here
|
||||
# too. The compare consumer opts out of marking, so nothing else promotes
|
||||
# or retires that request: skipping it left this one recorded as the
|
||||
# executor, ignoring its Stop and letting a late reset cancel it.
|
||||
if owner is not None:
|
||||
if resp.get("type", "") in ("gen_done", "gen_error"):
|
||||
self._release_worker(owner)
|
||||
else:
|
||||
self._mark_worker_started(owner)
|
||||
other.put(resp)
|
||||
return None
|
||||
return resp
|
||||
|
||||
def drain(timeout: float = 5.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
resp = read_one(timeout = min(0.5, deadline - time.monotonic()))
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
return
|
||||
continue
|
||||
if resp.get("type", "") in ("gen_done", "gen_error"):
|
||||
return
|
||||
logger.warning("Timed out waiting for gen_done after cancel")
|
||||
|
||||
def release() -> None:
|
||||
with self._mailbox_lock:
|
||||
self._direct_mailboxes.pop(request_id, None)
|
||||
|
||||
return read_one, drain, release
|
||||
|
||||
def _drain_until_gen_done(self, timeout: float = 5.0) -> None:
|
||||
"""Consume resp_queue events until gen_done/gen_error, discarding them.
|
||||
|
||||
|
|
@ -542,6 +643,7 @@ class InferenceOrchestrator:
|
|||
cancel_event = None,
|
||||
stats_holder: Optional[dict] = None,
|
||||
read_timeout: float = 30.0,
|
||||
mark_started: bool = True,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Yield tokens from a response stream until gen_done/gen_error.
|
||||
|
||||
|
|
@ -578,6 +680,11 @@ class InferenceOrchestrator:
|
|||
rtype = resp.get("type", "")
|
||||
if rtype == "status":
|
||||
continue
|
||||
# The worker is answering THIS request, so it is the one executing: only now may its
|
||||
# cancel event speak for the shared worker one. The dispatched path opts out: its
|
||||
# dispatcher already did this in worker order, which a mailbox read can lag behind.
|
||||
if mark_started:
|
||||
self._mark_worker_started(cancel_event)
|
||||
# Subprocess-level error (no request_id); request-scoped failures
|
||||
# arrive as gen_error below.
|
||||
if rtype == "error" and not resp.get("request_id"):
|
||||
|
|
@ -587,7 +694,13 @@ class InferenceOrchestrator:
|
|||
if rtype == "token":
|
||||
# Cancel from route (e.g. SSE connection closed).
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
self._cancel_generation()
|
||||
# Same rule as reset_generation_state: the shared worker event may only be set by
|
||||
# the generation the worker is running. A dispatched request can still be draining
|
||||
# stale mailbox tokens after the dispatcher retired it, and signalling from here
|
||||
# would end the next one instead. Tearing this stream down is always safe, so the
|
||||
# local drain happens either way.
|
||||
if self._owns_worker(cancel_event):
|
||||
self._cancel_generation()
|
||||
drain_on_cancel()
|
||||
return
|
||||
yield resp.get("text", "")
|
||||
|
|
@ -681,8 +794,17 @@ class InferenceOrchestrator:
|
|||
# Route to mailbox if a matching request_id exists
|
||||
if rid:
|
||||
with self._mailbox_lock:
|
||||
mbox = self._mailboxes.get(rid)
|
||||
mbox = self._mailboxes.get(rid) or self._direct_mailboxes.get(rid)
|
||||
owner = self._request_cancel_events.get(rid)
|
||||
if mbox is not None:
|
||||
# Worker order, not consumer order: retire a request the moment its last response
|
||||
# is routed. Waiting for the consumer's finally left it owning the worker after
|
||||
# the worker moved on, so a late Stop for it cancelled whichever request started next.
|
||||
if owner is not None:
|
||||
if rtype in ("gen_done", "gen_error"):
|
||||
self._release_worker(owner)
|
||||
else:
|
||||
self._mark_worker_started(owner)
|
||||
mbox.put(resp)
|
||||
continue
|
||||
|
||||
|
|
@ -798,6 +920,8 @@ class InferenceOrchestrator:
|
|||
)
|
||||
if not unloading:
|
||||
self._mailboxes[request_id] = mailbox
|
||||
if cancel_event is not None:
|
||||
self._request_cancel_events[request_id] = cancel_event
|
||||
# When bailing without a mailbox, note whether any OTHER compare request still
|
||||
# routes through the dispatcher; if none and this call started it, stop it below.
|
||||
orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes
|
||||
|
|
@ -813,11 +937,19 @@ class InferenceOrchestrator:
|
|||
yield GenStreamError("Error: model is being unloaded", public = True)
|
||||
return
|
||||
|
||||
# Claim before sending, like the locked path: dispatched runs are concurrent by design,
|
||||
# so without this a Stop on one saw no owner and reset the worker, ending its siblings.
|
||||
# Claim and enqueue under one lock, or two dispatcher threads interleave and claim order
|
||||
# stops matching the subprocess's command order, which _owns_worker reads.
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
with self._send_order_lock:
|
||||
self._claim_worker(cancel_event)
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
self._release_worker(cancel_event)
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes.pop(request_id, None)
|
||||
self._request_cancel_events.pop(request_id, None)
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
|
|
@ -836,10 +968,15 @@ class InferenceOrchestrator:
|
|||
cancel_event = cancel_event,
|
||||
stats_holder = stats_holder,
|
||||
read_timeout = _DISPATCH_READ_TIMEOUT,
|
||||
mark_started = False,
|
||||
)
|
||||
finally:
|
||||
# Normally already retired by the dispatcher at gen_done; this covers streams that
|
||||
# end without one (cancel, disconnect, a dead subprocess).
|
||||
self._release_worker(cancel_event)
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes.pop(request_id, None)
|
||||
self._request_cancel_events.pop(request_id, None)
|
||||
|
||||
def _drain_mailbox(
|
||||
self,
|
||||
|
|
@ -1578,6 +1715,11 @@ class InferenceOrchestrator:
|
|||
# Won the lock handoff during a switch; don't start on the outgoing model.
|
||||
yield GenStreamError("Error: model is being unloaded", public = True)
|
||||
return
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
# Stopped while queued on the lock. Sending anyway occupied the worker with a
|
||||
# run the user ended: the cancel is only seen on a token, so a long prefill
|
||||
# (or a generation that reaches gen_done without one) held up its siblings.
|
||||
return
|
||||
request_id = str(uuid.uuid4())
|
||||
image_b64 = self._pil_to_base64(image) if image is not None else None
|
||||
cmd = self._build_generate_cmd(
|
||||
|
|
@ -1599,22 +1741,95 @@ class InferenceOrchestrator:
|
|||
preserve_thinking = preserve_thinking,
|
||||
)
|
||||
|
||||
# Claim the worker BEFORE sending, so a Stop on some OTHER chat -- still queued on the
|
||||
# lock above, having generated nothing -- cannot reset the generation this is starting.
|
||||
# Claiming after the send left the command running unclaimed. Released in the finally.
|
||||
# Own mailbox: a compare request can start the dispatcher while this is streaming,
|
||||
# and it would otherwise consume our responses and drop them.
|
||||
read_one, drain, release_mailbox = self._direct_reader(request_id)
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
try:
|
||||
with self._send_order_lock:
|
||||
self._claim_worker(cancel_event)
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
yield from self._consume_token_stream(
|
||||
self._read_resp,
|
||||
lambda: self._drain_until_gen_done(timeout = 5.0),
|
||||
crash_context = "generation",
|
||||
cancel_event = cancel_event,
|
||||
stats_holder = stats_holder,
|
||||
)
|
||||
yield from self._consume_token_stream(
|
||||
read_one,
|
||||
lambda: drain(timeout = 5.0),
|
||||
crash_context = "generation",
|
||||
cancel_event = cancel_event,
|
||||
stats_holder = stats_holder,
|
||||
)
|
||||
finally:
|
||||
self._release_worker(cancel_event)
|
||||
release_mailbox()
|
||||
|
||||
def reset_generation_state(self):
|
||||
"""Cancel any ongoing generation and reset state."""
|
||||
def _claim_worker(self, cancel_event) -> None:
|
||||
"""Record this request as one the worker will run.
|
||||
|
||||
Admission only. The subprocess executes generations one at a time, so a
|
||||
dispatched request sitting behind another in the command queue is claimed
|
||||
but not executing, and must not be able to signal the shared cancel event
|
||||
(that would end whichever request IS executing). _mark_worker_started
|
||||
promotes it once the worker answers it.
|
||||
"""
|
||||
with self._active_cancel_lock:
|
||||
self._active_cancel_events.append(cancel_event)
|
||||
|
||||
def _mark_worker_started(self, cancel_event) -> None:
|
||||
"""Promote a claimed request to executing, on its first worker response.
|
||||
|
||||
Sole executor: the subprocess runs one generation at a time, so answering
|
||||
this one means it has left the previous one behind.
|
||||
"""
|
||||
if cancel_event is None:
|
||||
return
|
||||
with self._active_cancel_lock:
|
||||
if self._executing_cancel_events[:1] != [cancel_event]:
|
||||
self._executing_cancel_events[:] = [cancel_event]
|
||||
|
||||
def _release_worker(self, cancel_event) -> None:
|
||||
with self._active_cancel_lock:
|
||||
for bucket in (self._active_cancel_events, self._executing_cancel_events):
|
||||
try:
|
||||
bucket.remove(cancel_event)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
def _owns_worker(self, cancel_event) -> bool:
|
||||
"""Whether a reset from this request may signal the shared cancel event.
|
||||
|
||||
True when it is one of the EXECUTING generations, and when nothing is in
|
||||
flight at all: an error path that resets before anything started has no
|
||||
one else to interrupt, so it must not become a silent no-op. Claimed but
|
||||
queued does not count, or a Stop on a queued request would end the
|
||||
running one, including during the prefill before any response arrives.
|
||||
"""
|
||||
with self._active_cancel_lock:
|
||||
if not self._active_cancel_events:
|
||||
# Nothing in flight at all, so there is no one to protect.
|
||||
return True
|
||||
if self._executing_cancel_events:
|
||||
return any(ev is cancel_event for ev in self._executing_cancel_events)
|
||||
# Claimed but nothing has answered yet (A is in prefill). The worker takes commands
|
||||
# in order, so the oldest claim is the executor; anyone else here is queued behind it.
|
||||
return self._active_cancel_events[0] is cancel_event
|
||||
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
"""Cancel any ongoing generation and reset state.
|
||||
|
||||
``caller_cancel_event`` scopes the reset to one request. The worker has a
|
||||
single cancel event and generation is serialized on _gen_lock, so a chat
|
||||
that is still queued has no generation of its own to reset: calling this
|
||||
from its Stop handler would kill whichever chat currently holds the lock.
|
||||
Pass the request's own event and the reset is dropped unless that request
|
||||
is the one running. Omit it for genuinely global resets (unload, switch).
|
||||
"""
|
||||
if caller_cancel_event is not None and not self._owns_worker(caller_cancel_event):
|
||||
return
|
||||
self._cancel_generation()
|
||||
if not self._ensure_subprocess_alive():
|
||||
return
|
||||
|
|
@ -1673,35 +1888,40 @@ class InferenceOrchestrator:
|
|||
if use_adapter is not None:
|
||||
cmd["use_adapter"] = use_adapter
|
||||
|
||||
self._send_cmd(cmd)
|
||||
# Same shared-queue hazard as _generate_inner: see _direct_reader.
|
||||
read_one, _drain, release_mailbox = self._direct_reader(request_id)
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
|
||||
deadline = time.monotonic() + 120.0
|
||||
while time.monotonic() < deadline:
|
||||
remaining = max(0.1, deadline - time.monotonic())
|
||||
resp = self._read_resp(timeout = min(remaining, 1.0))
|
||||
deadline = time.monotonic() + 120.0
|
||||
while time.monotonic() < deadline:
|
||||
remaining = max(0.1, deadline - time.monotonic())
|
||||
resp = read_one(timeout = min(remaining, 1.0))
|
||||
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError(self._subprocess_crash_message("audio generation"))
|
||||
continue
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError(self._subprocess_crash_message("audio generation"))
|
||||
continue
|
||||
|
||||
rtype = resp.get("type", "")
|
||||
rtype = resp.get("type", "")
|
||||
|
||||
if rtype == "audio_done":
|
||||
wav_bytes = base64.b64decode(resp["wav_base64"])
|
||||
sample_rate = resp["sample_rate"]
|
||||
return wav_bytes, sample_rate
|
||||
if rtype == "audio_done":
|
||||
wav_bytes = base64.b64decode(resp["wav_base64"])
|
||||
sample_rate = resp["sample_rate"]
|
||||
return wav_bytes, sample_rate
|
||||
|
||||
if rtype == "audio_error":
|
||||
raise RuntimeError(resp.get("error", "Audio generation failed"))
|
||||
if rtype == "audio_error":
|
||||
raise RuntimeError(resp.get("error", "Audio generation failed"))
|
||||
|
||||
if rtype == "error":
|
||||
raise RuntimeError(resp.get("error", "Unknown error"))
|
||||
if rtype == "error":
|
||||
raise RuntimeError(resp.get("error", "Unknown error"))
|
||||
|
||||
if rtype == "status":
|
||||
continue
|
||||
if rtype == "status":
|
||||
continue
|
||||
|
||||
raise RuntimeError("Timeout waiting for audio generation (120s)")
|
||||
raise RuntimeError("Timeout waiting for audio generation (120s)")
|
||||
finally:
|
||||
release_mailbox()
|
||||
|
||||
def generate_whisper_response(
|
||||
self,
|
||||
|
|
@ -1775,6 +1995,9 @@ class InferenceOrchestrator:
|
|||
# Won the lock handoff during a switch; don't start on the outgoing model.
|
||||
yield GenStreamError("Error: model is being unloaded", public = True)
|
||||
return
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
# Stopped while queued on the lock, same as _generate_inner.
|
||||
return
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
# numpy array -> list for mp.Queue serialization
|
||||
|
|
@ -1797,18 +2020,28 @@ class InferenceOrchestrator:
|
|||
"repetition_penalty": repetition_penalty,
|
||||
}
|
||||
|
||||
# Same shared-queue hazard as _generate_inner: see _direct_reader.
|
||||
read_one, drain, release_mailbox = self._direct_reader(request_id)
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
try:
|
||||
# Claim under the send lock, like _generate_inner: unclaimed, a compare request queued
|
||||
# behind this looked like the oldest owner, so stopping it killed this one.
|
||||
with self._send_order_lock:
|
||||
self._claim_worker(cancel_event)
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
yield from self._consume_token_stream(
|
||||
self._read_resp,
|
||||
lambda: self._drain_until_gen_done(timeout = 5.0),
|
||||
crash_context = "audio input generation",
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
yield from self._consume_token_stream(
|
||||
read_one,
|
||||
lambda: drain(timeout = 5.0),
|
||||
crash_context = "audio input generation",
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
finally:
|
||||
self._release_worker(cancel_event)
|
||||
release_mailbox()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Local helpers (no subprocess needed)
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ from core.tool_healing import (
|
|||
from core.inference.tool_loop_controller import (
|
||||
ToolLoopController,
|
||||
append_deferred_nudges,
|
||||
awaiting_approval_status,
|
||||
coerce_tool_arguments,
|
||||
status_for_tool,
|
||||
tool_event_provenance,
|
||||
|
|
@ -1209,18 +1210,30 @@ def run_safetensors_tool_loop(
|
|||
start_event["awaiting_confirmation"] = needs_confirm
|
||||
|
||||
try:
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
# A gated call has not started: say waiting, not "Running" (GGUF parity).
|
||||
yield {
|
||||
"type": "status",
|
||||
"text": (
|
||||
awaiting_approval_status(decision.tool_name)
|
||||
if needs_confirm
|
||||
else decision.status_text
|
||||
),
|
||||
}
|
||||
yield start_event
|
||||
|
||||
if (
|
||||
decision_slot is not None
|
||||
and wait_tool_decision(
|
||||
_decision = (
|
||||
wait_tool_decision(
|
||||
decision_slot,
|
||||
approval_id,
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
== "deny"
|
||||
):
|
||||
if decision_slot is not None
|
||||
else None
|
||||
)
|
||||
if _decision is not None and _decision != "deny":
|
||||
# Approved: now it really is running.
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
if _decision == "deny":
|
||||
decision_slot = None
|
||||
if provisional_match:
|
||||
provisional_resolved = True
|
||||
|
|
|
|||
|
|
@ -238,6 +238,19 @@ def status_for_tool(tool_name: str, arguments: Mapping[str, Any]) -> str:
|
|||
return f"Calling: {tool_name}"
|
||||
|
||||
|
||||
def awaiting_approval_status(tool_name: str) -> str:
|
||||
"""Status text for a call parked on the approval prompt.
|
||||
|
||||
It has not started, so reporting "Running ..." with a climbing timer reads
|
||||
as a hang.
|
||||
"""
|
||||
if tool_name == "python":
|
||||
return "Waiting for approval: Python"
|
||||
if tool_name == "terminal":
|
||||
return "Waiting for approval: command"
|
||||
return f"Waiting for approval: {tool_name}"
|
||||
|
||||
|
||||
def is_tool_error(result: str) -> bool:
|
||||
return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
|
||||
|
||||
|
|
|
|||
|
|
@ -3105,6 +3105,22 @@ def is_always_safe_tool(name: str) -> bool:
|
|||
return name in _ALWAYS_SAFE_TOOLS
|
||||
|
||||
|
||||
# Tools whose provisional card is only a text preview of the arguments, so it can stream
|
||||
# while awaiting approval.
|
||||
_TEXT_PREVIEW_TOOLS = frozenset({"python", "terminal"})
|
||||
|
||||
|
||||
def has_text_only_provisional_card(name: str) -> bool:
|
||||
"""True when streaming this tool's arguments before approval shows only text.
|
||||
|
||||
A large code payload takes a minute or more to write, and suppressing the
|
||||
card until the call completes leaves the chat blank the whole time. Nothing
|
||||
runs before the decision either way, and you have to read the code to make
|
||||
it.
|
||||
"""
|
||||
return name in _TEXT_PREVIEW_TOOLS
|
||||
|
||||
|
||||
def is_potentially_unsafe_tool_call(name: str, arguments: dict) -> bool:
|
||||
"""Whether a tool call must still pause for approval in auto mode.
|
||||
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ from pathlib import Path
|
|||
from typing import Any
|
||||
|
||||
logger = get_logger(__name__)
|
||||
from utils.hardware import apply_gpu_ids
|
||||
from utils.hardware import apply_gpu_ids, is_apple_silicon
|
||||
|
||||
_SHARE_OBJECT_MAX_BYTES = 1 << 20
|
||||
_SHARE_OBJECT_ERROR_SIZE = -1
|
||||
|
|
@ -801,10 +801,7 @@ def run_inference_process(
|
|||
# ── 0. MLX fast-path — skip torch/transformers ──
|
||||
_ensure_backend_on_path()
|
||||
|
||||
from utils.hardware import hardware as _hw
|
||||
|
||||
_hw.detect_hardware()
|
||||
if _hw.DEVICE == _hw.DeviceType.MLX:
|
||||
if is_apple_silicon():
|
||||
# Non-fatal: fall through with the installed version, but log the cause
|
||||
# instead of swallowing it (issue #6103).
|
||||
try:
|
||||
|
|
@ -816,6 +813,11 @@ def run_inference_process(
|
|||
model_name,
|
||||
exc,
|
||||
)
|
||||
|
||||
from utils.hardware import hardware as _hw
|
||||
|
||||
_hw.detect_hardware()
|
||||
if _hw.DEVICE == _hw.DeviceType.MLX:
|
||||
try:
|
||||
from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed
|
||||
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ def spawn_worker(
|
|||
use_xet: bool,
|
||||
protected_blob_hashes: Optional[frozenset[str]] = None,
|
||||
cache_env: Optional[Mapping[str, str]] = None,
|
||||
allow_ambient_token: bool = True,
|
||||
) -> subprocess.Popen:
|
||||
"""Spawn the download worker.
|
||||
|
||||
|
|
@ -83,7 +84,8 @@ def spawn_worker(
|
|||
env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1"
|
||||
# No token in Unsloth settings: fall back to the backend's own HF_TOKEN so
|
||||
# private repos stay downloadable (needed while inkling repos are private).
|
||||
if not hf_token:
|
||||
# Not for a repo an API caller named: that would lend them the owner's identity.
|
||||
if not hf_token and allow_ambient_token:
|
||||
hf_token = os.environ.get("HF_TOKEN") or None
|
||||
env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0" if hf_token else "1"
|
||||
# hf_transfer's parallel Range chunks can leave sparse partials even in
|
||||
|
|
@ -239,13 +241,31 @@ def finalize_worker_exit(
|
|||
state = classify_exit(rc, cancel_requested = cancel_requested)
|
||||
if state == "complete":
|
||||
registry.set_job(key, "complete")
|
||||
# Where /v1 learns a new model exists: its resolver answers from a cached scan
|
||||
# with no watcher, so it would report the model absent and serve whatever is
|
||||
# resident. Models only: noting a dataset id as a local model would refuse a
|
||||
# bare request naming it instead of letting a foreign id fall through.
|
||||
if repo_type == "model":
|
||||
try:
|
||||
from core.inference.local_model_resolver import (
|
||||
invalidate_index,
|
||||
note_downloaded,
|
||||
warm_index_soon,
|
||||
)
|
||||
|
||||
note_downloaded(repo_id)
|
||||
invalidate_index()
|
||||
# Rebuild here, not on the first request, to keep the scan off the
|
||||
# request path.
|
||||
warm_index_soon()
|
||||
except Exception:
|
||||
pass
|
||||
if transport == download_registry.TRANSPORT_HTTP:
|
||||
registry.update_job_transport(key, download_registry.TRANSPORT_HTTP)
|
||||
if stderr_text:
|
||||
if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text:
|
||||
logger.warning(
|
||||
f"{log_prefix} complete with degraded diagnostics for "
|
||||
f"{label}: {stderr_text}"
|
||||
f"{log_prefix} complete with degraded diagnostics for {label}: {stderr_text}"
|
||||
)
|
||||
else:
|
||||
logger.info(f"{log_prefix} worker diagnostics for {label}: {stderr_text}")
|
||||
|
|
|
|||
|
|
@ -91,6 +91,7 @@ def _spawn_download_worker(
|
|||
use_xet: bool = True,
|
||||
protected_blob_hashes: Optional[frozenset[str]] = None,
|
||||
cache_env: Optional[dict[str, str]] = None,
|
||||
allow_ambient_token: bool = True,
|
||||
) -> subprocess.Popen:
|
||||
args = ["--repo-id", repo_id]
|
||||
if variant:
|
||||
|
|
@ -101,11 +102,21 @@ def _spawn_download_worker(
|
|||
use_xet = use_xet,
|
||||
protected_blob_hashes = protected_blob_hashes,
|
||||
cache_env = cache_env,
|
||||
allow_ambient_token = allow_ambient_token,
|
||||
)
|
||||
|
||||
|
||||
async def download_model_response(body: DownloadModelRequest, hf_token: Optional[str] = None):
|
||||
"""Start a background download for a HuggingFace model."""
|
||||
async def download_model_response(
|
||||
body: DownloadModelRequest,
|
||||
hf_token: Optional[str] = None,
|
||||
*,
|
||||
allow_ambient_token: bool = True,
|
||||
):
|
||||
"""Start a background download for a HuggingFace model.
|
||||
|
||||
``allow_ambient_token=False`` keeps the worker anonymous when the caller sent
|
||||
no token, for repos named over the API rather than chosen here.
|
||||
"""
|
||||
repo_id = body.repo_id.strip()
|
||||
if not _is_valid_repo_id(repo_id):
|
||||
raise HTTPException(
|
||||
|
|
@ -218,6 +229,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
|
|||
use_xet = use_xet,
|
||||
protected_blob_hashes = protected_blob_hashes,
|
||||
cache_env = cache_env,
|
||||
allow_ambient_token = allow_ambient_token,
|
||||
),
|
||||
hf_token = hf_token,
|
||||
label = label,
|
||||
|
|
|
|||
|
|
@ -8,12 +8,19 @@ filter_sensitive_data (structlog processor for sanitization), and
|
|||
get_logger (factory for structured loggers).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import structlog
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
# Annotations only: a runtime import makes the ASGI stack a hard dependency of
|
||||
# every CLI command.
|
||||
if TYPE_CHECKING:
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from utils.native_path_leases import redact_native_paths
|
||||
|
||||
|
|
|
|||
|
|
@ -1075,7 +1075,9 @@ async def liveness_check():
|
|||
"status": "alive",
|
||||
"service": "Unsloth UI Backend",
|
||||
"desktop_protocol_version": 1,
|
||||
"desktop_manageability_version": 1,
|
||||
# Lockstep with DESKTOP_MANAGEABILITY_VERSION in
|
||||
# studio/src-tauri/src/preflight/version.rs and `desktop-capabilities`.
|
||||
"desktop_manageability_version": 2,
|
||||
"supports_desktop_auth": True,
|
||||
"supports_desktop_backend_ownership": True,
|
||||
"studio_root_id": _studio_root_id(),
|
||||
|
|
@ -1098,7 +1100,8 @@ async def health_check(request: Request):
|
|||
"service": "Unsloth UI Backend",
|
||||
"chat_only": _hw_module.CHAT_ONLY,
|
||||
"desktop_protocol_version": 1,
|
||||
"desktop_manageability_version": 1,
|
||||
# Lockstep: see the note in /api/liveness above.
|
||||
"desktop_manageability_version": 2,
|
||||
"supports_desktop_auth": True,
|
||||
"supports_desktop_backend_ownership": True,
|
||||
# Opaque per-install id; launchers reject sibling Studios on the same port.
|
||||
|
|
@ -1249,16 +1252,18 @@ def _get_cached_system_gpu_info(logger) -> tuple[dict[str, Any], dict[str, Any]]
|
|||
)
|
||||
enriched_devices.append(enriched_dev)
|
||||
|
||||
# Whether GGUF loads accept an explicit gpu_ids pick: /load and
|
||||
# /validate 400 picks on XPU hosts (no visibility mask speaks torch-xpu
|
||||
# ordinals) and on Vulkan-only builds (--device pins ggml's own
|
||||
# ordinals), so the picker must not offer them.
|
||||
# Whether GGUF loads accept an explicit gpu_ids pick. /load and /validate
|
||||
# 400 picks on XPU hosts, where no visibility mask speaks torch-xpu
|
||||
# ordinals. A Vulkan build IS pinnable: its picks are ggml ordinals, the
|
||||
# same space `--device Vulkan<i>` uses, so check it first and let it
|
||||
# through even on an XPU host (the XPU ban is about torch ordinals).
|
||||
is_vulkan_build = False
|
||||
try:
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
from utils.hardware import DeviceType, get_device
|
||||
gpu_ids_supported = (
|
||||
get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend()
|
||||
)
|
||||
|
||||
is_vulkan_build = LlamaCppBackend._is_vulkan_backend()
|
||||
gpu_ids_supported = is_vulkan_build or get_device() != DeviceType.XPU
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not resolve gpu_ids support: {e}")
|
||||
gpu_ids_supported = True
|
||||
|
|
@ -1284,7 +1289,9 @@ def _get_cached_system_gpu_info(logger) -> tuple[dict[str, Any], dict[str, Any]]
|
|||
inference_gpu_info = (
|
||||
{
|
||||
**vulkan_info,
|
||||
"gguf_gpu_ids_supported": False,
|
||||
# Pinnable only once the probe actually enumerated devices:
|
||||
# without ordinals the frontend has nothing valid to offer.
|
||||
"gguf_gpu_ids_supported": bool(vulkan_info.get("devices")),
|
||||
}
|
||||
if vulkan_info is not None
|
||||
else gpu_info
|
||||
|
|
|
|||
|
|
@ -204,12 +204,26 @@ class LoadRequest(BaseModel):
|
|||
"auth, UI/server mode) are rejected. Ignored for non-GGUF models."
|
||||
),
|
||||
)
|
||||
force_cancel_active: bool = Field(
|
||||
False,
|
||||
description = (
|
||||
"Stop chats still generating instead of refusing with 409. A load "
|
||||
"replaces the llama-server every open conversation decodes on."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class UnloadRequest(BaseModel):
|
||||
"""Request to unload a model"""
|
||||
|
||||
model_path: str = Field(..., description = "Model identifier to unload")
|
||||
force_cancel_active: bool = Field(
|
||||
False,
|
||||
description = (
|
||||
"Stop chats still generating instead of refusing with 409. An "
|
||||
"unload takes away the llama-server they are decoding on."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TranscribeRequest(BaseModel):
|
||||
|
|
@ -373,6 +387,14 @@ class InstallLatestTransformersRequest(BaseModel):
|
|||
description = "Exact transformers version to install; must match the current "
|
||||
"latest PyPI release reported by /validate.",
|
||||
)
|
||||
force_cancel_active: bool = Field(
|
||||
False,
|
||||
description = (
|
||||
"Stop chats still generating instead of refusing with 409. The install "
|
||||
"is a step of the model swap that raised the same prompt, so a client "
|
||||
"that already got consent for that swap can carry it through here."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class InstallLatestTransformersResponse(BaseModel):
|
||||
|
|
@ -2088,7 +2110,8 @@ class AnthropicMessage(BaseModel):
|
|||
|
||||
|
||||
class AnthropicTool(BaseModel):
|
||||
# Client tools have input_schema; server tools may only have type/name.
|
||||
# User-defined client tools have input_schema; Anthropic-schema client tools
|
||||
# and server tools use type/name.
|
||||
type: Optional[str] = None
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -37,12 +37,14 @@ from utils.helper_precache_settings import (
|
|||
from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents
|
||||
from utils.openai_auto_switch_settings import (
|
||||
DEFAULT_AUTO_UNLOAD_KEEP_KV,
|
||||
DEFAULT_OPENAI_AUTO_DOWNLOAD_ENABLED,
|
||||
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
|
||||
get_auto_unload_idle_seconds,
|
||||
get_auto_unload_keep_kv,
|
||||
get_model_overrides,
|
||||
get_openai_auto_switch_enabled,
|
||||
get_stored_auto_unload_idle_seconds,
|
||||
get_stored_openai_auto_download_enabled,
|
||||
set_model_override,
|
||||
set_openai_auto_switch,
|
||||
)
|
||||
|
|
@ -112,6 +114,7 @@ class OpenAIAutoSwitchPayload(BaseModel):
|
|||
# None leaves the stored value untouched (partial updates can't clobber it).
|
||||
auto_unload_idle_seconds: Optional[int] = Field(default = None, ge = 0)
|
||||
auto_unload_keep_kv: Optional[bool] = None
|
||||
auto_download_model: Optional[bool] = None
|
||||
|
||||
|
||||
class OpenAIAutoSwitchResponse(BaseModel):
|
||||
|
|
@ -123,6 +126,8 @@ class OpenAIAutoSwitchResponse(BaseModel):
|
|||
# is false, so the UI can show idle-unload as active instead of "needs enable".
|
||||
idle_unload_active: bool = False
|
||||
auto_unload_keep_kv: bool = DEFAULT_AUTO_UNLOAD_KEEP_KV
|
||||
# Stored, not effective: the UI must round-trip the saved value across an auto-switch toggle.
|
||||
auto_download_model: bool = DEFAULT_OPENAI_AUTO_DOWNLOAD_ENABLED
|
||||
|
||||
|
||||
class ModelOverridePayload(BaseModel):
|
||||
|
|
@ -245,6 +250,7 @@ def get_openai_auto_switch(
|
|||
auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(),
|
||||
idle_unload_active = get_auto_unload_idle_seconds() > 0,
|
||||
auto_unload_keep_kv = get_auto_unload_keep_kv(),
|
||||
auto_download_model = get_stored_openai_auto_download_enabled(),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -253,8 +259,11 @@ def update_openai_auto_switch(
|
|||
payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject)
|
||||
) -> OpenAIAutoSwitchResponse:
|
||||
try:
|
||||
enabled, idle_seconds, keep_kv = set_openai_auto_switch(
|
||||
payload.enabled, payload.auto_unload_idle_seconds, payload.auto_unload_keep_kv
|
||||
enabled, idle_seconds, keep_kv, auto_download = set_openai_auto_switch(
|
||||
payload.enabled,
|
||||
payload.auto_unload_idle_seconds,
|
||||
payload.auto_unload_keep_kv,
|
||||
payload.auto_download_model,
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise log_and_http_error(
|
||||
|
|
@ -274,6 +283,7 @@ def update_openai_auto_switch(
|
|||
auto_unload_idle_seconds = idle_seconds,
|
||||
idle_unload_active = idle_unload_active,
|
||||
auto_unload_keep_kv = keep_kv,
|
||||
auto_download_model = auto_download,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1377,13 +1377,21 @@ def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
|
|||
set_tool_policy(enable_tools)
|
||||
|
||||
|
||||
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*: the admission queue caps concurrent
|
||||
# chats at the slot count, so a direct launch matches the CLI (VRAM fit may still cut it
|
||||
# back). Defined above run_server() so embedders that omit it do not serialise every chat.
|
||||
_PARALLEL_MIN = 1
|
||||
_PARALLEL_MAX = 64
|
||||
_PARALLEL_DEFAULT_PLAIN = 4
|
||||
|
||||
|
||||
def run_server(
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8888,
|
||||
frontend_path: Path = _DEFAULT_FRONTEND_PATH,
|
||||
silent: bool = False,
|
||||
api_only: bool = False,
|
||||
llama_parallel_slots: int = 1,
|
||||
llama_parallel_slots: int = _PARALLEL_DEFAULT_PLAIN,
|
||||
cloudflare: "Optional[bool]" = None,
|
||||
secure: bool = False,
|
||||
enable_tools: "Optional[bool]" = None,
|
||||
|
|
@ -1399,7 +1407,8 @@ def run_server(
|
|||
frontend_path: Path to frontend build directory (optional)
|
||||
silent: Suppress startup messages
|
||||
api_only: API server only, no frontend (for Tauri desktop app)
|
||||
llama_parallel_slots: parallel slots for llama-server
|
||||
llama_parallel_slots: parallel slots for llama-server (default
|
||||
_PARALLEL_DEFAULT_PLAIN, matching the CLI entry points)
|
||||
cloudflare: opt in to the public Cloudflare HTTPS tunnel for a wildcard
|
||||
bind. Tri-state: None (unset) and False both mean off; True enables it.
|
||||
--secure implies it (True) and rejects an explicit False.
|
||||
|
|
@ -1817,15 +1826,6 @@ def run_server(
|
|||
return app
|
||||
|
||||
|
||||
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_* and the shared
|
||||
# core/inference/llama_server_args.py PARALLEL_* (the per-load
|
||||
# LoadRequest.n_parallel bounds). Default 1 is for direct backend launches;
|
||||
# `unsloth studio run` always passes its own value (4).
|
||||
_PARALLEL_MIN = 1
|
||||
_PARALLEL_MAX = 64
|
||||
_PARALLEL_DEFAULT_PLAIN = 1
|
||||
|
||||
|
||||
def _build_arg_parser():
|
||||
"""Build the backend CLI argument parser.
|
||||
|
||||
|
|
@ -1920,8 +1920,8 @@ def _build_arg_parser():
|
|||
default = _PARALLEL_DEFAULT_PLAIN,
|
||||
help = (
|
||||
f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). "
|
||||
f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4. "
|
||||
"The Studio run settings (Parallel Slots) can override it per load."
|
||||
f"Default {_PARALLEL_DEFAULT_PLAIN}. The Studio run settings "
|
||||
"(Parallel Slots) override it per load."
|
||||
),
|
||||
)
|
||||
return parser
|
||||
|
|
|
|||
146
studio/backend/state/active_generations.py
Normal file
146
studio/backend/state/active_generations.py
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Registry of in-flight chat generations, keyed by conversation.
|
||||
|
||||
New Chat leaves the previous conversation streaming, so /load and /unload need
|
||||
to know which chats a reload would interrupt: they refuse with 409 unless the
|
||||
caller opts in to cancelling them, and GET /inference/active-generations lets
|
||||
the UI name them. A frontend guard alone would miss a second tab or a REST call.
|
||||
|
||||
Entries hold the same threading.Event as the per-run cancel registry in
|
||||
routes/inference.py, so cancel_all() closes each generation's own upstream
|
||||
stream and never signals llama-server itself.
|
||||
|
||||
A plain dict plus a threading.Lock: no signals, no process groups, no event loop
|
||||
affinity, so it behaves identically on Linux, macOS, Windows and WSL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Optional
|
||||
|
||||
# handle id -> entry. Keyed by handle, not thread_id: a tool continuation can register
|
||||
# before the previous leg unregisters, and one key would drop the other.
|
||||
_ACTIVE: dict[str, dict[str, Any]] = {}
|
||||
_LOCK = threading.Lock()
|
||||
|
||||
|
||||
class ActiveGeneration:
|
||||
"""Registers one in-flight generation for the duration of the block.
|
||||
|
||||
Each __enter__ mints its own handle, so overlapping uses never clobber.
|
||||
"""
|
||||
|
||||
__slots__ = ("thread_id", "cancel_event", "model", "kind", "_handle")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cancel_event: threading.Event,
|
||||
*,
|
||||
thread_id: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
kind: str = "chat",
|
||||
):
|
||||
self.thread_id = thread_id or None
|
||||
self.cancel_event = cancel_event
|
||||
self.model = model or None
|
||||
self.kind = kind
|
||||
self._handle: Optional[str] = None
|
||||
|
||||
def __enter__(self) -> "ActiveGeneration":
|
||||
self._handle = uuid.uuid4().hex
|
||||
with _LOCK:
|
||||
_ACTIVE[self._handle] = {
|
||||
"handle": self._handle,
|
||||
"thread_id": self.thread_id,
|
||||
"model": self.model,
|
||||
"kind": self.kind,
|
||||
"started_at": time.time(),
|
||||
"event": self.cancel_event,
|
||||
}
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc) -> bool:
|
||||
handle, self._handle = self._handle, None
|
||||
if handle is not None:
|
||||
with _LOCK:
|
||||
_ACTIVE.pop(handle, None)
|
||||
return False
|
||||
|
||||
|
||||
def snapshot() -> list[dict[str, Any]]:
|
||||
"""In-flight generations, newest last. Drops the Event: this is a response."""
|
||||
with _LOCK:
|
||||
entries = list(_ACTIVE.values())
|
||||
entries.sort(key = lambda e: e["started_at"])
|
||||
return [
|
||||
{
|
||||
"handle": e["handle"],
|
||||
"thread_id": e["thread_id"],
|
||||
"model": e["model"],
|
||||
"kind": e["kind"],
|
||||
"started_at": e["started_at"],
|
||||
}
|
||||
for e in entries
|
||||
]
|
||||
|
||||
|
||||
def active_thread_ids() -> list[str]:
|
||||
"""Distinct conversation ids with a generation in flight, in start order.
|
||||
|
||||
A first turn that races persistence has no thread id yet: count() sees it,
|
||||
this cannot name it.
|
||||
"""
|
||||
seen: list[str] = []
|
||||
for e in snapshot():
|
||||
tid = e["thread_id"]
|
||||
if tid and tid not in seen:
|
||||
seen.append(tid)
|
||||
return seen
|
||||
|
||||
|
||||
def count() -> int:
|
||||
"""Number of generations currently in flight."""
|
||||
with _LOCK:
|
||||
return len(_ACTIVE)
|
||||
|
||||
|
||||
def cancel_all() -> int:
|
||||
"""Signal every in-flight generation to stop. Returns how many were signalled.
|
||||
|
||||
Only sets the cancel events; each stream tears itself down. Entries are
|
||||
removed by their own __exit__, so one mid-cleanup is neither lost nor double
|
||||
counted.
|
||||
"""
|
||||
with _LOCK:
|
||||
events = [e["event"] for e in _ACTIVE.values()]
|
||||
for ev in events:
|
||||
try:
|
||||
ev.set()
|
||||
except Exception:
|
||||
pass
|
||||
return len(events)
|
||||
|
||||
|
||||
def cancel_thread(thread_id: str) -> int:
|
||||
"""Signal only the generations belonging to ``thread_id``."""
|
||||
if not thread_id:
|
||||
return 0
|
||||
with _LOCK:
|
||||
events = [e["event"] for e in _ACTIVE.values() if e["thread_id"] == thread_id]
|
||||
for ev in events:
|
||||
try:
|
||||
ev.set()
|
||||
except Exception:
|
||||
pass
|
||||
return len(events)
|
||||
|
||||
|
||||
def reset_for_tests() -> None:
|
||||
"""Drop every entry. Test-only; never called from request paths."""
|
||||
with _LOCK:
|
||||
_ACTIVE.clear()
|
||||
|
|
@ -58,6 +58,26 @@ def pytest_addoption(parser):
|
|||
# E2E server fixtures
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _no_background_model_scan(monkeypatch):
|
||||
"""Keep the /v1 admission hook from scanning the real HF cache during tests.
|
||||
|
||||
The hook warms the local-model index on a background thread: right in a server,
|
||||
wrong here, since it walks the developer's actual caches and the I/O starves the
|
||||
loop under timing-sensitive streaming tests. Warm tests patch it back.
|
||||
"""
|
||||
import time
|
||||
|
||||
from core.inference import local_model_resolver
|
||||
|
||||
monkeypatch.setattr(local_model_resolver, "warm_index_soon", lambda: None)
|
||||
# Start from a built, empty index: stubbing only the warm left the cold path
|
||||
# walking those caches inside the admission wait, so on a large install the
|
||||
# assertion became a 503 "still indexing". Cold-path tests reset _scan themselves;
|
||||
# _build_index is untouched so tests calling it directly still walk for real.
|
||||
monkeypatch.setattr(local_model_resolver, "_scan", (time.monotonic(), {}))
|
||||
|
||||
|
||||
@pytest.fixture(scope = "session")
|
||||
def studio_server(request):
|
||||
"""Yield ``(base_url, api_key)`` for e2e tests.
|
||||
|
|
|
|||
2635
studio/backend/tests/test_active_generations.py
Normal file
2635
studio/backend/tests/test_active_generations.py
Normal file
File diff suppressed because it is too large
Load diff
975
studio/backend/tests/test_anthropic_admission.py
Normal file
975
studio/backend/tests/test_anthropic_admission.py
Normal file
|
|
@ -0,0 +1,975 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""Admission-control wiring for the Anthropic /v1/messages endpoint.
|
||||
|
||||
The FIFO queue itself is unit-tested in test_llama_admission.py; here we exercise
|
||||
how anthropic_messages reserves a slot, queues when the backend is saturated,
|
||||
streams keep-alives while waiting, releases on completion, and maps rejects to
|
||||
429/503. Slot occupancy is driven directly through the shared queue (keyed by the
|
||||
backend base_url) so generation stays fast and no thread has to block.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import gc
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import warnings
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
_backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
sys.path.insert(0, _backend)
|
||||
|
||||
import routes.inference as inf_mod
|
||||
from routes.inference import (
|
||||
_anthropic_passthrough_retry_url,
|
||||
_anthropic_passthrough_stream,
|
||||
anthropic_messages,
|
||||
)
|
||||
from models.inference import AnthropicMessagesRequest
|
||||
from core.inference.api_monitor import ApiMonitor
|
||||
from core.inference.llama_admission import (
|
||||
ADMISSION_CONTROL_ENV,
|
||||
ADMISSION_KEEPALIVE_INTERVAL_ENV,
|
||||
ADMISSION_MAX_QUEUE_ENV,
|
||||
ADMISSION_QUEUE_PER_SLOT_ENV,
|
||||
ADMISSION_QUEUE_TIMEOUT_ENV,
|
||||
LlamaAdmissionConfig,
|
||||
get_llama_admission_queue,
|
||||
reset_llama_admission_queues,
|
||||
)
|
||||
from fastapi import HTTPException
|
||||
|
||||
_KEY = "http://llama.admission.test:9999"
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _isolate(monkeypatch):
|
||||
reset_llama_admission_queues()
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", ApiMonitor(max_entries = 64))
|
||||
monkeypatch.setattr(inf_mod, "_CANCEL_REGISTRY", {})
|
||||
for name in (
|
||||
ADMISSION_CONTROL_ENV,
|
||||
ADMISSION_QUEUE_TIMEOUT_ENV,
|
||||
ADMISSION_KEEPALIVE_INTERVAL_ENV,
|
||||
ADMISSION_MAX_QUEUE_ENV,
|
||||
ADMISSION_QUEUE_PER_SLOT_ENV,
|
||||
# Legacy spellings resolve too, so clear both for isolation.
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL",
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT",
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL",
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE",
|
||||
):
|
||||
monkeypatch.delenv(name, raising = False)
|
||||
yield
|
||||
reset_llama_admission_queues()
|
||||
|
||||
|
||||
class _Request:
|
||||
def __init__(self, disconnected = False):
|
||||
self.state = SimpleNamespace()
|
||||
self.url = SimpleNamespace(path = "/v1/messages")
|
||||
self.method = "POST"
|
||||
self._disconnected = disconnected
|
||||
|
||||
async def is_disconnected(self):
|
||||
return self._disconnected
|
||||
|
||||
|
||||
def _install_backend(
|
||||
monkeypatch,
|
||||
*,
|
||||
slots = 1,
|
||||
base_url = _KEY,
|
||||
count_tokens = None,
|
||||
):
|
||||
def _gen_plain(**_kwargs):
|
||||
yield "ok"
|
||||
|
||||
def _gen_tools(**_kwargs):
|
||||
yield {"type": "content", "text": "ok"}
|
||||
|
||||
backend = SimpleNamespace(
|
||||
is_loaded = True,
|
||||
is_vision = False,
|
||||
supports_tools = True,
|
||||
supports_tool_passthrough = False,
|
||||
model_identifier = "test-model",
|
||||
context_length = 2048,
|
||||
count_chat_tokens = count_tokens or (lambda *a, **k: 2),
|
||||
generate_chat_completion = _gen_plain,
|
||||
generate_chat_completion_with_tools = _gen_tools,
|
||||
effective_parallel_slots = slots,
|
||||
base_url = base_url,
|
||||
)
|
||||
monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
|
||||
return backend
|
||||
|
||||
|
||||
def _payload(**fields) -> AnthropicMessagesRequest:
|
||||
base = {"max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}
|
||||
base.update(fields)
|
||||
return AnthropicMessagesRequest(**base)
|
||||
|
||||
|
||||
def _record_admission_logs(monkeypatch):
|
||||
"""Capture _llama_admission_log output.
|
||||
|
||||
Through the logger rather than caplog: this one is a structlog bound logger,
|
||||
so it never reaches the stdlib handlers caplog installs.
|
||||
"""
|
||||
records = []
|
||||
|
||||
def _record(level):
|
||||
return lambda fmt, *args: records.append((level, fmt % args))
|
||||
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"logger",
|
||||
SimpleNamespace(
|
||||
debug = _record("debug"),
|
||||
info = _record("info"),
|
||||
warning = _record("warning"),
|
||||
),
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
def _snapshot(key = _KEY):
|
||||
return get_llama_admission_queue(key).snapshot()
|
||||
|
||||
|
||||
def _occupy(key, capacity, n):
|
||||
"""Hold ``n`` slots on the queue so the next reserve must wait; returns leases."""
|
||||
leases = []
|
||||
for _ in range(n):
|
||||
reservation = get_llama_admission_queue(key).reserve(
|
||||
capacity = capacity, config = LlamaAdmissionConfig()
|
||||
)
|
||||
lease = reservation.lease_nowait()
|
||||
assert lease is not None
|
||||
leases.append(lease)
|
||||
return leases
|
||||
|
||||
|
||||
async def _consume(response):
|
||||
chunks = []
|
||||
async for chunk in response.body_iterator:
|
||||
chunks.append(chunk.decode() if isinstance(chunk, (bytes, bytearray)) else chunk)
|
||||
return "".join(chunks)
|
||||
|
||||
|
||||
# ── Non-streaming ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_non_streaming_completes_and_releases_slot(monkeypatch):
|
||||
_install_backend(monkeypatch, slots = 2)
|
||||
|
||||
async def _run():
|
||||
response = await anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
assert response.status_code == 200
|
||||
snap = _snapshot()
|
||||
assert snap.active == 0 and snap.queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_non_streaming_queue_full_returns_429(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1) # slot busy
|
||||
# One waiter fills the max_queue=1; the next reserve rejects.
|
||||
get_llama_admission_queue(_KEY).reserve(
|
||||
capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
assert exc.value.status_code == 429
|
||||
# rate_limit_error is what Anthropic SDKs back off on; overloaded_error is 529.
|
||||
# The type string alone does not pin the envelope, since OpenAI's 429 uses the
|
||||
# same word. Assert the shape too, or emitting an OpenAI body still passes.
|
||||
detail = exc.value.detail
|
||||
assert detail["type"] == "error"
|
||||
assert "request_id" in detail
|
||||
assert set(detail["error"]) == {"type", "message"}
|
||||
assert detail["error"]["type"] == "rate_limit_error"
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_admission_events_are_logged_on_the_anthropic_surface(monkeypatch):
|
||||
# The OpenAI passthrough logs these with a mode; without the same on /v1/messages
|
||||
# an operator debugging a slow Anthropic client has nothing to look at, and the
|
||||
# pool is shared, so it is the same triage.
|
||||
records = _record_admission_logs(monkeypatch)
|
||||
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
get_llama_admission_queue(_KEY).reserve(
|
||||
capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)
|
||||
)
|
||||
with pytest.raises(HTTPException):
|
||||
await anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
full = [msg for _level, msg in records if "queue-full" in msg]
|
||||
assert full, records
|
||||
assert "llama admission queue-full" in full[0]
|
||||
assert "mode=anthropic_nonstream" in full[0]
|
||||
|
||||
|
||||
def test_streaming_admission_waiting_is_logged(monkeypatch):
|
||||
# queued and granted-after-wait were both emitted with nothing asserting them.
|
||||
records = _record_admission_logs(monkeypatch)
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
task = asyncio.create_task(_consume(response))
|
||||
await asyncio.sleep(0.15)
|
||||
for lease in held:
|
||||
lease.release()
|
||||
await asyncio.wait_for(task, timeout = 5)
|
||||
|
||||
asyncio.run(_run())
|
||||
events = [msg for _level, msg in records if "llama admission" in msg]
|
||||
# "llama admission queued", not "queued": every line carries a queued=N field,
|
||||
# so the bare substring matches any admission log at all.
|
||||
assert any(
|
||||
"llama admission queued" in m and "mode=anthropic_stream" in m for m in events
|
||||
), events
|
||||
granted = [m for m in events if "granted-after-wait" in m]
|
||||
assert granted, events
|
||||
# wait_ms is the point of the event: a grant that reports nothing is useless.
|
||||
assert re.search(r"wait_ms=\d+", granted[0]), granted
|
||||
|
||||
|
||||
def test_streaming_admission_timeout_is_logged(monkeypatch):
|
||||
records = _record_admission_logs(monkeypatch)
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15")
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1) # never released, so the waiter times out
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
await _consume(response)
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
timeouts = [msg for level, msg in records if "timeout" in msg and level == "warning"]
|
||||
assert timeouts, records
|
||||
assert "mode=anthropic_stream" in timeouts[0]
|
||||
|
||||
|
||||
def test_streaming_give_up_while_queued_is_logged(monkeypatch):
|
||||
# cancelled-before-upstream is the one that tells an operator a client walked
|
||||
# away rather than the backend being slow.
|
||||
records = _record_admission_logs(monkeypatch)
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True),
|
||||
request = _Request(disconnected = True),
|
||||
current_subject = "t",
|
||||
)
|
||||
await _consume(response)
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
events = [msg for _level, msg in records if "llama admission" in msg]
|
||||
assert any("llama admission cancelled-before-upstream" in m for m in events), events
|
||||
|
||||
|
||||
def test_non_streaming_times_out_returns_503(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1) # never released -> waiter times out
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
assert exc.value.status_code == 503
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_non_streaming_queued_then_admitted(monkeypatch):
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
task = asyncio.create_task(
|
||||
anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
assert _snapshot().queued == 1 # waiting on the busy slot
|
||||
held[0].release() # free it
|
||||
response = await asyncio.wait_for(task, timeout = 2)
|
||||
assert response.status_code == 200
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_capacity_enforced_from_effective_parallel_slots(monkeypatch):
|
||||
_install_backend(monkeypatch, slots = 3)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 3, 3) # all 3 slots busy
|
||||
task = asyncio.create_task(
|
||||
anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
snap = _snapshot()
|
||||
assert snap.capacity == 3 and snap.active == 3 and snap.queued == 1
|
||||
for lease in held:
|
||||
lease.release()
|
||||
response = await asyncio.wait_for(task, timeout = 2)
|
||||
assert response.status_code == 200
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_disabled_admission_bypasses_limit(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1) # would block if admission were on
|
||||
response = await asyncio.wait_for(
|
||||
anthropic_messages(_payload(), request = _Request(), current_subject = "t"),
|
||||
timeout = 2,
|
||||
)
|
||||
assert response.status_code == 200
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ── Streaming ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_streaming_completes_and_releases_slot(monkeypatch):
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
blob = await _consume(response)
|
||||
assert "event: message_start" in blob
|
||||
assert "event: message_stop" in blob
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_streaming_emits_keepalives_while_queued_then_streams(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
body = response.body_iterator
|
||||
# First chunk must be a keep-alive comment (slot still busy).
|
||||
first = await asyncio.wait_for(body.__anext__(), timeout = 2)
|
||||
first = first.decode() if isinstance(first, (bytes, bytearray)) else first
|
||||
assert first.startswith(":") # SSE comment keep-alive
|
||||
held[0].release() # free the slot -> real stream follows
|
||||
rest = await asyncio.wait_for(_drain(body), timeout = 2)
|
||||
assert "event: message_start" in rest
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_streaming_queue_full_returns_429(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
get_llama_admission_queue(_KEY).reserve(
|
||||
capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await anthropic_messages(_payload(stream = True), request = _Request(), current_subject = "t")
|
||||
assert exc.value.status_code == 429
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_streaming_disconnect_while_queued_frees_slot(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
body = response.body_iterator
|
||||
await asyncio.wait_for(body.__anext__(), timeout = 2) # one keep-alive
|
||||
assert _snapshot().queued == 1
|
||||
await body.aclose() # client goes away mid-wait
|
||||
held[0].release()
|
||||
await asyncio.sleep(0.05)
|
||||
snap = _snapshot()
|
||||
assert snap.queued == 0 and snap.active == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ── Shared queue + fairness + speed ───────────────────────────
|
||||
|
||||
|
||||
def test_shares_queue_with_openai_by_base_url(monkeypatch):
|
||||
"""The two API surfaces must land on one pool of the same llama-server slots.
|
||||
|
||||
Reserves through the OpenAI helper the /v1/chat/completions path uses, rather
|
||||
than poking the queue directly, so this fails if either side ever derives a
|
||||
different key.
|
||||
"""
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
openai_reservation, _ = inf_mod._openai_llama_admission_reserve(
|
||||
request = _Request(), llama_backend = inf_mod.get_llama_cpp_backend()
|
||||
)
|
||||
openai_lease = openai_reservation.lease_nowait()
|
||||
assert openai_lease is not None
|
||||
assert _snapshot().active == 1 # same key the Anthropic side will use
|
||||
|
||||
task = asyncio.create_task(
|
||||
anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
assert _snapshot().queued == 1 # queued behind the OpenAI generation
|
||||
openai_lease.release()
|
||||
assert (await asyncio.wait_for(task, timeout = 2)).status_code == 200
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_non_streaming_client_gone_while_queued_returns_499(monkeypatch):
|
||||
# The disconnect-while-queued branch; nothing else exercised 499.
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await anthropic_messages(
|
||||
_payload(), request = _Request(disconnected = True), current_subject = "t"
|
||||
)
|
||||
assert exc.value.status_code == 499
|
||||
assert _snapshot().queued == 0 # waiter cleaned up, not left parked
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_streaming_timeout_emits_an_error_event_and_frees_the_slot(monkeypatch):
|
||||
# Only the non-streaming 503 was covered; streaming reports in-band instead.
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15")
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1) # never released, so the waiter times out
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
body = await _consume(response)
|
||||
assert "event: error" in body
|
||||
assert "message_start" not in body # never reached the model
|
||||
for lease in held:
|
||||
lease.release()
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_fifo_fairness_across_many_waiters(monkeypatch):
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
order = []
|
||||
|
||||
async def _one(i):
|
||||
resp = await anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
order.append(i)
|
||||
return resp
|
||||
|
||||
tasks = [asyncio.create_task(_one(i)) for i in range(8)]
|
||||
await asyncio.sleep(0.2)
|
||||
assert _snapshot().queued == 8
|
||||
held[0].release()
|
||||
await asyncio.wait_for(asyncio.gather(*tasks), timeout = 5)
|
||||
assert order == list(range(8)) # granted in arrival order
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_uncontended_hot_path_is_fast(monkeypatch):
|
||||
_install_backend(monkeypatch, slots = 4)
|
||||
|
||||
async def _run():
|
||||
start = time.perf_counter()
|
||||
for _ in range(50):
|
||||
resp = await anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
assert resp.status_code == 200
|
||||
elapsed = time.perf_counter() - start
|
||||
# Generous ceiling on purpose: this guards against admission accidentally
|
||||
# serialising or sleeping on the uncontended path, not against a slow
|
||||
# runner, so it must not flake on a loaded CI box.
|
||||
assert elapsed < 10.0, f"50 uncontended round-trips took {elapsed:.2f}s"
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
async def _drain(body):
|
||||
chunks = []
|
||||
async for chunk in body:
|
||||
chunks.append(chunk.decode() if isinstance(chunk, (bytes, bytearray)) else chunk)
|
||||
return "".join(chunks)
|
||||
|
||||
|
||||
def test_streaming_midstream_cancel_finalizes_the_monitor(monkeypatch):
|
||||
# A mid-stream disconnect is delivered as CancelledError so the monitored body
|
||||
# can finalize its entry. Closing the inner iterator with aclose() instead
|
||||
# delivers GeneratorExit, and the entry stays "running" for the process life.
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
body = response.body_iterator
|
||||
await asyncio.wait_for(body.__anext__(), timeout = 2) # stream started
|
||||
assert inf_mod.api_monitor.active_count() == 1
|
||||
|
||||
# Propagates back out, as the un-admitted path did; what matters is that
|
||||
# the monitored body saw it on the way through.
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await body.athrow(asyncio.CancelledError()) # client vanished
|
||||
|
||||
assert inf_mod.api_monitor.active_count() == 0
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_streaming_give_up_while_queued_finalizes_the_monitor(monkeypatch):
|
||||
# Cancelled before the body ever ran, so nothing downstream can close the
|
||||
# entry out; the wrapper has to do it.
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
body = response.body_iterator
|
||||
await asyncio.wait_for(body.__anext__(), timeout = 2) # keep-alive, still queued
|
||||
assert inf_mod.api_monitor.active_count() == 1
|
||||
|
||||
await body.aclose() # give up while waiting
|
||||
|
||||
assert inf_mod.api_monitor.active_count() == 0
|
||||
for lease in held:
|
||||
lease.release()
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_every_dispatch_site_goes_through_admission():
|
||||
"""All six generation returns in anthropic_messages are admission-wrapped.
|
||||
|
||||
The tool paths need a passthrough-capable backend and a tools payload to reach
|
||||
at runtime, so guard them structurally instead: a new dispatch site added
|
||||
without admission (or one reverted to _monitored_anthropic) fails here.
|
||||
"""
|
||||
import ast
|
||||
import inspect
|
||||
|
||||
tree = ast.parse(inspect.getsource(inf_mod).replace("\t", " "))
|
||||
handler = next(
|
||||
node
|
||||
for node in ast.walk(tree)
|
||||
if isinstance(node, ast.AsyncFunctionDef) and node.name == "anthropic_messages"
|
||||
)
|
||||
# The wrappers themselves call _monitored_anthropic (the non-streaming one
|
||||
# through the swap-gate tracker); only the dispatch sites count.
|
||||
nested = {
|
||||
node
|
||||
for node in ast.walk(handler)
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
and node.name.startswith(("_admitted_anthropic", "_tracked_anthropic"))
|
||||
}
|
||||
inner = {id(n) for wrapper in nested for n in ast.walk(wrapper)}
|
||||
|
||||
called = []
|
||||
for node in ast.walk(handler):
|
||||
if id(node) in inner or not isinstance(node, ast.Call):
|
||||
continue
|
||||
if isinstance(node.func, ast.Name):
|
||||
called.append(node.func.id)
|
||||
|
||||
assert called.count("_admitted_anthropic") == 6
|
||||
assert called.count("_monitored_anthropic") == 0
|
||||
|
||||
|
||||
def test_queued_give_up_runs_the_response_pre_start_cleanup(monkeypatch):
|
||||
"""A stream abandoned while queued must run the builder's eager cleanup.
|
||||
|
||||
The passthrough enters a _TrackedCancel before returning its response and
|
||||
relies on the stream's finally to exit it. That finally never runs for a
|
||||
generator that never started, so the response carries a pre-start hook and
|
||||
the admission wrapper has to chain to it instead of replacing it.
|
||||
"""
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
ran = []
|
||||
|
||||
async def _hook():
|
||||
ran.append(True)
|
||||
|
||||
real = inf_mod._sse_streaming_response
|
||||
|
||||
def _tagged(content, *, unstarted_cleanup = None):
|
||||
return real(content, unstarted_cleanup = _hook)
|
||||
|
||||
monkeypatch.setattr(inf_mod, "_sse_streaming_response", _tagged)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
body = response.body_iterator
|
||||
await asyncio.wait_for(body.__anext__(), timeout = 2) # keep-alive, still queued
|
||||
await body.aclose() # give up before the body ran
|
||||
|
||||
assert ran == [True]
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_passthrough_stream_registers_a_pre_start_cleanup():
|
||||
# Structural guard: the tracker is entered eagerly, so the response must
|
||||
# carry the hook that exits it when the body never starts.
|
||||
import ast
|
||||
import inspect
|
||||
|
||||
src = inspect.getsource(inf_mod._anthropic_passthrough_stream)
|
||||
tree = ast.parse(src.replace("\t", " ").lstrip())
|
||||
returns = [n for n in ast.walk(tree) if isinstance(n, ast.Return) and n.value is not None]
|
||||
call = next(
|
||||
n.value
|
||||
for n in returns
|
||||
if isinstance(n.value, ast.Call)
|
||||
and getattr(n.value.func, "id", "") == "_sse_streaming_response"
|
||||
)
|
||||
hook = next(kw.value for kw in call.keywords if kw.arg == "unstarted_cleanup")
|
||||
# Not just present: a literal None passes the keyword check and still leaks.
|
||||
assert isinstance(hook, ast.Call)
|
||||
assert getattr(hook.func, "id", None) == "_tracked_cancel_unstarted_cleanup"
|
||||
|
||||
|
||||
def test_slot_is_released_even_if_closing_the_body_raises(monkeypatch):
|
||||
# A slot lost here never comes back: with no queue timeout the pool silently
|
||||
# shrinks and later callers wait forever, so the release must not sit behind
|
||||
# anything that can throw.
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _boom(iterator, *, cancelled):
|
||||
raise RuntimeError("close failed")
|
||||
|
||||
monkeypatch.setattr(inf_mod, "_close_openai_admitted_stream_iterator", _boom)
|
||||
|
||||
async def _run():
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
body = response.body_iterator
|
||||
await asyncio.wait_for(body.__anext__(), timeout = 2) # stream started
|
||||
assert _snapshot().active == 1
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await body.aclose()
|
||||
|
||||
assert _snapshot().active == 0 # slot returned despite the failure
|
||||
# And the pool still serves the next caller.
|
||||
again = get_llama_admission_queue(_KEY).reserve(capacity = 1, config = LlamaAdmissionConfig())
|
||||
lease = again.lease_nowait()
|
||||
assert lease is not None
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
_CLIENT_TOOLS = [
|
||||
{"name": "get_time", "description": "t", "input_schema": {"type": "object", "properties": {}}}
|
||||
]
|
||||
|
||||
|
||||
def _passthrough_payload(**fields):
|
||||
# server_tools off + declared tools + a passthrough-capable backend routes
|
||||
# anthropic_messages down the client-tool passthrough dispatch site.
|
||||
return _payload(tools = _CLIENT_TOOLS, enable_tools = False, **fields)
|
||||
|
||||
|
||||
def test_response_pre_start_cleanup_leaves_no_passthrough_tracker(monkeypatch):
|
||||
"""A disconnect before the body starts must leave no tracker and no slot.
|
||||
|
||||
The passthrough registers from inside its body rather than eagerly, so a
|
||||
generator that never runs registers nothing; the hook still has to hand the
|
||||
admission slot back. Asserting through _CANCEL_REGISTRY and the pool rather
|
||||
than the wiring, because the hook can be present and still be a no-op.
|
||||
"""
|
||||
backend = _install_backend(monkeypatch, slots = 1)
|
||||
backend.supports_tool_passthrough = True
|
||||
monkeypatch.setattr(inf_mod, "_CANCEL_REGISTRY", {})
|
||||
|
||||
async def _run():
|
||||
response = await anthropic_messages(
|
||||
_passthrough_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
assert inf_mod._CANCEL_REGISTRY == {}, "nothing runs the body's exit for it yet"
|
||||
|
||||
cleanup = getattr(response, "_unstarted_cleanup", None)
|
||||
assert cleanup is not None
|
||||
await cleanup() # what _SameTaskStreamingResponse runs on a pre-start disconnect
|
||||
|
||||
assert inf_mod._CANCEL_REGISTRY == {}
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_passthrough_dispatch_site_reserves_and_releases(monkeypatch):
|
||||
# Behavioural cover for a dispatch site the other tests never reach.
|
||||
backend = _install_backend(monkeypatch, slots = 1)
|
||||
backend.supports_tool_passthrough = True
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
task = asyncio.create_task(
|
||||
anthropic_messages(_passthrough_payload(), request = _Request(), current_subject = "t")
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
assert _snapshot().queued == 1 # queued behind the busy slot, not bypassing
|
||||
for lease in held:
|
||||
lease.release()
|
||||
with contextlib.suppress(Exception):
|
||||
await asyncio.wait_for(task, timeout = 2) # upstream is not mocked
|
||||
assert _snapshot().active == 0 and _snapshot().queued == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_stream_setup_failure_returns_the_slot(monkeypatch):
|
||||
# count_chat_tokens makes a blocking HTTP call to llama-server, so a dead
|
||||
# server raises here: after lease_nowait() took the slot, before a body
|
||||
# exists to release it. Nothing else can hand the slot back.
|
||||
def _boom(*_a, **_k):
|
||||
raise RuntimeError("tokenizer unreachable")
|
||||
|
||||
_install_backend(monkeypatch, slots = 1, count_tokens = _boom)
|
||||
|
||||
async def _run():
|
||||
with pytest.raises(RuntimeError):
|
||||
await anthropic_messages(_payload(stream = True), request = _Request(), current_subject = "t")
|
||||
snap = _snapshot()
|
||||
assert snap.active == 0, f"slot leaked after stream setup failed: {snap}"
|
||||
# And the pool still serves the next caller.
|
||||
again = get_llama_admission_queue(_KEY).reserve(capacity = 1, config = LlamaAdmissionConfig())
|
||||
assert again.lease_nowait() is not None
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_queued_non_stream_cancel_does_not_leak_a_coroutine(monkeypatch):
|
||||
# The non-stream path builds the generation coroutine before reserving and
|
||||
# only awaits it once admitted. Giving up while queued must close it.
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1)
|
||||
task = asyncio.create_task(
|
||||
anthropic_messages(_payload(), request = _Request(), current_subject = "t")
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
assert _snapshot().queued == 1
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
with warnings.catch_warnings(record = True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
asyncio.run(_run())
|
||||
gc.collect()
|
||||
leaked = [w for w in caught if "never awaited" in str(w.message)]
|
||||
assert not leaked, [str(w.message) for w in leaked]
|
||||
|
||||
|
||||
def test_stream_timeout_marks_the_monitor_entry_as_error(monkeypatch):
|
||||
# The finally finishes the entry as "cancelled"; without the fail() first, a
|
||||
# timed-out request is indistinguishable from a client hang-up in the
|
||||
# monitor. api_monitor.finish is a no-op on an already terminal entry.
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.15")
|
||||
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.05")
|
||||
_install_backend(monkeypatch, slots = 1)
|
||||
|
||||
async def _run():
|
||||
held = _occupy(_KEY, 1, 1) # never released, so the waiter times out
|
||||
response = await anthropic_messages(
|
||||
_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
async for _ in response.body_iterator:
|
||||
pass
|
||||
entries = inf_mod.api_monitor.snapshot()
|
||||
assert entries and entries[0]["status"] == "error", entries
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
class _RespawnBackend:
|
||||
"""Backend whose base_url moves to a new port once respawned."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
mtp_handled = False,
|
||||
fallback_in_progress = False,
|
||||
):
|
||||
self.base_url = "http://127.0.0.1:57953"
|
||||
self.context_length = 4096
|
||||
self.respawn_calls = 0
|
||||
self._mtp_handled = mtp_handled
|
||||
self._mtp_runtime_fallback_in_progress = fallback_in_progress
|
||||
|
||||
def count_chat_tokens(self, *_a, **_k):
|
||||
return 2
|
||||
|
||||
def _maybe_recover_from_mtp_crash(self, _exc):
|
||||
return self._mtp_handled
|
||||
|
||||
def _respawn_if_dead(self):
|
||||
self.respawn_calls += 1
|
||||
self.base_url = "http://127.0.0.1:62933"
|
||||
return True
|
||||
|
||||
|
||||
def test_retry_url_stands_down_while_an_mtp_fallback_is_reloading():
|
||||
# Only the first caller gets True from _maybe_recover_from_mtp_crash; the rest
|
||||
# see False and must still stand down, or they respawn the same MTP config
|
||||
# underneath the fallback already reloading without it.
|
||||
backend = _RespawnBackend(mtp_handled = False, fallback_in_progress = True)
|
||||
|
||||
url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x")))
|
||||
|
||||
assert url is None
|
||||
assert backend.respawn_calls == 0
|
||||
|
||||
|
||||
class _PtRequest:
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
|
||||
async def _passthrough_response(backend):
|
||||
return await _anthropic_passthrough_stream(
|
||||
_PtRequest(),
|
||||
threading.Event(),
|
||||
backend,
|
||||
[{"role": "user", "content": "hi"}],
|
||||
[],
|
||||
0.7,
|
||||
0.95,
|
||||
20,
|
||||
16,
|
||||
"msg_tracker_probe",
|
||||
"test-model",
|
||||
)
|
||||
|
||||
|
||||
def test_disconnect_during_the_opening_lines_exits_the_tracker():
|
||||
# Suspended inside emitter.start()'s yields the generator has not reached the
|
||||
# try/finally that exits the tracker, so those yields need their own.
|
||||
backend = _RespawnBackend()
|
||||
|
||||
async def _run():
|
||||
response = await _passthrough_response(backend)
|
||||
body = response.body_iterator
|
||||
await asyncio.wait_for(body.__anext__(), timeout = 2) # first start line
|
||||
assert inf_mod._CANCEL_REGISTRY, "tracker should be registered"
|
||||
await body.aclose()
|
||||
assert inf_mod._CANCEL_REGISTRY == {}, "tracker leaked"
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_cancel_during_the_opening_lines_exits_the_tracker():
|
||||
# Same window, delivered the way _SameTaskStreamingResponse delivers it.
|
||||
backend = _RespawnBackend()
|
||||
|
||||
async def _run():
|
||||
response = await _passthrough_response(backend)
|
||||
body = response.body_iterator
|
||||
await asyncio.wait_for(body.__anext__(), timeout = 2)
|
||||
assert inf_mod._CANCEL_REGISTRY, "tracker should be registered"
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await body.athrow(asyncio.CancelledError())
|
||||
assert inf_mod._CANCEL_REGISTRY == {}, "tracker leaked"
|
||||
|
||||
asyncio.run(_run())
|
||||
|
|
@ -28,6 +28,7 @@ from models.inference import (
|
|||
)
|
||||
from core.inference.anthropic_compat import (
|
||||
anthropic_messages_to_openai,
|
||||
anthropic_schema_client_tool_kind,
|
||||
anthropic_tools_to_openai,
|
||||
build_anthropic_sse_event,
|
||||
AnthropicStreamEmitter,
|
||||
|
|
@ -626,6 +627,41 @@ class TestAnthropicToolsToOpenAI:
|
|||
]
|
||||
assert anthropic_tools_to_openai(tools) == []
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("type_", "name", "kind"),
|
||||
[
|
||||
("bash_20250124", "bash", "bash"),
|
||||
("text_editor_20250728", "str_replace_based_edit_tool", "text_editor"),
|
||||
("computer_20251124", "computer", "computer"),
|
||||
("memory_20250818", "memory", "memory"),
|
||||
],
|
||||
)
|
||||
def test_schema_client_tools_are_converted_to_openai_functions(self, type_, name, kind):
|
||||
tool = {"type": type_, "name": name}
|
||||
|
||||
[result] = anthropic_tools_to_openai([tool])
|
||||
|
||||
assert anthropic_schema_client_tool_kind(tool) == kind
|
||||
assert result["function"]["name"] == name
|
||||
assert result["function"]["parameters"]["type"] == "object"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("type_", "supports_undo"),
|
||||
[
|
||||
("text_editor_20241022", True),
|
||||
("text_editor_20250124", True),
|
||||
("text_editor_20250429", False),
|
||||
("text_editor_20250728", False),
|
||||
],
|
||||
)
|
||||
def test_text_editor_commands_follow_tool_version(self, type_, supports_undo):
|
||||
[result] = anthropic_tools_to_openai(
|
||||
[{"type": type_, "name": "str_replace_based_edit_tool"}]
|
||||
)
|
||||
|
||||
commands = result["function"]["parameters"]["properties"]["command"]["enum"]
|
||||
assert ("undo_edit" in commands) is supports_undo
|
||||
|
||||
def test_server_tool_selection_merges_enabled_tools_extension(self):
|
||||
all_tools = [
|
||||
{"type": "function", "function": {"name": "web_search"}},
|
||||
|
|
@ -1523,6 +1559,17 @@ def _reset_policy():
|
|||
reset_tool_policy()
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _reset_admission_queues():
|
||||
# The admission queue is process-global; isolate the shared "llama-server" key
|
||||
# so one test's leftover reservation can't stall the next.
|
||||
from core.inference.llama_admission import reset_llama_admission_queues
|
||||
|
||||
reset_llama_admission_queues()
|
||||
yield
|
||||
reset_llama_admission_queues()
|
||||
|
||||
|
||||
class TestAnthropicMessagesToolRouting:
|
||||
class _Request:
|
||||
state = SimpleNamespace()
|
||||
|
|
@ -1724,6 +1771,116 @@ class TestAnthropicMessagesToolRouting:
|
|||
assert exc.value.status_code == 400
|
||||
assert "Mixing Anthropic server tools" in exc.value.detail
|
||||
|
||||
def test_explicit_server_loop_and_client_tools_rejected_with_400(self, monkeypatch):
|
||||
_mock_backend(monkeypatch)
|
||||
payload = _basic_payload(
|
||||
enable_tools = True,
|
||||
tools = [{"name": "Write", "input_schema": {"type": "object"}}],
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
assert exc.value.status_code == 400
|
||||
assert "Mixing Anthropic server tools" in exc.value.detail
|
||||
|
||||
def test_explicit_server_loop_and_schema_client_tools_rejected_with_400(self, monkeypatch):
|
||||
_mock_backend(monkeypatch)
|
||||
payload = _basic_payload(
|
||||
enable_tools = True,
|
||||
tools = [{"type": "bash_20250124", "name": "bash"}],
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
assert exc.value.status_code == 400
|
||||
assert "Mixing Anthropic server tools" in exc.value.detail
|
||||
|
||||
def test_process_tool_policy_does_not_steal_schema_client_tools(self, monkeypatch):
|
||||
import routes.inference as inf_mod
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
backend = _mock_backend(monkeypatch)
|
||||
captured = {}
|
||||
|
||||
async def _passthrough(*args, **kwargs):
|
||||
captured["tools"] = args[2]
|
||||
return JSONResponse(
|
||||
{
|
||||
"id": "msg_test",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"model": "test-model",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(inf_mod, "_anthropic_passthrough_non_streaming", _passthrough)
|
||||
set_tool_policy(True)
|
||||
payload = _basic_payload(tools = [{"type": "bash_20250124", "name": "bash"}])
|
||||
|
||||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
|
||||
assert backend.calls == []
|
||||
assert captured["tools"][0]["function"]["name"] == "bash"
|
||||
|
||||
@pytest.mark.parametrize("permission_mode", [None, "ask"])
|
||||
@pytest.mark.parametrize(
|
||||
("tool_policy", "enable_tools"),
|
||||
[(True, None), (False, True)],
|
||||
)
|
||||
def test_process_tool_policy_does_not_steal_client_tools(
|
||||
self, monkeypatch, permission_mode, tool_policy, enable_tools
|
||||
):
|
||||
"""A server-wide tool default must not replace Claude Code's own tools."""
|
||||
import routes.inference as inf_mod
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
backend = _mock_backend(monkeypatch)
|
||||
captured = {}
|
||||
|
||||
async def _passthrough(*args, **kwargs):
|
||||
captured["tools"] = args[2]
|
||||
return JSONResponse(
|
||||
{
|
||||
"id": "msg_test",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"model": "test-model",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(inf_mod, "_anthropic_passthrough_non_streaming", _passthrough)
|
||||
set_tool_policy(tool_policy)
|
||||
fields = {
|
||||
"tools": [
|
||||
{
|
||||
"name": "Write",
|
||||
"description": "Write a file",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"path": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
if enable_tools is not None:
|
||||
fields["enable_tools"] = enable_tools
|
||||
if permission_mode is not None:
|
||||
fields["permission_mode"] = permission_mode
|
||||
payload = _basic_payload(**fields)
|
||||
|
||||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
|
||||
assert backend.calls == []
|
||||
assert captured["tools"][0]["function"]["name"] == "Write"
|
||||
|
||||
def test_mixed_rejected_when_client_tool_name_collides_with_server_alias(self, monkeypatch):
|
||||
# Regression: a client tool sharing a name with a mapped server tool
|
||||
# (e.g. a custom "web_search") must still trigger the mixed-mode 400;
|
||||
|
|
@ -1769,6 +1926,15 @@ class TestAnthropicMessagesToolRouting:
|
|||
assert exc.value.status_code == 400
|
||||
assert "name" in exc.value.detail
|
||||
|
||||
def test_schema_client_tool_missing_name_rejected_with_400(self, monkeypatch):
|
||||
_mock_backend(monkeypatch)
|
||||
payload = _basic_payload(tools = [{"type": "bash_20250124"}])
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
|
||||
assert exc.value.status_code == 400
|
||||
assert "name" in exc.value.detail
|
||||
|
||||
def test_client_tool_empty_name_rejected_with_400(self, monkeypatch):
|
||||
# Same silent-disable class as missing-name: `name: ""` passes the
|
||||
# isinstance check but is dropped by anthropic_tools_to_openai's
|
||||
|
|
|
|||
266
studio/backend/tests/test_anthropic_passthrough_respawn.py
Normal file
266
studio/backend/tests/test_anthropic_passthrough_respawn.py
Normal file
|
|
@ -0,0 +1,266 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""Restart survival for the Anthropic /v1/messages passthrough.
|
||||
|
||||
A crashed llama-server relaunches on a NEW ephemeral port. Before the retry the
|
||||
passthrough kept posting to the dead port, so a Claude Code session stayed broken
|
||||
until the next explicit load. These cover the respawn-and-retry on both the
|
||||
streaming and non-streaming passthroughs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
_backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
sys.path.insert(0, _backend)
|
||||
|
||||
import routes.inference as inf_mod
|
||||
from routes.inference import (
|
||||
_anthropic_passthrough_non_streaming,
|
||||
_anthropic_passthrough_retry_url,
|
||||
_anthropic_passthrough_stream,
|
||||
)
|
||||
|
||||
_DEAD = "http://127.0.0.1:57953"
|
||||
_FRESH = "http://127.0.0.1:62933"
|
||||
|
||||
|
||||
class _Backend:
|
||||
"""Stub llama backend whose base_url moves to a new port once respawned."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
respawn_ok = True,
|
||||
mtp_handled = False,
|
||||
):
|
||||
self.base_url = _DEAD
|
||||
self.context_length = 4096
|
||||
self.respawn_calls = 0
|
||||
self.mtp_calls = 0
|
||||
self._respawn_ok = respawn_ok
|
||||
self._mtp_handled = mtp_handled
|
||||
|
||||
def count_chat_tokens(self, *_args, **_kwargs):
|
||||
return 2
|
||||
|
||||
def _maybe_recover_from_mtp_crash(self, _exc):
|
||||
self.mtp_calls += 1
|
||||
return self._mtp_handled
|
||||
|
||||
def _respawn_if_dead(self):
|
||||
self.respawn_calls += 1
|
||||
if not self._respawn_ok:
|
||||
return False
|
||||
self.base_url = _FRESH
|
||||
return True
|
||||
|
||||
|
||||
class _Request:
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
|
||||
class _FakeNonStreamingClient:
|
||||
def __init__(self):
|
||||
self.urls = []
|
||||
self.closed = False
|
||||
|
||||
async def aclose(self):
|
||||
self.closed = True
|
||||
|
||||
async def post(self, url, **_kwargs):
|
||||
self.urls.append(url)
|
||||
if url.startswith(_DEAD):
|
||||
raise httpx.ConnectError("connection refused")
|
||||
return httpx.Response(
|
||||
200,
|
||||
json = {
|
||||
"choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 2, "completion_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _install_stream_transport(monkeypatch, calls):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(str(request.url))
|
||||
if str(request.url).startswith(_DEAD):
|
||||
raise httpx.ConnectError("connection refused")
|
||||
content = (
|
||||
f"data: {json.dumps({'choices': [{'delta': {'content': 'hi'}}]})}\n\n"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
content = content.encode(),
|
||||
headers = {"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
real_client = httpx.AsyncClient
|
||||
|
||||
def _client(*_args, **kwargs):
|
||||
return real_client(transport = transport, timeout = kwargs.get("timeout", 600))
|
||||
|
||||
monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client)
|
||||
|
||||
|
||||
async def _run_stream(backend):
|
||||
response = await _anthropic_passthrough_stream(
|
||||
_Request(),
|
||||
threading.Event(),
|
||||
backend,
|
||||
[{"role": "user", "content": "hi"}],
|
||||
[],
|
||||
0.7,
|
||||
0.95,
|
||||
20,
|
||||
16,
|
||||
"msg_1",
|
||||
"test-model",
|
||||
)
|
||||
chunks = []
|
||||
async for chunk in response.body_iterator:
|
||||
chunks.append(chunk.decode() if isinstance(chunk, (bytes, bytearray)) else chunk)
|
||||
return "".join(chunks)
|
||||
|
||||
|
||||
async def _run_non_streaming(backend):
|
||||
return await _anthropic_passthrough_non_streaming(
|
||||
backend,
|
||||
[{"role": "user", "content": "hi"}],
|
||||
[],
|
||||
0.7,
|
||||
0.95,
|
||||
20,
|
||||
16,
|
||||
"msg_1",
|
||||
"test-model",
|
||||
)
|
||||
|
||||
|
||||
# ── Helper ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_retry_url_rebuilds_from_the_respawned_base_url():
|
||||
backend = _Backend()
|
||||
|
||||
url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x")))
|
||||
|
||||
assert url == f"{_FRESH}/v1/chat/completions"
|
||||
assert backend.respawn_calls == 1
|
||||
|
||||
|
||||
def test_retry_url_is_none_when_nothing_respawned():
|
||||
backend = _Backend(respawn_ok = False)
|
||||
|
||||
url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x")))
|
||||
|
||||
assert url is None
|
||||
|
||||
|
||||
def test_retry_url_defers_to_the_mtp_crash_recovery():
|
||||
# An MTP+tensor crash schedules its own reload; retrying would race it.
|
||||
backend = _Backend(mtp_handled = True)
|
||||
|
||||
url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x")))
|
||||
|
||||
assert url is None
|
||||
assert backend.respawn_calls == 0
|
||||
|
||||
|
||||
def test_retry_url_tolerates_a_backend_without_respawn_hooks():
|
||||
backend = SimpleNamespace(base_url = _DEAD)
|
||||
|
||||
url = asyncio.run(_anthropic_passthrough_retry_url(backend, httpx.ConnectError("x")))
|
||||
|
||||
assert url is None
|
||||
|
||||
|
||||
# ── Non-streaming ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_non_streaming_retries_against_the_new_port(monkeypatch):
|
||||
client = _FakeNonStreamingClient()
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
backend = _Backend()
|
||||
|
||||
response = asyncio.run(_run_non_streaming(backend))
|
||||
|
||||
assert response.status_code == 200
|
||||
assert backend.respawn_calls == 1
|
||||
assert client.urls == [f"{_DEAD}/v1/chat/completions", f"{_FRESH}/v1/chat/completions"]
|
||||
|
||||
|
||||
def test_non_streaming_raises_when_the_server_stays_dead(monkeypatch):
|
||||
client = _FakeNonStreamingClient()
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
backend = _Backend(respawn_ok = False)
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
asyncio.run(_run_non_streaming(backend))
|
||||
|
||||
assert client.urls == [f"{_DEAD}/v1/chat/completions"] # no blind retry
|
||||
|
||||
|
||||
def test_non_streaming_does_not_retry_an_mtp_crash(monkeypatch):
|
||||
client = _FakeNonStreamingClient()
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
backend = _Backend(mtp_handled = True)
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
asyncio.run(_run_non_streaming(backend))
|
||||
|
||||
assert backend.respawn_calls == 0
|
||||
|
||||
|
||||
# ── Streaming ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_streaming_retries_against_the_new_port(monkeypatch):
|
||||
calls = []
|
||||
_install_stream_transport(monkeypatch, calls)
|
||||
backend = _Backend()
|
||||
|
||||
blob = asyncio.run(_run_stream(backend))
|
||||
|
||||
assert backend.respawn_calls == 1
|
||||
assert calls == [f"{_DEAD}/v1/chat/completions", f"{_FRESH}/v1/chat/completions"]
|
||||
# The retried stream really produced the turn, not just a clean-looking stop.
|
||||
assert "event: message_start" in blob
|
||||
assert "event: message_stop" in blob
|
||||
assert "hi" in blob
|
||||
|
||||
|
||||
def test_streaming_emits_an_error_event_when_the_server_stays_dead(monkeypatch):
|
||||
calls = []
|
||||
_install_stream_transport(monkeypatch, calls)
|
||||
backend = _Backend(respawn_ok = False)
|
||||
|
||||
blob = asyncio.run(_run_stream(backend))
|
||||
|
||||
assert calls == [f"{_DEAD}/v1/chat/completions"] # no blind retry
|
||||
assert "event: error" in blob
|
||||
|
||||
|
||||
def test_streaming_does_not_retry_an_mtp_crash(monkeypatch):
|
||||
calls = []
|
||||
_install_stream_transport(monkeypatch, calls)
|
||||
backend = _Backend(mtp_handled = True)
|
||||
|
||||
blob = asyncio.run(_run_stream(backend))
|
||||
|
||||
assert backend.respawn_calls == 0
|
||||
assert calls == [f"{_DEAD}/v1/chat/completions"]
|
||||
assert "event: error" in blob
|
||||
|
|
@ -258,3 +258,157 @@ def test_api_monitor_append_reply_exact_cap_then_more_marks_truncated():
|
|||
monitor.append_reply(entry_id, "y")
|
||||
reply = monitor.snapshot()[0]["reply"]
|
||||
assert len(reply) == m._MAX_REPLY_CHARS and reply.endswith("...")
|
||||
|
||||
|
||||
def test_api_monitor_disabled_is_noop():
|
||||
monitor = ApiMonitor(max_entries = 3, enabled = False)
|
||||
|
||||
request_id = monitor.start(
|
||||
endpoint = "/v1/chat/completions",
|
||||
method = "POST",
|
||||
model = "local-model",
|
||||
prompt = "user: hello",
|
||||
context_length = 100,
|
||||
)
|
||||
load_id = monitor.record_lifecycle(
|
||||
event = "load",
|
||||
model = "local-model",
|
||||
running = True,
|
||||
)
|
||||
unload_id = monitor.record_lifecycle(
|
||||
event = "unload",
|
||||
model = "local-model",
|
||||
)
|
||||
assert request_id == load_id == unload_id == ""
|
||||
|
||||
# Every mutator must be a safe no-op on the falsy id.
|
||||
monitor.append_reply(request_id, "hi")
|
||||
monitor.set_reply(request_id, "hi")
|
||||
monitor.set_usage(request_id, prompt_tokens = 4, completion_tokens = 6)
|
||||
monitor.relabel(load_id, "renamed-model")
|
||||
monitor.set_progress(load_id, 50)
|
||||
monitor.finish(load_id)
|
||||
monitor.fail_open(load_id, "boom")
|
||||
monitor.fail(request_id, "boom")
|
||||
monitor.discard(unload_id)
|
||||
|
||||
assert monitor.snapshot() == []
|
||||
assert monitor.active_count() == 0
|
||||
assert monitor.get(request_id) is None
|
||||
|
||||
|
||||
def test_api_monitor_disable_env_var_truthy(monkeypatch):
|
||||
import core.inference.api_monitor as m
|
||||
for value in ("1", "true", "yes", "on", "TRUE", "On", " yes "):
|
||||
monkeypatch.setenv(m._DISABLE_ENV, value)
|
||||
assert m._api_monitor_disabled() is True, value
|
||||
|
||||
|
||||
def test_api_monitor_disable_env_var_falsy(monkeypatch):
|
||||
import core.inference.api_monitor as m
|
||||
for value in ("", "0", "false", "no", "off", "disabled"):
|
||||
monkeypatch.setenv(m._DISABLE_ENV, value)
|
||||
assert m._api_monitor_disabled() is False, value
|
||||
|
||||
|
||||
def test_api_monitor_disable_env_var_unset(monkeypatch):
|
||||
import core.inference.api_monitor as m
|
||||
monkeypatch.delenv(m._DISABLE_ENV, raising = False)
|
||||
assert m._api_monitor_disabled() is False
|
||||
|
||||
|
||||
# ── model lifecycle rows (load / unload) ────────────────────────────
|
||||
|
||||
|
||||
def test_lifecycle_load_row_opens_running_then_closes():
|
||||
monitor = ApiMonitor(max_entries = 5)
|
||||
event_id = monitor.record_lifecycle(event = "load", model = "org/A-GGUF", running = True)
|
||||
row = monitor.snapshot()[0]
|
||||
assert row["kind"] == "lifecycle" and row["event"] == "load"
|
||||
assert row["status"] == "running" and row["duration_ms"] is None
|
||||
# A load in progress is not an in-flight API request.
|
||||
assert monitor.active_count() == 0
|
||||
|
||||
monitor.relabel(event_id, "org/A-GGUF:Q4_K_M")
|
||||
monitor.finish(event_id)
|
||||
row = monitor.snapshot()[0]
|
||||
assert row["status"] == "completed"
|
||||
assert row["model"] == "org/A-GGUF:Q4_K_M"
|
||||
assert row["duration_ms"] is not None
|
||||
|
||||
|
||||
def test_lifecycle_unload_row_is_terminal_on_arrival():
|
||||
monitor = ApiMonitor(max_entries = 5)
|
||||
monitor.record_lifecycle(event = "unload", model = "org/A-GGUF", reason = "idle")
|
||||
row = monitor.snapshot()[0]
|
||||
assert row["status"] == "completed"
|
||||
assert (row["event"], row["reason"]) == ("unload", "idle")
|
||||
assert monitor.active_count() == 0
|
||||
|
||||
|
||||
def test_lifecycle_rows_are_visible_to_every_subject():
|
||||
# A load is server-wide, so it must not vanish for other API keys like a request does.
|
||||
monitor = ApiMonitor(max_entries = 5)
|
||||
monitor.start(
|
||||
endpoint = "/v1/chat/completions",
|
||||
method = "POST",
|
||||
model = "m",
|
||||
prompt = "hi",
|
||||
subject = "alice",
|
||||
)
|
||||
event_id = monitor.record_lifecycle(event = "unload", model = "org/A-GGUF")
|
||||
|
||||
bob = monitor.snapshot(subject = "bob")
|
||||
assert [r["kind"] for r in bob] == ["lifecycle"]
|
||||
assert monitor.get(event_id, subject = "bob") is not None
|
||||
assert len(monitor.snapshot(subject = "alice")) == 2
|
||||
|
||||
|
||||
def test_request_rows_stay_private_to_their_subject():
|
||||
monitor = ApiMonitor(max_entries = 5)
|
||||
rid = monitor.start(
|
||||
endpoint = "/v1/chat/completions",
|
||||
method = "POST",
|
||||
model = "m",
|
||||
prompt = "hi",
|
||||
subject = "alice",
|
||||
)
|
||||
assert monitor.snapshot(subject = "bob") == []
|
||||
assert monitor.get(rid, subject = "bob") is None
|
||||
|
||||
|
||||
def test_discard_drops_a_row_that_never_happened():
|
||||
# A load that found the model already resident must leave no trace.
|
||||
monitor = ApiMonitor(max_entries = 5)
|
||||
event_id = monitor.record_lifecycle(event = "load", model = "org/A-GGUF", running = True)
|
||||
monitor.discard(event_id)
|
||||
assert monitor.snapshot() == []
|
||||
monitor.discard(event_id) # idempotent
|
||||
|
||||
|
||||
def test_fail_open_never_touches_a_finished_row():
|
||||
# Called from a finally, so it must not stamp an error onto a load that succeeded.
|
||||
monitor = ApiMonitor(max_entries = 5)
|
||||
event_id = monitor.record_lifecycle(event = "load", model = "org/A-GGUF", running = True)
|
||||
monitor.finish(event_id)
|
||||
monitor.fail_open(event_id, "Load did not complete")
|
||||
row = monitor.snapshot()[0]
|
||||
assert row["status"] == "completed" and row["error"] is None
|
||||
|
||||
still_open = monitor.record_lifecycle(event = "load", model = "org/B-GGUF", running = True)
|
||||
monitor.fail_open(still_open, "Load did not complete")
|
||||
assert monitor.snapshot()[0]["status"] == "error"
|
||||
|
||||
|
||||
def test_lifecycle_rows_share_the_retention_budget():
|
||||
monitor = ApiMonitor(max_entries = 2)
|
||||
for i in range(4):
|
||||
monitor.record_lifecycle(event = "unload", model = f"org/M{i}")
|
||||
models = [r["model"] for r in monitor.snapshot()]
|
||||
assert models == ["org/M3", "org/M2"]
|
||||
|
||||
|
||||
def test_request_rows_report_kind_request():
|
||||
monitor = ApiMonitor(max_entries = 2)
|
||||
monitor.start(endpoint = "/v1/chat/completions", method = "POST", model = "m", prompt = "hi")
|
||||
assert monitor.snapshot()[0]["kind"] == "request"
|
||||
|
|
|
|||
|
|
@ -6,10 +6,14 @@ from the OpenAI JSON-string form to a dict before rendering. Strict tool
|
|||
templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and
|
||||
raise "Can only get item pairs from a mapping." on the string form when a prior
|
||||
tool call is re-rendered on the next turn (MLX + transformers paths).
|
||||
|
||||
It must likewise split parallel tool calls for templates that render only one
|
||||
call per message (Llama 3.x).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -21,6 +25,7 @@ if str(_BACKEND) not in sys.path:
|
|||
|
||||
from core.inference.chat_template_helpers import ( # noqa: E402
|
||||
_normalize_tool_call_arguments,
|
||||
_split_parallel_tool_calls,
|
||||
apply_chat_template_for_generation,
|
||||
)
|
||||
|
||||
|
|
@ -155,3 +160,152 @@ def test_unrelated_template_error_still_propagates_with_dict_args():
|
|||
|
||||
with pytest.raises(ValueError, match = "broken"):
|
||||
apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"}))
|
||||
|
||||
|
||||
def _parallel_conv(
|
||||
*,
|
||||
ids = ("c1", "c2"),
|
||||
results_have_ids = True,
|
||||
content = "sure",
|
||||
):
|
||||
a, b = ids
|
||||
return [
|
||||
{"role": "user", "content": "search then render"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
"tool_calls": [
|
||||
{
|
||||
"type": "function",
|
||||
"id": a,
|
||||
"function": {"name": "web_search", "arguments": {"query": "x"}},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"id": b,
|
||||
"function": {"name": "render_html", "arguments": {"html": "<canvas>"}},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"name": "web_search",
|
||||
**({"tool_call_id": a} if results_have_ids else {}),
|
||||
"content": "no text",
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"name": "render_html",
|
||||
**({"tool_call_id": b} if results_have_ids else {}),
|
||||
"content": "ok",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class _SingleToolCallTokenizer:
|
||||
"""Mimics the Llama 3.x template: rejects >1 call per message."""
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
*,
|
||||
tokenize = False,
|
||||
add_generation_prompt = True,
|
||||
**kw,
|
||||
):
|
||||
for msg in messages:
|
||||
if len(msg.get("tool_calls") or ()) > 1:
|
||||
raise ValueError("This model only supports single tool-calls at once!")
|
||||
return "RENDERED"
|
||||
|
||||
|
||||
def test_parallel_calls_split_into_sequential_single_call_turns():
|
||||
out = _split_parallel_tool_calls(_parallel_conv())
|
||||
assert [(m["role"], m.get("name")) for m in out] == [
|
||||
("user", None),
|
||||
("assistant", None),
|
||||
("tool", "web_search"),
|
||||
("assistant", None),
|
||||
("tool", "render_html"),
|
||||
]
|
||||
assert [len(m["tool_calls"]) for m in out if m.get("tool_calls")] == [1, 1]
|
||||
assert out[1]["tool_calls"][0]["function"]["name"] == "web_search"
|
||||
assert out[3]["tool_calls"][0]["function"]["name"] == "render_html"
|
||||
|
||||
|
||||
def test_split_pairs_results_by_tool_call_id_not_position():
|
||||
conv = _parallel_conv()
|
||||
conv[2], conv[3] = conv[3], conv[2] # results arrive out of order
|
||||
out = _split_parallel_tool_calls(conv)
|
||||
assert out[1]["tool_calls"][0]["id"] == "c1" and out[2]["tool_call_id"] == "c1"
|
||||
assert out[3]["tool_calls"][0]["id"] == "c2" and out[4]["tool_call_id"] == "c2"
|
||||
|
||||
|
||||
def test_split_falls_back_to_order_when_results_have_no_ids():
|
||||
out = _split_parallel_tool_calls(_parallel_conv(results_have_ids = False))
|
||||
assert [m["role"] for m in out] == ["user", "assistant", "tool", "assistant", "tool"]
|
||||
assert out[2]["name"] == "web_search" and out[4]["name"] == "render_html"
|
||||
|
||||
|
||||
def test_split_keeps_content_on_first_piece_only():
|
||||
out = _split_parallel_tool_calls(_parallel_conv(content = "sure"))
|
||||
assert out[1]["content"] == "sure"
|
||||
assert out[3]["content"] == ""
|
||||
|
||||
|
||||
def test_split_keeps_unmatched_results_after_the_split():
|
||||
conv = _parallel_conv()
|
||||
del conv[3] # second call never returned a result
|
||||
out = _split_parallel_tool_calls(conv)
|
||||
assert [m["role"] for m in out] == ["user", "assistant", "tool", "assistant"]
|
||||
|
||||
|
||||
def test_split_leaves_later_turns_intact():
|
||||
conv = _parallel_conv() + [
|
||||
{"role": "assistant", "content": "done"},
|
||||
{"role": "user", "content": "thanks"},
|
||||
]
|
||||
out = _split_parallel_tool_calls(conv)
|
||||
assert [m["role"] for m in out[-2:]] == ["assistant", "user"]
|
||||
assert out[-2]["content"] == "done"
|
||||
|
||||
|
||||
def test_single_call_and_plain_conversations_pass_through_unchanged():
|
||||
conv = _conv({"query": "x"})
|
||||
assert _split_parallel_tool_calls(conv) is conv
|
||||
plain = [{"role": "user", "content": "hi"}]
|
||||
assert _split_parallel_tool_calls(plain) is plain
|
||||
|
||||
|
||||
def test_render_succeeds_on_single_call_template_with_parallel_calls():
|
||||
# Regression: two calls in one turn used to break every later render.
|
||||
result = apply_chat_template_for_generation(_SingleToolCallTokenizer(), _parallel_conv())
|
||||
assert result == "RENDERED"
|
||||
|
||||
|
||||
def test_string_arguments_and_parallel_calls_are_repaired_together():
|
||||
conv = _parallel_conv()
|
||||
for call in conv[1]["tool_calls"]:
|
||||
call["function"]["arguments"] = json.dumps(call["function"]["arguments"])
|
||||
|
||||
class _StrictAndSingleCall(_SingleToolCallTokenizer):
|
||||
def apply_chat_template(self, messages, **kw):
|
||||
for msg in messages:
|
||||
for call in msg.get("tool_calls", []) or []:
|
||||
if isinstance(call.get("function", {}).get("arguments"), str):
|
||||
raise TypeError("Can only get item pairs from a mapping.")
|
||||
return super().apply_chat_template(messages, **kw)
|
||||
|
||||
assert apply_chat_template_for_generation(_StrictAndSingleCall(), conv) == "RENDERED"
|
||||
|
||||
|
||||
def test_lenient_template_never_sees_a_split_conversation():
|
||||
seen = {}
|
||||
|
||||
class _Lenient:
|
||||
def apply_chat_template(self, messages, **kw):
|
||||
seen["n"] = len(messages)
|
||||
return "RENDERED"
|
||||
|
||||
apply_chat_template_for_generation(_Lenient(), _parallel_conv())
|
||||
assert seen["n"] == 4 # unsplit
|
||||
|
|
|
|||
|
|
@ -304,12 +304,16 @@ def test_load_request_accepts_valid_tensor_split(good):
|
|||
def test_route_normalizes_explicit_extras_before_reload_dedupe():
|
||||
route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
|
||||
load_impl = route_src[route_src.index("async def _load_model_impl") :]
|
||||
preserve = load_impl.index("_gpu_layers_override = parse_gpu_layers_override")
|
||||
translate = load_impl.index(
|
||||
'request = request.model_copy(update = {"gpu_layers": _gpu_layers_override})'
|
||||
)
|
||||
strip = load_impl.index("_stripped_explicit = strip_shadowing_flags")
|
||||
normalize = load_impl.index(
|
||||
'request = request.model_copy(update = {"llama_extra_args": extra_llama_args})'
|
||||
)
|
||||
dedupe = load_impl.index("and _request_matches_loaded_settings(")
|
||||
assert strip < normalize < dedupe
|
||||
assert preserve < translate < strip < normalize < dedupe
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
|
||||
|
|
@ -1048,7 +1052,7 @@ def _rocm_torch_stub(monkeypatch):
|
|||
def test_subset_pin_masks_via_rocr_on_rocm(monkeypatch):
|
||||
# A GPU-subset pin must exclude the rest at the ROCr/HSA layer: HIP masking
|
||||
# still enumerates every agent first, which segfaults the build on an
|
||||
# unsupported deselected GPU (e.g. a gfx1103 iGPU under a gfx110X prebuilt).
|
||||
# unsupported deselected GPU (e.g. a gfx1036 iGPU under a gfx103X prebuilt).
|
||||
# ROCR drops it at the driver layer; only one mask is set (HIP cleared).
|
||||
_rocm_torch_stub(monkeypatch)
|
||||
env = {"HIP_VISIBLE_DEVICES": "9"} # stale/inherited HIP mask must not survive
|
||||
|
|
|
|||
|
|
@ -427,14 +427,16 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
|
|||
|
||||
self.assertTrue(result["available"])
|
||||
self.assertEqual(result["backend"], "vulkan")
|
||||
self.assertEqual(result["index_kind"], "relative")
|
||||
# ggml Vulkan ordinals are the space `--device Vulkan<i>` pins, so they
|
||||
# are selectable, unlike a torch-xpu relative ordinal.
|
||||
self.assertEqual(result["index_kind"], "vulkan")
|
||||
self.assertEqual(result["parent_visible_gpu_ids"], [])
|
||||
self.assertEqual(
|
||||
result["devices"],
|
||||
[
|
||||
{
|
||||
"index": 0,
|
||||
"index_kind": "relative",
|
||||
"index_kind": "vulkan",
|
||||
"visible_ordinal": 0,
|
||||
"name": "Vulkan0",
|
||||
"memory_total_gb": 8.0,
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ def _dispatcher():
|
|||
o._dispatcher_stop = threading.Event()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
return o
|
||||
|
||||
|
||||
|
|
@ -118,3 +119,68 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
|
|||
kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False
|
||||
for kw in call.keywords
|
||||
), f"httpx.AsyncClient at line {call.lineno} must set trust_env=False"
|
||||
|
||||
|
||||
def _direct_reader_host():
|
||||
"""Orchestrator with only what _direct_reader and the ownership helpers touch."""
|
||||
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._active_cancel_lock = threading.Lock()
|
||||
o._active_cancel_events = []
|
||||
o._executing_cancel_events = []
|
||||
o._dispatcher_thread = None
|
||||
return o
|
||||
|
||||
|
||||
def test_rerouting_a_foreign_response_moves_worker_ownership():
|
||||
# A _gen_lock reader already blocked on resp_queue can beat the compare dispatcher to
|
||||
# that request's first response. The compare consumer passes mark_started=False, so if
|
||||
# this path does not promote it nothing does: the direct request stays recorded as the
|
||||
# executor, so the compare chat's Stop is ignored and a late reset from the direct one
|
||||
# cancels the compare generation instead.
|
||||
o = _direct_reader_host()
|
||||
mine, theirs = threading.Event(), threading.Event()
|
||||
o._request_cancel_events = {"mine": mine, "theirs": theirs}
|
||||
o._claim_worker(mine)
|
||||
o._mark_worker_started(mine)
|
||||
o._claim_worker(theirs)
|
||||
compare_mailbox = queue.Queue()
|
||||
o._mailboxes["theirs"] = compare_mailbox
|
||||
|
||||
read_one, _drain, release = _direct_reader_calls(o, "mine")
|
||||
o._scripted = [{"request_id": "theirs", "type": "token", "text": "hi"}]
|
||||
|
||||
assert read_one(timeout = 0.1) is None, "a foreign response is routed, not returned"
|
||||
assert compare_mailbox.get_nowait()["text"] == "hi"
|
||||
assert o._owns_worker(theirs), "the compare request is the one the worker answered"
|
||||
assert not o._owns_worker(mine), "so a late reset from the direct request must not fire"
|
||||
release()
|
||||
|
||||
|
||||
def test_rerouting_a_foreign_gen_done_retires_that_request():
|
||||
# The other half of the dispatcher's move: once its last response is routed, the
|
||||
# request no longer owns the worker, or a Stop for it would end whatever starts next.
|
||||
o = _direct_reader_host()
|
||||
mine, theirs = threading.Event(), threading.Event()
|
||||
o._request_cancel_events = {"mine": mine, "theirs": theirs}
|
||||
o._claim_worker(theirs)
|
||||
o._mark_worker_started(theirs)
|
||||
o._claim_worker(mine)
|
||||
o._mailboxes["theirs"] = queue.Queue()
|
||||
|
||||
read_one, _drain, release = _direct_reader_calls(o, "mine")
|
||||
o._scripted = [{"request_id": "theirs", "type": "gen_done"}]
|
||||
|
||||
assert read_one(timeout = 0.1) is None
|
||||
assert not o._owns_worker(theirs), "retired once its last response was routed"
|
||||
assert o._owns_worker(mine), "the next claim takes over"
|
||||
release()
|
||||
|
||||
|
||||
def _direct_reader_calls(o, request_id):
|
||||
"""_direct_reader wired to a scripted _read_resp (o._scripted, popped in order)."""
|
||||
o._read_resp = lambda timeout = 1.0: o._scripted.pop(0) if o._scripted else None
|
||||
return o._direct_reader(request_id)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,9 @@ by default; --published-repo overrides).
|
|||
|
||||
These back the in-app update for source-build (markerless) installs: the backend
|
||||
asks the installer whether an official prebuilt exists for this host without
|
||||
downloading. Network and host detection are stubbed; no GPU or internet needed.
|
||||
downloading. Network and host detection are stubbed; no GPU or internet needed. The one
|
||||
exception is the windows-rocm floor guard, which reads the fork's published manifest
|
||||
because nothing in-tree mirrors it, and skips when that release is unreachable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -32,6 +34,18 @@ FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp
|
|||
UPSTREAM = ilp.UPSTREAM_REPO # ggml-org/llama.cpp
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _no_ambient_hip_device_mask(monkeypatch):
|
||||
"""These tests describe hosts through HostInfo, not through the environment.
|
||||
|
||||
A mask inherited from the shell (ML boxes commonly export CUDA_VISIBLE_DEVICES) means
|
||||
the arch probe saw only part of the GPUs, which the Windows auto-Vulkan guard treats as
|
||||
an unknown physical inventory. Clear all three so a host is described by its fields
|
||||
alone; the tests that are about the mask set it explicitly."""
|
||||
for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
|
||||
monkeypatch.delenv(_env, raising = False)
|
||||
|
||||
|
||||
def _host(**kw):
|
||||
base = dict(
|
||||
system = "Linux",
|
||||
|
|
@ -407,7 +421,9 @@ 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)
|
||||
routed, repo, tag, _persist = 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
|
||||
|
|
@ -416,7 +432,9 @@ def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
|
|||
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)
|
||||
_routed, repo, tag, _persist = ilp._route_to_vulkan_prebuilt(
|
||||
host, UPSTREAM, "b9596", force_cpu = False
|
||||
)
|
||||
assert repo == UPSTREAM
|
||||
assert tag == "b9596"
|
||||
|
||||
|
|
@ -424,7 +442,9 @@ def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
|
|||
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)
|
||||
routed, repo, tag, _persist = 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
|
||||
|
|
@ -536,20 +556,20 @@ def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
|
|||
has_physical_nvidia = True,
|
||||
has_usable_nvidia = False,
|
||||
)
|
||||
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
|
||||
_routed, repo, _tag, _persist = 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)
|
||||
_routed, repo, _tag, _persist = 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)
|
||||
routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
|
||||
assert repo == FORK
|
||||
assert routed is host
|
||||
|
||||
|
|
@ -797,3 +817,800 @@ def test_detect_host_cim_rescues_exploding_registry(monkeypatch):
|
|||
)
|
||||
assert host.has_intel_gpu is True
|
||||
assert "powershell" in captured
|
||||
|
||||
|
||||
def _windows_amd_host(**overrides):
|
||||
defaults = dict(
|
||||
system = "Windows",
|
||||
machine = "amd64",
|
||||
is_windows = True,
|
||||
is_linux = False,
|
||||
is_macos = False,
|
||||
is_x86_64 = True,
|
||||
is_arm64 = False,
|
||||
nvidia_smi = None,
|
||||
driver_cuda_version = None,
|
||||
compute_caps = [],
|
||||
visible_cuda_devices = None,
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
has_rocm = True,
|
||||
has_intel_gpu = False,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return ilp.HostInfo(**defaults)
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_auto_fallback_for_legacy_amd_gfx():
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
assert routed.has_intel_gpu is True
|
||||
assert routed.has_rocm is False
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_keeps_hip_when_one_gpu_is_supported():
|
||||
host = _windows_amd_host(
|
||||
rocm_gfx_target = "gfx1201",
|
||||
rocm_gfx_targets = ["gfx1201", "gfx803"],
|
||||
)
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_auto_fallback_skips_hip_masked_hosts():
|
||||
# A HIP mask can hide a HIP-capable dGPU, but the Vulkan runtime honours none of them,
|
||||
# so auto-routing would let the installed backend grab the gfx1201 the user masked
|
||||
# off.
|
||||
host = _windows_amd_host(
|
||||
rocm_gfx_target = "gfx803",
|
||||
rocm_gfx_targets = ["gfx1201", "gfx803"],
|
||||
)
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
assert routed is host
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_auto_fallback_when_no_amd_gpu_reaches_floor():
|
||||
# Every physical AMD device is below the floor, so no card can be exposed to HIP and
|
||||
# the #7357 auto-Vulkan fallback still fires.
|
||||
host = _windows_amd_host(
|
||||
rocm_gfx_target = "gfx900",
|
||||
rocm_gfx_targets = ["gfx803", "gfx900"],
|
||||
)
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
assert routed.has_rocm is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mask_env", ["HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"]
|
||||
)
|
||||
def test_auto_vulkan_declines_when_a_hip_device_mask_filtered_the_probe(mask_env, monkeypatch):
|
||||
# hipinfo is a HIP application, so under a mask rocm_gfx_targets is the VISIBLE set and
|
||||
# a HIP-capable card can be hidden entirely. "No AMD GPU here reaches the floor" is then
|
||||
# unprovable, and Vulkan honours none of these masks, so the auto fallback must decline
|
||||
# rather than hand it the reserved card.
|
||||
monkeypatch.setenv(mask_env, "1")
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mask_value", ["", " ", "-1"])
|
||||
def test_auto_vulkan_declines_when_the_mask_hides_every_amd_gpu(mask_value, monkeypatch):
|
||||
# An all-hiding mask is the strongest form of the same signal, not an exemption:
|
||||
# detect_host() resolves no arch under it, but a forwarded --rocm-gfx still reconstructs
|
||||
# one (setup infers it from the display-adapter name, which no HIP mask touches), so
|
||||
# auto-routing would hand Vulkan every AMD GPU the user hid from HIP.
|
||||
monkeypatch.setenv("HIP_VISIBLE_DEVICES", mask_value)
|
||||
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
|
||||
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
|
||||
assert ilp._active_rocm_gfx_target(host) == "gfx803"
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_hip_device_mask_check_is_presence_not_value(monkeypatch):
|
||||
# Presence is the whole test: any value means the HIP view is not the physical one, and
|
||||
# no value can be read as "the probe saw everything".
|
||||
assert ilp._hip_visible_device_mask_set() is False
|
||||
for value in ("", " ", "-1", "0", "1", "0,1"):
|
||||
monkeypatch.setenv("HIP_VISIBLE_DEVICES", value)
|
||||
assert ilp._hip_visible_device_mask_set() is True, value
|
||||
monkeypatch.delenv("HIP_VISIBLE_DEVICES")
|
||||
monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "0")
|
||||
assert ilp._hip_visible_device_mask_set() is True
|
||||
monkeypatch.delenv("ROCR_VISIBLE_DEVICES")
|
||||
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
|
||||
assert ilp._hip_visible_device_mask_set() is True
|
||||
|
||||
|
||||
def test_masked_probe_suppression_does_not_touch_non_amd_auto_paths(monkeypatch):
|
||||
# The mask says nothing about an Intel iGPU, whose Vulkan auto path is unrelated.
|
||||
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
|
||||
host = _host(
|
||||
system = "Windows",
|
||||
is_windows = True,
|
||||
has_intel_gpu = True,
|
||||
has_rocm = False,
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
)
|
||||
_routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(
|
||||
host, FORK, "pin", force_cpu = False
|
||||
)
|
||||
assert repo == UPSTREAM
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_hip_masked_host_still_honours_explicit_optin(monkeypatch):
|
||||
# The mask guard only suppresses the AUTOMATIC fallback; an explicit opt-in is the user
|
||||
# taking responsibility for the Vulkan device mask themselves.
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
host = _windows_amd_host(
|
||||
rocm_gfx_target = "gfx803",
|
||||
rocm_gfx_targets = ["gfx1201", "gfx803"],
|
||||
)
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(
|
||||
host, FORK, "pin", force_cpu = False, llama_backend = "vulkan"
|
||||
)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
|
||||
|
||||
def test_auto_vulkan_is_repository_specific_for_fork_only_gfx():
|
||||
# gfx1034 is served only by the fork's gfx103X bundle: ggml-org's windows-hip radeon
|
||||
# build does not target it and direct_upstream_release_plan() offers win-hip then CPU
|
||||
# with no Vulkan branch, so the predicate must answer per repo.
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, UPSTREAM) is True
|
||||
# An arch upstream really does build stays on HIP for both repos.
|
||||
supported = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(supported, FORK) is False
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(supported, UPSTREAM) is False
|
||||
# A family label is a bundle name, not an arch: upstream builds every member but
|
||||
# gfx1034 / gfx1103, and the label cannot say which card this is, so it stays on HIP
|
||||
# rather than moving the covered members onto Vulkan.
|
||||
family = _windows_amd_host(rocm_gfx_target = "gfx110X", rocm_gfx_targets = ["gfx110X"])
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(family, UPSTREAM) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"repo", ["acme/llama.cpp-mirror", "GGML-ORG/llama.cpp", "unslothAI/llama.cpp"]
|
||||
)
|
||||
def test_fork_only_gfx_coverage_is_not_granted_to_other_repos(repo):
|
||||
# Only the fork is planned from a manifest: resolve_simple_install_release_plans()
|
||||
# compares == DEFAULT_PUBLISHED_REPO and sends everything else, mirrors and differently
|
||||
# cased spellings alike, to direct_upstream_release_plan(). Granting a fork-only arch
|
||||
# coverage there lands it on win-hip-radeon or CPU instead of Vulkan, so the predicate
|
||||
# must gate on the fork rather than exempt one name.
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, repo) is True
|
||||
supported = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(supported, repo) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("repo", [None, ""])
|
||||
def test_empty_published_repo_gets_fork_coverage(repo):
|
||||
# Negative control: the resolver defaults an empty repo to the fork, so the predicate
|
||||
# must too, or the default install path loses its fork-only archs.
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, repo) is False
|
||||
|
||||
|
||||
def test_upstream_windows_hip_targets_are_a_subset_of_the_combined_floor():
|
||||
# The floor must stay a superset, else auto-Vulkan steals a host upstream builds for.
|
||||
assert ilp.UPSTREAM_WINDOWS_HIP_GFX_TARGETS <= ilp.WINDOWS_HIP_PREBUILT_GFX_TARGETS
|
||||
# The fork-only extras are exactly the archs that must route to Vulkan upstream.
|
||||
assert ilp.WINDOWS_HIP_PREBUILT_GFX_TARGETS - ilp.UPSTREAM_WINDOWS_HIP_GFX_TARGETS == {
|
||||
"gfx908",
|
||||
"gfx90a",
|
||||
"gfx1034",
|
||||
"gfx1103",
|
||||
}
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_unknown_gfx_does_not_auto_fallback():
|
||||
host = _windows_amd_host(
|
||||
has_rocm = True,
|
||||
rocm_gfx_target = None,
|
||||
rocm_gfx_targets = [],
|
||||
)
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_family_gfx_token_keeps_rocm():
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx110X", rocm_gfx_targets = ["gfx110X"])
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_gfx1103_keeps_rocm():
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1103", rocm_gfx_targets = ["gfx1103"])
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_gfx1034_keeps_rocm():
|
||||
# gfx1034 (RX 6500/6400-class) is covered by the fork's gfx103X bundle.
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_explicit_opt_in_on_mixed_amd(monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
|
||||
host = _windows_amd_host(
|
||||
rocm_gfx_target = "gfx1201",
|
||||
rocm_gfx_targets = ["gfx1201", "gfx803"],
|
||||
)
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
assert routed.has_rocm is False
|
||||
|
||||
|
||||
def test_direct_upstream_windows_amd_legacy_gfx_routes_to_vulkan():
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
rel = _upstream_release(
|
||||
"b9925",
|
||||
[
|
||||
"llama-b9925-bin-win-hip-radeon-x64.zip",
|
||||
"llama-b9925-bin-win-vulkan-x64.zip",
|
||||
"llama-b9925-bin-win-cpu-x64.zip",
|
||||
],
|
||||
)
|
||||
plan = ilp.direct_upstream_release_plan(rel, routed, repo, "latest")
|
||||
assert persist == "vulkan"
|
||||
assert plan.attempts[0].install_kind == "windows-vulkan"
|
||||
|
||||
|
||||
def test_llama_backend_env_requests_vulkan(monkeypatch):
|
||||
assert ilp.llama_backend_from_env() is None
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
|
||||
assert ilp.llama_backend_from_env() == "vulkan"
|
||||
assert ilp.force_vulkan_requested() is True
|
||||
|
||||
|
||||
def test_llama_cpp_backend_env_does_not_trigger_vulkan(monkeypatch):
|
||||
# UNSLOTH_LLAMA_CPP_BACKEND is a separate setup variable (auto/cpu) whose other values
|
||||
# setup warns about and ignores, so reading it here would opt in behind that warning.
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "vulkan")
|
||||
assert ilp.llama_backend_from_env() is None
|
||||
assert ilp.force_vulkan_requested() is False
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_hidden_physical_nvidia_amd_not_rerouted():
|
||||
# Vulkan ignores CUDA_VISIBLE_DEVICES, so a CUDA-masked NVIDIA card next to a legacy
|
||||
# AMD gfx must not auto-route: Vulkan could grab the reserved NVIDIA GPU.
|
||||
host = _windows_amd_host(
|
||||
rocm_gfx_target = "gfx803",
|
||||
rocm_gfx_targets = ["gfx803"],
|
||||
has_physical_nvidia = True,
|
||||
has_usable_nvidia = False,
|
||||
)
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_route_to_vulkan_prebuilt_explicit_opt_in_overrides_hidden_nvidia(monkeypatch):
|
||||
# The physical-NVIDIA guard only gates the AMD auto path; an explicit opt-in wins.
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
|
||||
host = _windows_amd_host(
|
||||
rocm_gfx_target = "gfx803",
|
||||
rocm_gfx_targets = ["gfx803"],
|
||||
has_physical_nvidia = True,
|
||||
has_usable_nvidia = False,
|
||||
)
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
|
||||
|
||||
# The gfx archs the fork's llama-prebuilt-manifest.json maps to a windows-rocm bundle.
|
||||
# Static because parametrisation happens at import time and the routing tests below must
|
||||
# stay offline; the guard further down re-derives it from the published manifest and fails
|
||||
# on drift, so this is a checked mirror, not a second source of truth.
|
||||
_FORK_WINDOWS_ROCM_GFX = (
|
||||
"gfx908",
|
||||
"gfx90a",
|
||||
"gfx1030",
|
||||
"gfx1031",
|
||||
"gfx1032",
|
||||
"gfx1034",
|
||||
"gfx1100",
|
||||
"gfx1101",
|
||||
"gfx1102",
|
||||
"gfx1103",
|
||||
"gfx1150",
|
||||
"gfx1151",
|
||||
"gfx1200",
|
||||
"gfx1201",
|
||||
)
|
||||
|
||||
|
||||
def _published_fork_windows_rocm_artifacts():
|
||||
"""The fork's windows-rocm artifact records, read the way an install reads them.
|
||||
|
||||
_download_host_resolved_release is the path a default fork install takes first: it
|
||||
resolves the latest release off the download host and hands llama-prebuilt-manifest.json
|
||||
to parse_published_release_bundle, so these are the very records
|
||||
published_rocm_choice_for_host later matches a host gfx against. No api.github.com call,
|
||||
hence no shared rate-limit bucket to exhaust.
|
||||
|
||||
The manifest ships only as a release asset and nothing in-tree mirrors it, so this is
|
||||
the one honest source. Only OSError and the release-side PrebuiltFallback become a skip,
|
||||
so an offline run stays quiet while a manifest that fetches but no longer parses still
|
||||
fails loudly."""
|
||||
try:
|
||||
resolved = ilp._download_host_resolved_release(FORK)
|
||||
except OSError as exc:
|
||||
pytest.skip(f"{FORK} release manifest unreachable: {exc}")
|
||||
except ilp.PrebuiltFallback as exc:
|
||||
pytest.skip(f"{FORK} latest release was rejected before its manifest parsed: {exc}")
|
||||
if resolved is None:
|
||||
pytest.skip(f"{FORK} published no resolvable latest release")
|
||||
tag = resolved.bundle.release_tag
|
||||
artifacts = [
|
||||
artifact
|
||||
for artifact in resolved.bundle.artifacts
|
||||
if artifact.install_kind == "windows-rocm"
|
||||
]
|
||||
assert artifacts, f"{FORK}@{tag} manifest listed no windows-rocm artifacts"
|
||||
return tag, artifacts
|
||||
|
||||
|
||||
def test_windows_hip_gfx_floor_covers_every_fork_windows_rocm_bundle():
|
||||
# Derived from the published manifest, not a second literal: a gfx the fork builds but
|
||||
# the floor omits bypasses the fork manifest, downgrading a hash-approved windows-rocm
|
||||
# bundle to an unhashed upstream Vulkan build. A newly published arch must redden here.
|
||||
tag, artifacts = _published_fork_windows_rocm_artifacts()
|
||||
# published_rocm_choice_for_host serves a bundle on a concrete mapped_targets entry or on
|
||||
# the umbrella gfx_target itself, so both spellings must clear a floor. A gfx_target
|
||||
# absent from its own mapped_targets is the family label (gfx110X); one present in it is
|
||||
# a standalone bundle (gfx908) already counted as concrete.
|
||||
concrete = {target.lower() for artifact in artifacts for target in artifact.mapped_targets}
|
||||
labels = {
|
||||
artifact.gfx_target.lower()
|
||||
for artifact in artifacts
|
||||
if artifact.gfx_target and artifact.gfx_target.lower() not in concrete
|
||||
}
|
||||
unfloored = sorted(concrete - ilp.WINDOWS_HIP_PREBUILT_GFX_TARGETS)
|
||||
assert (
|
||||
not unfloored
|
||||
), f"auto-Vulkan would steal windows-rocm archs published in {FORK}@{tag}: {unfloored}"
|
||||
unlabelled = sorted(labels - ilp.WINDOWS_ROCM_FAMILY_GFX_LABELS)
|
||||
assert not unlabelled, (
|
||||
f"update markers forward family labels {FORK}@{tag} publishes but "
|
||||
f"WINDOWS_ROCM_FAMILY_GFX_LABELS omits: {unlabelled}"
|
||||
)
|
||||
# Keep the import-time tuple the offline routing tests parametrise on an exact mirror.
|
||||
assert set(_FORK_WINDOWS_ROCM_GFX) == concrete, (
|
||||
f"_FORK_WINDOWS_ROCM_GFX drifted from {FORK}@{tag}: "
|
||||
f"gained {sorted(concrete - set(_FORK_WINDOWS_ROCM_GFX))}, "
|
||||
f"lost {sorted(set(_FORK_WINDOWS_ROCM_GFX) - concrete)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gfx", _FORK_WINDOWS_ROCM_GFX)
|
||||
def test_route_to_vulkan_prebuilt_keeps_every_fork_windows_rocm_arch(gfx, monkeypatch):
|
||||
# No ambient opt-in: this asserts the AUTO path leaves covered archs alone.
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
host = _windows_amd_host(rocm_gfx_target = gfx, rocm_gfx_targets = [gfx])
|
||||
routed, repo, tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert (repo, tag) == (FORK, "pin")
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_forwarded_gfx_does_not_undo_visible_device_auto_vulkan(monkeypatch):
|
||||
# Mixed-AMD Windows host: GPU 0 = gfx1100 (HIP prebuilt exists), GPU 1 = gfx1010 (none).
|
||||
# Under CUDA_VISIBLE_DEVICES=1 setup.ps1 still resolves GPU 0 and forwards gfx1100, but
|
||||
# detect_host() resolved the visible gfx1010, so folding the forward in must not
|
||||
# reinstate gfx1100 and install a HIP bundle the visible GPU cannot run.
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1010", rocm_gfx_targets = ["gfx1100", "gfx1010"])
|
||||
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx1100")
|
||||
assert ilp._active_rocm_gfx_target(host) == "gfx1010"
|
||||
assert host.rocm_gfx_targets == ["gfx1100", "gfx1010"]
|
||||
# gfx1100 is masked off, not absent, and Vulkan does not honour the HIP mask, so the
|
||||
# automatic fallback stays off and the HIP / fork path is kept.
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_forwarded_gfx_absent_from_probe_keeps_the_physical_hip_card(monkeypatch):
|
||||
# Mixed-AMD Windows host: GPU 0 = gfx1100 (HIP prebuilt exists), GPU 1 = gfx803 (below
|
||||
# the floor). CUDA_VISIBLE_DEVICES=1 reserves the gfx1100, so detect_host() picks gfx803
|
||||
# as active but still reports both cards, and setup forwards a third arch the probe never
|
||||
# saw (a stale env var, or name inference reading the other card). That forward selects
|
||||
# the HIP target but must not delete the probe's inventory, or the floor check concludes
|
||||
# no AMD GPU here reaches HIP and auto-routes to Vulkan, which ignores the HIP mask and
|
||||
# enumerates the reserved gfx1100.
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx1100", "gfx803"])
|
||||
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx900")
|
||||
assert ilp._active_rocm_gfx_target(host) == "gfx900"
|
||||
assert host.rocm_gfx_targets == ["gfx1100", "gfx803", "gfx900"]
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_forwarded_gfx_absent_from_probe_keeps_a_single_probed_hip_card(monkeypatch):
|
||||
# Same rule on a single-GPU box: a stale below-floor forward over a probe-confirmed
|
||||
# gfx1100 must not auto-route that machine to Vulkan.
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
|
||||
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
|
||||
assert ilp._active_rocm_gfx_target(host) == "gfx803"
|
||||
assert host.rocm_gfx_targets == ["gfx1100", "gfx803"]
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
|
||||
|
||||
|
||||
def test_forwarded_gfx_absent_from_probe_still_allows_explicit_vulkan(monkeypatch):
|
||||
# The physical-inventory rule gates the AUTO path only; naming the backend wins.
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
|
||||
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
|
||||
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
|
||||
|
||||
def test_forwarded_gfx_on_unprobed_host_still_auto_vulkans(monkeypatch):
|
||||
# Negative control: a driver-only AMD host runs no successful probe (no hipinfo, amd-smi
|
||||
# suppressed), so --rocm-gfx is the ONLY source of the arch and there is no inventory to
|
||||
# preserve. This is the #7357 path the feature exists for; it must still reach Vulkan.
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
|
||||
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
|
||||
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
|
||||
assert host.rocm_gfx_targets == ["gfx803"]
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is True
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
|
||||
|
||||
def test_forwarded_gfx_still_fills_an_unprobed_arch(monkeypatch):
|
||||
# Negative control: on an amd-smi-only host detect_host() reports no arch, so the
|
||||
# forward is the only source and must still apply.
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
|
||||
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
|
||||
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx1151")
|
||||
assert ilp._active_rocm_gfx_target(host) == "gfx1151"
|
||||
assert ilp._should_auto_vulkan_for_amd_windows(host) is False
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_llama_backend_hip_opts_out_of_auto_vulkan(monkeypatch):
|
||||
# hip names a backend, so it keeps the fork path even on an auto-fallback arch.
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "hip")
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
|
||||
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert routed is host
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
assert ilp.force_vulkan_requested() is False
|
||||
|
||||
|
||||
def test_explicit_backend_beats_legacy_force_vulkan(monkeypatch):
|
||||
# A stale UNSLOTH_FORCE_VULKAN must not overrule UNSLOTH_LLAMA_BACKEND=rocm (== hip).
|
||||
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "rocm")
|
||||
assert ilp.resolved_llama_backend() == "hip"
|
||||
assert ilp.force_vulkan_requested() is False
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == FORK
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_unknown_llama_backend_value_falls_through_to_legacy_flag(monkeypatch):
|
||||
# An unrecognised value is ignored, not an error, so the legacy flag still works.
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "banana")
|
||||
assert ilp.resolved_llama_backend() is None
|
||||
assert ilp.force_vulkan_requested() is False
|
||||
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
|
||||
assert ilp.force_vulkan_requested() is True
|
||||
|
||||
|
||||
def test_llama_backend_flag_beats_conflicting_env(monkeypatch):
|
||||
# --llama-backend is the caller's explicit request and outranks the env.
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "hip")
|
||||
assert ilp.force_vulkan_requested("vulkan") is True
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(
|
||||
host, FORK, "pin", force_cpu = False, llama_backend = "vulkan"
|
||||
)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
|
||||
|
||||
def _windows_arm64_host(**overrides):
|
||||
defaults = dict(
|
||||
system = "Windows",
|
||||
machine = "ARM64",
|
||||
is_windows = True,
|
||||
is_linux = False,
|
||||
is_macos = False,
|
||||
is_x86_64 = False,
|
||||
is_arm64 = True,
|
||||
nvidia_smi = None,
|
||||
driver_cuda_version = None,
|
||||
compute_caps = [],
|
||||
visible_cuda_devices = None,
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
has_rocm = False,
|
||||
has_intel_gpu = False,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return ilp.HostInfo(**defaults)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"env, flag",
|
||||
[
|
||||
({"UNSLOTH_LLAMA_BACKEND": "vulkan"}, None),
|
||||
({"UNSLOTH_FORCE_VULKAN": "1"}, None),
|
||||
({}, "vulkan"),
|
||||
],
|
||||
)
|
||||
def test_vulkan_opt_in_ignored_on_windows_arm64(monkeypatch, env, flag):
|
||||
# Upstream builds win-vulkan for x64 only (arm64 gets CPU + opencl-adreno), so rewriting
|
||||
# the host would only swap the published arm64 bundle for the upstream CPU one.
|
||||
for name, value in env.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
host = _windows_arm64_host()
|
||||
routed, repo, tag, persist = ilp._route_to_vulkan_prebuilt(
|
||||
host, FORK, "pin", force_cpu = False, llama_backend = flag
|
||||
)
|
||||
assert routed is host
|
||||
assert (repo, tag) == (FORK, "pin")
|
||||
assert persist is None
|
||||
|
||||
|
||||
def test_vulkan_opt_in_still_routes_on_windows_x64(monkeypatch):
|
||||
# Negative control for the arm64 guard: x64 keeps its Vulkan routing.
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
|
||||
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
|
||||
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
|
||||
assert repo == UPSTREAM
|
||||
assert persist == "vulkan"
|
||||
|
||||
|
||||
def _choice(install_kind, name = "asset.zip"):
|
||||
return ilp.AssetChoice(
|
||||
repo = UPSTREAM,
|
||||
tag = "b9925",
|
||||
name = name,
|
||||
url = f"https://example/{name}",
|
||||
source_label = "upstream",
|
||||
install_kind = install_kind,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["windows-vulkan", "linux-vulkan"])
|
||||
def test_persisted_llama_backend_keeps_vulkan_for_a_vulkan_bundle(kind):
|
||||
assert ilp.persisted_llama_backend("vulkan", _choice(kind)) == "vulkan"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", ["windows-arm64", "windows-cpu", "linux-cpu", "windows-rocm"])
|
||||
def test_persisted_llama_backend_drops_vulkan_for_a_non_vulkan_bundle(kind):
|
||||
# _plan_llama_phase re-asserts the marker's backend on every later update, so a Vulkan
|
||||
# request that fell through to CPU must not leave a marker claiming Vulkan.
|
||||
assert ilp.persisted_llama_backend("vulkan", _choice(kind)) is None
|
||||
|
||||
|
||||
def test_persisted_llama_backend_passes_none_through():
|
||||
assert ilp.persisted_llama_backend(None, _choice("windows-vulkan")) is None
|
||||
|
||||
|
||||
def test_marker_records_no_backend_when_vulkan_fell_back_to_cpu(tmp_path):
|
||||
# End to end over write_prebuilt_metadata: describe the CPU attempt that actually won,
|
||||
# so the next update re-detects instead of re-asserting Vulkan forever.
|
||||
checksums = ilp.ApprovedReleaseChecksums(
|
||||
repo = UPSTREAM,
|
||||
release_tag = "b9925",
|
||||
upstream_tag = "b9925",
|
||||
source_repo = UPSTREAM,
|
||||
source_repo_url = f"https://github.com/{UPSTREAM}",
|
||||
)
|
||||
cpu = _choice("windows-arm64", "llama-b9925-bin-win-cpu-arm64.zip")
|
||||
ilp.write_prebuilt_metadata(
|
||||
tmp_path,
|
||||
requested_tag = "latest",
|
||||
llama_tag = "b9925",
|
||||
release_tag = "b9925",
|
||||
choice = cpu,
|
||||
approved_checksums = checksums,
|
||||
prebuilt_fallback_used = False,
|
||||
llama_backend = "vulkan",
|
||||
)
|
||||
marker = json.loads((tmp_path / "UNSLOTH_PREBUILT_INFO.json").read_text())
|
||||
assert marker["asset"] == "llama-b9925-bin-win-cpu-arm64.zip"
|
||||
assert marker["llama_backend"] is None
|
||||
|
||||
vulkan = _choice("windows-vulkan", "llama-b9925-bin-win-vulkan-x64.zip")
|
||||
ilp.write_prebuilt_metadata(
|
||||
tmp_path,
|
||||
requested_tag = "latest",
|
||||
llama_tag = "b9925",
|
||||
release_tag = "b9925",
|
||||
choice = vulkan,
|
||||
approved_checksums = checksums,
|
||||
prebuilt_fallback_used = False,
|
||||
llama_backend = "vulkan",
|
||||
)
|
||||
marker = json.loads((tmp_path / "UNSLOTH_PREBUILT_INFO.json").read_text())
|
||||
assert marker["llama_backend"] == "vulkan"
|
||||
|
||||
|
||||
# UNSLOTH_LLAMA_CPP_BACKEND (setup.sh/setup.ps1, "auto"|"cpu") and
|
||||
# UNSLOTH_LLAMA_BACKEND (this module, a backend name) are different variables at
|
||||
# different layers, and both accept "cpu". setup translates its own =cpu into
|
||||
# --force-cpu to pin the CPU-only bundle on a GPU host, which is what keeps Intel
|
||||
# iGPU Vulkan crashes away (#7213). Vulkan is opt-in here, so no trigger it adds
|
||||
# may outrank that flag on any host.
|
||||
_SIM_PLATFORMS = {
|
||||
# WSL presents as Linux to this resolver, so it rides the Linux row.
|
||||
"Linux": dict(
|
||||
system = "Linux",
|
||||
is_windows = False,
|
||||
is_linux = True,
|
||||
is_macos = False,
|
||||
machine = "x86_64",
|
||||
is_x86_64 = True,
|
||||
is_arm64 = False,
|
||||
),
|
||||
"Windows": dict(
|
||||
system = "Windows",
|
||||
is_windows = True,
|
||||
is_linux = False,
|
||||
is_macos = False,
|
||||
machine = "amd64",
|
||||
is_x86_64 = True,
|
||||
is_arm64 = False,
|
||||
),
|
||||
"macOS": dict(
|
||||
system = "Darwin",
|
||||
is_windows = False,
|
||||
is_linux = False,
|
||||
is_macos = True,
|
||||
machine = "arm64",
|
||||
is_x86_64 = False,
|
||||
is_arm64 = True,
|
||||
),
|
||||
}
|
||||
_SIM_GPUS = {
|
||||
"nvidia": dict(
|
||||
has_physical_nvidia = True,
|
||||
has_usable_nvidia = True,
|
||||
has_rocm = False,
|
||||
has_intel_gpu = False,
|
||||
nvidia_smi = "/usr/bin/nvidia-smi",
|
||||
driver_cuda_version = "12.4",
|
||||
compute_caps = ["8.9"],
|
||||
),
|
||||
"amd": dict(
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
has_rocm = True,
|
||||
has_intel_gpu = False,
|
||||
nvidia_smi = None,
|
||||
driver_cuda_version = None,
|
||||
compute_caps = [],
|
||||
rocm_gfx_target = "gfx803",
|
||||
rocm_gfx_targets = ["gfx803"],
|
||||
),
|
||||
"intel": dict(
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
has_rocm = False,
|
||||
has_intel_gpu = True,
|
||||
nvidia_smi = None,
|
||||
driver_cuda_version = None,
|
||||
compute_caps = [],
|
||||
),
|
||||
"cpu_only": dict(
|
||||
has_physical_nvidia = False,
|
||||
has_usable_nvidia = False,
|
||||
has_rocm = False,
|
||||
has_intel_gpu = False,
|
||||
nvidia_smi = None,
|
||||
driver_cuda_version = None,
|
||||
compute_caps = [],
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _sim_host(platform_name, gpu_name):
|
||||
base = dict(visible_cuda_devices = None)
|
||||
base.update(_SIM_PLATFORMS[platform_name])
|
||||
base.update(_SIM_GPUS[gpu_name])
|
||||
return ilp.HostInfo(**base)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("platform_name", sorted(_SIM_PLATFORMS))
|
||||
@pytest.mark.parametrize("gpu_name", sorted(_SIM_GPUS))
|
||||
@pytest.mark.parametrize("backend_env", [None, "vulkan", "hip", "rocm", "cpu"])
|
||||
def test_forced_cpu_outranks_every_vulkan_trigger(
|
||||
monkeypatch, platform_name, gpu_name, backend_env
|
||||
):
|
||||
"""A deliberate CPU install stays CPU on every host, whatever asks for Vulkan."""
|
||||
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
|
||||
if backend_env is None:
|
||||
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
|
||||
else:
|
||||
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", backend_env)
|
||||
# The legacy switch too, so a stale one cannot smuggle Vulkan past --force-cpu.
|
||||
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
|
||||
|
||||
repo, tag = "unslothai/llama.cpp-prebuilt", "latest"
|
||||
_, out_repo, _, persist = ilp._route_to_vulkan_prebuilt(
|
||||
_sim_host(platform_name, gpu_name),
|
||||
repo,
|
||||
tag,
|
||||
force_cpu = True,
|
||||
llama_backend = "vulkan",
|
||||
)
|
||||
assert out_repo == repo, (platform_name, gpu_name, backend_env)
|
||||
assert persist is None, (platform_name, gpu_name, backend_env)
|
||||
|
||||
|
||||
def test_the_forced_cpu_guard_is_not_vacuous():
|
||||
"""The same host DOES take Vulkan once the CPU pin is gone, or the check above
|
||||
would pass on a resolver that had stopped routing to Vulkan entirely."""
|
||||
repo, tag = "unslothai/llama.cpp-prebuilt", "latest"
|
||||
_, out_repo, _, persist = ilp._route_to_vulkan_prebuilt(
|
||||
_sim_host("Linux", "amd"),
|
||||
repo,
|
||||
tag,
|
||||
force_cpu = False,
|
||||
llama_backend = "vulkan",
|
||||
)
|
||||
assert out_repo != repo or persist == "vulkan"
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from core.inference.llama_admission import (
|
|||
ADMISSION_CONTROL_ENV,
|
||||
ADMISSION_KEEPALIVE_INTERVAL_ENV,
|
||||
ADMISSION_MAX_QUEUE_ENV,
|
||||
ADMISSION_QUEUE_PER_SLOT_ENV,
|
||||
ADMISSION_QUEUE_TIMEOUT_ENV,
|
||||
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
|
||||
DEFAULT_ADMISSION_MAX_QUEUE,
|
||||
|
|
@ -28,8 +29,23 @@ from core.inference.llama_admission import (
|
|||
)
|
||||
|
||||
|
||||
_ADMISSION_ENV = (
|
||||
ADMISSION_CONTROL_ENV,
|
||||
ADMISSION_QUEUE_TIMEOUT_ENV,
|
||||
ADMISSION_KEEPALIVE_INTERVAL_ENV,
|
||||
ADMISSION_MAX_QUEUE_ENV,
|
||||
ADMISSION_QUEUE_PER_SLOT_ENV,
|
||||
*llama_admission._LEGACY_ENV.values(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _reset_queues():
|
||||
def _reset_queues(monkeypatch):
|
||||
# Clear ambient settings for every test, not just the ones that remember to:
|
||||
# a canonical name set on the machine silently beats the legacy name a test
|
||||
# is exercising, and the queue registry is process-global.
|
||||
for name in _ADMISSION_ENV:
|
||||
monkeypatch.delenv(name, raising = False)
|
||||
reset_llama_admission_queues()
|
||||
yield
|
||||
reset_llama_admission_queues()
|
||||
|
|
@ -41,15 +57,25 @@ def test_admission_config_defaults(monkeypatch):
|
|||
ADMISSION_QUEUE_TIMEOUT_ENV,
|
||||
ADMISSION_KEEPALIVE_INTERVAL_ENV,
|
||||
ADMISSION_MAX_QUEUE_ENV,
|
||||
ADMISSION_QUEUE_PER_SLOT_ENV,
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL",
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT",
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL",
|
||||
"UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE",
|
||||
):
|
||||
monkeypatch.delenv(name, raising = False)
|
||||
|
||||
config = llama_admission_config_from_env()
|
||||
|
||||
# Literals, not the module constants: comparing a default to itself would let
|
||||
# any future value change through silently.
|
||||
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
|
||||
assert config.queue_timeout_s is None # wait forever
|
||||
assert config.keepalive_interval_s == 5.0
|
||||
assert config.max_queue is None # no absolute cap
|
||||
assert config.queue_per_slot == 16
|
||||
assert (DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, DEFAULT_ADMISSION_MAX_QUEUE) == (None, None)
|
||||
assert DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S == 5.0
|
||||
|
||||
|
||||
def test_admission_config_env_overrides(monkeypatch):
|
||||
|
|
@ -66,6 +92,25 @@ def test_admission_config_env_overrides(monkeypatch):
|
|||
assert config.max_queue is None
|
||||
|
||||
|
||||
def test_admission_config_honors_legacy_openai_compat_env(monkeypatch):
|
||||
# The queue is shared with /v1/messages now, but existing OPENAI_COMPAT
|
||||
# settings must keep working.
|
||||
monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", "7")
|
||||
monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL", "off")
|
||||
|
||||
config = llama_admission_config_from_env()
|
||||
|
||||
assert config.max_queue == 7
|
||||
assert config.enabled is False
|
||||
|
||||
|
||||
def test_admission_config_prefers_neutral_env_over_legacy(monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", "7")
|
||||
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "3")
|
||||
|
||||
assert llama_admission_config_from_env().max_queue == 3
|
||||
|
||||
|
||||
def test_admission_config_positive_queue_timeout_env(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600")
|
||||
|
||||
|
|
@ -106,6 +151,160 @@ def test_fifo_capacity_one_grants_next_waiter_on_release():
|
|||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_pool_hands_out_distinct_slots_and_reuses_them():
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
leases = [queue.reserve(capacity = 3, config = config).lease_nowait() for _ in range(3)]
|
||||
assert sorted(lease.slot for lease in leases) == [0, 1, 2] # one slot each
|
||||
snapshot = queue.snapshot()
|
||||
assert (snapshot.active, snapshot.free, snapshot.capacity) == (3, 0, 3)
|
||||
|
||||
# A freed slot returns to the pool and is handed to the next caller.
|
||||
freed = leases[1].slot
|
||||
leases[1].release()
|
||||
assert queue.snapshot().free == 1
|
||||
reused = queue.reserve(capacity = 3, config = config).lease_nowait()
|
||||
assert reused.slot == freed
|
||||
|
||||
reused.release()
|
||||
leases[0].release()
|
||||
leases[2].release()
|
||||
snapshot = queue.snapshot()
|
||||
assert (snapshot.active, snapshot.free) == (0, 3)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_pool_waiter_is_handed_a_real_slot():
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
held = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
waiting = queue.reserve(capacity = 1, config = config)
|
||||
assert waiting.lease_nowait() is None
|
||||
assert queue.snapshot().free == 0
|
||||
|
||||
held.release()
|
||||
granted = await waiting.wait(0.1)
|
||||
assert granted is not None and granted.slot == 0 # the slot just freed
|
||||
granted.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_shrinking_capacity_retires_slots_beyond_the_new_pool():
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
leases = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
|
||||
assert queue.snapshot().capacity == 4
|
||||
|
||||
# llama-server reloaded with fewer --parallel slots; in-flight holders keep
|
||||
# running and their slots retire instead of returning to the smaller pool.
|
||||
shrunk = queue.reserve(capacity = 2, config = config)
|
||||
assert shrunk.lease_nowait() is None # all 4 still held, nothing free
|
||||
for lease in leases:
|
||||
lease.release()
|
||||
|
||||
granted = await shrunk.wait(0.1)
|
||||
assert granted is not None and granted.slot < 2
|
||||
granted.release()
|
||||
snapshot = queue.snapshot()
|
||||
assert (snapshot.capacity, snapshot.active, snapshot.free) == (2, 0, 2)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_queue_limit_scales_with_the_serving_slots():
|
||||
# The wait line follows --parallel: 16 per slot, floored at 64 so a 1-slot
|
||||
# backend keeps the depth it had before scaling existed.
|
||||
config = LlamaAdmissionConfig()
|
||||
assert config.queue_limit(4) == 64 # --parallel 4 (the default)
|
||||
assert config.queue_limit(8) == 128 # --parallel 8
|
||||
assert config.queue_limit(16) == 256
|
||||
assert config.queue_limit(1) == 64 # floor, not 16
|
||||
assert config.queue_limit(2) == 64 # floor, not 32
|
||||
# An explicit cap wins, and a None multiplier means an unbounded line.
|
||||
assert LlamaAdmissionConfig(max_queue = 5).queue_limit(8) == 5
|
||||
assert LlamaAdmissionConfig(queue_per_slot = None).queue_limit(8) is None
|
||||
# Non-positive settings mean unbounded, never "reject everything".
|
||||
assert LlamaAdmissionConfig(max_queue = 0).queue_limit(4) is None
|
||||
assert LlamaAdmissionConfig(max_queue = -1).queue_limit(4) is None
|
||||
assert LlamaAdmissionConfig(queue_per_slot = 0).queue_limit(4) is None
|
||||
assert LlamaAdmissionConfig(queue_per_slot = -3).queue_limit(4) is None
|
||||
|
||||
|
||||
def test_queue_limit_rejects_only_once_the_line_is_full():
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
# Explicit cap, so the test drives rejection without standing up the 64
|
||||
# waiters the scaled floor would otherwise require.
|
||||
config = LlamaAdmissionConfig(max_queue = 4)
|
||||
|
||||
held = [queue.reserve(capacity = 2, config = config).lease_nowait() for _ in range(2)]
|
||||
parked = [queue.reserve(capacity = 2, config = config) for _ in range(4)]
|
||||
assert queue.snapshot().queued == 4
|
||||
|
||||
with pytest.raises(LlamaAdmissionQueueFull):
|
||||
queue.reserve(capacity = 2, config = config)
|
||||
|
||||
for reservation in parked:
|
||||
reservation.cancel()
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_waiting_is_never_timed_out_by_default():
|
||||
# "Wait forever": the default config sets no queue timeout at all.
|
||||
assert llama_admission_config_from_env().queue_timeout_s is None
|
||||
assert LlamaAdmissionConfig().queue_timeout_s is None
|
||||
|
||||
|
||||
def test_single_request_at_a_time_never_queues_or_allocates_waiters():
|
||||
# The common serving case: one request in flight at a time must take a slot
|
||||
# straight away and never touch the wait line.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
for _ in range(50):
|
||||
reservation = queue.reserve(capacity = 4, config = config)
|
||||
lease = reservation.lease_nowait()
|
||||
assert lease is not None # admitted immediately
|
||||
assert queue.snapshot().queued == 0 # nobody ever lined up
|
||||
lease.release()
|
||||
snapshot = queue.snapshot()
|
||||
assert (snapshot.active, snapshot.free, snapshot.queued) == (0, 4, 0)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_unbounded_queue_keeps_waiting_instead_of_rejecting():
|
||||
# queue_per_slot None is the "pool + unbounded wait line" mode: nothing is
|
||||
# ever rejected, callers just line up for the next free slot.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig(max_queue = None, queue_per_slot = None)
|
||||
|
||||
held = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
waiters = [queue.reserve(capacity = 1, config = config) for _ in range(200)]
|
||||
assert queue.snapshot().queued == 200 # no LlamaAdmissionQueueFull
|
||||
|
||||
held.release()
|
||||
first = await waiters[0].wait(0.1)
|
||||
assert first is not None
|
||||
first.release()
|
||||
for waiter in waiters[1:]:
|
||||
waiter.cancel()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_queue_full_rejects_excess_waiter():
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
|
|
@ -288,6 +487,105 @@ def test_lease_release_is_idempotent_under_concurrent_calls():
|
|||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_releasing_a_stale_lease_does_not_free_someone_elses_slot():
|
||||
# The concurrent test above passes without the _released guard: the racing
|
||||
# calls all target a still-live slot, which the bitmask already absorbs. The
|
||||
# case the guard exists for is a slot released twice with a reuse in between.
|
||||
# It is live: _wait_for_openai_admission_non_streaming releases and re-raises,
|
||||
# then the caller's finally cancels the reservation and releases the same
|
||||
# lease again, by which point the slot can belong to another request.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
stale = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
stale.release()
|
||||
other = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
assert other.slot == stale.slot # the slot got reused
|
||||
|
||||
stale.release()
|
||||
assert queue.snapshot().active == 1, "stale release handed back a live slot"
|
||||
other.release()
|
||||
assert queue.snapshot().active == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_grant_reclaims_the_slot_when_the_waiters_loop_is_gone():
|
||||
# _grant_waiters_locked takes the slot before scheduling delivery, so if the
|
||||
# schedule fails the bit is already set. Leaving it set strands the slot for
|
||||
# good, because _free is rebuilt from the bitmask.
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
held = None
|
||||
|
||||
dead = asyncio.new_event_loop()
|
||||
try:
|
||||
|
||||
async def _fill_and_queue():
|
||||
nonlocal held
|
||||
held = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
assert queue.reserve(capacity = 1, config = config).lease_nowait() is None
|
||||
|
||||
dead.run_until_complete(_fill_and_queue())
|
||||
finally:
|
||||
dead.close()
|
||||
|
||||
held.release() # grant path now hits the closed loop
|
||||
assert queue.snapshot().active == 0
|
||||
assert queue.is_idle()
|
||||
|
||||
|
||||
def test_cancel_returns_the_granted_slot_when_the_waiters_loop_is_gone():
|
||||
# Routes cancel() from finally blocks, so a raise here would mask their
|
||||
# exception and skip the release that hands the granted slot back.
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
held = reservation = None
|
||||
|
||||
dead = asyncio.new_event_loop()
|
||||
try:
|
||||
|
||||
async def _fill_and_queue():
|
||||
nonlocal held, reservation
|
||||
held = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
reservation = queue.reserve(capacity = 1, config = config)
|
||||
|
||||
dead.run_until_complete(_fill_and_queue())
|
||||
held.release() # promotes the waiter, so cancel() has a lease to return
|
||||
finally:
|
||||
dead.close()
|
||||
|
||||
reservation.cancel()
|
||||
assert queue.snapshot().active == 0
|
||||
assert queue.is_idle()
|
||||
|
||||
|
||||
def test_delivery_to_an_already_finished_waiter_releases_the_slot():
|
||||
# A slot is taken before delivery is scheduled, so if the waiter finishes in
|
||||
# that window someone has to hand it back. _deliver_lease does it twice over,
|
||||
# in the dead-waiter branch and in the InvalidStateError backstop; this pins
|
||||
# the outcome, not which one. Reaches into the waiter because no public call
|
||||
# leaves that window open: queue.cancel() reclaims granted_lease itself.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
held = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
reservation = queue.reserve(capacity = 1, config = config)
|
||||
waiter = reservation._waiter
|
||||
|
||||
held.release() # schedules _deliver_lease, sets granted_lease
|
||||
waiter.future.cancel() # finishes the future before the callback runs
|
||||
assert waiter.granted_lease is not None
|
||||
await asyncio.sleep(0) # let the callback run
|
||||
|
||||
assert queue.snapshot().active == 0
|
||||
assert queue.is_idle()
|
||||
|
||||
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.
|
||||
|
|
@ -318,3 +616,453 @@ def test_new_key_retains_in_flight_prior_load_queue():
|
|||
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"}
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_capacity_shrink_never_admits_past_the_new_ceiling():
|
||||
# A load that downshifts --parallel (or an unload resetting it to 1) shrinks the
|
||||
# pool while slots are still held. Those holdovers keep occupying the backend, so
|
||||
# they must count against the ceiling; sizing on free ids alone over-admits.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
held = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
|
||||
assert all(lease is not None for lease in held)
|
||||
waiter = queue.reserve(capacity = 4, config = config)
|
||||
|
||||
queue.reserve(capacity = 1, config = config) # capacity collapses to 1
|
||||
# Release the one id that still falls inside the shrunk pool, so it goes
|
||||
# back on the free list; ids at or above capacity retire instead.
|
||||
low = min(held, key = lambda lease: lease.slot)
|
||||
assert low.slot == 0
|
||||
low.release()
|
||||
|
||||
# The other 3 holdovers are still generating, which already meets the new
|
||||
# ceiling, so the freed id must not be handed on. Gating on "is an id free"
|
||||
# alone grants it here and puts 4 generations on a 1-slot backend.
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await waiter.wait(0.2)
|
||||
assert queue.snapshot().active == 3
|
||||
|
||||
waiter.cancel()
|
||||
for lease in held:
|
||||
if lease is not low:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_queue_per_slot_env_is_parsed(monkeypatch):
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "4")
|
||||
assert llama_admission_config_from_env().queue_limit(32) == 128
|
||||
# Non-positive asks for an unbounded line rather than rejecting everything.
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "0")
|
||||
assert llama_admission_config_from_env().queue_limit(32) is None
|
||||
|
||||
|
||||
def test_max_queue_zero_from_env_is_unbounded_end_to_end(monkeypatch):
|
||||
# Guards the whole env path, not just the parsed field: a regression that let
|
||||
# queue_per_slot survive MAX_QUEUE=0 would silently re-bound the line.
|
||||
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0")
|
||||
config = llama_admission_config_from_env()
|
||||
assert config.max_queue is None and config.queue_per_slot is None
|
||||
assert config.queue_limit(1) is None and config.queue_limit(64) is None
|
||||
|
||||
|
||||
def test_legacy_env_fallback_covers_every_setting(monkeypatch):
|
||||
for canonical, legacy in llama_admission._LEGACY_ENV.items():
|
||||
monkeypatch.delenv(canonical, raising = False)
|
||||
monkeypatch.setenv(legacy, "0" if "CONTROL" in canonical else "7")
|
||||
config = llama_admission_config_from_env()
|
||||
assert config.enabled is False
|
||||
assert config.queue_timeout_s == 7.0
|
||||
assert config.keepalive_interval_s == 7.0
|
||||
assert config.max_queue == 7
|
||||
|
||||
|
||||
def test_empty_canonical_env_falls_through_to_legacy(monkeypatch):
|
||||
# The branch _raw_env exists for: set but blank must not mask the legacy name.
|
||||
monkeypatch.setenv(ADMISSION_CONTROL_ENV, " ")
|
||||
monkeypatch.setenv(llama_admission._LEGACY_ENV[ADMISSION_CONTROL_ENV], "0")
|
||||
assert llama_admission_config_from_env().enabled is False
|
||||
|
||||
|
||||
def test_explicit_queue_per_slot_is_not_floored(monkeypatch):
|
||||
# The floor exists so a 1-slot backend keeps its old depth by default, not to
|
||||
# override an operator who asked for a shallow line.
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "2")
|
||||
config = llama_admission_config_from_env()
|
||||
assert config.queue_limit(1) == 2
|
||||
assert config.queue_limit(8) == 16
|
||||
|
||||
# Unset, the default multiplier is floored instead.
|
||||
monkeypatch.delenv(ADMISSION_QUEUE_PER_SLOT_ENV, raising = False)
|
||||
assert llama_admission_config_from_env().queue_limit(1) == 64
|
||||
|
||||
# A value that does not parse falls back to the default multiplier, so it has
|
||||
# to keep the default's floor. Otherwise a typo quietly shrinks the line 4x.
|
||||
for garbage in ("abc", "1e3", "16.0"):
|
||||
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, garbage)
|
||||
assert llama_admission_config_from_env().queue_limit(1) == 64, garbage
|
||||
|
||||
|
||||
def test_module_imports_on_python_39(monkeypatch):
|
||||
"""No 3.10+ API on an import path. The package declares >=3.9 but CI only
|
||||
runs 3.12, so a regression here would ship broken."""
|
||||
import ast
|
||||
import pathlib
|
||||
|
||||
src = pathlib.Path(llama_admission.__file__).read_text(encoding = "utf-8")
|
||||
tree = ast.parse(src)
|
||||
|
||||
# int.bit_count() (3.10+)
|
||||
assert not [
|
||||
n
|
||||
for n in ast.walk(tree)
|
||||
if isinstance(n, ast.Call)
|
||||
and isinstance(n.func, ast.Attribute)
|
||||
and n.func.attr == "bit_count"
|
||||
]
|
||||
# dataclass(slots = ...) is 3.10+, so every dataclass must take it through
|
||||
# the version gate instead of naming it. A new one that forgets the gate
|
||||
# loses slots silently, so require the **_SLOTS unpack rather than allow it.
|
||||
seen = 0
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.Call):
|
||||
continue
|
||||
name = getattr(node.func, "id", None) or getattr(node.func, "attr", None)
|
||||
if name != "dataclass":
|
||||
continue
|
||||
seen += 1
|
||||
assert "slots" not in {kw.arg for kw in node.keywords}
|
||||
assert [
|
||||
kw
|
||||
for kw in node.keywords
|
||||
if kw.arg is None and getattr(kw.value, "id", None) == "_SLOTS"
|
||||
], ast.dump(node)
|
||||
assert seen
|
||||
|
||||
|
||||
def test_slots_gate_matches_the_running_interpreter():
|
||||
"""The gate is only worth having if it actually applies where it can."""
|
||||
import sys
|
||||
|
||||
gated = (LlamaAdmissionConfig, llama_admission.LlamaAdmissionSnapshot, llama_admission._Waiter)
|
||||
if sys.version_info >= (3, 10):
|
||||
assert llama_admission._SLOTS == {"slots": True}
|
||||
for cls in gated:
|
||||
assert getattr(cls, "__slots__", None), cls
|
||||
else:
|
||||
assert llama_admission._SLOTS == {}
|
||||
|
||||
# Construct through the gate either way: slots=True rebuilds the class, so a
|
||||
# field it cannot carry over would only show up on instantiation.
|
||||
config = LlamaAdmissionConfig(max_queue = 7)
|
||||
assert config.max_queue == 7 and config.queue_limit(4) == 7
|
||||
assert llama_admission.LlamaAdmissionSnapshot("k", 1, 1, 0).capacity == 1
|
||||
|
||||
|
||||
def test_held_count_tracks_the_bitmask():
|
||||
# _held replaces int.bit_count(); the two must never drift apart.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
popcount = lambda: bin(queue._in_use).count("1")
|
||||
|
||||
leases = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
|
||||
assert queue._held == popcount() == 4
|
||||
leases[1].release()
|
||||
assert queue._held == popcount() == 3
|
||||
shrunk = queue.reserve(capacity = 2, config = config) # shrink with slots held
|
||||
assert queue._held == popcount() == 3
|
||||
shrunk.cancel() # else it is granted a slot as the others drain
|
||||
for lease in leases:
|
||||
lease.release()
|
||||
assert queue._held == popcount() == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_snapshot_free_never_exceeds_what_can_be_admitted():
|
||||
# After a shrink, low ids can sit in _free while holdovers fill the ceiling.
|
||||
# Reporting them as free made the admission log contradict itself.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
held = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
|
||||
queue.reserve(capacity = 1, config = config) # capacity collapses to 1
|
||||
min(held, key = lambda lease: lease.slot).release()
|
||||
|
||||
snapshot = queue.snapshot()
|
||||
assert snapshot.free == 0, snapshot # nothing is actually takeable
|
||||
assert snapshot.active == 3
|
||||
for lease in held:
|
||||
lease.release()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_a_newcomer_does_not_barge_past_a_parked_waiter():
|
||||
# Anti-starvation, pinned as behaviour rather than as the `if not self._waiters`
|
||||
# check: _take_slot_locked consults _can_admit_locked anyway, so either alone
|
||||
# refuses the newcomer. This fails if both ever go.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
held = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
parked = queue.reserve(capacity = 1, config = config)
|
||||
assert parked.lease_nowait() is None
|
||||
|
||||
held.release()
|
||||
newcomer = queue.reserve(capacity = 1, config = config)
|
||||
assert newcomer.lease_nowait() is None, "newcomer barged past the parked waiter"
|
||||
assert (await parked.wait(0.1)) is not None
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_dead_waiters_stop_counting_against_the_queue_limit():
|
||||
# A future cancelled out of band leaves the entry in the deque: cancel() is not
|
||||
# called, so only the prune drops it. Without that, depth, is_idle() and the
|
||||
# queue-full limit all drift for the life of the queue.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig(max_queue = 2)
|
||||
|
||||
held = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
first = queue.reserve(capacity = 1, config = config)
|
||||
second = queue.reserve(capacity = 1, config = config)
|
||||
assert queue.snapshot().queued == 2
|
||||
with pytest.raises(LlamaAdmissionQueueFull):
|
||||
queue.reserve(capacity = 1, config = config)
|
||||
|
||||
first._waiter.future.cancel()
|
||||
second._waiter.future.cancel()
|
||||
assert queue.snapshot().queued == 0, "dead waiters still occupy the line"
|
||||
# The freed depth is usable again, and an idle queue is evictable.
|
||||
queue.reserve(capacity = 1, config = config).cancel()
|
||||
held.release()
|
||||
assert queue.is_idle()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_parking_frees_the_slot_for_a_waiter():
|
||||
"""A holder waiting on a tool approval must not hold a decode slot.
|
||||
|
||||
It is not generating, and with several prompts unanswered every slot would
|
||||
be held by a run parked on a human while llama-server sits idle.
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
first_lease.park()
|
||||
assert first_lease.slot is None, "the slot went back to the pool"
|
||||
second_lease = await second.wait(0.1)
|
||||
assert second_lease is not None, "parking did not free the slot"
|
||||
|
||||
# The parked holder keeps its lease, so releasing it is still correct.
|
||||
first_lease.unpark()
|
||||
first_lease.release()
|
||||
second_lease.release()
|
||||
assert queue.snapshot().active == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_unpark_without_park_is_a_no_op():
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
first = queue.reserve(capacity = 1, config = config)
|
||||
first_lease = first.lease_nowait()
|
||||
assert first_lease is not None
|
||||
first_lease.unpark()
|
||||
first_lease.unpark()
|
||||
|
||||
second = queue.reserve(capacity = 1, config = config)
|
||||
assert second.lease_nowait() is None, "capacity leaked past the limit"
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_releasing_a_parked_lease_leaves_the_queue_evictable():
|
||||
# is_idle() drives registry eviction, and a parked holder owns no slot, so
|
||||
# nothing but the parked count keeps its queue alive. A stuck count would
|
||||
# pin every dead queue for the life of the process.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
lease = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
lease.park()
|
||||
assert not queue.is_idle(), "a parked holder is coming back to this queue"
|
||||
lease.release()
|
||||
assert queue.is_idle()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_unpark_waits_instead_of_putting_two_holders_on_one_slot():
|
||||
# park() hands the freed slot to a waiter, so by the time the user answers an approval
|
||||
# prompt someone else may be decoding in it. Resuming regardless left two holders
|
||||
# against capacity 1, and the resumed tool loop went past the admission limit.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None, "A takes the only slot"
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
assert b.lease_nowait() is None, "B waits behind A"
|
||||
|
||||
a_lease.park() # A parks on an approval prompt; its slot goes to B
|
||||
b_lease = await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2)
|
||||
assert b_lease is not None, "B was granted the parked slot"
|
||||
|
||||
# A answers the prompt while B is still decoding: it must WAIT.
|
||||
resumed = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.05)
|
||||
assert not resumed.done(), "A must not resume while B holds the slot"
|
||||
assert queue.snapshot().active <= 1, "never over capacity while waiting"
|
||||
|
||||
b_lease.release()
|
||||
await asyncio.wait_for(resumed, timeout = 2)
|
||||
assert a_lease.slot is not None, "A took a real slot back"
|
||||
assert queue.snapshot().active <= 1, "still within capacity after resuming"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_unpark_gives_up_when_the_caller_is_cancelled():
|
||||
# A holder being torn down must not sit in the wait loop.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
a_lease.park()
|
||||
assert await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2) is not None
|
||||
|
||||
ev = threading.Event()
|
||||
waiting = asyncio.ensure_future(a_lease.unpark_async(cancel_event = ev, poll_s = 0.01))
|
||||
await asyncio.sleep(0.03)
|
||||
assert not waiting.done()
|
||||
ev.set()
|
||||
await asyncio.wait_for(waiting, timeout = 2)
|
||||
assert a_lease.slot is None, "gave up without a slot rather than over-admitting"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_an_approved_chat_is_not_overtaken_by_later_arrivals():
|
||||
# A parks on an approval prompt, B takes the slot, C arrives afterwards. release() grants
|
||||
# under the same lock, so a plain poll in unpark_async never saw a free slot: A waited
|
||||
# behind every later arrival and starved.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
a_lease.park() # A's slot goes to B
|
||||
b_lease = await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2)
|
||||
assert b_lease is not None
|
||||
|
||||
# A is approved and starts waiting; C arrives only after that.
|
||||
resumed = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.03)
|
||||
c = queue.reserve(capacity = 1, config = config)
|
||||
assert c.lease_nowait() is None
|
||||
|
||||
b_lease.release() # the slot frees exactly once
|
||||
await asyncio.wait_for(resumed, timeout = 2)
|
||||
# A resumed; C is still queued behind it rather than having overtaken it.
|
||||
assert c.lease_nowait() is None
|
||||
assert queue.snapshot().active <= 1
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_two_approved_chats_do_not_block_each_other():
|
||||
# A bare pending-count made every approved holder count against every other: park A, admit
|
||||
# and park B, admit C, approve both, and once C released the predicate stayed false forever.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
a_lease.park() # A parks; B is admitted
|
||||
b_lease = await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2)
|
||||
assert b_lease is not None
|
||||
|
||||
c = queue.reserve(capacity = 1, config = config)
|
||||
b_lease.park() # B parks too; C is admitted
|
||||
c_lease = await asyncio.wait_for(c.wait(timeout_s = 1), timeout = 2)
|
||||
assert c_lease is not None
|
||||
|
||||
# Both approvals come back while C is still decoding.
|
||||
first = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.02)
|
||||
second = asyncio.ensure_future(b_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.02)
|
||||
assert not first.done() and not second.done()
|
||||
|
||||
c_lease.release()
|
||||
# The earlier approval goes first; the other follows once it releases.
|
||||
await asyncio.wait_for(first, timeout = 2)
|
||||
assert not second.done(), "the second approval waits its turn, not forever"
|
||||
a_lease.release()
|
||||
await asyncio.wait_for(second, timeout = 2)
|
||||
assert queue.snapshot().active <= 1
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_an_immediate_arrival_cannot_take_an_approved_chats_slot():
|
||||
# The fairness reservation lived only in _grant_waiters_locked. reserve()'s fast path
|
||||
# ignored it, so a request arriving in the window between the slot freeing and the
|
||||
# approved chat's next poll took the slot straight off the top.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
a_lease.park() # A is on an approval prompt; its slot is up for grabs
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
b_lease = b.lease_nowait()
|
||||
assert b_lease is not None
|
||||
|
||||
resumed = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.03) # A is approved and now holds a ticket
|
||||
|
||||
# No await between these two: C arrives before A's poll can run again.
|
||||
b_lease.release()
|
||||
c = queue.reserve(capacity = 1, config = config)
|
||||
assert c.lease_nowait() is None, "the freed slot is reserved for the approved chat"
|
||||
|
||||
await asyncio.wait_for(resumed, timeout = 2)
|
||||
assert queue.snapshot().active <= 1
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
|
|
|||
|
|
@ -2076,6 +2076,50 @@ def test_large_python_tool_call_emits_early_provisional_start(monkeypatch):
|
|||
assert any(e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events)
|
||||
|
||||
|
||||
def test_gated_python_call_still_streams_its_arguments(monkeypatch):
|
||||
"""A call awaiting approval still streams its code into the card.
|
||||
|
||||
Suppressing it left the chat completely blank for as long as the model took
|
||||
to write the payload, which for a large file is minutes. Nothing runs before
|
||||
the decision either way, and the code is what the user is approving.
|
||||
"""
|
||||
|
||||
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
|
||||
assert len(json.dumps({"code": big_code})) > _PROVISIONAL_ARGS_MIN_CHARS
|
||||
|
||||
first_stream = _streamed_structured_tool_call("python", {"code": big_code}, "call_gated")
|
||||
final_stream = [_sse({"content": "Done."}), _done()]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", lambda name, arguments, **_k: "OK")
|
||||
monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow")
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "write code"}],
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
confirm_tool_calls = True,
|
||||
permission_mode = "ask",
|
||||
max_tool_iterations = 1,
|
||||
)
|
||||
)
|
||||
|
||||
tool_starts = [e for e in events if e.get("type") == "tool_start"]
|
||||
provisional = [e for e in tool_starts if not e.get("arguments")]
|
||||
assert len(provisional) == 1, tool_starts
|
||||
assert provisional[0]["tool_call_id"] == "call_gated"
|
||||
|
||||
args_events = [e for e in events if e.get("type") == "tool_args"]
|
||||
assert args_events, "gated call streamed no arguments"
|
||||
assert "total += 119" in "".join(e["text"] for e in args_events)
|
||||
|
||||
# The approval prompt still fires, and it comes after the code is on screen.
|
||||
gated = [e for e in tool_starts if e.get("awaiting_confirmation")]
|
||||
assert gated, tool_starts
|
||||
assert events.index(provisional[0]) < events.index(gated[0])
|
||||
|
||||
|
||||
def test_auto_mode_render_html_suppresses_provisional_card_under_confirm(monkeypatch):
|
||||
"""render_html is no longer unconditionally safe (a networked canvas asks), so
|
||||
with confirm_tool_calls set under permission_mode="auto" its early provisional
|
||||
|
|
|
|||
|
|
@ -473,6 +473,7 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
|
|||
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
|
||||
|
||||
def _on_start(cmd):
|
||||
captured["cmd"] = cmd
|
||||
_write_install(
|
||||
install_dir,
|
||||
"b9518",
|
||||
|
|
@ -480,6 +481,7 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
|
|||
asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
|
||||
)
|
||||
|
||||
captured: dict = {}
|
||||
popen_kwargs: dict = {}
|
||||
_patch_installer_popen(
|
||||
monkeypatch,
|
||||
|
|
@ -497,6 +499,8 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
|
|||
time.sleep(0.05)
|
||||
assert job["state"] == "success", job
|
||||
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
|
||||
assert popen_kwargs["env"]["UNSLOTH_LLAMA_BACKEND"] == "vulkan"
|
||||
assert "--llama-backend" in captured["cmd"] and "vulkan" in captured["cmd"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ is_managed_flag = _lsa.is_managed_flag
|
|||
parse_cache_override = _lsa.parse_cache_override
|
||||
parse_cache_override_per_axis = _lsa.parse_cache_override_per_axis
|
||||
parse_ctx_override = _lsa.parse_ctx_override
|
||||
parse_gpu_layers_override = _lsa.parse_gpu_layers_override
|
||||
parse_split_mode_override = _lsa.parse_split_mode_override
|
||||
resolve_cache_type_kv = _lsa.resolve_cache_type_kv
|
||||
resolve_tensor_parallel = _lsa.resolve_tensor_parallel
|
||||
|
|
@ -453,6 +454,45 @@ def test_validate_extra_args_rejects_malformed_ctx_override():
|
|||
validate_extra_args(["--ctx-size", "abc"])
|
||||
|
||||
|
||||
# ── parse_gpu_layers_override ────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args,expected",
|
||||
[
|
||||
(None, None),
|
||||
([], None),
|
||||
(["--top-k", "20"], None),
|
||||
(["--gpu-layers", "20"], 20),
|
||||
(["--gpu-layers=20"], 20),
|
||||
(["--n-gpu-layers", "0"], 0),
|
||||
(["-ngl", "-1"], -1),
|
||||
(["-ngl", "12", "--gpu-layers", "20"], 20),
|
||||
],
|
||||
)
|
||||
def test_parse_gpu_layers_override(args, expected):
|
||||
assert parse_gpu_layers_override(args) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args",
|
||||
[
|
||||
["--gpu-layers"],
|
||||
["--gpu-layers", "--top-k"],
|
||||
["--gpu-layers", "abc"],
|
||||
["--gpu-layers=-2"],
|
||||
],
|
||||
)
|
||||
def test_parse_gpu_layers_override_rejects_malformed_values(args):
|
||||
with pytest.raises(ValueError, match = "gpu-layers|GPU layers"):
|
||||
parse_gpu_layers_override(args)
|
||||
|
||||
|
||||
def test_validate_extra_args_rejects_malformed_gpu_layers_override():
|
||||
with pytest.raises(ValueError, match = "GPU layers"):
|
||||
validate_extra_args(["-ngl", "abc"])
|
||||
|
||||
|
||||
# ── parse_cache_override ─────────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
|
@ -376,6 +379,128 @@ def test_worker_share_object_receives_distributed_payload(monkeypatch):
|
|||
assert response["object"] == shared_obj
|
||||
|
||||
|
||||
def test_worker_activates_mlx_sidecar_before_hardware_detection(tmp_path):
|
||||
backend_dir = Path(__file__).resolve().parent.parent
|
||||
fake_modules = tmp_path / "base"
|
||||
sidecar = tmp_path / ".venv_t5_530"
|
||||
packages = {
|
||||
fake_modules / "transformers" / "__init__.py": '__version__ = "4.57.6"\n',
|
||||
fake_modules / "mlx" / "__init__.py": "",
|
||||
fake_modules / "mlx" / "core.py": "",
|
||||
fake_modules / "mlx_lm" / "__init__.py": "import transformers\n",
|
||||
fake_modules / "mlx_lm" / "sample_utils.py": "",
|
||||
fake_modules / "mlx_vlm" / "__init__.py": "",
|
||||
sidecar / "transformers" / "__init__.py": '__version__ = "5.3.0"\n',
|
||||
}
|
||||
for path, contents in packages.items():
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_text(contents)
|
||||
|
||||
script = r"""
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.environ["FAKE_MODULES"])
|
||||
from core.inference import worker
|
||||
from utils.hardware import hardware
|
||||
import utils.mlx_repair as mlx_repair
|
||||
import utils.transformers_version as transformers_version
|
||||
|
||||
bootstrap_roots = sorted(
|
||||
{
|
||||
name.split(".", 1)[0]
|
||||
for name in sys.modules
|
||||
if name.split(".", 1)[0]
|
||||
in {
|
||||
"huggingface_hub",
|
||||
"mlx",
|
||||
"mlx_lm",
|
||||
"mlx_vlm",
|
||||
"torch",
|
||||
"transformers",
|
||||
"unsloth",
|
||||
"unsloth_zoo",
|
||||
}
|
||||
}
|
||||
)
|
||||
assert not bootstrap_roots, f"worker bootstrap imported ML modules: {bootstrap_roots}"
|
||||
|
||||
worker.is_apple_silicon = lambda: True
|
||||
hardware.is_apple_silicon = lambda: True
|
||||
hardware._has_torch = lambda: False
|
||||
mlx_repair._mlx_versions_satisfy_minimums = lambda: True
|
||||
transformers_version._VENV_T5_530_DIR = os.environ["SIDECAR"]
|
||||
transformers_version._ensure_venv_t5_530_exists = lambda: True
|
||||
|
||||
observed = {"bootstrap_roots": bootstrap_roots}
|
||||
|
||||
def capture_active_version(_backend, _config, _responses):
|
||||
module = sys.modules["transformers"]
|
||||
observed["active"] = module.__version__
|
||||
observed["file"] = module.__file__
|
||||
observed["device"] = hardware.DEVICE.value
|
||||
|
||||
class CommandQueue:
|
||||
def get(self, timeout):
|
||||
return {"type": "shutdown"}
|
||||
|
||||
class ResponseQueue:
|
||||
def put(self, _response):
|
||||
pass
|
||||
|
||||
worker._handle_load = capture_active_version
|
||||
worker.run_inference_process(
|
||||
cmd_queue = CommandQueue(),
|
||||
resp_queue = ResponseQueue(),
|
||||
cancel_event = None,
|
||||
config = {
|
||||
"model_name": "Ministral-3-regression",
|
||||
"hf_token": "",
|
||||
"resolved_gpu_ids": None,
|
||||
"device_backend": "mlx",
|
||||
},
|
||||
)
|
||||
observed["tier"] = transformers_version.get_transformers_tier(
|
||||
"Ministral-3-regression"
|
||||
)
|
||||
print("RESULT " + json.dumps(observed, sort_keys = True))
|
||||
"""
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
cwd = backend_dir,
|
||||
env = {
|
||||
**__import__("os").environ,
|
||||
"FAKE_MODULES": str(fake_modules),
|
||||
"SIDECAR": str(sidecar),
|
||||
"UNSLOTH_STUDIO_HOME": str(tmp_path),
|
||||
"HF_HOME": str(tmp_path / "hf"),
|
||||
"HF_HUB_CACHE": str(tmp_path / "hf" / "hub"),
|
||||
"HF_HUB_OFFLINE": "1",
|
||||
"TRANSFORMERS_OFFLINE": "1",
|
||||
},
|
||||
capture_output = True,
|
||||
text = True,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
result_line = next(
|
||||
(
|
||||
line.removeprefix("RESULT ")
|
||||
for line in result.stdout.splitlines()
|
||||
if line.startswith("RESULT ")
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert result_line is not None, result.stdout + result.stderr
|
||||
observed = json.loads(result_line)
|
||||
assert observed["bootstrap_roots"] == []
|
||||
assert observed["tier"] == "530"
|
||||
assert observed["device"] == "mlx"
|
||||
assert observed["active"] == "5.3.0"
|
||||
assert observed["file"] == str(sidecar / "transformers" / "__init__.py")
|
||||
|
||||
|
||||
def test_worker_share_object_oversize_notifies_peers(monkeypatch):
|
||||
from core.inference import worker
|
||||
|
||||
|
|
|
|||
|
|
@ -37,6 +37,23 @@ def test_directory_path_uses_basename():
|
|||
assert public_model_id("a/b/c") == "c"
|
||||
|
||||
|
||||
def test_hf_cache_snapshot_recovers_the_repo_id():
|
||||
from core.inference.model_ids import hf_cache_repo_id
|
||||
|
||||
# The snapshot basename is a commit sha, so recover org/name instead.
|
||||
snapshot = (
|
||||
"/home/u/.cache/huggingface/hub/models--unsloth--gemma-4-31B-it-GGUF"
|
||||
"/snapshots/c1ac76e99d5513b141e8adde7288b85c3f9c32ec"
|
||||
)
|
||||
assert public_model_id(snapshot) == "unsloth/gemma-4-31B-it-GGUF"
|
||||
# A file inside the snapshot resolves the same way, not to the file stem.
|
||||
assert public_model_id(snapshot + "/gemma-4-31B-it-UD-Q5_K_XL.gguf") == (
|
||||
"unsloth/gemma-4-31B-it-GGUF"
|
||||
)
|
||||
assert hf_cache_repo_id("/opt/models/plain.gguf") is None
|
||||
assert hf_cache_repo_id(None) is None
|
||||
|
||||
|
||||
def test_relative_and_home_paths_are_sanitized():
|
||||
# ./ ../ ~ prefixed paths are local and must not be echoed raw.
|
||||
assert public_model_id("./model.gguf") == "model"
|
||||
|
|
|
|||
1782
studio/backend/tests/test_openai_auto_download.py
Normal file
1782
studio/backend/tests/test_openai_auto_download.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -11,6 +11,7 @@ import asyncio
|
|||
import os
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import routes.inference as inference_route
|
||||
from models.inference import LoadRequest
|
||||
|
|
@ -18,6 +19,18 @@ from core.inference import local_model_resolver as resolver
|
|||
from utils import openai_auto_switch_settings as settings
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _clean_resolver_index():
|
||||
"""Drop the scan cache around every test.
|
||||
|
||||
The /v1 admission hook warms the index in the background, so a test exercising it
|
||||
can publish its fixture's scan and, inside the TTL, hand it to the next test.
|
||||
"""
|
||||
resolver.invalidate_index()
|
||||
yield
|
||||
resolver.invalidate_index()
|
||||
|
||||
|
||||
class _FakeBackend:
|
||||
effective_parallel_slots = 1
|
||||
_slot_save_binary = None
|
||||
|
|
@ -94,7 +107,7 @@ class _LoadRecorder:
|
|||
|
||||
def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder):
|
||||
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled)
|
||||
monkeypatch.setattr(resolver, "resolve_local_gguf", lambda _m: resolves_to)
|
||||
monkeypatch.setattr(resolver, "resolve_local_gguf", lambda _m, **_kw: resolves_to)
|
||||
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
|
||||
# Auto-switch loads via _load_model_impl (the /load route holds the lifecycle
|
||||
# gate that auto-switch already owns, so it calls the impl directly).
|
||||
|
|
@ -116,7 +129,11 @@ def test_flag_off_never_loads(monkeypatch):
|
|||
backend = backend,
|
||||
recorder = rec,
|
||||
)
|
||||
_run_hook("unsloth/B-GGUF")
|
||||
# Off means no load, but A must not answer as B either: say why instead.
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
_run_hook("unsloth/B-GGUF")
|
||||
assert excinfo.value.status_code == 404
|
||||
assert "Switch model by request" in str(excinfo.value.detail)
|
||||
assert rec.calls == []
|
||||
|
||||
|
||||
|
|
@ -387,6 +404,45 @@ def test_resolver_nonstring_model_is_failsafe():
|
|||
assert resolver.resolve_local_gguf(None) is None
|
||||
|
||||
|
||||
def test_describe_local_miss_separates_missing_repo_from_missing_quant(monkeypatch):
|
||||
# Two different misses: the repo isn't downloaded, or only that quant is absent.
|
||||
monkeypatch.setattr(
|
||||
resolver,
|
||||
"_build_index",
|
||||
lambda: {"unsloth/b-gguf": _entry("unsloth/B-GGUF", "UD-Q5_K_XL", "Q4_K_M")},
|
||||
)
|
||||
resolver._scan = (0.0, {})
|
||||
assert resolver.describe_local_miss("unsloth/B-GGUF:Q8_0") == (
|
||||
resolver.MISS_VARIANT_NOT_FOUND,
|
||||
("UD-Q5_K_XL", "Q4_K_M"),
|
||||
)
|
||||
# Split the same way resolve_local_gguf does, so the two never disagree.
|
||||
assert resolver.describe_local_miss("unsloth/b-gguf:q8_0")[0] == (
|
||||
resolver.MISS_VARIANT_NOT_FOUND
|
||||
)
|
||||
# Unknown repo, and a bare id with no ":VARIANT" to blame.
|
||||
assert resolver.describe_local_miss("totally/unknown:Q8_0") == (
|
||||
resolver.MISS_MODEL_NOT_FOUND,
|
||||
(),
|
||||
)
|
||||
assert resolver.describe_local_miss("unsloth/B-GGUF") == (resolver.MISS_MODEL_NOT_FOUND, ())
|
||||
|
||||
|
||||
def test_describe_local_miss_is_failsafe(monkeypatch):
|
||||
# Runs inside an error path, so a broken scan must degrade, not turn a 4xx into a 500.
|
||||
def boom():
|
||||
raise RuntimeError("scan blew up")
|
||||
|
||||
monkeypatch.setattr(resolver, "_build_index", boom)
|
||||
resolver._scan = (0.0, {})
|
||||
assert resolver.describe_local_miss("unsloth/B-GGUF:Q8_0") == (
|
||||
resolver.MISS_MODEL_NOT_FOUND,
|
||||
(),
|
||||
)
|
||||
assert resolver.describe_local_miss(123) == (resolver.MISS_MODEL_NOT_FOUND, ())
|
||||
assert resolver.describe_local_miss("") == (resolver.MISS_MODEL_NOT_FOUND, ())
|
||||
|
||||
|
||||
def test_resolver_exact_id_with_colon_wins(monkeypatch):
|
||||
# A local id that itself contains a colon (e.g. a Windows path) must match
|
||||
# exactly rather than being split at the drive-letter colon.
|
||||
|
|
@ -537,7 +593,9 @@ def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path):
|
|||
"dir": str(tmp_path),
|
||||
"slots": [{"id": 0, "filename": saved.name}],
|
||||
}
|
||||
monkeypatch.setattr(settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True))
|
||||
monkeypatch.setattr(
|
||||
settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True, False)
|
||||
)
|
||||
monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0)
|
||||
|
||||
payload = settings_route.OpenAIAutoSwitchPayload(enabled = False)
|
||||
|
|
@ -1548,22 +1606,28 @@ def test_load_route_holds_lifecycle_gate(monkeypatch):
|
|||
|
||||
|
||||
def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded():
|
||||
# Both replacement directions drain active inference, then recheck whether a
|
||||
# sidecar install reserved the lifecycle gate during that wait. Exact-model
|
||||
# reuse exits earlier, so an already-loaded model never waits on unrelated inference.
|
||||
# Both replacement directions drain, then recheck whether a sidecar install reserved the
|
||||
# gate meanwhile. That recheck is the last thing that can reject the load, so the
|
||||
# destructive cancel must follow it. Exact-model reuse exits earlier and never waits.
|
||||
import inspect
|
||||
|
||||
src = inspect.getsource(inference_route._load_model_impl)
|
||||
already_loaded = src.index('status = "already_loaded"')
|
||||
standard_branch = src.index("# ── Standard path")
|
||||
|
||||
gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:"))
|
||||
gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait)
|
||||
gguf_cancel = src.index("on_reload_confirmed(cancel = True)", gguf_wait)
|
||||
unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait)
|
||||
standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1)
|
||||
standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait)
|
||||
unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
|
||||
already_loaded = src.index('status = "already_loaded"')
|
||||
|
||||
assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth
|
||||
assert standard_wait < standard_sidecar_check < unload_gguf
|
||||
standard_wait = src.index("await _wait_for_model_switch_idle", standard_branch)
|
||||
standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait)
|
||||
standard_cancel = src.index("on_reload_confirmed(cancel = True)", standard_wait)
|
||||
unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
|
||||
|
||||
assert already_loaded < gguf_wait < gguf_sidecar_check < gguf_cancel < unload_unsloth
|
||||
assert standard_branch < standard_wait < standard_sidecar_check
|
||||
assert standard_sidecar_check < standard_cancel < unload_gguf
|
||||
|
||||
|
||||
def test_switch_waiter_deregisters_before_swap_gate_release():
|
||||
|
|
@ -1877,7 +1941,10 @@ def test_env_idle_standalone_reloads_freed_model_with_auto_switch_off(monkeypatc
|
|||
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) # standalone env TTL
|
||||
monkeypatch.setattr(kw, "_inflight", 0)
|
||||
monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF"))
|
||||
_run_hook("org/B-GGUF")
|
||||
# A is restored, but the request named B, so it is told so rather than served A.
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
_run_hook("org/B-GGUF")
|
||||
assert excinfo.value.status_code == 404
|
||||
# Resolver skipped (auto-switch off), so only the stash reload runs: the freed A
|
||||
# is restored, not the resolves_to target B.
|
||||
assert len(rec.calls) == 1
|
||||
|
|
@ -2947,9 +3014,13 @@ def test_require_vision_ignores_reload_stash(monkeypatch):
|
|||
monkeypatch.setattr(
|
||||
inference_route, "_target_is_vision", lambda _p: False
|
||||
) # would reject if used
|
||||
asyncio.run(
|
||||
inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True)
|
||||
)
|
||||
# 404 because the restored A is not the requested B, whose quant makes it a real reference.
|
||||
with pytest.raises(HTTPException):
|
||||
asyncio.run(
|
||||
inference_route._maybe_auto_switch_model(
|
||||
"org/B-GGUF:UD-Q6_K_XL", object(), "t", require_vision = True
|
||||
)
|
||||
)
|
||||
assert len(rec.calls) == 1
|
||||
assert rec.calls[0].model_path == "/cache/snap/A" # restored despite require_vision
|
||||
|
||||
|
|
@ -3290,13 +3361,19 @@ def test_no_model_loaded_detail_appends_hint_only_when_off(monkeypatch):
|
|||
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.
|
||||
def _run_responses_stream_no_model(
|
||||
monkeypatch,
|
||||
*,
|
||||
enabled,
|
||||
active_model_name,
|
||||
resolves_to = None,
|
||||
):
|
||||
# Drive _responses_stream's GGUF-not-loaded guard. Returns (status, detail).
|
||||
from fastapi import HTTPException
|
||||
from models.inference import ResponsesRequest, ChatMessage
|
||||
|
||||
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled)
|
||||
monkeypatch.setattr(resolver, "resolve_local_gguf", lambda name: resolves_to)
|
||||
monkeypatch.setattr(
|
||||
inference_route, "get_llama_cpp_backend", lambda: _FakeBackend(loaded_id = None)
|
||||
)
|
||||
|
|
@ -3309,29 +3386,230 @@ def _run_responses_stream_no_model(monkeypatch, *, enabled, active_model_name):
|
|||
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
|
||||
return exc.value.status_code, 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)
|
||||
# The hint attaches whenever the toggle is off, whatever is active. With it on the name
|
||||
# resolved to nothing local, so 404 rather than 400.
|
||||
off_status, hinted = _run_responses_stream_no_model(
|
||||
monkeypatch, enabled = False, active_model_name = None
|
||||
)
|
||||
assert off_status == 400
|
||||
assert "Model auto-switch" in hinted
|
||||
|
||||
on = _run_responses_stream_no_model(monkeypatch, enabled = True, active_model_name = None)
|
||||
on_status, on = _run_responses_stream_no_model(
|
||||
monkeypatch, enabled = True, active_model_name = None
|
||||
)
|
||||
assert on_status == 404
|
||||
assert "Model auto-switch" not in on
|
||||
assert "unsloth/Qwen3.5-4B-GGUF" in on
|
||||
|
||||
non_gguf_loaded = _run_responses_stream_no_model(
|
||||
non_gguf_status, non_gguf_loaded = _run_responses_stream_no_model(
|
||||
monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct"
|
||||
)
|
||||
assert non_gguf_status == 400
|
||||
assert "Model auto-switch" in non_gguf_loaded
|
||||
|
||||
|
||||
def _wire_unloaded_chat(
|
||||
monkeypatch,
|
||||
*,
|
||||
enabled,
|
||||
catalog = ("org/A-GGUF", "org/B-GGUF"),
|
||||
):
|
||||
# Nothing loaded, so a chat request hits "no model loaded". Pin the catalog for determinism.
|
||||
async def _catalog():
|
||||
return [{"id": mid} for mid in catalog]
|
||||
|
||||
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled)
|
||||
monkeypatch.setattr(resolver, "resolve_local_gguf", lambda _m, **_kw: None)
|
||||
monkeypatch.setattr(
|
||||
resolver, "describe_local_miss", lambda _m: (resolver.MISS_MODEL_NOT_FOUND, ())
|
||||
)
|
||||
monkeypatch.setattr(inference_route, "_openai_catalog_objects", _catalog)
|
||||
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": None, "models": {}})(),
|
||||
)
|
||||
|
||||
|
||||
def _chat_error(payload):
|
||||
from fastapi import HTTPException
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester"))
|
||||
return exc.value.status_code, exc.value.detail
|
||||
|
||||
|
||||
def test_chat_names_undownloaded_model_404s_with_available_ids(monkeypatch):
|
||||
# The reported bug: the model is not here, so the switch did nothing and /inference/load
|
||||
# cannot fix it. Name it and list what can serve.
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True)
|
||||
status, detail = _chat_error(_chat_request(model = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL"))
|
||||
assert status == 404
|
||||
assert "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL" in detail
|
||||
assert "org/A-GGUF, org/B-GGUF" in detail
|
||||
assert "GET /v1/models" in detail
|
||||
assert "POST /inference/load" not in detail
|
||||
|
||||
|
||||
def test_chat_undownloaded_model_with_empty_catalog(monkeypatch):
|
||||
# Nothing downloaded: an empty list would read as a bug, so say so plainly.
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True, catalog = ())
|
||||
status, detail = _chat_error(_chat_request(model = "org/nope-GGUF"))
|
||||
assert status == 404
|
||||
assert "no models are downloaded yet" in detail
|
||||
|
||||
|
||||
def test_chat_wrong_quant_lists_the_local_quants(monkeypatch):
|
||||
# Repo downloaded, only the quant missing: sibling quants, not the catalog.
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True)
|
||||
monkeypatch.setattr(
|
||||
resolver,
|
||||
"describe_local_miss",
|
||||
lambda _m: (resolver.MISS_VARIANT_NOT_FOUND, ("Q4_K_M", "Q8_0")),
|
||||
)
|
||||
status, detail = _chat_error(_chat_request(model = "org/A-GGUF:UD-Q5_K_XL"))
|
||||
assert status == 404
|
||||
assert "'org/A-GGUF' is downloaded, but the quant 'UD-Q5_K_XL' is not" in detail
|
||||
assert "Q4_K_M, Q8_0" in detail
|
||||
|
||||
|
||||
def test_chat_error_unchanged_when_auto_switch_off(monkeypatch):
|
||||
# Toggle off: nothing resolved, so keep the pre-existing status and text, hint included.
|
||||
_wire_unloaded_chat(monkeypatch, enabled = False)
|
||||
status, detail = _chat_error(_chat_request(model = "org/nope-GGUF"))
|
||||
assert status == 400
|
||||
assert detail.startswith("No model loaded. Call POST /inference/load first.")
|
||||
assert "Model auto-switch" in detail
|
||||
|
||||
|
||||
def test_chat_error_unchanged_when_no_model_named(monkeypatch):
|
||||
# An omitted model means "serve whatever is loaded", so there is no name to report.
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True)
|
||||
status, detail = _chat_error(_chat_request())
|
||||
assert status == 400
|
||||
assert detail == "No model loaded. Call POST /inference/load first."
|
||||
|
||||
|
||||
def test_chat_not_downloaded_error_survives_a_broken_catalog_scan(monkeypatch):
|
||||
# Layered onto an already-failing path, so a broken scan must not make it a 500.
|
||||
async def _boom():
|
||||
raise RuntimeError("catalog scan blew up")
|
||||
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True)
|
||||
monkeypatch.setattr(inference_route, "_openai_catalog_objects", _boom)
|
||||
status, detail = _chat_error(_chat_request(model = "org/nope-GGUF"))
|
||||
assert status == 400
|
||||
assert detail.startswith("No model loaded. Call POST /inference/load first.")
|
||||
|
||||
|
||||
def test_chat_available_id_list_is_capped(monkeypatch):
|
||||
# A machine with 40 GGUFs must not print all 40 into a terminal error.
|
||||
_wire_unloaded_chat(
|
||||
monkeypatch, enabled = True, catalog = tuple(f"org/m{i:02d}-GGUF" for i in range(20))
|
||||
)
|
||||
status, detail = _chat_error(_chat_request(model = "org/nope-GGUF"))
|
||||
assert status == 404
|
||||
assert "and 12 more" in detail
|
||||
assert "org/m08-GGUF" not in detail
|
||||
|
||||
|
||||
def test_anthropic_undownloaded_model_uses_the_anthropic_envelope(monkeypatch):
|
||||
# Shared with /v1/messages, so the 404 must not leak an OpenAI-shaped body.
|
||||
from fastapi import HTTPException
|
||||
|
||||
async def _noop_switch(*a, **k):
|
||||
return None
|
||||
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True)
|
||||
monkeypatch.setattr(inference_route, "_automatic_model_load_may_run", lambda: True)
|
||||
monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _noop_switch)
|
||||
|
||||
request = type("_R", (), {"url": type("_U", (), {"path": "/v1/messages"})()})()
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(inference_route.anthropic_messages(_anthropic_payload(64), request, "tester"))
|
||||
assert exc.value.status_code == 404
|
||||
body = exc.value.detail
|
||||
assert body["type"] == "error"
|
||||
assert body["error"]["type"] == "not_found_error"
|
||||
assert "claude-x" in body["error"]["message"]
|
||||
|
||||
|
||||
def test_chat_undownloaded_model_uses_the_openai_envelope(monkeypatch):
|
||||
# The OpenAI surface carries param/code so SDK clients can branch on it.
|
||||
from fastapi import HTTPException
|
||||
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True)
|
||||
request = type("_R", (), {"url": type("_U", (), {"path": "/v1/chat/completions"})()})()
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(
|
||||
inference_route.openai_chat_completions(
|
||||
_chat_request(model = "org/nope-GGUF"), request, "tester"
|
||||
)
|
||||
)
|
||||
assert exc.value.status_code == 404
|
||||
err = exc.value.detail["error"]
|
||||
assert err["type"] == "not_found_error"
|
||||
assert err["code"] == "model_not_found"
|
||||
assert err["param"] == "model"
|
||||
|
||||
|
||||
def test_gguf_only_paths_keep_the_generic_error_for_the_resident_non_gguf_model(monkeypatch):
|
||||
# resolve_local_gguf misses a resident Transformers model the catalog does list, so
|
||||
# "not downloaded" would contradict itself.
|
||||
resident = "unsloth/Qwen3.5-4B-GGUF" # the id _run_responses_stream_no_model asks for
|
||||
|
||||
async def _catalog():
|
||||
return [{"id": resident}]
|
||||
|
||||
monkeypatch.setattr(inference_route, "_openai_catalog_objects", _catalog)
|
||||
status, detail = _run_responses_stream_no_model(
|
||||
monkeypatch, enabled = True, active_model_name = resident
|
||||
)
|
||||
assert status == 400
|
||||
assert "requires a GGUF model" in detail
|
||||
assert "not downloaded" not in detail
|
||||
|
||||
|
||||
def test_completions_keeps_the_generic_error_for_the_resident_non_gguf_model(monkeypatch):
|
||||
# Same contradiction on the raw-body surface, via _auto_switch_from_request_body.
|
||||
from fastapi import HTTPException
|
||||
|
||||
resident = "unsloth/Llama-3.2-1B-Instruct"
|
||||
_wire_unloaded_chat(monkeypatch, enabled = True, catalog = (resident,))
|
||||
monkeypatch.setattr(
|
||||
inference_route,
|
||||
"get_inference_backend",
|
||||
lambda: type("_B", (), {"active_model_name": resident, "models": {}})(),
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
asyncio.run(
|
||||
inference_route.openai_completions(
|
||||
_json_body_request({"model": resident, "prompt": "hi"}), "tester"
|
||||
)
|
||||
)
|
||||
assert exc.value.status_code == 503
|
||||
assert exc.value.detail.startswith("No GGUF model loaded.")
|
||||
assert "not downloaded" not in exc.value.detail
|
||||
|
||||
|
||||
def test_responses_stream_keeps_generic_error_when_target_is_local(monkeypatch):
|
||||
# Resolves locally yet nothing is loaded: the switch failed, so keep the generic 400.
|
||||
status, detail = _run_responses_stream_no_model(
|
||||
monkeypatch,
|
||||
enabled = True,
|
||||
active_model_name = None,
|
||||
resolves_to = ("/p/A", "Q4_K_M", "unsloth/Qwen3.5-4B-GGUF"),
|
||||
)
|
||||
assert status == 400
|
||||
assert "not downloaded" not in detail
|
||||
|
||||
|
||||
# ── idle-unload KV persistence (slot save/restore) ──────────────────
|
||||
|
||||
|
||||
|
|
@ -3784,10 +4062,11 @@ def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch):
|
|||
monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600")
|
||||
|
||||
assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None
|
||||
enabled, idle, keep_kv = settings.set_openai_auto_switch(False, None, False)
|
||||
enabled, idle, keep_kv, auto_dl = settings.set_openai_auto_switch(False, None, False)
|
||||
assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched
|
||||
assert settings.OPENAI_AUTO_DOWNLOAD_SETTING_KEY not in store # nor auto-download
|
||||
assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active
|
||||
assert (enabled, idle, keep_kv) == (False, 600, False)
|
||||
assert (enabled, idle, keep_kv, auto_dl) == (False, 600, False, False)
|
||||
|
||||
|
||||
def test_load_impl_notes_loaded_with_backend_off_loop():
|
||||
|
|
@ -3869,3 +4148,233 @@ def test_env_idle_below_floor_is_clamped(monkeypatch):
|
|||
assert settings.get_auto_unload_idle_seconds() == 600
|
||||
monkeypatch.delenv(settings.MODEL_IDLE_TTL_ENV_VAR)
|
||||
assert settings.get_auto_unload_idle_seconds() == 0
|
||||
|
||||
|
||||
def test_a_tag_that_names_no_quant_resolves_to_the_repo(monkeypatch):
|
||||
# A downloaded but unloaded GGUF asked for as org/model:latest missed the resolver,
|
||||
# so the switch could not load it (404ing on a quant that was never a quant with
|
||||
# auto-download on, refusing with it off). A real quant that is not on disk must
|
||||
# still miss, or a swap would serve the wrong weights under the right name.
|
||||
from core.inference.local_model_resolver import _LocalGgufEntry
|
||||
|
||||
import time
|
||||
|
||||
entry = _LocalGgufEntry("org/model", "/srv/models/org--model", ("Q4_K_M",))
|
||||
# Fresh stamp so _index serves this instead of rescanning over it.
|
||||
monkeypatch.setattr(resolver, "_scan", (time.monotonic(), {"org/model": entry}))
|
||||
for tag in ("org/model:latest", "org/model:8b", "org/model"):
|
||||
assert resolver.resolve_local_gguf(tag) == (
|
||||
"/srv/models/org--model",
|
||||
"Q4_K_M",
|
||||
"org/model",
|
||||
)
|
||||
assert resolver.resolve_local_gguf("org/model:Q8_0") is None
|
||||
assert resolver.resolve_local_gguf("org/model:Q4_K_M") == (
|
||||
"/srv/models/org--model",
|
||||
"Q4_K_M",
|
||||
"org/model",
|
||||
)
|
||||
|
||||
|
||||
def test_any_finished_download_drops_the_resolver_cache(monkeypatch):
|
||||
# Only the API auto-download watcher invalidated, so a GGUF fetched in the Hub UI
|
||||
# stayed absent to the cache-only request path and the resident model answered.
|
||||
# Every worker exits through here.
|
||||
import logging
|
||||
|
||||
from hub.services import download_lifecycle
|
||||
|
||||
class _Proc:
|
||||
stderr = None
|
||||
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
class _Registry:
|
||||
def cancel_requested(self, key):
|
||||
return False
|
||||
|
||||
def drop_process(self, key, proc):
|
||||
return True
|
||||
|
||||
def get_job_metadata(self, key):
|
||||
return None
|
||||
|
||||
def set_job(self, key, state):
|
||||
self.state = state
|
||||
|
||||
resolver._scan = (1234.0, {"already-here": "entry"})
|
||||
assert (
|
||||
download_lifecycle.finalize_worker_exit(
|
||||
_Registry(),
|
||||
"org/model:Q4_K_M",
|
||||
_Proc(),
|
||||
hf_token = None,
|
||||
label = "org/model",
|
||||
log_prefix = "[test]",
|
||||
logger = logging.getLogger(__name__),
|
||||
repo_type = "model",
|
||||
repo_id = "org/model",
|
||||
)
|
||||
== "complete"
|
||||
)
|
||||
stamp, entries = resolver._scan
|
||||
assert stamp == 0.0, "a finished download left the scan looking fresh"
|
||||
# Evidence for models already indexed has to survive, or a bare request for one
|
||||
# of them during the rebuild is answered by whatever is resident.
|
||||
assert entries == {"already-here": "entry"}
|
||||
|
||||
|
||||
def test_invalidating_keeps_the_entries_it_already_had(monkeypatch):
|
||||
# The request path reads this cache without scanning, so emptying it leaves no
|
||||
# evidence until the rebuild lands. Only a completed download invalidates, and
|
||||
# that only adds, so the entries stay true.
|
||||
import time
|
||||
|
||||
entry = resolver._LocalGgufEntry("org/old", "/srv/models/org--old", ("Q4_K_M",))
|
||||
monkeypatch.setattr(resolver, "_scan", (time.monotonic(), {"org/old": entry}))
|
||||
resolver.invalidate_index()
|
||||
assert resolver._scan[0] == 0.0
|
||||
assert resolver.resolve_local_gguf("org/old", allow_scan = False) == (
|
||||
"/srv/models/org--old",
|
||||
"Q4_K_M",
|
||||
"org/old",
|
||||
)
|
||||
|
||||
|
||||
def test_a_bare_local_id_takes_the_quant_a_plain_load_would(monkeypatch, tmp_path):
|
||||
# list_local_gguf_variants orders by descending size, so the head is the biggest
|
||||
# quant. Resolving a bare id to that could evict a working model and then OOM on an
|
||||
# F16 next to a fitting Q4, and /v1/models advertised the same head for pinning.
|
||||
from core.inference.local_model_resolver import _local_gguf_entry
|
||||
|
||||
for name, size in (("model-F16.gguf", 900), ("model-Q4_K_M.gguf", 100)):
|
||||
(tmp_path / name).write_bytes(b"\0" * size)
|
||||
entry = _local_gguf_entry("org/model", type("I", (), {"path": str(tmp_path)})())
|
||||
assert entry is not None
|
||||
assert set(entry.variants) == {"F16", "Q4_K_M"}
|
||||
assert entry.variants[0] == "Q4_K_M", "a bare id would have resolved to F16"
|
||||
|
||||
|
||||
def test_local_and_remote_agree_on_the_preferred_quant():
|
||||
# A bare id must mean the same quant whichever side answered it.
|
||||
from core.inference.openai_auto_download import _match_variant, preferred_quant
|
||||
|
||||
labels = ("F16", "Q8_0", "UD-Q4_K_XL", "Q4_K_M")
|
||||
assert preferred_quant(labels) == _match_variant(None, dict.fromkeys(labels, 1))
|
||||
assert preferred_quant(labels) not in ("F16",)
|
||||
|
||||
|
||||
def test_a_just_downloaded_model_is_evidence_before_the_scan_indexes_it(monkeypatch):
|
||||
# The retained index covers what was known, but nothing covers the model that just
|
||||
# landed until the next scan: a bare request for it was answered by the resident one.
|
||||
import logging
|
||||
|
||||
from hub.services import download_lifecycle
|
||||
|
||||
class _Proc:
|
||||
stderr = None
|
||||
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
class _Registry:
|
||||
def cancel_requested(self, key):
|
||||
return False
|
||||
|
||||
def drop_process(self, key, proc):
|
||||
return True
|
||||
|
||||
def get_job_metadata(self, key):
|
||||
return None
|
||||
|
||||
def set_job(self, key, state):
|
||||
pass
|
||||
|
||||
assert not resolver.recently_downloaded("org/fresh")
|
||||
download_lifecycle.finalize_worker_exit(
|
||||
_Registry(),
|
||||
"org/fresh:Q4_K_M",
|
||||
_Proc(),
|
||||
hf_token = None,
|
||||
label = "org/fresh",
|
||||
log_prefix = "[test]",
|
||||
logger = logging.getLogger(__name__),
|
||||
repo_type = "model",
|
||||
repo_id = "org/fresh",
|
||||
)
|
||||
assert resolver.recently_downloaded("org/fresh"), "no evidence for the new model"
|
||||
assert resolver.recently_downloaded("ORG/Fresh"), "evidence must be case-insensitive"
|
||||
assert not resolver.recently_downloaded("org/other")
|
||||
|
||||
# The scan that indexes it supersedes the note.
|
||||
monkeypatch.setattr(resolver, "_build_index", dict)
|
||||
resolver._index()
|
||||
assert not resolver.recently_downloaded("org/fresh")
|
||||
|
||||
|
||||
def test_a_finished_dataset_is_not_recorded_as_a_local_model(monkeypatch):
|
||||
# finalize_worker_exit is shared with dataset downloads. Noting one as a local model
|
||||
# would refuse a bare /v1 request naming that id instead of letting a foreign id
|
||||
# fall through, and would kick off a multi-directory scan for nothing.
|
||||
import logging
|
||||
import time
|
||||
|
||||
from hub.services import download_lifecycle
|
||||
|
||||
class _Proc:
|
||||
stderr = None
|
||||
|
||||
def wait(self):
|
||||
return 0
|
||||
|
||||
class _Registry:
|
||||
def cancel_requested(self, key):
|
||||
return False
|
||||
|
||||
def drop_process(self, key, proc):
|
||||
return True
|
||||
|
||||
def get_job_metadata(self, key):
|
||||
return None
|
||||
|
||||
def set_job(self, key, state):
|
||||
pass
|
||||
|
||||
stamp = time.monotonic()
|
||||
monkeypatch.setattr(resolver, "_scan", (stamp, {"kept": "entry"}))
|
||||
download_lifecycle.finalize_worker_exit(
|
||||
_Registry(),
|
||||
"org/corpus",
|
||||
_Proc(),
|
||||
hf_token = None,
|
||||
label = "org/corpus",
|
||||
log_prefix = "[test]",
|
||||
logger = logging.getLogger(__name__),
|
||||
repo_type = "dataset",
|
||||
repo_id = "org/corpus",
|
||||
)
|
||||
assert not resolver.recently_downloaded("org/corpus")
|
||||
assert resolver._scan == (stamp, {"kept": "entry"}), "a dataset invalidated the index"
|
||||
|
||||
|
||||
def test_two_local_paths_differing_only_in_case_are_not_the_same_model(monkeypatch):
|
||||
# _loaded_satisfies lowercased the request and every backend identifier, so on a
|
||||
# case-sensitive filesystem /srv/models/foo.gguf read as satisfied by a resident
|
||||
# /srv/models/Foo.gguf. A repo alias must still stay case-insensitive.
|
||||
import os
|
||||
|
||||
loaded = _FakeBackend(loaded_id = "/srv/models/Foo.gguf")
|
||||
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: loaded)
|
||||
monkeypatch.setattr(
|
||||
inference_route,
|
||||
"get_inference_backend",
|
||||
lambda: type("B", (), {"active_model_name": None})(),
|
||||
)
|
||||
assert inference_route._loaded_satisfies("/srv/models/Foo.gguf") is True
|
||||
same = os.path.normcase("A") == os.path.normcase("a")
|
||||
assert inference_route._loaded_satisfies("/srv/models/foo.gguf") is same
|
||||
|
||||
alias = _FakeBackend(loaded_id = "unsloth/Qwen3-4B-GGUF")
|
||||
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: alias)
|
||||
assert inference_route._loaded_satisfies("unsloth/qwen3-4b-gguf") is True
|
||||
|
|
|
|||
|
|
@ -64,8 +64,10 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
|
|||
]
|
||||
|
||||
monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog)
|
||||
# GGUF-ness is read from the on-disk files; drive it off each info's flag here.
|
||||
monkeypatch.setattr(resolver, "info_has_local_gguf", lambda info: info.is_gguf)
|
||||
# GGUF-ness and the quant labels come from one on-disk scan; drive both off the flag.
|
||||
monkeypatch.setattr(
|
||||
resolver, "local_gguf_quants", lambda info: ("Q8_0",) if info.is_gguf else None
|
||||
)
|
||||
|
||||
data = asyncio.run(inf._openai_catalog_objects())
|
||||
ids = {m["id"]: m for m in data}
|
||||
|
|
@ -73,8 +75,9 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
|
|||
# Loaded model is present, marked loaded, and keeps context fields.
|
||||
assert ids["Qwen3-Q4"]["loaded"] is True
|
||||
assert ids["Qwen3-Q4"]["context_length"] == 4096
|
||||
# Available-but-not-loaded GGUF models are listed too.
|
||||
# Not-loaded GGUFs are listed too, with the quant a client appends to pin them.
|
||||
assert ids["Llama-8B-Q8"]["loaded"] is False
|
||||
assert ids["Llama-8B-Q8"]["quant"] == "Q8_0"
|
||||
# The HF-cache GGUF is listed despite model_format being unset.
|
||||
assert ids["org/Foo"]["loaded"] is False
|
||||
# The non-GGUF model is filtered out (/v1 can never serve it).
|
||||
|
|
@ -205,3 +208,156 @@ def test_cached_local_catalog_offloads_and_caches(monkeypatch):
|
|||
assert second is first or [i.id for i in second] == [i.id for i in first]
|
||||
assert calls["scan"] == 1 # cached: scanned once for two calls
|
||||
assert calls["threaded"] == 1 # offloaded to a worker thread
|
||||
|
||||
|
||||
def test_monitor_active_model_is_a_public_id_not_a_host_path(monkeypatch):
|
||||
# The settings UI renders this and --secure serves it publicly, so never a load path.
|
||||
class _Llama:
|
||||
is_loaded = True
|
||||
model_identifier = "/home/me/.cache/huggingface/hub/models--org--A-GGUF/snapshots/abc"
|
||||
hf_variant = "UD-Q4_K_XL"
|
||||
_openai_advertised_id = "org/A-GGUF"
|
||||
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _Llama())
|
||||
assert inf._monitor_active_model() == "org/A-GGUF:UD-Q4_K_XL"
|
||||
|
||||
|
||||
def test_monitor_active_model_cleans_a_path_with_no_advertised_id(monkeypatch):
|
||||
class _Llama:
|
||||
is_loaded = True
|
||||
model_identifier = "/data/models/Llama-8B-Q8.gguf"
|
||||
hf_variant = None
|
||||
_openai_advertised_id = None
|
||||
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _Llama())
|
||||
label = inf._monitor_active_model()
|
||||
assert "/" not in label and ".gguf" not in label
|
||||
|
||||
|
||||
def test_lifecycle_label_recovers_the_repo_id_from_an_hf_cache_path():
|
||||
# An auto-switch load gets the snapshot dir, whose basename is a commit sha.
|
||||
snap = "/home/me/.cache/huggingface/hub/models--unsloth--gemma-4-E4B-it-GGUF/snapshots/bfc15c3"
|
||||
assert (
|
||||
inf._lifecycle_model_label(snap, "UD-Q4_K_XL") == "unsloth/gemma-4-E4B-it-GGUF:UD-Q4_K_XL"
|
||||
)
|
||||
|
||||
|
||||
def test_lifecycle_model_label_is_path_free():
|
||||
label = inf._lifecycle_model_label("/data/models/Llama-8B-Q8.gguf", "Q8_0")
|
||||
assert "/" not in label and ".gguf" not in label
|
||||
assert inf._lifecycle_model_label("org/A-GGUF", "Q4_K_M") == "org/A-GGUF:Q4_K_M"
|
||||
# An id that already carries a quant is not double-suffixed.
|
||||
assert inf._lifecycle_model_label("org/A-GGUF:Q4_K_M", "Q8_0") == "org/A-GGUF:Q4_K_M"
|
||||
|
||||
|
||||
def test_a_standalone_gguf_does_not_advertise_a_quant_that_stops_resolving(monkeypatch):
|
||||
# llama.cpp reads hf_variant off the filename, but the resolver stores standalone files
|
||||
# with no quants, so a pinned "<stem>:<quant>" would 404 once it is not resident.
|
||||
from core.inference.local_model_resolver import _LocalGgufEntry
|
||||
|
||||
standalone = _LocalGgufEntry("Qwen3-Q4", "/srv/models/Qwen3-Q4.gguf", ())
|
||||
repo = _LocalGgufEntry("org/Foo", "/hf/models--org--Foo/snapshots/a", ("Q4_K_M",))
|
||||
monkeypatch.setattr(resolver, "_scan", (1.0, {"qwen3-q4": standalone, "org/foo": repo}))
|
||||
monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth())
|
||||
|
||||
llama = _FakeLlama()
|
||||
llama.hf_variant = "Q4_K_M"
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: llama)
|
||||
assert "quant" not in inf._openai_model_objects()[0]
|
||||
|
||||
# The same quant on a repo the resolver does list stays advertised.
|
||||
llama.model_identifier = "org/Foo"
|
||||
assert inf._openai_model_objects()[0]["quant"] == "Q4_K_M"
|
||||
|
||||
# A cold index cannot prove the reference either, and publishing on no proof is
|
||||
# exactly what hands out the pin that later fails to resolve.
|
||||
monkeypatch.setattr(resolver, "_scan", (0.0, {}))
|
||||
# Stub the walk: a real multi-root scan inside the cold-wait budget makes this
|
||||
# test time out into a 503 under load instead of asserting what it is here for.
|
||||
monkeypatch.setattr(resolver, "_build_index", lambda: {})
|
||||
monkeypatch.setattr(resolver, "warm_index_soon", lambda: None)
|
||||
assert "quant" not in inf._openai_model_objects()[0]
|
||||
|
||||
|
||||
def test_a_loaded_alias_advertises_the_quant_that_is_actually_loaded(monkeypatch):
|
||||
# Marking the alias loaded while still publishing the preferred on-disk quant said
|
||||
# alias:Q4 was loaded while Q8 was serving, and pinning that 404s with switching off.
|
||||
monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth())
|
||||
llama = _FakeLlama()
|
||||
llama.hf_variant = "Q8_0"
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: llama)
|
||||
|
||||
alias = _Info("/srv/models", "Qwen3", model_id = "publisher/Qwen3")
|
||||
alias.path = "/srv/models" # holds the resident /srv/models/Qwen3-Q4.gguf
|
||||
|
||||
async def _fake_catalog():
|
||||
return [alias]
|
||||
|
||||
monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog)
|
||||
monkeypatch.setattr(resolver, "local_gguf_quants", lambda info: ("Q4_K_M", "Q8_0"))
|
||||
ids = {m["id"]: m for m in asyncio.run(inf._openai_catalog_objects())}
|
||||
assert ids["publisher/Qwen3"]["loaded"] is True
|
||||
assert ids["publisher/Qwen3"]["quant"] == "Q8_0"
|
||||
|
||||
|
||||
def test_a_nested_model_directory_is_not_the_resident_one(monkeypatch):
|
||||
# Two indexed models can nest (/models/A holding A, /models/A/sub/B holding B). A
|
||||
# plain prefix test made loading B mark A resident, so a request for A was answered
|
||||
# with B's weights. The innermost indexed model owns the file.
|
||||
outer = _Info("/models/A", "A", model_id = "publisher/A")
|
||||
outer.path = "/models/A"
|
||||
inner = _Info("/models/A/sub/B", "B", model_id = "publisher/B")
|
||||
inner.path = "/models/A/sub/B"
|
||||
monkeypatch.setitem(inf._CATALOG_CACHE, "models", [outer, inner])
|
||||
|
||||
llama = _FakeLlama()
|
||||
llama.gguf_path = "/models/A/sub/B/model-Q4_K_M.gguf"
|
||||
llama.model_identifier = llama.gguf_path
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: llama)
|
||||
monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth())
|
||||
|
||||
assert inf._resolves_to_resident("/models/A/sub/B") is True
|
||||
assert inf._resolves_to_resident("/models/A") is False
|
||||
# With nothing indexed there is no nesting to tell apart, so the directory-to-file
|
||||
# match this exists for must still hold.
|
||||
monkeypatch.setitem(inf._CATALOG_CACHE, "models", [])
|
||||
assert inf._resolves_to_resident("/models/A") is True
|
||||
|
||||
|
||||
def test_a_transformers_model_does_not_mark_a_gguf_alias_loaded(monkeypatch):
|
||||
# Every entry in this loop is advertised as GGUF with a GGUF quant. A Transformers
|
||||
# model live from a directory that also holds GGUF exports is not one, and marking
|
||||
# the alias loaded had the examples pin a quant nothing can serve with switching off.
|
||||
unsloth = _FakeUnsloth()
|
||||
unsloth.active_model_name = "/srv/models"
|
||||
monkeypatch.setattr(inf, "get_inference_backend", lambda: unsloth)
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama(loaded = False))
|
||||
|
||||
alias = _Info("/srv/models", "Qwen3", model_id = "publisher/Qwen3")
|
||||
alias.path = "/srv/models" # also holds /srv/models/Qwen3-Q4.gguf
|
||||
|
||||
async def _fake_catalog():
|
||||
return [alias]
|
||||
|
||||
monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog)
|
||||
monkeypatch.setattr(resolver, "local_gguf_quants", lambda info: ("Q4_K_M",))
|
||||
ids = {m["id"]: m for m in asyncio.run(inf._openai_catalog_objects())}
|
||||
assert ids["publisher/Qwen3"]["loaded"] is False
|
||||
|
||||
|
||||
def test_an_alias_for_the_resident_weights_is_not_listed_as_unloaded(monkeypatch):
|
||||
# A GGUF loaded by absolute path keys the resident entry by basename, so an id-only dedup
|
||||
# would emit the alias again marked not loaded.
|
||||
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama())
|
||||
monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth())
|
||||
|
||||
alias = _Info("/srv/models", "Qwen3", model_id = "publisher/Qwen3")
|
||||
alias.path = "/srv/models" # holds the resident /srv/models/Qwen3-Q4.gguf
|
||||
|
||||
async def _fake_catalog():
|
||||
return [alias]
|
||||
|
||||
monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog)
|
||||
monkeypatch.setattr(resolver, "local_gguf_quants", lambda info: ("Q4_K_M",))
|
||||
ids = {m["id"]: m for m in asyncio.run(inf._openai_catalog_objects())}
|
||||
assert ids["publisher/Qwen3"]["loaded"] is True
|
||||
|
|
|
|||
|
|
@ -1191,6 +1191,41 @@ class TestBuildPassthroughPayloadToolChoice:
|
|||
body = _build_passthrough_payload(**self._args(), tool_choice = tc)
|
||||
assert body["tool_choice"] == tc
|
||||
|
||||
def test_llama_incompatible_tool_constraints_are_omitted(self):
|
||||
args = self._args()
|
||||
schema = args["openai_tools"][0]["function"]["parameters"]
|
||||
schema["properties"] = {
|
||||
"declarationKey": {"type": "string", "pattern": r"\S"},
|
||||
"exactKey": {"type": "string", "pattern": r"^[A-Z]+$"},
|
||||
"nested": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"anyOf": [
|
||||
{"type": "string", "pattern": "token"},
|
||||
{"type": "string", "pattern": "^fixed$"},
|
||||
],
|
||||
"default": {"pattern": "annotation data"},
|
||||
},
|
||||
},
|
||||
"largeScript": {"type": "string", "minLength": 1, "maxLength": 65536},
|
||||
"boundedScript": {"type": "string", "maxLength": 2000},
|
||||
}
|
||||
|
||||
body = _build_passthrough_payload(**args)
|
||||
forwarded = body["tools"][0]["function"]["parameters"]["properties"]
|
||||
|
||||
assert forwarded["declarationKey"] == {"type": "string"}
|
||||
assert forwarded["exactKey"]["pattern"] == r"^[A-Z]+$"
|
||||
nested = forwarded["nested"]["items"]
|
||||
assert nested["anyOf"][0] == {"type": "string"}
|
||||
assert nested["anyOf"][1]["pattern"] == "^fixed$"
|
||||
assert nested["default"] == {"pattern": "annotation data"}
|
||||
assert forwarded["largeScript"] == {"type": "string", "minLength": 1}
|
||||
assert forwarded["boundedScript"]["maxLength"] == 2000
|
||||
assert schema["properties"]["declarationKey"]["pattern"] == r"\S"
|
||||
assert schema["properties"]["nested"]["items"]["anyOf"][0]["pattern"] == "token"
|
||||
assert schema["properties"]["largeScript"]["maxLength"] == 65536
|
||||
|
||||
def test_stream_omits_usage_options_when_client_did_not_request_them(self):
|
||||
args = self._args()
|
||||
args["stream"] = True
|
||||
|
|
@ -1611,7 +1646,7 @@ class TestOpenAICompatibilityHelpers:
|
|||
def test_openai_stream_error_sse_closes_with_done(self):
|
||||
error = {"error": {"message": "boom"}}
|
||||
assert _openai_stream_error_sse(error) == (
|
||||
'data: {"error": {"message": "boom"}}\n\n' "data: [DONE]\n\n"
|
||||
'data: {"error": {"message": "boom"}}\n\ndata: [DONE]\n\n'
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -4580,6 +4615,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"prompt": "hi", "stream": False}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
class FailingAsyncClient:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
|
@ -4587,14 +4625,18 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, *_args, **_kwargs):
|
||||
raise httpx.ConnectError("llama down")
|
||||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
# Per-request client so a forced swap can close it mid-call; the pooled one is shared.
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"nonstreaming_client",
|
||||
"_cancelable_nonstreaming_client",
|
||||
lambda: FailingAsyncClient(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -4632,9 +4674,15 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"prompt": "hi", "stream": False}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
captured = []
|
||||
|
||||
class CapturingClient:
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, _url, *, json, **_kwargs):
|
||||
captured.append(dict(json))
|
||||
return httpx.Response(
|
||||
|
|
@ -4652,7 +4700,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient())
|
||||
monkeypatch.setattr(
|
||||
inf_mod, "_cancelable_nonstreaming_client", lambda: CapturingClient()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"get_llama_cpp_backend",
|
||||
|
|
@ -4683,9 +4733,15 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"prompt": "hi", "stream": False, "max_tokens": 0}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
captured = []
|
||||
|
||||
class CapturingClient:
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, _url, *, json, **_kwargs):
|
||||
captured.append(dict(json))
|
||||
return httpx.Response(
|
||||
|
|
@ -4703,7 +4759,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient())
|
||||
monkeypatch.setattr(
|
||||
inf_mod, "_cancelable_nonstreaming_client", lambda: CapturingClient()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"get_llama_cpp_backend",
|
||||
|
|
@ -4741,6 +4799,7 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: UnusedClient())
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: UnusedClient())
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"get_llama_cpp_backend",
|
||||
|
|
@ -4845,6 +4904,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"input": ["alpha", "beta"], "model": "embed"}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
class FakeAsyncClient:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
|
@ -4852,6 +4914,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, *_args, **_kwargs):
|
||||
assert monitor.active_count() == 1
|
||||
return httpx.Response(
|
||||
|
|
@ -4864,9 +4929,10 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
# Per-request client so a forced swap can close it mid-call; the pooled one is shared.
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"nonstreaming_client",
|
||||
"_cancelable_nonstreaming_client",
|
||||
lambda: FakeAsyncClient(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -6337,7 +6403,7 @@ class TestApiMonitorSafetensorsUsage:
|
|||
}
|
||||
yield "safe reply"
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
pass
|
||||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
|
|
@ -6408,7 +6474,7 @@ class TestApiMonitorSafetensorsUsage:
|
|||
cancel_event.set()
|
||||
yield {"type": "content", "text": "ignored"}
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
pass
|
||||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
|
|
@ -6469,11 +6535,18 @@ class TestApiMonitorSafetensorsUsage:
|
|||
def generate_chat_completion_with_tools(self, **_kwargs):
|
||||
yield {"type": "content", "text": "unused"}
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
nonlocal reset_called
|
||||
reset_called = True
|
||||
|
||||
async def fake_to_thread(*_args, **_kwargs):
|
||||
async def fake_to_thread(
|
||||
func = None,
|
||||
*_args,
|
||||
**_kwargs,
|
||||
):
|
||||
# Only the generation hop should cancel; resolution runs before the row opens.
|
||||
if getattr(func, "__name__", "") == "resolve_local_gguf":
|
||||
return None
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
|
|
|
|||
|
|
@ -19,6 +19,10 @@ def _bare_orchestrator():
|
|||
"""An orchestrator without the real __init__ subprocess/network."""
|
||||
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
|
||||
o._gen_lock = threading.Lock()
|
||||
o._send_order_lock = threading.Lock()
|
||||
o._active_cancel_lock = threading.Lock()
|
||||
o._active_cancel_events = []
|
||||
o._executing_cancel_events = []
|
||||
o._cancel_event = threading.Event() # stands in for the mp.Event
|
||||
o._drain_event = threading.Event() # stands in for the unload-drain mp.Event
|
||||
o._proc = object() # truthy so _ensure_subprocess_alive reports alive
|
||||
|
|
@ -775,6 +779,7 @@ def test_dispatched_bails_when_unload_flips_before_mailbox_registration(monkeypa
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
monkeypatch.setattr(o, "_start_dispatcher", lambda: None)
|
||||
|
|
@ -817,6 +822,7 @@ def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeyp
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -846,6 +852,7 @@ def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(mo
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -872,6 +879,7 @@ def test_dispatched_happy_path_registers_and_sends(monkeypatch):
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1338,6 +1346,7 @@ def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch):
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = None # none running -> this call starts it
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1382,6 +1391,7 @@ def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch)
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = None
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1419,6 +1429,7 @@ def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch):
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher() # already running
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1545,6 +1556,7 @@ def test_concurrent_start_dispatcher_spawns_exactly_one():
|
|||
o._resp_queue = _queue.Queue() # real queue so the dispatcher loop blocks and stays alive
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._dispatcher_thread = None
|
||||
o._dispatcher_stop = threading.Event()
|
||||
o._dispatcher_lifecycle_lock = threading.Lock()
|
||||
|
|
@ -1660,6 +1672,7 @@ def test_queued_start_behind_unload_stop_spawns_no_dispatcher():
|
|||
o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._dispatcher_stop = threading.Event()
|
||||
o._dispatcher_lifecycle_lock = threading.Lock()
|
||||
o._unload_pending = False
|
||||
|
|
@ -1713,3 +1726,310 @@ def test_queued_start_behind_unload_stop_spawns_no_dispatcher():
|
|||
assert o._dispatcher_thread is None, "the stop cleared it and the queued start spawned nothing"
|
||||
live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()]
|
||||
assert live == [], "no fresh dispatcher may be left to consume the unloaded reply"
|
||||
|
||||
|
||||
def _dispatch(o, resps):
|
||||
"""Run the dispatcher over a fixed response list and stop it."""
|
||||
import queue as _queue
|
||||
|
||||
o._resp_queue = _queue.Queue()
|
||||
for r in resps:
|
||||
o._resp_queue.put(r)
|
||||
o._dispatcher_stop = threading.Event()
|
||||
t = threading.Thread(target = o._dispatcher_loop, daemon = True)
|
||||
t.start()
|
||||
deadline = time.monotonic() + 5.0
|
||||
while not o._resp_queue.empty() and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
o._dispatcher_stop.set()
|
||||
t.join(timeout = 5.0)
|
||||
|
||||
|
||||
def test_worker_ownership_follows_the_worker_not_the_consumer():
|
||||
# The subprocess runs one generation at a time and can start B while A's consumer has yet to
|
||||
# drain its mailbox. A must stop owning the worker the moment its gen_done is routed, else
|
||||
# a late Stop for A cancels B.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._mailboxes = {"a": _queue.Queue(), "b": _queue.Queue()}
|
||||
o._request_cancel_events = {"a": a_cancel, "b": b_cancel}
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
|
||||
_dispatch(o, [{"type": "token", "request_id": "a", "token": "hi"}])
|
||||
assert o._owns_worker(a_cancel), "the request the worker is answering owns it"
|
||||
assert not o._owns_worker(b_cancel), "a queued request does not"
|
||||
|
||||
# A finishes. B has been sent but has not answered yet (it is prefilling), so the gap
|
||||
# between the two is the window a late Stop for A used to fire into.
|
||||
_dispatch(o, [{"type": "gen_done", "request_id": "a"}])
|
||||
assert not o._owns_worker(a_cancel), "a finished request stops owning the worker"
|
||||
assert o._owns_worker(b_cancel), "the next queued request is the one prefilling"
|
||||
|
||||
# Worker moves on to B, still before A's consumer reads anything.
|
||||
_dispatch(o, [{"type": "token", "request_id": "b", "token": "yo"}])
|
||||
assert not o._owns_worker(a_cancel), "a finished request must not cancel its successor"
|
||||
assert o._owns_worker(b_cancel), "the worker moved on to B, so B owns it"
|
||||
|
||||
# A's own stream unwinding afterwards must not disturb B.
|
||||
o._release_worker(a_cancel)
|
||||
assert o._owns_worker(b_cancel)
|
||||
|
||||
|
||||
def test_status_responses_do_not_transfer_worker_ownership():
|
||||
# Status lines are not an answer to any request; the dispatcher drops them before routing.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._mailboxes = {"a": _queue.Queue(), "b": _queue.Queue()}
|
||||
o._request_cancel_events = {"a": a_cancel, "b": b_cancel}
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
|
||||
_dispatch(o, [{"type": "status", "request_id": "b", "message": "loading"}])
|
||||
# Nothing has answered, so the oldest claim is still the one prefilling.
|
||||
assert o._owns_worker(a_cancel)
|
||||
assert not o._owns_worker(b_cancel)
|
||||
|
||||
|
||||
def test_only_the_latest_responder_executes():
|
||||
# The subprocess runs one generation at a time, so answering B means it has left A.
|
||||
# _generate_inner promotes from its own consumer and can share the worker with a
|
||||
# dispatched request, so the two must not both count as executing.
|
||||
o = _bare_orchestrator()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
|
||||
o._mark_worker_started(a_cancel)
|
||||
assert o._owns_worker(a_cancel)
|
||||
o._mark_worker_started(b_cancel)
|
||||
assert o._owns_worker(b_cancel), "the latest responder is the one executing"
|
||||
assert not o._owns_worker(a_cancel), "and it is the only one"
|
||||
# Idempotent: more of B's own tokens must not disturb it.
|
||||
o._mark_worker_started(b_cancel)
|
||||
assert o._owns_worker(b_cancel)
|
||||
|
||||
|
||||
def test_a_stale_mailbox_read_does_not_cancel_the_running_generation():
|
||||
# A dispatched consumer can still be draining tokens after the dispatcher retired its request
|
||||
# and started the next one. Stopping it then must tear down only its own stream: signalling
|
||||
# the shared worker event would end its successor.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._mailboxes = {"a": _queue.Queue(), "b": _queue.Queue()}
|
||||
o._request_cancel_events = {"a": a_cancel, "b": b_cancel}
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
# Worker finished A and moved on to B.
|
||||
_dispatch(
|
||||
o,
|
||||
[
|
||||
{"type": "gen_done", "request_id": "a"},
|
||||
{"type": "token", "request_id": "b", "token": "yo"},
|
||||
],
|
||||
)
|
||||
assert o._owns_worker(b_cancel) and not o._owns_worker(a_cancel)
|
||||
|
||||
# A's consumer now reads a token buffered before that, with A stopped.
|
||||
a_cancel.set()
|
||||
stale = [{"type": "token", "request_id": "a", "text": "late"}]
|
||||
drained = []
|
||||
list(
|
||||
o._consume_token_stream(
|
||||
lambda timeout: stale.pop(0) if stale else None,
|
||||
lambda: drained.append(True),
|
||||
crash_context = "generation",
|
||||
cancel_event = a_cancel,
|
||||
mark_started = False,
|
||||
)
|
||||
)
|
||||
assert drained, "the stopped stream still tears itself down"
|
||||
assert not o._cancel_event.is_set(), "a retired request must not signal the shared worker event"
|
||||
|
||||
# The generation that does own the worker still can.
|
||||
b_cancel.set()
|
||||
stale_b = [{"type": "token", "request_id": "b", "text": "live"}]
|
||||
list(
|
||||
o._consume_token_stream(
|
||||
lambda timeout: stale_b.pop(0) if stale_b else None,
|
||||
lambda: None,
|
||||
crash_context = "generation",
|
||||
cancel_event = b_cancel,
|
||||
mark_started = False,
|
||||
)
|
||||
)
|
||||
assert o._cancel_event.is_set(), "the running generation's own Stop must reach the worker"
|
||||
|
||||
|
||||
def test_a_dispatcher_started_mid_stream_still_reaches_the_direct_reader():
|
||||
# A compare request can start the dispatcher while an ordinary chat is streaming. The
|
||||
# dispatcher then owns resp_queue, and without a mailbox for the direct reader it dropped
|
||||
# that chat's tokens and its gen_done as unaddressed, hanging it.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
|
||||
read_one, _drain, release = o._direct_reader("direct-1")
|
||||
try:
|
||||
_dispatch(
|
||||
o,
|
||||
[
|
||||
{"type": "token", "request_id": "direct-1", "text": "hi"},
|
||||
{"type": "gen_done", "request_id": "direct-1"},
|
||||
],
|
||||
)
|
||||
assert read_one(timeout = 0.1) == {
|
||||
"type": "token",
|
||||
"request_id": "direct-1",
|
||||
"text": "hi",
|
||||
}, "the dispatcher must route to the direct reader, not drop"
|
||||
assert read_one(timeout = 0.1)["type"] == "gen_done"
|
||||
finally:
|
||||
release()
|
||||
assert o._direct_mailboxes == {}, "the mailbox is dropped when the stream ends"
|
||||
|
||||
|
||||
def test_the_direct_reader_hands_back_a_compare_response_it_took():
|
||||
# The mirror race: this reader is already blocked on resp_queue when a compare request's
|
||||
# dispatcher starts, so it can take that request's response first. Consuming it would
|
||||
# corrupt this chat and hang the compare pane.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
compare_box: _queue.Queue = _queue.Queue()
|
||||
o._mailboxes = {"compare-1": compare_box}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._resp_queue = _queue.Queue()
|
||||
o._dispatcher_thread = None # no dispatcher yet: this reader owns the queue
|
||||
|
||||
read_one, _drain, release = o._direct_reader("direct-1")
|
||||
try:
|
||||
o._resp_queue.put({"type": "token", "request_id": "compare-1", "text": "theirs"})
|
||||
o._resp_queue.put({"type": "token", "request_id": "direct-1", "text": "mine"})
|
||||
assert read_one(timeout = 0.1) is None, "a foreign response is not ours to yield"
|
||||
assert compare_box.get_nowait()["text"] == "theirs", "it goes to its own mailbox"
|
||||
assert read_one(timeout = 0.1)["text"] == "mine"
|
||||
finally:
|
||||
release()
|
||||
|
||||
|
||||
def test_a_direct_mailbox_is_not_mistaken_for_compare_activity():
|
||||
# _mailboxes means "compare requests are in flight" to the unload and distributed paths,
|
||||
# so an ordinary chat's mailbox must live somewhere else.
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
_read_one, _drain, release = o._direct_reader("direct-1")
|
||||
try:
|
||||
assert o._mailboxes == {}
|
||||
assert "direct-1" in o._direct_mailboxes
|
||||
finally:
|
||||
release()
|
||||
|
||||
|
||||
def test_replacing_the_subprocess_clears_worker_scoped_state():
|
||||
# Ownership is keyed only by cancel-event identity, so a consumer still blocked on its
|
||||
# mailbox when the worker was replaced stayed recorded as the executor. A generation on
|
||||
# the fresh worker then failed _owns_worker and could not be stopped.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
dead = threading.Event()
|
||||
o._mailboxes = {"compare-1": _queue.Queue()}
|
||||
o._direct_mailboxes = {"direct-1": _queue.Queue()}
|
||||
o._request_cancel_events = {"compare-1": dead}
|
||||
o._claim_worker(dead)
|
||||
o._mark_worker_started(dead)
|
||||
assert o._owns_worker(dead)
|
||||
|
||||
o._reset_worker_scoped_state()
|
||||
|
||||
assert o._mailboxes == {} and o._direct_mailboxes == {}
|
||||
assert o._request_cancel_events == {}
|
||||
assert o._active_cancel_events == [] and o._executing_cancel_events == []
|
||||
# A generation on the fresh worker owns it rather than being refused by a ghost.
|
||||
fresh = threading.Event()
|
||||
o._claim_worker(fresh)
|
||||
assert o._owns_worker(fresh), "the dead worker's request must not outrank a live one"
|
||||
|
||||
|
||||
def test_audio_input_claims_the_worker_before_sending():
|
||||
# Unclaimed, a compare request queued behind an audio-input generation looked like the
|
||||
# oldest owner, so stopping that queued request signalled the worker and killed this.
|
||||
import ast
|
||||
import pathlib
|
||||
|
||||
src = pathlib.Path(orch_mod.__file__).read_text(encoding = "utf-8")
|
||||
tree = ast.parse(src)
|
||||
fn = next(
|
||||
n
|
||||
for n in ast.walk(tree)
|
||||
if isinstance(n, ast.FunctionDef) and n.name == "_generate_audio_input_inner"
|
||||
)
|
||||
body = ast.get_source_segment(src, fn) or ""
|
||||
claim = body.find("self._claim_worker(cancel_event)")
|
||||
send = body.find("self._send_cmd(cmd)")
|
||||
assert claim != -1, "_generate_audio_input_inner must claim the worker"
|
||||
assert send != -1
|
||||
assert claim < send, "the claim has to happen before the command is enqueued"
|
||||
assert "with self._send_order_lock:" in body, "claim and send must be one critical section"
|
||||
assert "self._release_worker(cancel_event)" in body
|
||||
|
||||
|
||||
def test_generation_stopped_while_queued_is_never_sent(monkeypatch):
|
||||
# Two chats on the serialized backend: the second blocks on _gen_lock, and Stop sets its
|
||||
# event while it waits. Sending anyway occupied the worker with a run the user ended --
|
||||
# the cancel is only checked on a token, so a long prefill (or a generation that reaches
|
||||
# gen_done without one) still held up its siblings.
|
||||
o = _bare_orchestrator()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda *a, **k: None)
|
||||
monkeypatch.setattr(
|
||||
o, "_send_cmd", lambda cmd: pytest.fail("must not send a generation already stopped")
|
||||
)
|
||||
stopped = threading.Event()
|
||||
stopped.set()
|
||||
|
||||
out = list(
|
||||
o._generate_inner(messages = [{"role": "user", "content": "hi"}], cancel_event = stopped)
|
||||
)
|
||||
|
||||
assert out == [], "a stopped request yields nothing rather than an error banner"
|
||||
assert o._active_cancel_events == [], "it must not claim the worker either"
|
||||
assert o._gen_lock.acquire(blocking = False)
|
||||
o._gen_lock.release()
|
||||
|
||||
|
||||
def test_audio_input_stopped_while_queued_is_never_sent(monkeypatch):
|
||||
# Same lock, same hole.
|
||||
o = _bare_orchestrator()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
o, "_send_cmd", lambda cmd: pytest.fail("must not send a generation already stopped")
|
||||
)
|
||||
stopped = threading.Event()
|
||||
stopped.set()
|
||||
|
||||
out = list(o._generate_audio_input_inner(audio_array = [0.0, 0.1], cancel_event = stopped))
|
||||
|
||||
assert out == []
|
||||
assert o._active_cancel_events == []
|
||||
assert o._gen_lock.acquire(blocking = False)
|
||||
o._gen_lock.release()
|
||||
|
|
|
|||
|
|
@ -312,7 +312,7 @@ def _load_impl_source() -> str:
|
|||
def test_route_resolves_slots_once_before_dedupe_guard_and_load():
|
||||
load_impl = _load_impl_source()
|
||||
resolve = load_impl.index("request.n_parallel")
|
||||
fallback = load_impl.index('getattr(fastapi_request.app.state, "llama_parallel_slots", 1)')
|
||||
fallback = load_impl.index('getattr(_app_state, "llama_parallel_slots", 1)')
|
||||
dedupe = load_impl.index("requested_parallel_slots = _n_parallel")
|
||||
guard = load_impl.index("_guard_chat_load_against_training")
|
||||
# The GGUF launch kwargs, not the guard's own kwarg (which shares the spelling).
|
||||
|
|
@ -324,7 +324,10 @@ def test_route_resolves_slots_once_before_dedupe_guard_and_load():
|
|||
# nothing re-reads app.state after the single resolution point.
|
||||
assert load_impl.count("n_parallel = _n_parallel") == 2
|
||||
assert "n_parallel = _n_parallel" in load_impl[load_kwargs : load_kwargs + 800]
|
||||
assert load_impl.count('getattr(fastapi_request.app.state, "llama_parallel_slots", 1)') == 1
|
||||
assert load_impl.count('getattr(_app_state, "llama_parallel_slots", 1)') == 1
|
||||
# Reached by getattr, so a direct caller passing a request without an app
|
||||
# cannot raise; a bare attribute read would also be a second resolution.
|
||||
assert "fastapi_request.app.state" not in load_impl
|
||||
|
||||
|
||||
def test_route_dedupe_compares_requested_slots_and_skips_diffusion():
|
||||
|
|
@ -361,7 +364,7 @@ def test_parallel_slot_echo_reports_none_for_diffusion():
|
|||
def test_validate_route_prefers_request_n_parallel():
|
||||
validate_impl = _route_source()[_route_source().index("async def validate_model") :]
|
||||
resolve = validate_impl.index("request.n_parallel")
|
||||
fallback = validate_impl.index('getattr(fastapi_request.app.state, "llama_parallel_slots", 1)')
|
||||
fallback = validate_impl.index('"llama_parallel_slots",')
|
||||
guard = validate_impl.index("_guard_chat_load_against_training")
|
||||
assert guard < resolve and guard < fallback, "the guard call resolves the slots inline"
|
||||
|
||||
|
|
|
|||
|
|
@ -504,11 +504,12 @@ def _upstream_message(
|
|||
|
||||
|
||||
class ScriptedClient:
|
||||
"""Fake nonstreaming_client() returning scripted JSON bodies, counting POSTs."""
|
||||
"""Fake upstream client returning scripted JSON bodies, counting POSTs."""
|
||||
|
||||
def __init__(self, bodies):
|
||||
self.bodies = list(bodies)
|
||||
self.posts = []
|
||||
self.closed = False
|
||||
|
||||
async def post(
|
||||
self,
|
||||
|
|
@ -520,6 +521,10 @@ class ScriptedClient:
|
|||
self.posts.append(json)
|
||||
return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])
|
||||
|
||||
async def aclose(self):
|
||||
# The Anthropic pass-through owns its client and closes it in a finally.
|
||||
self.closed = True
|
||||
|
||||
|
||||
async def _drive_non_streaming(monkeypatch, payload, bodies):
|
||||
import routes.inference as inf_mod
|
||||
|
|
@ -867,7 +872,7 @@ class TestNudgeRetryAnthropic:
|
|||
from routes.inference import _anthropic_passthrough_non_streaming
|
||||
|
||||
client = ScriptedClient(bodies)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
response = await _anthropic_passthrough_non_streaming(
|
||||
_llama_backend(),
|
||||
[{"role": "user", "content": "hi"}],
|
||||
|
|
@ -925,7 +930,7 @@ class TestAnthropicPassthroughHealingText:
|
|||
from routes.inference import _anthropic_passthrough_non_streaming
|
||||
|
||||
client = ScriptedClient([upstream])
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
response = await _anthropic_passthrough_non_streaming(
|
||||
_llama_backend(),
|
||||
[{"role": "user", "content": "hi"}],
|
||||
|
|
@ -1171,7 +1176,7 @@ class TestAnthropicNonStreamingRoute:
|
|||
from routes.inference import _anthropic_passthrough_non_streaming
|
||||
|
||||
client = ScriptedClient(bodies)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
response = await _anthropic_passthrough_non_streaming(
|
||||
_llama_backend(),
|
||||
[{"role": "user", "content": "hi"}],
|
||||
|
|
|
|||
|
|
@ -5051,3 +5051,27 @@ class TestFalseAlarmMarkerProse:
|
|||
assert [c[0] for c in exec_fn.calls] == ["web_search", "python"]
|
||||
assistant = next(m for m in convs[1] if m["role"] == "assistant")
|
||||
assert '"python"' not in (assistant.get("content") or "")
|
||||
|
||||
|
||||
def test_both_tool_loops_say_they_are_waiting_for_approval():
|
||||
"""A gated call must not report "Running" in either loop.
|
||||
|
||||
The GGUF loop was fixed first and the safetensors one was missed, so the
|
||||
badge counted up "Running ..." against a prompt nobody had answered yet.
|
||||
Asserted on the source so the two paths cannot drift apart again.
|
||||
"""
|
||||
import ast
|
||||
import os
|
||||
|
||||
backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
for name in ("core/inference/safetensors_agentic.py", "core/inference/llama_cpp.py"):
|
||||
with open(os.path.join(backend, name), encoding = "utf-8") as f:
|
||||
tree = ast.parse(f.read())
|
||||
calls = [
|
||||
node
|
||||
for node in ast.walk(tree)
|
||||
if isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Name)
|
||||
and node.func.id == "awaiting_approval_status"
|
||||
]
|
||||
assert calls, f"{name} still announces a gated tool call as running"
|
||||
|
|
|
|||
|
|
@ -95,7 +95,7 @@ class _ScriptedBackend:
|
|||
for snap in snapshots:
|
||||
yield snap
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
self.reset_count += 1
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ holds sidecar transformers modules (breaking the rename on Windows). The methods
|
|||
the handle and return False so callers can refuse the swap.
|
||||
"""
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from core.export.orchestrator import ExportOrchestrator
|
||||
|
|
@ -52,6 +54,14 @@ def _bare_inference():
|
|||
o._resp_queue = _Q()
|
||||
o._cancel_event = None
|
||||
o._drain_event = None
|
||||
# Worker-scoped bookkeeping the teardown clears (see _reset_worker_scoped_state).
|
||||
o._active_cancel_lock = threading.Lock()
|
||||
o._active_cancel_events = []
|
||||
o._executing_cancel_events = []
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
return o
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -56,7 +56,9 @@ def test_system_gpu_info_preserves_vulkan_visibility_metrics(monkeypatch):
|
|||
assert gpu["available"] is False
|
||||
assert gpu["backend"] == "cpu"
|
||||
assert gpu["index_kind"] == "relative"
|
||||
assert gpu["gguf_gpu_ids_supported"] is False
|
||||
# A Vulkan llama.cpp build accepts gpu_ids even when torch training is
|
||||
# CPU-only: the pick is a ggml ordinal, not a torch device index.
|
||||
assert gpu["gguf_gpu_ids_supported"] is True
|
||||
assert gpu["devices"] == []
|
||||
assert inference_gpu["backend"] == "vulkan"
|
||||
assert inference_gpu["devices"] == [vulkan_device]
|
||||
|
|
@ -124,7 +126,8 @@ def test_system_gpu_info_keeps_forced_vulkan_separate_from_training_metrics(monk
|
|||
assert gpu["devices"][0]["vram_used_gb"] == 6.0
|
||||
assert inference_gpu["backend"] == "vulkan"
|
||||
assert inference_gpu["devices"][0]["vram_used_gb"] == 1.0
|
||||
assert inference_gpu["gguf_gpu_ids_supported"] is False
|
||||
# Probed devices exist, so the ordinals are known and picks are offered.
|
||||
assert inference_gpu["gguf_gpu_ids_supported"] is True
|
||||
|
||||
|
||||
def test_system_gpu_info_does_not_merge_metrics_across_backend_index_spaces(monkeypatch):
|
||||
|
|
@ -169,3 +172,85 @@ def test_system_gpu_info_does_not_merge_metrics_across_backend_index_spaces(monk
|
|||
|
||||
assert gpu["devices"] == [vulkan_device]
|
||||
assert inference_gpu == gpu
|
||||
|
||||
|
||||
def test_vulkan_inference_gpu_uses_real_device_names_and_igpu_flag(monkeypatch):
|
||||
"""The picker and the GPU labels need ggml's real device description, not a
|
||||
Vulkan<i> placeholder, and an explicit iGPU flag rather than inferring one
|
||||
from a zero total. Memory still comes from _get_gpu_memory so the iGPU host
|
||||
reserve is applied; budgeting off the raw shared total would hand out the
|
||||
whole machine's RAM with no OS headroom.
|
||||
"""
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
from utils.hardware.hardware import get_vulkan_inference_gpu_info
|
||||
|
||||
monkeypatch.setattr(
|
||||
LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda binary = None: True)
|
||||
)
|
||||
# Fit view: discrete card keeps its total, iGPU reports 0 with capped free.
|
||||
monkeypatch.setattr(
|
||||
LlamaCppBackend,
|
||||
"_get_gpu_memory",
|
||||
staticmethod(lambda binary = None: [(0, 15 * 1024, 16 * 1024), (1, 12 * 1024, 0)]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LlamaCppBackend,
|
||||
"vulkan_device_inventory",
|
||||
staticmethod(
|
||||
lambda binary = None: [
|
||||
{
|
||||
"index": 0,
|
||||
"name": "AMD Radeon RX 9070 XT",
|
||||
"free_mib": 15 * 1024,
|
||||
"total_mib": 16 * 1024,
|
||||
"is_igpu": False,
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"name": "AMD Radeon(TM) 8060S Graphics",
|
||||
"free_mib": 89 * 1024,
|
||||
"total_mib": 91 * 1024,
|
||||
"is_igpu": True,
|
||||
},
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
info = get_vulkan_inference_gpu_info()
|
||||
assert info is not None and info["index_kind"] == "vulkan"
|
||||
dgpu, igpu = info["devices"]
|
||||
|
||||
assert dgpu["name"] == "AMD Radeon RX 9070 XT"
|
||||
assert dgpu["index_kind"] == "vulkan"
|
||||
assert dgpu["shared_memory"] is False
|
||||
assert dgpu["memory_total_gb"] == 16.0
|
||||
|
||||
assert igpu["name"] == "AMD Radeon(TM) 8060S Graphics"
|
||||
assert igpu["shared_memory"] is True
|
||||
# The capped free budget from _get_gpu_memory, NOT the 91 GiB raw total.
|
||||
assert igpu["memory_total_gb"] == 12.0
|
||||
|
||||
|
||||
def test_vulkan_inference_gpu_falls_back_to_ordinal_names(monkeypatch):
|
||||
"""A probe that cannot resolve descriptions must not lose the device list:
|
||||
names degrade to Vulkan<i> and the memory readings still get through."""
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
from utils.hardware.hardware import get_vulkan_inference_gpu_info
|
||||
|
||||
monkeypatch.setattr(
|
||||
LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda binary = None: True)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LlamaCppBackend,
|
||||
"_get_gpu_memory",
|
||||
staticmethod(lambda binary = None: [(0, 15 * 1024, 16 * 1024)]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LlamaCppBackend,
|
||||
"vulkan_device_inventory",
|
||||
staticmethod(lambda binary = None: (_ for _ in ()).throw(RuntimeError("probe failed"))),
|
||||
)
|
||||
|
||||
info = get_vulkan_inference_gpu_info()
|
||||
assert info["devices"][0]["name"] == "Vulkan0"
|
||||
assert info["devices"][0]["memory_total_gb"] == 16.0
|
||||
|
|
|
|||
80
studio/backend/tests/test_tool_sandbox_per_thread.py
Normal file
80
studio/backend/tests/test_tool_sandbox_per_thread.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Every conversation runs its tools in its own sandbox directory.
|
||||
|
||||
Parallel chats lean on this: two conversations can be mid tool call at the same
|
||||
time, so a shared working directory would let one overwrite the other's files.
|
||||
The session id is the chat's thread id (or project-<id> for project chats), and
|
||||
the dir is derived from it here.
|
||||
|
||||
HOME is redirected at import time, so nothing touches the real ~/studio_sandbox.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
_backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
sys.path.insert(0, _backend)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def workdir(tmp_path, monkeypatch):
|
||||
"""_get_workdir with HOME pointed at tmp_path and its cache cleared."""
|
||||
from core.inference import tools
|
||||
|
||||
monkeypatch.setattr(os.path, "expanduser", lambda path: str(tmp_path))
|
||||
monkeypatch.setattr(tools, "_workdirs", {})
|
||||
return tools._get_workdir
|
||||
|
||||
|
||||
def test_two_conversations_get_two_directories(workdir, tmp_path):
|
||||
a = workdir("thread-alpha")
|
||||
b = workdir("thread-beta")
|
||||
assert a != b
|
||||
assert os.path.basename(a) == "thread-alpha"
|
||||
assert os.path.basename(b) == "thread-beta"
|
||||
assert os.path.isdir(a) and os.path.isdir(b)
|
||||
assert os.path.dirname(a) == os.path.dirname(b) == str(tmp_path / "studio_sandbox")
|
||||
|
||||
|
||||
def test_the_same_conversation_keeps_its_directory(workdir):
|
||||
# A later turn, or a tool continuation, must land back in the same place.
|
||||
assert workdir("thread-alpha") == workdir("thread-alpha")
|
||||
|
||||
|
||||
def test_a_directory_is_private_to_its_conversation(workdir):
|
||||
a = workdir("thread-alpha")
|
||||
b = workdir("thread-beta")
|
||||
with open(os.path.join(a, "secret.txt"), "w", encoding = "utf-8") as f:
|
||||
f.write("alpha")
|
||||
assert os.listdir(b) == []
|
||||
|
||||
|
||||
def test_project_chats_deliberately_share_one_workspace(workdir, monkeypatch):
|
||||
# Chats in a project are meant to see each other's files.
|
||||
from core.inference import tools
|
||||
monkeypatch.setattr(tools, "_get_project_workdir", lambda sid: "/tmp/project-ws")
|
||||
assert tools._get_workdir("project-abc") == "/tmp/project-ws"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"session_id",
|
||||
["../escape", "a/b", "", " ", "x" * 65],
|
||||
)
|
||||
def test_a_session_id_cannot_escape_the_sandbox_root(workdir, tmp_path, session_id):
|
||||
resolved = workdir(session_id) if session_id else workdir(None)
|
||||
root = os.path.realpath(str(tmp_path / "studio_sandbox"))
|
||||
assert os.path.realpath(resolved).startswith(root + os.sep)
|
||||
assert os.path.basename(resolved) in {"_invalid", "_default"}
|
||||
|
||||
|
||||
def test_no_session_id_falls_back_to_default(workdir):
|
||||
assert os.path.basename(workdir(None)) == "_default"
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform == "win32", reason = "POSIX permission bits")
|
||||
def test_directories_are_private_to_the_user(workdir):
|
||||
assert os.stat(workdir("thread-alpha")).st_mode & 0o777 == 0o700
|
||||
|
|
@ -668,7 +668,8 @@ def test_route_history_and_passthrough_forward_the_display_gate():
|
|||
blocks = {
|
||||
"safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)",
|
||||
"anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)",
|
||||
"anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)",
|
||||
# Anchored on the code, not the comment above it, so rewrapping prose cannot break this.
|
||||
"anthropic passthrough": r"if not healing_active:.*?\.strip\(\)",
|
||||
}
|
||||
for label, pat in blocks.items():
|
||||
m = _re.search(pat, _src, _re.DOTALL)
|
||||
|
|
|
|||
|
|
@ -209,7 +209,24 @@ def _force_missing_fla_imports(monkeypatch):
|
|||
monkeypatch.setattr(builtins, "__import__", fake_import)
|
||||
|
||||
|
||||
def _pin_fla_model_types(monkeypatch):
|
||||
"""Pin the auto-discovered FLA allowlist to the Qwen GDN families.
|
||||
|
||||
`_discover_fla_model_types` scans the *installed* transformers, and
|
||||
`models/qwen3_5/` only exists from 5.x. The backend supports
|
||||
`transformers>=4.51`, so on a 4.x install the gate returns False and every
|
||||
Qwen3.5 assertion below silently no-ops. Pinning keeps these tests hermetic
|
||||
across the supported range.
|
||||
"""
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"_discover_fla_model_types",
|
||||
lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_6", "qwen3_next"}),
|
||||
)
|
||||
|
||||
|
||||
def test_flash_linear_attention_installs_pinned_pair_for_qwen3_5(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
||||
monkeypatch.setattr(worker._sp, "run", run_mock)
|
||||
|
|
@ -315,6 +332,7 @@ def test_flash_linear_attention_skipped_via_env(monkeypatch):
|
|||
|
||||
|
||||
def test_flash_linear_attention_skipped_below_torch_2_7(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 5))
|
||||
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
||||
|
|
@ -332,6 +350,7 @@ def test_flash_linear_attention_skipped_below_torch_2_7(monkeypatch):
|
|||
|
||||
|
||||
def test_flash_linear_attention_install_includes_einops(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9))
|
||||
|
|
@ -358,6 +377,7 @@ def test_flash_linear_attention_install_includes_einops(monkeypatch):
|
|||
|
||||
def test_flash_linear_attention_logs_post_install_import_failure(monkeypatch):
|
||||
"""pip exits 0 but `import fla.modules` still fails (missing transitive)."""
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9))
|
||||
|
|
@ -402,6 +422,7 @@ def test_tilelang_backend_skipped_on_unsupported_linux_arch(monkeypatch):
|
|||
|
||||
|
||||
def test_tilelang_backend_pins_only_binary(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
||||
|
|
@ -442,6 +463,7 @@ def _force_missing_tilelang_imports(monkeypatch):
|
|||
|
||||
|
||||
def test_tilelang_backend_installs_pinned_pair_for_qwen3_5(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
||||
|
|
@ -472,6 +494,7 @@ def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch):
|
|||
2 (install): plain apache-tvm-ffi + tilelang -- resolves missing transitive
|
||||
deps without --force-reinstall, so it never replaces correct packages.
|
||||
"""
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11")
|
||||
|
|
@ -533,6 +556,7 @@ def test_tilelang_backend_skipped_on_windows(monkeypatch):
|
|||
|
||||
|
||||
def test_tilelang_backend_swallows_install_timeout(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
||||
|
|
@ -586,6 +610,7 @@ def test_tilelang_backend_skipped_via_env(monkeypatch):
|
|||
|
||||
|
||||
def test_tilelang_backend_swallows_install_failure(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: None)
|
||||
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
||||
|
|
@ -649,6 +674,7 @@ def _patch_iu_gates(monkeypatch, fla_gate, conv_gate):
|
|||
|
||||
|
||||
def test_hook_installs_when_gate_returns_false(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
fla_gate = _make_fake_gate(initial_return = False)
|
||||
conv_gate = _make_fake_gate(initial_return = False)
|
||||
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
||||
|
|
@ -716,6 +742,7 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch):
|
|||
|
||||
|
||||
def test_hook_idempotent_on_repeat_call(monkeypatch):
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
fla_gate = _make_fake_gate(initial_return = False)
|
||||
conv_gate = _make_fake_gate(initial_return = False)
|
||||
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
||||
|
|
@ -924,6 +951,7 @@ def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch)
|
|||
|
||||
def test_hook_does_install_tilelang_for_qwen35(monkeypatch):
|
||||
"""Positive control for finding #1: Qwen3.5 still gets tilelang."""
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
fla_gate = _make_fake_gate(initial_return = False)
|
||||
conv_gate = _make_fake_gate(initial_return = True)
|
||||
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
||||
|
|
@ -953,6 +981,7 @@ def test_tilelang_repair_does_not_touch_torch_cuda_stack(monkeypatch):
|
|||
forced step so --force-reinstall doesn't cascade through
|
||||
apache-tvm-ffi's dep graph and pull a different torch wheel.
|
||||
"""
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
||||
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
||||
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.10")
|
||||
|
|
@ -1065,6 +1094,7 @@ def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch):
|
|||
probe) but tilelang is missing or apache-tvm-ffi is on the broken
|
||||
list, the post-available action must still run tilelang.
|
||||
"""
|
||||
_pin_fla_model_types(monkeypatch)
|
||||
fla_gate = _make_fake_gate(initial_return = True)
|
||||
conv_gate = _make_fake_gate(initial_return = True)
|
||||
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
||||
|
|
|
|||
|
|
@ -1706,7 +1706,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
|
|||
"backend": _backend_label(device),
|
||||
"parent_visible_gpu_ids": [],
|
||||
"devices": [],
|
||||
"index_kind": "relative",
|
||||
"index_kind": "vulkan",
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -2600,22 +2600,35 @@ def get_vulkan_inference_gpu_info() -> Optional[Dict[str, Any]]:
|
|||
"backend_cuda_visible_devices": None,
|
||||
"parent_visible_gpu_ids": [],
|
||||
"devices": [],
|
||||
"index_kind": "relative",
|
||||
"index_kind": "vulkan",
|
||||
}
|
||||
# Identity (real device description, explicit iGPU flag) comes from the
|
||||
# inventory; the memory numbers stay on _get_gpu_memory, which applies the
|
||||
# iGPU host reserve and zeroes a shared total. Budgeting an APU off the raw
|
||||
# shared total instead would hand out the whole machine's RAM with no OS
|
||||
# headroom. Join by ordinal; a probe failure just leaves names unresolved.
|
||||
identity: Dict[int, Dict[str, Any]] = {}
|
||||
try:
|
||||
identity = {row["index"]: row for row in LlamaCppBackend.vulkan_device_inventory()}
|
||||
except Exception as e:
|
||||
logger.debug("Vulkan device inventory failed, falling back to ordinals: %s", e)
|
||||
|
||||
try:
|
||||
for ordinal, free_mib, total_mib in LlamaCppBackend._get_gpu_memory():
|
||||
# Integrated Vulkan GPUs report total=0 because their memory is
|
||||
# shared. Publish the capped free value as their usable inference
|
||||
# budget and mark it so clients do not add system RAM again.
|
||||
shared_memory = total_mib == 0
|
||||
info = identity.get(ordinal, {})
|
||||
# _get_gpu_memory reports total 0 for a shared pool; prefer the
|
||||
# explicit flag when the inventory resolved this ordinal.
|
||||
shared_memory = bool(info["is_igpu"]) if "is_igpu" in info else total_mib == 0
|
||||
budget_mib = total_mib or free_mib
|
||||
used_mib = max(0, total_mib - free_mib) if total_mib else None
|
||||
result["devices"].append(
|
||||
{
|
||||
"index": ordinal,
|
||||
"index_kind": "relative",
|
||||
# ggml Vulkan ordinals are the space `--device Vulkan<i>` pins,
|
||||
# so unlike a torch-xpu relative ordinal these are selectable.
|
||||
"index_kind": "vulkan",
|
||||
"visible_ordinal": ordinal,
|
||||
"name": f"Vulkan{ordinal}",
|
||||
"name": info.get("name") or f"Vulkan{ordinal}",
|
||||
"memory_total_gb": round(budget_mib / 1024, 2),
|
||||
"vram_used_gb": round(used_mib / 1024, 2) if used_mib is not None else None,
|
||||
"vram_free_gb": round(free_mib / 1024, 2),
|
||||
|
|
@ -2727,7 +2740,7 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]:
|
|||
"backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
|
||||
"parent_visible_gpu_ids": [],
|
||||
"devices": [],
|
||||
"index_kind": "relative",
|
||||
"index_kind": "vulkan",
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -403,6 +403,7 @@ def _run_llama_phase(
|
|||
pin_release_tag: Optional[str],
|
||||
set_progress,
|
||||
force_cpu: bool = False,
|
||||
llama_backend: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""The llama phase of a chained update: put the backend into a maintenance
|
||||
state, run the installer for the latest prebuilt, then refresh caches so the
|
||||
|
|
@ -454,14 +455,15 @@ def _run_llama_phase(
|
|||
# updates. A natural fallback (or a legacy marker without the flag) heals to GPU (#6097).
|
||||
if force_cpu:
|
||||
cmd.append("--force-cpu")
|
||||
if llama_backend == "vulkan":
|
||||
cmd.extend(["--llama-backend", "vulkan"])
|
||||
logger.info("llama update: installing", cmd = " ".join(cmd))
|
||||
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():
|
||||
# Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm box would
|
||||
# otherwise re-route and silently replace it. Re-assert via setup's env/CLI flags.
|
||||
if llama_backend == "vulkan" or (asset and "vulkan" in asset.lower()):
|
||||
env["UNSLOTH_FORCE_VULKAN"] = "1"
|
||||
env["UNSLOTH_LLAMA_BACKEND"] = "vulkan"
|
||||
_flow.stream_installer(
|
||||
cmd,
|
||||
env,
|
||||
|
|
@ -578,6 +580,9 @@ def _plan_llama_phase() -> dict:
|
|||
from_tag = marker.get("tag") or marker.get("release_tag")
|
||||
asset = marker.get("asset")
|
||||
force_cpu = bool(marker.get("force_cpu"))
|
||||
llama_backend = marker.get("llama_backend")
|
||||
if llama_backend == "vulkan" or (asset and "vulkan" in str(asset).lower()):
|
||||
llama_backend = "vulkan"
|
||||
# Install exactly the release the banner offered: the installer's own
|
||||
# "latest" is commit-date ordered and can lag the published_at pick
|
||||
# above, reinstalling the current build in a loop (the #6219 class).
|
||||
|
|
@ -621,6 +626,7 @@ def _plan_llama_phase() -> dict:
|
|||
asset = (res or {}).get("asset")
|
||||
# Source builds carry no forced-CPU marker, so nothing to preserve here.
|
||||
force_cpu = False
|
||||
llama_backend = None
|
||||
# No pin: source-build detection resolves via --resolve-prebuilt latest,
|
||||
# the same resolver the unpinned apply uses, so the two already agree.
|
||||
pin_release_tag = None
|
||||
|
|
@ -643,6 +649,7 @@ def _plan_llama_phase() -> dict:
|
|||
"pin_release_tag": pin_release_tag,
|
||||
"from_tag": from_tag,
|
||||
"force_cpu": force_cpu,
|
||||
"llama_backend": llama_backend,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -695,6 +702,7 @@ def start_update() -> dict:
|
|||
llama_spec["pin_release_tag"],
|
||||
set_progress,
|
||||
force_cpu = llama_spec.get("force_cpu", False),
|
||||
llama_backend = llama_spec.get("llama_backend"),
|
||||
)
|
||||
)
|
||||
if llama_spec
|
||||
|
|
|
|||
|
|
@ -3,10 +3,13 @@
|
|||
|
||||
"""Persisted opt-in controls for OpenAI-compatible model auto-switching.
|
||||
|
||||
Two settings, both off by default so existing API behavior is unchanged:
|
||||
All off by default so existing API behavior is unchanged:
|
||||
- ``openai_api_auto_switch_model``: when on, a ``/v1`` request whose ``model``
|
||||
names a downloaded local GGUF different from the loaded one transparently
|
||||
loads it before serving (llama-swap-style). Unknown names pass through.
|
||||
- ``openai_api_auto_download_model``: when on, a ``/v1`` request naming an
|
||||
undownloaded GGUF repo starts a background download instead of failing.
|
||||
Gated on auto-switch, which is what serves the model once it lands.
|
||||
- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is
|
||||
unloaded after this many idle seconds to free VRAM. Enabled values have a
|
||||
60s floor (0 stays "off"): a tiny TTL tears the model down between turns of
|
||||
|
|
@ -29,12 +32,14 @@ import time
|
|||
from typing import Any, Optional
|
||||
|
||||
OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model"
|
||||
OPENAI_AUTO_DOWNLOAD_SETTING_KEY = "openai_api_auto_download_model"
|
||||
AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds"
|
||||
AUTO_UNLOAD_KEEP_KV_SETTING_KEY = "openai_api_auto_unload_keep_kv"
|
||||
MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides"
|
||||
MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
|
||||
|
||||
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
|
||||
DEFAULT_OPENAI_AUTO_DOWNLOAD_ENABLED = False
|
||||
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
|
||||
DEFAULT_AUTO_UNLOAD_KEEP_KV = True
|
||||
MIN_AUTO_UNLOAD_IDLE_SECONDS = 60
|
||||
|
|
@ -95,6 +100,22 @@ def get_openai_auto_switch_enabled() -> bool:
|
|||
return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED
|
||||
|
||||
|
||||
def get_stored_openai_auto_download_enabled() -> bool:
|
||||
"""The persisted auto-download flag, independent of auto-switch, so the UI
|
||||
round-trips the saved value across an auto-switch toggle instead of erasing it."""
|
||||
parsed = _coerce_bool(_cached_setting(OPENAI_AUTO_DOWNLOAD_SETTING_KEY, None))
|
||||
return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_DOWNLOAD_ENABLED
|
||||
|
||||
|
||||
def get_openai_auto_download_enabled() -> bool:
|
||||
"""Whether a /v1 request may download a GGUF repo it names but doesn't have.
|
||||
|
||||
Gated on auto-switch: that is what loads the model once it lands, so without
|
||||
it we would fetch gigabytes nothing can serve.
|
||||
"""
|
||||
return get_stored_openai_auto_download_enabled() and get_openai_auto_switch_enabled()
|
||||
|
||||
|
||||
def _stored_idle_seconds() -> Optional[int]:
|
||||
"""The persisted idle TTL as an int, or None when never set."""
|
||||
return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None))
|
||||
|
|
@ -170,7 +191,8 @@ def set_openai_auto_switch(
|
|||
enabled: Any,
|
||||
idle_seconds: Any,
|
||||
keep_kv: Any = None,
|
||||
) -> tuple[bool, int, bool]:
|
||||
auto_download: Any = None,
|
||||
) -> tuple[bool, int, bool, bool]:
|
||||
"""One-transaction write; ``None`` leaves a stored value untouched."""
|
||||
parsed_enabled = _coerce_bool(enabled)
|
||||
if parsed_enabled is None:
|
||||
|
|
@ -190,6 +212,11 @@ def set_openai_auto_switch(
|
|||
parsed_keep_kv = _coerce_bool(keep_kv)
|
||||
if parsed_keep_kv is None:
|
||||
raise ValueError("Keep KV on idle unload must be true or false.")
|
||||
parsed_auto_download = None
|
||||
if auto_download is not None:
|
||||
parsed_auto_download = _coerce_bool(auto_download)
|
||||
if parsed_auto_download is None:
|
||||
raise ValueError("Auto-download missing models must be true or false.")
|
||||
from storage.studio_db import upsert_app_settings
|
||||
|
||||
updates: dict[str, Any] = {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled}
|
||||
|
|
@ -197,16 +224,25 @@ def set_openai_auto_switch(
|
|||
updates[AUTO_UNLOAD_IDLE_SETTING_KEY] = parsed_idle
|
||||
if parsed_keep_kv is not None:
|
||||
updates[AUTO_UNLOAD_KEEP_KV_SETTING_KEY] = parsed_keep_kv
|
||||
if parsed_auto_download is not None:
|
||||
updates[OPENAI_AUTO_DOWNLOAD_SETTING_KEY] = parsed_auto_download
|
||||
upsert_app_settings(updates)
|
||||
_invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY)
|
||||
if parsed_idle is not None:
|
||||
_invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
|
||||
if parsed_keep_kv is not None:
|
||||
_invalidate(AUTO_UNLOAD_KEEP_KV_SETTING_KEY)
|
||||
if parsed_auto_download is not None:
|
||||
_invalidate(OPENAI_AUTO_DOWNLOAD_SETTING_KEY)
|
||||
return (
|
||||
parsed_enabled,
|
||||
parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(),
|
||||
parsed_keep_kv if parsed_keep_kv is not None else get_auto_unload_keep_kv(),
|
||||
(
|
||||
parsed_auto_download
|
||||
if parsed_auto_download is not None
|
||||
else get_stored_openai_auto_download_enabled()
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -395,9 +395,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
|
|||
);
|
||||
}
|
||||
|
||||
const showApp = status === "running";
|
||||
const desktopBooting = status === "running" && !desktopAuthReady;
|
||||
const showInteractiveApp = showApp && desktopAuthReady;
|
||||
const showApp = status === "running" && desktopAuthReady;
|
||||
const startupStatus = status === "running" ? "starting" : status;
|
||||
const startupProgressDetail = progressDetail;
|
||||
const usesCustomTitlebar = shouldUseCustomWindowTitlebar();
|
||||
|
|
@ -409,28 +407,12 @@ function TauriWrapper({ children }: { children: ReactNode }) {
|
|||
<TauriUpdateLayer isExternalServer={isExternalServer}>
|
||||
<LlamaUpdateBanner
|
||||
positioned={false}
|
||||
enabled={showInteractiveApp && !hidesTitlebarSidebar}
|
||||
enabled={!hidesTitlebarSidebar}
|
||||
/>
|
||||
{showInteractiveApp ? (
|
||||
<DownloadManagerPanel positioned={false} />
|
||||
) : null}
|
||||
<DownloadManagerPanel positioned={false} />
|
||||
</TauriUpdateLayer>
|
||||
{showInteractiveApp ? <NativeIntentDrain /> : null}
|
||||
{showInteractiveApp ? children : null}
|
||||
{desktopBooting ? (
|
||||
<div className="pointer-events-none fixed inset-x-0 bottom-5 z-[9999] flex justify-center px-4">
|
||||
<div className="absolute inset-x-4 bottom-16 mx-auto flex max-w-[520px] flex-col items-center gap-2 rounded-2xl border border-border/70 bg-background/95 px-6 py-5 text-center shadow-xl">
|
||||
<div className="font-medium text-sm">Preparing Unsloth</div>
|
||||
<div className="text-muted-foreground text-xs">
|
||||
The local backend is ready. Signing in to your desktop session
|
||||
before loading chats.
|
||||
</div>
|
||||
</div>
|
||||
<div className="rounded-full border border-border/70 bg-background/95 px-4 py-2 text-xs text-muted-foreground shadow-lg">
|
||||
Signing in to desktop session...
|
||||
</div>
|
||||
</div>
|
||||
) : null}
|
||||
<NativeIntentDrain />
|
||||
{children}
|
||||
</>
|
||||
) : (
|
||||
<StartupScreen
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import {
|
|||
import {
|
||||
ChatPage,
|
||||
clearNewChatDraft,
|
||||
StopRunningChatsDialog,
|
||||
useChatRuntimeStore,
|
||||
type ChatSearch,
|
||||
} from "@/features/chat";
|
||||
|
|
@ -227,6 +228,8 @@ function RootLayout() {
|
|||
<HfTokenWarningDialog />
|
||||
<RemoteCodeConsentDialog />
|
||||
<TransformersUpgradeDialog />
|
||||
{/* At the root, not under /chat: a swap can start from the Hub too. */}
|
||||
<StopRunningChatsDialog />
|
||||
{hideNavbar ? (
|
||||
<main className="flex-1 pt-[var(--studio-hidden-route-top-inset,0px)] [--studio-titlebar-height:var(--studio-hidden-route-top-inset,0px)]">
|
||||
<Suspense fallback={<RouteFallback />}>
|
||||
|
|
|
|||
|
|
@ -520,15 +520,46 @@ export function AppSidebar() {
|
|||
});
|
||||
const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
|
||||
const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId);
|
||||
const anyChatRunning = useChatRuntimeStore((s) =>
|
||||
Object.values(s.runningByThreadId).some(Boolean),
|
||||
);
|
||||
// The thread currently generating (if any), so "Return to Chat" lands on the
|
||||
// live chat rather than an empty new-chat draft left active after New Chat.
|
||||
const runningThreadId = useChatRuntimeStore((s) => {
|
||||
const entry = Object.entries(s.runningByThreadId).find(([, on]) => on);
|
||||
return entry ? entry[0] : null;
|
||||
});
|
||||
// The whole map, so each row can show its own spinner.
|
||||
const runningThreadIds = useChatRuntimeStore((s) => s.runningByThreadId);
|
||||
// Rows, not raw thread ids: a compare conversation runs two pane threads but is one chat
|
||||
// in the sidebar, so counting the map said "2 Chats" for a single compare row.
|
||||
const runningChatCount = useMemo(() => {
|
||||
const running = new Set(
|
||||
Object.entries(runningThreadIds)
|
||||
.filter(([, on]) => on)
|
||||
.map(([id]) => id),
|
||||
);
|
||||
if (running.size === 0) return 0;
|
||||
let rows = 0;
|
||||
for (const item of allChatItems) {
|
||||
const ids = item.type === "compare" ? (item.threadIds ?? []) : [item.id];
|
||||
let claimed = false;
|
||||
for (const id of ids) {
|
||||
if (running.delete(id)) claimed = true;
|
||||
}
|
||||
if (claimed) rows += 1;
|
||||
}
|
||||
// Anything left belongs to no known row (a first turn mid-persist); count it as one.
|
||||
return rows + running.size;
|
||||
}, [runningThreadIds, allChatItems]);
|
||||
const anyChatRunning = runningChatCount > 0;
|
||||
// Where "Return to Chat" lands: the newest running chat, not the empty draft New Chat left
|
||||
// active (map insertion order is start order). A compare row runs pane threads that /chat
|
||||
// cannot address, so resolve those back to the pair id the route expects.
|
||||
const runningTarget = useMemo(() => {
|
||||
const ids = Object.entries(runningThreadIds)
|
||||
.filter(([, on]) => on)
|
||||
.map(([id]) => id);
|
||||
const id = ids.length > 0 ? ids[ids.length - 1] : null;
|
||||
if (!id) return null;
|
||||
const pair = allChatItems.find(
|
||||
(item) => item.type === "compare" && (item.threadIds ?? []).includes(id),
|
||||
);
|
||||
return pair
|
||||
? { id: pair.id, compare: true as const }
|
||||
: { id, compare: false as const };
|
||||
}, [runningThreadIds, allChatItems]);
|
||||
const activeThreadId = isChatRoute
|
||||
? (search.thread as string | undefined) ??
|
||||
(search.compare as string | undefined) ??
|
||||
|
|
@ -892,6 +923,12 @@ export function AppSidebar() {
|
|||
variant: "project" | "recent",
|
||||
) {
|
||||
const isPinned = pinnedIdSet.has(item.id);
|
||||
// A compare row's id is the pair id while runningByThreadId is keyed per pane thread,
|
||||
// so aggregate its member threads instead.
|
||||
const isGenerating =
|
||||
item.type === "compare"
|
||||
? (item.threadIds ?? []).some((id) => Boolean(runningThreadIds[id]))
|
||||
: Boolean(runningThreadIds[item.id]);
|
||||
const itemClass =
|
||||
variant === "project"
|
||||
? "group/project-chat-item relative"
|
||||
|
|
@ -951,6 +988,8 @@ export function AppSidebar() {
|
|||
data-testid="recent-thread"
|
||||
data-thread-type={item.type}
|
||||
data-thread-id={item.id}
|
||||
data-generating={isGenerating ? "true" : undefined}
|
||||
aria-busy={isGenerating || undefined}
|
||||
isActive={activeThreadId === item.id}
|
||||
className={buttonClass}
|
||||
onClick={() => {
|
||||
|
|
@ -976,6 +1015,14 @@ export function AppSidebar() {
|
|||
<span className="truncate">
|
||||
{pendingRename?.id === item.id ? pendingRename.title : item.title}
|
||||
</span>
|
||||
{isGenerating && (
|
||||
<Spinner
|
||||
data-testid="chat-row-spinner"
|
||||
// role="status" + label: announced, not motion-only.
|
||||
label={translate("shell.navigation.chatGenerating")}
|
||||
className="ml-auto size-3.5 shrink-0 text-muted-foreground"
|
||||
/>
|
||||
)}
|
||||
</SidebarMenuButton>
|
||||
{variant === "project" && (
|
||||
<button
|
||||
|
|
@ -1288,9 +1335,16 @@ export function AppSidebar() {
|
|||
icon={PencilEdit02Icon}
|
||||
label={
|
||||
showReturnToChat
|
||||
? t("shell.navigation.returnToChat")
|
||||
? runningChatCount > 1
|
||||
// Name the count rather than imply a single live chat.
|
||||
? t("shell.navigation.returnToChats", {
|
||||
count: runningChatCount,
|
||||
})
|
||||
: t("shell.navigation.returnToChat")
|
||||
: t("shell.navigation.newChat")
|
||||
}
|
||||
// Off-route this row is the only sign chats are still running.
|
||||
spinner={anyChatRunning && !isChatRoute}
|
||||
active={
|
||||
isChatRoute &&
|
||||
!search.thread &&
|
||||
|
|
@ -1301,8 +1355,13 @@ export function AppSidebar() {
|
|||
if (showReturnToChat) {
|
||||
// Prefer the running thread so we return to the live generation,
|
||||
// not the empty new chat that became active after New Chat.
|
||||
if (runningThreadId && runningThreadId !== storeThreadId) {
|
||||
navigate({ to: "/chat", search: { thread: runningThreadId } });
|
||||
if (runningTarget && runningTarget.id !== storeThreadId) {
|
||||
navigate({
|
||||
to: "/chat",
|
||||
search: runningTarget.compare
|
||||
? { compare: runningTarget.id }
|
||||
: { thread: runningTarget.id },
|
||||
});
|
||||
} else {
|
||||
navigate({ to: "/chat" });
|
||||
}
|
||||
|
|
|
|||
|
|
@ -371,6 +371,9 @@ function useRafCoalescedText(text: string, isStreaming: boolean): string {
|
|||
|
||||
const MarkdownTextImpl = () => {
|
||||
const { text, status } = useMessagePartText();
|
||||
// Parts are keyed by index, so switching conversations hands this instance a different
|
||||
// message, and Streamdown only extends its parsed blocks: key it per message instead.
|
||||
const messageId = useAuiState(({ message }) => message.id);
|
||||
const displayText = useRafCoalescedText(text, status.type === "running");
|
||||
const processedText = useMemo(
|
||||
() => preprocessLaTeX(displayText),
|
||||
|
|
@ -385,6 +388,7 @@ const MarkdownTextImpl = () => {
|
|||
return (
|
||||
<div data-status={status.type} className="min-w-0 max-w-full">
|
||||
<Streamdown
|
||||
key={messageId}
|
||||
mode="streaming"
|
||||
isAnimating={status.type === "running"}
|
||||
plugins={{ code, math, mermaid }}
|
||||
|
|
|
|||
|
|
@ -724,10 +724,10 @@ function startPromptQueue(
|
|||
}
|
||||
}
|
||||
|
||||
function stopPromptQueueRun() {
|
||||
function stopPromptQueueRun(cancelActiveRun = true) {
|
||||
const activeItem = promptQueueItems[Math.max(promptQueueIndex, 0)];
|
||||
const activeTarget = activeItem?.target;
|
||||
const shouldCancelActiveRun = Boolean(activeItem?.dispatched);
|
||||
const shouldCancelActiveRun = cancelActiveRun && Boolean(activeItem?.dispatched);
|
||||
resetPromptQueue();
|
||||
if (!shouldCancelActiveRun) {
|
||||
return;
|
||||
|
|
@ -740,7 +740,11 @@ function stopPromptQueueRun() {
|
|||
}
|
||||
|
||||
if (typeof window !== "undefined") {
|
||||
window.addEventListener(PROMPT_QUEUE_STOP_EVENT, () => stopPromptQueueRun());
|
||||
window.addEventListener(PROMPT_QUEUE_STOP_EVENT, (event) => {
|
||||
// Navigation leaves the dispatched prompt streaming; an explicit stop cancels it too.
|
||||
const detail = (event as CustomEvent<{ cancelActiveRun?: boolean }>).detail;
|
||||
stopPromptQueueRun(detail?.cancelActiveRun ?? true);
|
||||
});
|
||||
}
|
||||
|
||||
interface PromptQueueCallbacks {
|
||||
|
|
@ -2760,9 +2764,28 @@ const ArtifactsToggle: FC = () => {
|
|||
};
|
||||
|
||||
const ToolStatusDisplay: FC = () => {
|
||||
const toolStatus = useChatRuntimeStore((s) => s.toolStatus);
|
||||
// This conversation's tool call only: a global status would put one chat's "Running
|
||||
// Python..." above every composer. remoteId, not id: the adapter keys this map by
|
||||
// unstable_threadId, so reading id lost the status of every restored chat.
|
||||
const threadListItemId = useAuiState(
|
||||
({ threadListItem }) => threadListItem.remoteId,
|
||||
);
|
||||
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const [elapsed, setElapsed] = useState(0);
|
||||
const entry = useChatRuntimeStore((s) => {
|
||||
// A first turn starts before its id is persisted, so the adapter files it under
|
||||
// "__default"; only this thread's own run may claim it. Two first turns share that key
|
||||
// with nothing to tell them apart, so claim it only when it holds one run.
|
||||
const unresolved = s.toolStatusByThreadId.__default;
|
||||
const own =
|
||||
s.toolStatusByThreadId[threadListItemId ?? ""] ??
|
||||
(isThreadRunning && unresolved?.length === 1 ? unresolved : undefined);
|
||||
// Newest of the runs behind this key: separate entries, so one finishing cannot blank
|
||||
// a sibling still running a tool.
|
||||
return own?.[own.length - 1];
|
||||
});
|
||||
const toolStatus = entry?.status ?? null;
|
||||
const startedAt = entry?.startedAt ?? null;
|
||||
const [now, setNow] = useState(() => Date.now());
|
||||
const [visible, setVisible] = useState(false);
|
||||
const visibleRef = useRef(false);
|
||||
|
||||
|
|
@ -2771,15 +2794,14 @@ const ToolStatusDisplay: FC = () => {
|
|||
}, [visible]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!toolStatus) {
|
||||
setElapsed(0);
|
||||
if (!startedAt) {
|
||||
if (!isThreadRunning) {
|
||||
setVisible(false);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
setElapsed(0);
|
||||
setNow(Date.now());
|
||||
|
||||
// Debounce visibility by 300ms when the badge isn't already on screen.
|
||||
// Once visible from a prior tool, later tools show immediately so it
|
||||
|
|
@ -2789,24 +2811,27 @@ const ToolStatusDisplay: FC = () => {
|
|||
showTimer = setTimeout(() => setVisible(true), 300);
|
||||
}
|
||||
|
||||
const interval = setInterval(() => {
|
||||
setElapsed((prev) => prev + 1);
|
||||
}, 1000);
|
||||
const interval = setInterval(() => setNow(Date.now()), 1000);
|
||||
return () => {
|
||||
clearInterval(interval);
|
||||
if (showTimer) {
|
||||
clearTimeout(showTimer);
|
||||
}
|
||||
};
|
||||
}, [toolStatus, isThreadRunning]);
|
||||
}, [startedAt, isThreadRunning]);
|
||||
|
||||
if (!(toolStatus && visible)) {
|
||||
if (!(toolStatus && startedAt && visible)) {
|
||||
return null;
|
||||
}
|
||||
// From the store's start time, so returning to the conversation resumes rather than restarting.
|
||||
const elapsed = Math.max(0, Math.floor((now - startedAt) / 1000));
|
||||
const isRunning = toolStatus.startsWith("Running");
|
||||
const StatusIcon = isRunning ? TerminalIcon : GlobeIcon;
|
||||
return (
|
||||
<div className="mb-2 flex w-full flex-row items-center gap-2 px-1.5 pt-0.5 pb-1">
|
||||
<div
|
||||
data-testid="composer-tool-status"
|
||||
className="mb-2 flex w-full flex-row items-center gap-2 px-1.5 pt-0.5 pb-1"
|
||||
>
|
||||
<div className="flex animate-pulse items-center gap-2 rounded-full border border-primary/20 bg-primary/5 px-3 py-1.5 text-xs text-primary">
|
||||
<StatusIcon className="size-3.5" />
|
||||
<span>{toolStatus}</span>
|
||||
|
|
@ -3769,9 +3794,15 @@ const DiffusionCanvas: FC = () => {
|
|||
const isRunning = useAuiState(
|
||||
({ message }) => message.status?.type === "running",
|
||||
);
|
||||
// A non-null canvas is set only by diffusion_frame events (diffusion models only),
|
||||
// so it is a sufficient gate; loadedIsDiffusion can lag the first frame on a fresh load.
|
||||
const canvas = useChatRuntimeStore((s) => s.activeDiffusionCanvas);
|
||||
// Only this conversation's own frames render here; a first turn has no id yet, so it reads
|
||||
// "__default", which is where its run files them until the thread persists.
|
||||
const threadKey =
|
||||
useAuiState(({ threadListItem }) => threadListItem.remoteId) ?? "__default";
|
||||
// A canvas is set only by diffusion_frame events, so its presence is a sufficient gate;
|
||||
// loadedIsDiffusion can lag the first frame on a fresh load.
|
||||
const canvas = useChatRuntimeStore(
|
||||
(s) => s.activeDiffusionCanvasByThreadId[threadKey],
|
||||
);
|
||||
if (!isRunning || !canvas) {
|
||||
return null;
|
||||
}
|
||||
|
|
|
|||
219
studio/frontend/src/components/assistant-ui/tool-code-cell.tsx
Normal file
219
studio/frontend/src/components/assistant-ui/tool-code-cell.tsx
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { code as codePlugin } from "@streamdown/code";
|
||||
import { CopyIcon, DownloadIcon } from "lucide-react";
|
||||
import { Tick02Icon } from "@/lib/tick-icon";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { Streamdown } from "streamdown";
|
||||
|
||||
const COPY_RESET_MS = 2000;
|
||||
const SHIKI_THEME = ["github-light", "github-dark"] as [
|
||||
"github-light",
|
||||
"github-dark",
|
||||
];
|
||||
/** Past this the block stays plain monospace: shiki is not worth the main-thread time. */
|
||||
const MAX_HIGHLIGHT_CHARS = 20_000;
|
||||
/** Within this many px of the bottom counts as following the stream. */
|
||||
const PIN_SLACK_PX = 40;
|
||||
|
||||
export function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
const copy = useCallback(async () => {
|
||||
if (await copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
function DownloadBtn({ code, name }: { code: string; name: string }) {
|
||||
const download = useCallback(() => {
|
||||
if (typeof document === "undefined") {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const blob = new Blob([code], { type: "text/plain;charset=utf-8" });
|
||||
const url = URL.createObjectURL(blob);
|
||||
const anchor = document.createElement("a");
|
||||
anchor.href = url;
|
||||
anchor.download = name;
|
||||
document.body.appendChild(anchor);
|
||||
anchor.click();
|
||||
anchor.remove();
|
||||
// Revoke next tick, after the click consumes the URL.
|
||||
setTimeout(() => URL.revokeObjectURL(url), 0);
|
||||
} catch {
|
||||
// Never break the transcript over a download.
|
||||
}
|
||||
}, [code, name]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={download}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Download"
|
||||
>
|
||||
<DownloadIcon className="size-3" />
|
||||
Download
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** A fence longer than any backtick run in the code, so a script containing ``` cannot end it early. */
|
||||
function fenceFor(source: string): string {
|
||||
const longest = (source.match(/`+/g) ?? []).reduce(
|
||||
(max, run) => Math.max(max, run.length),
|
||||
0,
|
||||
);
|
||||
return "`".repeat(Math.max(3, longest + 1));
|
||||
}
|
||||
|
||||
/** Syntax-highlighted code via Streamdown + shiki. Always in the DOM as plain monospace, but
|
||||
* shiki only tokenizes once the block nears the viewport, so a long transcript does not
|
||||
* highlight every script up front. Immediate where IntersectionObserver is missing. */
|
||||
function HighlightedCode({
|
||||
code: source,
|
||||
language,
|
||||
plain = false,
|
||||
}: {
|
||||
code: string;
|
||||
language: string;
|
||||
plain?: boolean;
|
||||
}) {
|
||||
const markdown = useMemo(() => {
|
||||
const fence = fenceFor(source);
|
||||
return `${fence}${language}\n${source}\n${fence}`;
|
||||
}, [source, language]);
|
||||
const containerRef = useRef<HTMLDivElement>(null);
|
||||
const [nearViewport, setNearViewport] = useState(
|
||||
() => typeof IntersectionObserver === "undefined",
|
||||
);
|
||||
// Pinned to the bottom until the reader scrolls up, so a streaming payload visibly grows.
|
||||
const pinnedToBottom = useRef(true);
|
||||
useEffect(() => {
|
||||
if (nearViewport) return;
|
||||
const el = containerRef.current;
|
||||
if (!el) return;
|
||||
const io = new IntersectionObserver(
|
||||
(entries) => {
|
||||
if (entries.some((entry) => entry.isIntersecting)) {
|
||||
setNearViewport(true);
|
||||
io.disconnect();
|
||||
}
|
||||
},
|
||||
// Highlight just before the block enters view, so it is ready on arrival.
|
||||
{ rootMargin: "200px" },
|
||||
);
|
||||
io.observe(el);
|
||||
return () => io.disconnect();
|
||||
}, [nearViewport]);
|
||||
|
||||
useEffect(() => {
|
||||
const el = containerRef.current;
|
||||
if (plain && el && pinnedToBottom.current) {
|
||||
el.scrollTop = el.scrollHeight;
|
||||
}
|
||||
}, [plain, source]);
|
||||
|
||||
const handleScroll = () => {
|
||||
const el = containerRef.current;
|
||||
if (el) {
|
||||
pinnedToBottom.current =
|
||||
el.scrollHeight - el.scrollTop - el.clientHeight < PIN_SLACK_PX;
|
||||
}
|
||||
};
|
||||
|
||||
// Skip shiki while the model is writing (it re-tokenizes every fragment) and on payloads too big.
|
||||
const highlight =
|
||||
nearViewport && !plain && source.length <= MAX_HIGHLIGHT_CHARS;
|
||||
|
||||
return (
|
||||
<div
|
||||
ref={containerRef}
|
||||
onScroll={handleScroll}
|
||||
className="max-h-48 overflow-auto text-xs [&_pre]:!m-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_pre]:!text-xs [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!p-3 [&_[data-streamdown=code-block]]:!border-0"
|
||||
>
|
||||
{highlight ? (
|
||||
<Streamdown
|
||||
mode="static"
|
||||
plugins={{ code: codePlugin }}
|
||||
controls={{ code: false }}
|
||||
shikiTheme={SHIKI_THEME}
|
||||
>
|
||||
{markdown}
|
||||
</Streamdown>
|
||||
) : (
|
||||
// A div, not a <pre>: the container's [&_pre]:!p-0 would strip the padding and shift
|
||||
// the content when shiki swaps in. whitespace-pre so long lines scroll.
|
||||
<div className="whitespace-pre p-3 font-mono text-xs text-muted-foreground">
|
||||
{source}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
/** The code a tool is about to run, in the card's collapsible content so the chevron hides code and output together. */
|
||||
export function ToolCodeCell({
|
||||
label,
|
||||
code,
|
||||
language,
|
||||
downloadName,
|
||||
streaming = false,
|
||||
}: {
|
||||
label: string;
|
||||
code: string;
|
||||
language: string;
|
||||
downloadName: string;
|
||||
streaming?: boolean;
|
||||
}) {
|
||||
return (
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">
|
||||
{label}
|
||||
</span>
|
||||
<div className="flex items-center gap-1">
|
||||
<CopyBtn text={code} />
|
||||
<DownloadBtn code={code} name={downloadName} />
|
||||
</div>
|
||||
</div>
|
||||
<HighlightedCode code={code} language={language} plain={streaming} />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -10,7 +10,11 @@ import {
|
|||
} from "react";
|
||||
import { useAuiState } from "@assistant-ui/react";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import { toolOutputKey, useToolPaneScope } from "@/features/chat";
|
||||
import {
|
||||
toolOutputKey,
|
||||
useToolPaneScope,
|
||||
useUnresolvedToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
import { ChevronDownIcon } from "lucide-react";
|
||||
import { Wrench01Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
|
|
@ -239,16 +243,23 @@ const ToolGroupImpl: FC<
|
|||
// Force the group open when any call is receiving tool_output events.
|
||||
const toolLiveOutput = useChatRuntimeStore((s) => s.toolLiveOutput);
|
||||
const paneScope = useToolPaneScope();
|
||||
const unresolvedScope = useUnresolvedToolPaneScope();
|
||||
const hasLiveOutput = useAuiState(({ message }) =>
|
||||
message.parts
|
||||
.slice(startIndex, endIndex + 1)
|
||||
.some(
|
||||
(part) =>
|
||||
part.type === "tool-call" &&
|
||||
Object.prototype.hasOwnProperty.call(
|
||||
// Either scope: a first turn writes under the unresolved one for its whole
|
||||
// life, even after the autosave assigns the id (see useToolOutputFor).
|
||||
(Object.prototype.hasOwnProperty.call(
|
||||
toolLiveOutput,
|
||||
toolOutputKey(paneScope, part.toolCallId),
|
||||
),
|
||||
) ||
|
||||
Object.prototype.hasOwnProperty.call(
|
||||
toolLiveOutput,
|
||||
toolOutputKey(unresolvedScope, part.toolCallId),
|
||||
)),
|
||||
),
|
||||
);
|
||||
// Keep the group open once a confirmation or live output forced it (so an
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
"use client";
|
||||
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import { toolOutputKey, useToolPaneScope } from "@/features/chat";
|
||||
import { useToolOutputFor, useToolPaneScope } from "@/features/chat";
|
||||
import { useEffect, useMemo, useRef } from "react";
|
||||
import { tailText } from "./tool-result-output";
|
||||
|
||||
|
|
@ -16,8 +16,10 @@ import { tailText } from "./tool-result-output";
|
|||
*/
|
||||
export function ToolLiveOutput({ toolCallId }: { toolCallId: string }) {
|
||||
const paneScope = useToolPaneScope();
|
||||
const output = useChatRuntimeStore(
|
||||
(s) => s.toolLiveOutput[toolOutputKey(paneScope, toolCallId)] ?? "",
|
||||
const output = useToolOutputFor(
|
||||
useChatRuntimeStore((s) => s.toolLiveOutput),
|
||||
paneScope,
|
||||
toolCallId,
|
||||
);
|
||||
const scrollRef = useRef<HTMLPreElement>(null);
|
||||
// Pinned to the bottom until the user scrolls up (handler below), so
|
||||
|
|
|
|||
|
|
@ -3,28 +3,25 @@
|
|||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { getAuthToken } from "@/features/auth/session";
|
||||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { useToolArgsStatus } from "@assistant-ui/react";
|
||||
import { code as codePlugin } from "@streamdown/code";
|
||||
import { CodeIcon, CopyIcon, DownloadIcon } from "lucide-react";
|
||||
import { Tick02Icon } from "@/lib/tick-icon";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { CodeIcon } from "lucide-react";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { Streamdown } from "streamdown";
|
||||
import { memo } from "react";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
import { CopyBtn, ToolCodeCell } from "./tool-code-cell";
|
||||
import { ToolLiveOutput } from "./tool-live-output";
|
||||
import { ToolResultOutput } from "./tool-result-output";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
useToolAwaitingApproval,
|
||||
useToolOutputFor,
|
||||
useToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
|
||||
|
|
@ -34,151 +31,6 @@ interface StructuredResult {
|
|||
sessionId: string;
|
||||
}
|
||||
|
||||
const MAX_DISPLAY = 10_000;
|
||||
const COPY_RESET_MS = 2000;
|
||||
const SHIKI_THEME = ["github-light", "github-dark"] as ["github-light", "github-dark"];
|
||||
|
||||
function truncate(text: string): string {
|
||||
return text.length <= MAX_DISPLAY
|
||||
? text
|
||||
: `${text.slice(0, MAX_DISPLAY)}\n... (truncated)`;
|
||||
}
|
||||
|
||||
function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
const copy = useCallback(async () => {
|
||||
if (await copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** Save the script as a .py file via a client-side Blob. */
|
||||
function DownloadBtn({ code, name = "script.py" }: { code: string; name?: string }) {
|
||||
const download = useCallback(() => {
|
||||
if (typeof document === "undefined") {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const blob = new Blob([code], { type: "text/x-python" });
|
||||
const url = URL.createObjectURL(blob);
|
||||
const anchor = document.createElement("a");
|
||||
anchor.href = url;
|
||||
anchor.download = name;
|
||||
document.body.appendChild(anchor);
|
||||
anchor.click();
|
||||
anchor.remove();
|
||||
// Revoke next tick, after the click consumes the URL.
|
||||
setTimeout(() => URL.revokeObjectURL(url), 0);
|
||||
} catch {
|
||||
// Best-effort: never break the transcript over a download.
|
||||
}
|
||||
}, [code, name]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={download}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Download script"
|
||||
>
|
||||
<DownloadIcon className="size-3" />
|
||||
Download
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** Syntax-highlighted code via Streamdown + shiki; inherits parent container.
|
||||
* The script is always in the DOM (a plain monospace placeholder), but shiki
|
||||
* only tokenizes once the block scrolls near the viewport, so a long transcript
|
||||
* with many scripts doesn't highlight every one up front. Falls back to
|
||||
* immediate highlight when IntersectionObserver is unavailable (SSR / tests). */
|
||||
function HighlightedCode({ code: source, language }: { code: string; language: string }) {
|
||||
const display = useMemo(() => truncate(source), [source]);
|
||||
const markdown = useMemo(
|
||||
() => `\`\`\`${language}\n${display}\n\`\`\``,
|
||||
[display, language],
|
||||
);
|
||||
const containerRef = useRef<HTMLDivElement>(null);
|
||||
const [highlight, setHighlight] = useState(
|
||||
() => typeof IntersectionObserver === "undefined",
|
||||
);
|
||||
useEffect(() => {
|
||||
if (highlight) return;
|
||||
const el = containerRef.current;
|
||||
if (!el) return;
|
||||
const io = new IntersectionObserver(
|
||||
(entries) => {
|
||||
if (entries.some((entry) => entry.isIntersecting)) {
|
||||
setHighlight(true);
|
||||
io.disconnect();
|
||||
}
|
||||
},
|
||||
// Highlight just before the block enters view so it's colorized by the
|
||||
// time the user reaches it, without tokenizing off-screen scripts.
|
||||
{ rootMargin: "200px" },
|
||||
);
|
||||
io.observe(el);
|
||||
return () => io.disconnect();
|
||||
}, [highlight]);
|
||||
return (
|
||||
<div
|
||||
ref={containerRef}
|
||||
className="max-h-48 overflow-auto text-xs [&_pre]:!m-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_pre]:!text-xs [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!p-3 [&_[data-streamdown=code-block]]:!border-0"
|
||||
>
|
||||
{highlight ? (
|
||||
<Streamdown
|
||||
mode="static"
|
||||
plugins={{ code: codePlugin }}
|
||||
controls={{ code: false }}
|
||||
shikiTheme={SHIKI_THEME}
|
||||
>
|
||||
{markdown}
|
||||
</Streamdown>
|
||||
) : (
|
||||
// A div, not a <pre>: the container's [&_pre]:!p-0 would override a
|
||||
// <pre>'s padding and shift the content by p-3 when shiki swaps in. Keep
|
||||
// the same p-3, and whitespace-pre (not pre-wrap) so long lines scroll in
|
||||
// the container's overflow-auto exactly like the highlighted <pre>, rather
|
||||
// than wrapping taller and then collapsing when shiki swaps in.
|
||||
<div className="whitespace-pre p-3 font-mono text-xs text-muted-foreground">
|
||||
{display}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function isStructuredResult(val: unknown): val is StructuredResult {
|
||||
return (
|
||||
typeof val === "object" &&
|
||||
|
|
@ -221,46 +73,50 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
// Show the fuller live stream over a truncated result, keeping its exit
|
||||
// status. Session-transient: after a reload only the result remains.
|
||||
const paneScope = useToolPaneScope();
|
||||
const fullOutput = useChatRuntimeStore(
|
||||
(s) => s.toolFullOutput[toolOutputKey(paneScope, toolCallId)] ?? "",
|
||||
const fullOutput = useToolOutputFor(
|
||||
useChatRuntimeStore((s) => s.toolFullOutput),
|
||||
paneScope,
|
||||
toolCallId,
|
||||
);
|
||||
const displayOutput = preferFullToolOutput(fullOutput, output);
|
||||
|
||||
const authToken = getAuthToken();
|
||||
// The gate only opens once the call parsed, so a pending approval means the script is
|
||||
// written even while the args status still reads as streaming.
|
||||
const awaitingApproval = useToolAwaitingApproval(toolCallId);
|
||||
const isWriting = isWritingCode && !awaitingApproval;
|
||||
|
||||
return (
|
||||
// Status/output collapse from history; the script source renders outside
|
||||
// ToolFallbackContent so it stays visible on reopen (#7165).
|
||||
// Script, status and output all collapse behind the one chevron.
|
||||
<ToolFallbackRoot defaultOpen={isRunning}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={firstLine ? `Python: ${firstLine}` : "Python"}
|
||||
status={status}
|
||||
icon={CodeIcon}
|
||||
/>
|
||||
{code && (
|
||||
<div className="mt-1 pl-5">
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">
|
||||
script
|
||||
</span>
|
||||
<div className="flex items-center gap-1">
|
||||
<CopyBtn text={code} />
|
||||
<DownloadBtn code={code} />
|
||||
</div>
|
||||
</div>
|
||||
<HighlightedCode code={code} language="python" />
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
<ToolFallbackContent>
|
||||
{code && (
|
||||
<ToolCodeCell
|
||||
label="script"
|
||||
code={code}
|
||||
language="python"
|
||||
downloadName="script.py"
|
||||
streaming={isWriting}
|
||||
/>
|
||||
)}
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
{/* Output */}
|
||||
{isRunning ? (
|
||||
<>
|
||||
<div className="mt-2 flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<Spinner className="size-3.5" />
|
||||
<span>{isWritingCode ? "Writing code…" : "Running…"}</span>
|
||||
<span>
|
||||
{awaitingApproval
|
||||
? "Waiting for approval…"
|
||||
: isWriting
|
||||
? "Writing code…"
|
||||
: "Running…"}
|
||||
</span>
|
||||
</div>
|
||||
{/* Live stdout streamed via tool_output SSE events. */}
|
||||
<ToolLiveOutput toolCallId={toolCallId} />
|
||||
|
|
|
|||
|
|
@ -3,69 +3,27 @@
|
|||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { useToolArgsStatus } from "@assistant-ui/react";
|
||||
import { CopyIcon, TerminalIcon } from "lucide-react";
|
||||
import { Tick02Icon } from "@/lib/tick-icon";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { TerminalIcon } from "lucide-react";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { memo, useCallback, useEffect, useRef, useState } from "react";
|
||||
import { memo } from "react";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
import { CopyBtn, ToolCodeCell } from "./tool-code-cell";
|
||||
import { ToolLiveOutput } from "./tool-live-output";
|
||||
import { ToolResultOutput } from "./tool-result-output";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
useToolAwaitingApproval,
|
||||
useToolOutputFor,
|
||||
useToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
|
||||
const COPY_RESET_MS = 2000;
|
||||
|
||||
function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
const copy = useCallback(async () => {
|
||||
if (await copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
toolCallId,
|
||||
args,
|
||||
|
|
@ -87,13 +45,19 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
// Show the fuller live stream over a truncated result, keeping its exit
|
||||
// status. Session-transient: after a reload only the result remains.
|
||||
const paneScope = useToolPaneScope();
|
||||
const fullOutput = useChatRuntimeStore(
|
||||
(s) => s.toolFullOutput[toolOutputKey(paneScope, toolCallId)] ?? "",
|
||||
const fullOutput = useToolOutputFor(
|
||||
useChatRuntimeStore((s) => s.toolFullOutput),
|
||||
paneScope,
|
||||
toolCallId,
|
||||
);
|
||||
const displayOutput = preferFullToolOutput(fullOutput, output);
|
||||
// The gate only opens once the call parsed, so a pending approval means the command is
|
||||
// written even while the args status still reads as streaming.
|
||||
const awaitingApproval = useToolAwaitingApproval(toolCallId);
|
||||
const isWriting = isWritingCommand && !awaitingApproval;
|
||||
|
||||
return (
|
||||
// Open when mounted mid-run so live output shows; collapsed from history.
|
||||
// Open mid-run so command and live output show, collapsed from history.
|
||||
<ToolFallbackRoot defaultOpen={isRunning}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={command ? `$ ${command.slice(0, 60)}` : "Terminal"}
|
||||
|
|
@ -101,12 +65,27 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
icon={TerminalIcon}
|
||||
/>
|
||||
<ToolFallbackContent>
|
||||
{command && (
|
||||
<ToolCodeCell
|
||||
label="command"
|
||||
code={command}
|
||||
language="bash"
|
||||
downloadName="command.sh"
|
||||
streaming={isWriting}
|
||||
/>
|
||||
)}
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
{isRunning ? (
|
||||
<>
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<Spinner className="size-3.5" />
|
||||
<span>{isWritingCommand ? "Writing command…" : "Running…"}</span>
|
||||
<span>
|
||||
{awaitingApproval
|
||||
? "Waiting for approval…"
|
||||
: isWriting
|
||||
? "Writing command…"
|
||||
: "Running…"}
|
||||
</span>
|
||||
</div>
|
||||
{/* Live stdout streamed via tool_output SSE events. */}
|
||||
<ToolLiveOutput toolCallId={toolCallId} />
|
||||
|
|
|
|||
|
|
@ -8,8 +8,8 @@ import {
|
|||
MultiplicationSignCircleIcon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { useTheme } from "@/features/settings/stores/theme-store";
|
||||
import { createLoadingToastIcon } from "@/lib/toast";
|
||||
import { Toaster as Sonner, type ToasterProps } from "sonner";
|
||||
|
||||
// Make toast text selectable. Sonner's onPointerDown calls setPointerCapture(),
|
||||
|
|
@ -78,7 +78,7 @@ const Toaster = ({ ...props }: ToasterProps) => {
|
|||
/>
|
||||
),
|
||||
// App-wide arc spinner so loading toasts match the "Downloading model" toast.
|
||||
loading: <Spinner className="size-4 text-muted-foreground" />,
|
||||
loading: createLoadingToastIcon(),
|
||||
}}
|
||||
style={
|
||||
{
|
||||
|
|
|
|||
|
|
@ -6,15 +6,22 @@
|
|||
import { Loader2Icon } from "lucide-react";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
/**
|
||||
* App-wide spinner: a clean circular arc with a rounded cap (lucide
|
||||
* Loader2 / LoaderCircle), animated, inheriting the current text color.
|
||||
*/
|
||||
function Spinner({ className }: { className?: string }) {
|
||||
/** App-wide spinner inheriting the current text color. `label` overrides the announcement
|
||||
* where "loading" is not what it means (a sidebar chat is generating). */
|
||||
function Spinner({
|
||||
className,
|
||||
label = "Loading",
|
||||
"data-testid": dataTestId,
|
||||
}: {
|
||||
className?: string;
|
||||
label?: string;
|
||||
"data-testid"?: string;
|
||||
}) {
|
||||
return (
|
||||
<Loader2Icon
|
||||
role="status"
|
||||
aria-label="Loading"
|
||||
aria-label={label}
|
||||
data-testid={dataTestId}
|
||||
className={cn("size-4 shrink-0 animate-spin", className)}
|
||||
/>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ 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";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { createLoadingToastIcon, toast } from "@/lib/toast";
|
||||
import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core";
|
||||
import type { ChatModelAdapter } from "@assistant-ui/react";
|
||||
import { parsePartialJsonObject } from "assistant-stream/utils";
|
||||
|
|
@ -59,6 +59,7 @@ import {
|
|||
shouldPreserveFullOutput,
|
||||
toolOutputKey,
|
||||
toolPaneScope,
|
||||
toolThreadScope,
|
||||
} from "../tool-output-scope";
|
||||
import type { ModelType } from "../types";
|
||||
import { isMultimodalResponse } from "../types/api";
|
||||
|
|
@ -1512,13 +1513,38 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
const trustRemoteCode = store.params.trustRemoteCode ?? false;
|
||||
const specSettings = resolveSpeculativeSettingsForLoad();
|
||||
const lastLoaded = readLastLocalModelLoad();
|
||||
const toastId = toast("Loading a model…", {
|
||||
let autoLoadToastDismissed = false;
|
||||
const toastId = toast.message("Loading a model…", {
|
||||
description: lastLoaded
|
||||
? "Loading last used model."
|
||||
: "Auto-selecting the smallest downloaded model.",
|
||||
duration: 5000,
|
||||
duration: Number.POSITIVE_INFINITY,
|
||||
closeButton: true,
|
||||
icon: createLoadingToastIcon(),
|
||||
onDismiss: () => {
|
||||
autoLoadToastDismissed = true;
|
||||
},
|
||||
});
|
||||
const updateAutoLoadToast = (message: string, description: string): void => {
|
||||
if (autoLoadToastDismissed) return;
|
||||
toast.message(message, {
|
||||
id: toastId,
|
||||
description,
|
||||
duration: Number.POSITIVE_INFINITY,
|
||||
});
|
||||
};
|
||||
const showAutoLoadSuccess = (message: string): void => {
|
||||
const options = {
|
||||
description: undefined,
|
||||
duration: 5000,
|
||||
icon: undefined,
|
||||
};
|
||||
if (autoLoadToastDismissed) {
|
||||
toast.success(message, options);
|
||||
return;
|
||||
}
|
||||
toast.success(message, { ...options, id: toastId });
|
||||
};
|
||||
let blockedByTrustRemoteCode = false;
|
||||
let hadNonTrustFailure = false;
|
||||
let loadAttempts = 0;
|
||||
|
|
@ -1781,7 +1807,7 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
ggufVariant: candidate.ggufVariant,
|
||||
});
|
||||
}
|
||||
toast.success(candidate.successLabel, { id: toastId });
|
||||
showAutoLoadSuccess(candidate.successLabel);
|
||||
return true;
|
||||
}
|
||||
try {
|
||||
|
|
@ -1807,11 +1833,10 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
isAutoLoadableGgufVariant(entry),
|
||||
);
|
||||
if (variant) {
|
||||
toast("Loading last used model…", {
|
||||
id: toastId,
|
||||
description: `${repo.repo_id} (${variant.quant})`,
|
||||
duration: 5000,
|
||||
});
|
||||
updateAutoLoadToast(
|
||||
"Loading last used model…",
|
||||
`${repo.repo_id} (${variant.quant})`,
|
||||
);
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
|
|
@ -1836,11 +1861,7 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
const repo = findCachedRepo(modelRepos, lastLoaded.id);
|
||||
if (repo) {
|
||||
try {
|
||||
toast("Loading last used model…", {
|
||||
id: toastId,
|
||||
description: repo.repo_id,
|
||||
duration: 5000,
|
||||
});
|
||||
updateAutoLoadToast("Loading last used model…", repo.repo_id);
|
||||
if (
|
||||
await loadAutoLoadCandidate({
|
||||
id: repo.repo_id,
|
||||
|
|
@ -1861,11 +1882,10 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
}
|
||||
}
|
||||
}
|
||||
toast("Loading a model…", {
|
||||
id: toastId,
|
||||
description: "Auto-selecting the smallest downloaded model.",
|
||||
duration: 5000,
|
||||
});
|
||||
updateAutoLoadToast(
|
||||
"Loading a model…",
|
||||
"Auto-selecting the smallest downloaded model.",
|
||||
);
|
||||
}
|
||||
|
||||
// GGUF first: smallest-total-size repo, then its smallest variant.
|
||||
|
|
@ -1956,12 +1976,10 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
}
|
||||
|
||||
// No cached models — try downloading a small default GGUF.
|
||||
toast("Downloading a small model…", {
|
||||
id: toastId,
|
||||
description:
|
||||
"No downloaded models found. Fetching Qwen3.5-4B-MTP (UD-Q4_K_XL).",
|
||||
duration: 30000,
|
||||
});
|
||||
updateAutoLoadToast(
|
||||
"Downloading a small model…",
|
||||
"No downloaded models found. Fetching Qwen3.5-4B-MTP (UD-Q4_K_XL).",
|
||||
);
|
||||
try {
|
||||
const rt = useChatRuntimeStore.getState();
|
||||
if (
|
||||
|
|
@ -2057,7 +2075,7 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
kind: "gguf",
|
||||
ggufVariant: "UD-Q4_K_XL",
|
||||
});
|
||||
toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId });
|
||||
showAutoLoadSuccess("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)");
|
||||
return { loaded: true, blockedByTrustRemoteCode: false };
|
||||
} catch {
|
||||
toast.dismiss(toastId);
|
||||
|
|
@ -2222,7 +2240,20 @@ export function createOpenAIStreamAdapter(
|
|||
: undefined;
|
||||
|
||||
const threadKey = resolvedThreadId;
|
||||
runtime.setThreadRunning(threadKey, true);
|
||||
// The run is durable on the server, but Stop, archive and delete reach a background
|
||||
// thread only through this map: without a handle the supervisor kept planning against
|
||||
// a deleted conversation. Registered before the run exists, since the thread can be
|
||||
// stopped while createResearchRun is still in flight.
|
||||
let researchRunId: string | null = null;
|
||||
let researchStopRequested = false;
|
||||
const researchServerCancel = () => {
|
||||
researchStopRequested = true;
|
||||
if (researchRunId) {
|
||||
void cancelResearchRun(researchRunId).catch(() => {});
|
||||
}
|
||||
};
|
||||
runtime.registerThreadServerCancel(threadKey, researchServerCancel);
|
||||
runtime.setThreadRunning(threadKey, true, { owner: researchServerCancel });
|
||||
let report = "";
|
||||
let releaseResearchFollow: (() => void) | null = null;
|
||||
const researchFollowController = new AbortController();
|
||||
|
|
@ -2262,6 +2293,13 @@ export function createOpenAIStreamAdapter(
|
|||
blockedDomains: [...runtime.researchWebsitePolicy.blockedDomains],
|
||||
},
|
||||
});
|
||||
researchRunId = createdRun.id;
|
||||
if (researchStopRequested) {
|
||||
// Stopped while createResearchRun was still in flight, so the handle had no
|
||||
// id to act on. Replay it rather than following a run the user already ended.
|
||||
void cancelResearchRun(createdRun.id).catch(() => {});
|
||||
return;
|
||||
}
|
||||
releaseResearchFollow = beginExternalResearchFollow(
|
||||
createdRun,
|
||||
detachResearchFollow,
|
||||
|
|
@ -2320,7 +2358,8 @@ export function createOpenAIStreamAdapter(
|
|||
} finally {
|
||||
abortSignal.removeEventListener("abort", forwardAdapterAbort);
|
||||
releaseResearchFollow?.();
|
||||
runtime.setThreadRunning(threadKey, false);
|
||||
runtime.clearThreadServerCancel(threadKey, researchServerCancel);
|
||||
runtime.setThreadRunning(threadKey, false, { owner: researchServerCancel });
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
|
@ -2329,17 +2368,21 @@ export function createOpenAIStreamAdapter(
|
|||
? `${sandboxSessionId || "_default"}:${resolvedThreadId}`
|
||||
: sandboxSessionId || "_default";
|
||||
const toolConfirmationIdsByBackendId = new Map<string, string>();
|
||||
// Store keys are pane-scoped since local tool ids ("call_0") repeat across
|
||||
// turns and concurrent panes (compare mode). Track this run's keys so
|
||||
// cleanup can't wipe another pane's.
|
||||
const toolOutputPaneScope = toolPaneScope(
|
||||
options.modelType,
|
||||
options.pairId,
|
||||
// Local tool ids ("call_0") repeat across turns, panes and conversations, so scope by pane
|
||||
// AND thread. unstable_threadId alone, no activeThreadId fallback: the reader has only
|
||||
// threadListItem.remoteId, which is exactly this value.
|
||||
const toolOutputPaneScope = toolThreadScope(
|
||||
toolPaneScope(options.modelType, options.pairId),
|
||||
unstable_threadId,
|
||||
);
|
||||
const scopedToolOutputKey = (id: string) =>
|
||||
toolOutputKey(toolOutputPaneScope, id);
|
||||
const runToolLiveOutputKeys = new Set<string>();
|
||||
const resolvedThreadKey = resolvedThreadId ?? null;
|
||||
// Which conversation was on screen when this run started. A first turn has no id yet, so
|
||||
// this is the only way to tell later whether the user has switched away from it.
|
||||
const activeThreadIdAtRunStart =
|
||||
useChatRuntimeStore.getState().activeThreadId ?? null;
|
||||
const pendingImageEditReferenceForRun = runtime.pendingImageEditReference;
|
||||
const selectedImageEditReference =
|
||||
(pendingImageEditReferenceForRun?.threadId ?? null) ===
|
||||
|
|
@ -2745,8 +2788,11 @@ export function createOpenAIStreamAdapter(
|
|||
// waitForRunEnd resolves instead of hanging: this gate fires
|
||||
// before the streaming path's setThreadRunning(true).
|
||||
const gatedThreadKey = resolvedThreadId || "__default";
|
||||
runtime.setThreadRunning(gatedThreadKey, true);
|
||||
runtime.setThreadRunning(gatedThreadKey, false);
|
||||
// Own token: siblings share "__default", so an ownerless clear would drop their
|
||||
// entries while they are still generating.
|
||||
const gateOwner = () => {};
|
||||
runtime.setThreadRunning(gatedThreadKey, true, { owner: gateOwner });
|
||||
runtime.setThreadRunning(gatedThreadKey, false, { owner: gateOwner });
|
||||
clearSelectedImageEditReference();
|
||||
throw new Error(imageGateReason);
|
||||
}
|
||||
|
|
@ -2764,13 +2810,44 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
const useAdapter = await resolveUseAdapter(resolvedThreadId, options);
|
||||
|
||||
const threadKey = resolvedThreadId || "__default";
|
||||
// A first turn files its handles under "__default"; autosave then assigns a real id and
|
||||
// adoptDefaultThreadRun re-keys them mid-run. Resolve per use so later writes and the
|
||||
// final clear follow the run instead of stranding entries behind.
|
||||
const liveThreadKey = (owner: () => void) =>
|
||||
threadKey === "__default"
|
||||
? useChatRuntimeStore.getState().runKeyForOwner(threadKey, owner)
|
||||
: threadKey;
|
||||
|
||||
// Per-run token so a delayed stop POST can't match the next run.
|
||||
const cancelId =
|
||||
typeof crypto !== "undefined" && "randomUUID" in crypto
|
||||
? crypto.randomUUID()
|
||||
: `${Date.now()}-${Math.random().toString(36).slice(2)}`;
|
||||
|
||||
// Per-run abort, chained to assistant-ui's signal. cancelByThreadId only holds the visible
|
||||
// thread's cancelRun(), so this controller is the only way to end a backgrounded chat's
|
||||
// request; the cancel POST below reaches llama-server only.
|
||||
const runAbort = new AbortController();
|
||||
const runSignal = runAbort.signal;
|
||||
const forwardAbort = () => runAbort.abort(abortSignal.reason);
|
||||
// Declared here, not at its registration below: it doubles as this run's identity token
|
||||
// on the per-thread maps (see registerThreadServerCancel).
|
||||
const serverCancel = () => runAbort.abort();
|
||||
if (abortSignal.aborted) {
|
||||
forwardAbort();
|
||||
} else {
|
||||
abortSignal.addEventListener("abort", forwardAbort, { once: true });
|
||||
}
|
||||
|
||||
// ── Audio model path (non-streaming) ─────────────────────
|
||||
const activeModel = runtime.models.find(
|
||||
(m) => m.id === params.checkpoint,
|
||||
);
|
||||
if (activeModel?.isAudio && !activeModel?.hasAudioInput) {
|
||||
const threadKey = resolvedThreadId || "__default";
|
||||
runtime.setThreadRunning(threadKey, true);
|
||||
const audioCancel = () => runAbort.abort();
|
||||
runtime.registerThreadServerCancel(threadKey, audioCancel);
|
||||
runtime.setThreadRunning(threadKey, true, { owner: audioCancel });
|
||||
try {
|
||||
yield {
|
||||
content: [{ type: "text" as const, text: "Generating audio..." }],
|
||||
|
|
@ -2780,6 +2857,10 @@ export function createOpenAIStreamAdapter(
|
|||
{
|
||||
model: params.checkpoint,
|
||||
messages: outboundMessages,
|
||||
// Same run in both registries: without it the backend files this under no
|
||||
// thread, and the stop-chats prompt counts the named local run and the
|
||||
// unnamed backend one as two.
|
||||
...(resolvedThreadId ? { thread_id: resolvedThreadId } : {}),
|
||||
stream: false,
|
||||
temperature: params.temperature,
|
||||
top_p: params.topP,
|
||||
|
|
@ -2790,7 +2871,7 @@ export function createOpenAIStreamAdapter(
|
|||
presence_penalty: params.presencePenalty,
|
||||
...(useAdapter === undefined ? {} : { use_adapter: useAdapter }),
|
||||
},
|
||||
abortSignal,
|
||||
runSignal,
|
||||
);
|
||||
|
||||
const audioUrl = `data:audio/wav;base64,${result.audio.data}`;
|
||||
|
|
@ -2803,19 +2884,21 @@ export function createOpenAIStreamAdapter(
|
|||
],
|
||||
};
|
||||
} catch (err) {
|
||||
if (!abortSignal.aborted) {
|
||||
if (!runSignal.aborted) {
|
||||
toast.error("Audio generation failed", {
|
||||
description: err instanceof Error ? err.message : "Unknown error",
|
||||
});
|
||||
}
|
||||
throw err;
|
||||
} finally {
|
||||
runtime.setThreadRunning(threadKey, false);
|
||||
abortSignal.removeEventListener("abort", forwardAbort);
|
||||
const audioKey = liveThreadKey(audioCancel);
|
||||
runtime.setThreadRunning(audioKey, false, { owner: audioCancel });
|
||||
runtime.clearThreadServerCancel(audioKey, audioCancel);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const threadKey = resolvedThreadId || "__default";
|
||||
let waitingFirstChunk = true;
|
||||
let firstTokenSettled = false;
|
||||
const streamStartTime = Date.now();
|
||||
|
|
@ -2846,10 +2929,15 @@ export function createOpenAIStreamAdapter(
|
|||
const warmupDelayMs = 450;
|
||||
const warmupTimer = setTimeout(() => {
|
||||
if (!waitingFirstChunk) return;
|
||||
if (abortSignal.aborted) return;
|
||||
if (runSignal.aborted) return;
|
||||
runtime.setGeneratingStatus("waiting");
|
||||
}, warmupDelayMs);
|
||||
runtime.setThreadRunning(threadKey, true);
|
||||
// Flagged local/external so the model-swap gate only counts the chats a reload ends; the
|
||||
// backend leaves external-provider runs out of active_generations for the same reason.
|
||||
runtime.setThreadRunning(threadKey, true, {
|
||||
local: !isExternalRequest,
|
||||
owner: serverCancel,
|
||||
});
|
||||
let cumulativeText = "";
|
||||
let reasoningStartAt: number | null = null;
|
||||
let reasoningDuration = 0;
|
||||
|
|
@ -3015,21 +3103,12 @@ export function createOpenAIStreamAdapter(
|
|||
timings?: ServerTimings;
|
||||
} | null = null;
|
||||
|
||||
// Per-run cancellation token so a delayed stop POST can't match
|
||||
// the next run on the same thread.
|
||||
const cancelId =
|
||||
typeof crypto !== "undefined" && "randomUUID" in crypto
|
||||
? crypto.randomUUID()
|
||||
: `${Date.now()}-${Math.random().toString(36).slice(2)}`;
|
||||
|
||||
// Colab-style proxies can swallow fetch aborts, so also POST
|
||||
// /inference/cancel explicitly on abort.
|
||||
const onAbortCancel = () => {
|
||||
// assistant-ui aborts with AbortError(detach=true) when a thread's runtime
|
||||
// unmounts (navigation / background thread switch) and detach=false for an
|
||||
// explicit Stop. Only a real Stop cancels the backend run; a detach must
|
||||
// leave a backgrounded generation streaming.
|
||||
if ((abortSignal.reason as { detach?: boolean } | undefined)?.detach) {
|
||||
// assistant-ui aborts with detach=true when a runtime unmounts and detach=false for an
|
||||
// explicit Stop. Only a real Stop cancels the backend run; runSignal forwards the reason.
|
||||
if ((runSignal.reason as { detach?: boolean } | undefined)?.detach) {
|
||||
return;
|
||||
}
|
||||
const body: Record<string, string> = { cancel_id: cancelId };
|
||||
|
|
@ -3050,11 +3129,17 @@ export function createOpenAIStreamAdapter(
|
|||
keepalive: true,
|
||||
}).catch(() => {});
|
||||
};
|
||||
|
||||
// Stop handle for when this conversation is not the visible one, which cancelByThreadId
|
||||
// cannot reach. Aborting this run's own controller closes just its request, and the
|
||||
// listener above posts its cancel_id so llama-server stops decoding too. For an
|
||||
// external provider the abort is the stop, since its cancel_id is never registered.
|
||||
runtime.registerThreadServerCancel(threadKey, serverCancel);
|
||||
try {
|
||||
if (abortSignal.aborted) {
|
||||
if (runSignal.aborted) {
|
||||
onAbortCancel();
|
||||
} else {
|
||||
abortSignal.addEventListener("abort", onAbortCancel, { once: true });
|
||||
runSignal.addEventListener("abort", onAbortCancel, { once: true });
|
||||
}
|
||||
|
||||
const {
|
||||
|
|
@ -3526,7 +3611,7 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
clearSelectedImageEditReference();
|
||||
await ThreadAutosaveHandle.awaitFirstSave(resolvedThreadId);
|
||||
const stream = streamChatCompletions(requestPayload, abortSignal);
|
||||
const stream = streamChatCompletions(requestPayload, runSignal);
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const chunkModel = (chunk as { model?: unknown }).model;
|
||||
|
|
@ -3539,7 +3624,11 @@ export function createOpenAIStreamAdapter(
|
|||
chunk as unknown as { _toolStatus?: string }
|
||||
)._toolStatus;
|
||||
if (toolStatusText !== undefined) {
|
||||
runtime.setToolStatus(toolStatusText || null);
|
||||
runtime.setToolStatus(
|
||||
liveThreadKey(serverCancel),
|
||||
toolStatusText || null,
|
||||
serverCancel,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
|
|
@ -3568,7 +3657,9 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
)._diffusionFrame;
|
||||
if (diffusionFrame !== undefined) {
|
||||
runtime.setActiveDiffusionCanvas({
|
||||
// Keyed by thread so a background run's frames stay out of the visible chat
|
||||
// instead of overwriting the frame it is painting.
|
||||
runtime.setActiveDiffusionCanvas(liveThreadKey(serverCancel), {
|
||||
block: diffusionFrame.block ?? 0,
|
||||
step: diffusionFrame.step ?? 0,
|
||||
total: diffusionFrame.total ?? 0,
|
||||
|
|
@ -3709,8 +3800,16 @@ export function createOpenAIStreamAdapter(
|
|||
const approvalId = (toolEvent.approval_id as string) || "";
|
||||
const awaitingConfirmation =
|
||||
toolEvent.awaiting_confirmation === true;
|
||||
// Reuse a provisional card's part id, else the confirmation-scoped id
|
||||
// opens a second card and the first spins "Running" forever.
|
||||
const openPartId = backendToolCallId
|
||||
? toolPartIdByBackendId.get(backendToolCallId)
|
||||
: undefined;
|
||||
const reuseOpenPart =
|
||||
!!openPartId &&
|
||||
toolCallParts.some((p) => p.toolCallId === openPartId);
|
||||
const id =
|
||||
awaitingConfirmation && approvalId
|
||||
awaitingConfirmation && approvalId && !reuseOpenPart
|
||||
? `${toolConfirmationScopeId}:${approvalId}`
|
||||
: backendToolCallId
|
||||
? resolveToolPartId(backendToolCallId)
|
||||
|
|
@ -4289,9 +4388,17 @@ export function createOpenAIStreamAdapter(
|
|||
// Anthropic-only (billed at the write premium).
|
||||
const cacheWriteTokens = meta?.usage?.cache_creation_input_tokens ?? 0;
|
||||
|
||||
// Gate on the captured checkpoint still being active so a late
|
||||
// completion from provider A doesn't populate the bar after a
|
||||
// mid-stream switch to provider B.
|
||||
// Gate on the captured checkpoint so a late completion from provider A cannot populate
|
||||
// the bar after a mid-stream switch to B, and on the captured thread so a background
|
||||
// run finishing after New Chat cannot repaint another chat's usage. An unresolved run
|
||||
// has no id to compare, so compare what was on screen when it started. A first turn is
|
||||
// adopted onto an id mid-run and autosave moves activeThreadId with it, so read the
|
||||
// adopted key, or the run stays "unresolved" for life and the bar stays blank.
|
||||
const usageKey = liveThreadKey(serverCancel);
|
||||
const usageThreadKey = usageKey === "__default" ? null : usageKey;
|
||||
const usageThreadIsVisible =
|
||||
useChatRuntimeStore.getState().activeThreadId ===
|
||||
(usageThreadKey ?? activeThreadIdAtRunStart);
|
||||
if (
|
||||
meta?.usage &&
|
||||
typeof meta.usage.prompt_tokens === "number" &&
|
||||
|
|
@ -4299,13 +4406,23 @@ export function createOpenAIStreamAdapter(
|
|||
typeof meta.usage.total_tokens === "number" &&
|
||||
useChatRuntimeStore.getState().params.checkpoint === params.checkpoint
|
||||
) {
|
||||
useChatRuntimeStore.getState().setContextUsage({
|
||||
const usage = {
|
||||
promptTokens: meta.usage.prompt_tokens,
|
||||
completionTokens: meta.usage.completion_tokens,
|
||||
totalTokens: meta.usage.total_tokens,
|
||||
cachedTokens,
|
||||
cacheWriteTokens,
|
||||
});
|
||||
};
|
||||
// File it under this run's own thread even when the gate below blocks the visible
|
||||
// write, so switching back re-applies it.
|
||||
if (usageThreadKey !== null) {
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.setThreadContextUsage(usageThreadKey, usage);
|
||||
}
|
||||
if (usageThreadIsVisible) {
|
||||
useChatRuntimeStore.getState().setContextUsage(usage);
|
||||
}
|
||||
}
|
||||
|
||||
const finishedAt = Date.now();
|
||||
|
|
@ -4358,7 +4475,7 @@ export function createOpenAIStreamAdapter(
|
|||
settleFirstTokenErr(
|
||||
err instanceof Error ? err : new Error("Generation failed"),
|
||||
);
|
||||
if (!abortSignal.aborted) {
|
||||
if (!runSignal.aborted) {
|
||||
const msg = err instanceof Error ? err.message : String(err);
|
||||
if (err instanceof GenerationLengthError) {
|
||||
toast.error("Response ran out of tokens", {
|
||||
|
|
@ -4396,13 +4513,18 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
throw err;
|
||||
} finally {
|
||||
abortSignal.removeEventListener("abort", onAbortCancel);
|
||||
runSignal.removeEventListener("abort", onAbortCancel);
|
||||
abortSignal.removeEventListener("abort", forwardAbort);
|
||||
// Resolve once: the clears below drop the owner the lookup keys on.
|
||||
const cleanupKey = liveThreadKey(serverCancel);
|
||||
const confirmStore = useChatRuntimeStore.getState();
|
||||
for (const part of toolCallParts) {
|
||||
confirmStore.clearToolConfirmation(part.toolCallId);
|
||||
}
|
||||
runtime.setGeneratingStatus(null);
|
||||
runtime.setToolStatus(null);
|
||||
// Scoped by thread AND by run: a global clear wiped every other running chat's badge,
|
||||
// and an unowned one wiped a concurrent run's badge behind the same key.
|
||||
runtime.setToolStatus(cleanupKey, null, serverCancel);
|
||||
// Clear only this run's live keys (a concurrent pane owns its own). A
|
||||
// key still here streamed stdout but never reached tool_end (SSE drop or
|
||||
// cancel), so promote it to full output first, else the partial
|
||||
|
|
@ -4416,20 +4538,23 @@ export function createOpenAIStreamAdapter(
|
|||
store.clearToolLiveOutput(liveKey);
|
||||
}
|
||||
runToolLiveOutputKeys.clear();
|
||||
// Drop the transient denoising canvas so the finished bubble shows only
|
||||
// the committed markdown answer (cancellation/error included).
|
||||
runtime.setActiveDiffusionCanvas(null);
|
||||
// Drop the transient denoising canvas so the finished bubble shows only the committed
|
||||
// answer. Scoped: a global clear wiped another denoising chat's frame.
|
||||
runtime.clearActiveDiffusionCanvasForThread(cleanupKey);
|
||||
clearTimeout(warmupTimer);
|
||||
if (waitingFirstChunk) {
|
||||
if (firstTokenSettled) {
|
||||
settleFirstTokenOk();
|
||||
} else if (abortSignal.aborted) {
|
||||
} else if (runSignal.aborted) {
|
||||
settleFirstTokenErr(new Error("Cancelled"));
|
||||
} else {
|
||||
settleFirstTokenErr(new Error("No tokens received"));
|
||||
}
|
||||
}
|
||||
runtime.setThreadRunning(threadKey, false);
|
||||
// serverCancel narrows both clears: runs with no resolved thread id share the "__default"
|
||||
// key, so a blind clear could drop a sibling's entry.
|
||||
runtime.setThreadRunning(cleanupKey, false, { owner: serverCancel });
|
||||
runtime.clearThreadServerCancel(cleanupKey, serverCancel);
|
||||
}
|
||||
},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -129,6 +129,27 @@ export async function getApiMonitorEntry(id: string): Promise<ApiMonitorEntry> {
|
|||
return parseJsonOrThrow<ApiMonitorEntry>(response);
|
||||
}
|
||||
|
||||
export interface ActiveGenerationsResponse {
|
||||
count: number;
|
||||
/** Conversations with a generation in flight. Shorter than `count` when a
|
||||
* first turn started before its thread id was persisted. */
|
||||
thread_ids: string[];
|
||||
/** One entry per in-flight request. `kind` is "chat" unless it is an
|
||||
* embeddings / completions / audio call, which has no conversation. */
|
||||
active?: { thread_id: string | null; kind?: string }[];
|
||||
parallel_slots: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Chats generating on the backend right now. Authoritative where `runningByThreadId` is not:
|
||||
* that map is per-tab, empty after a reload and blind to a second tab, and /load and /unload
|
||||
* 409 on these.
|
||||
*/
|
||||
export async function getActiveGenerations(): Promise<ActiveGenerationsResponse> {
|
||||
const response = await authFetch("/api/inference/active-generations");
|
||||
return parseJsonOrThrow<ActiveGenerationsResponse>(response);
|
||||
}
|
||||
|
||||
export async function loadModel(
|
||||
payload: LoadModelRequest,
|
||||
): Promise<LoadModelResponse> {
|
||||
|
|
|
|||
|
|
@ -2682,9 +2682,10 @@ export function ChatPage({
|
|||
ggufNativeContextLength: null,
|
||||
activeNativePathToken: null,
|
||||
activeNativePathExpiresAtMs: null,
|
||||
// Clear previous-model counters, else the relaxed external-provider
|
||||
// render gate shows stale stats until the next completion.
|
||||
// Clear previous-model counters, else the relaxed external-provider render gate shows
|
||||
// stale stats. The per-thread copies go too, so a switch back cannot re-apply.
|
||||
contextUsage: null,
|
||||
contextUsageByThreadId: {},
|
||||
supportsReasoning: reasoningCaps.supportsReasoning,
|
||||
reasoningAlwaysOn: reasoningCaps.reasoningAlwaysOn,
|
||||
reasoningStyle: reasoningCaps.reasoningStyle,
|
||||
|
|
@ -2906,7 +2907,13 @@ export function ChatPage({
|
|||
) {
|
||||
return;
|
||||
}
|
||||
store.setContextUsage(usage);
|
||||
// Key by the thread this restore read, like the history loader: the await above can
|
||||
// outlast a switch away, and an unkeyed write would file this thread's usage under
|
||||
// the incoming one.
|
||||
store.setThreadContextUsage(threadId, usage);
|
||||
if (store.activeThreadId === threadId) {
|
||||
store.setContextUsage(usage);
|
||||
}
|
||||
})
|
||||
.catch((error) => {
|
||||
if (!isExpectedBackgroundChatStorageError(error)) {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,92 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
AlertDialogCancel,
|
||||
AlertDialogContent,
|
||||
AlertDialogDescription,
|
||||
AlertDialogFooter,
|
||||
AlertDialogHeader,
|
||||
AlertDialogTitle,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
import { useStopRunningChatsDialogStore } from "../stores/stop-running-chats-dialog-store";
|
||||
|
||||
/**
|
||||
* Confirmation for applying a model or reload-required setting while chats are generating.
|
||||
* They share one llama-server, so the swap ends all of them: name them and make the user
|
||||
* opt in rather than truncating silently.
|
||||
*/
|
||||
export function StopRunningChatsDialog() {
|
||||
const open = useStopRunningChatsDialogStore((s) => s.open);
|
||||
const count = useStopRunningChatsDialogStore((s) => s.count);
|
||||
const titles = useStopRunningChatsDialogStore((s) => s.titles);
|
||||
const action = useStopRunningChatsDialogStore((s) => s.action);
|
||||
const hasNonChat = useStopRunningChatsDialogStore((s) => s.hasNonChat);
|
||||
const effect = useStopRunningChatsDialogStore((s) => s.effect);
|
||||
const resolve = useStopRunningChatsDialogStore((s) => s.resolve);
|
||||
|
||||
// Embeddings, raw completions and audio share the model but are not conversations,
|
||||
// so name them generically rather than offering to stop chats that do not exist.
|
||||
const noun = hasNonChat
|
||||
? count === 1
|
||||
? "request"
|
||||
: "requests"
|
||||
: count === 1
|
||||
? "chat"
|
||||
: "chats";
|
||||
const sharer = hasNonChat ? "request" : "conversation";
|
||||
// Ejecting leaves no model loaded. Saying it "reloads the model" and offering "Stop and
|
||||
// reload" promised the opposite of what confirming does, for the destructive one.
|
||||
const unloads = effect === "unload";
|
||||
const lead = unloads
|
||||
? `${action || "Unloading the model"} leaves no model loaded, and every open ${sharer} shares it, `
|
||||
: `${action ? `${action} reloads the model, ` : "Reloading the model "}which every open ${sharer} shares, `;
|
||||
const shown = titles.slice(0, 5);
|
||||
const remaining = Math.max(0, titles.length - shown.length);
|
||||
|
||||
return (
|
||||
<AlertDialog
|
||||
open={open}
|
||||
onOpenChange={(next) => {
|
||||
// Escape / overlay click must resolve, or the caller's await hangs.
|
||||
if (!next) resolve(false);
|
||||
}}
|
||||
>
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>
|
||||
Stop {count} running {noun}?
|
||||
</AlertDialogTitle>
|
||||
<AlertDialogDescription>
|
||||
{lead}so {count === 1 ? "this" : "these"} {noun} will stop
|
||||
{hasNonChat ? "" : " generating"}. Work produced so far is kept.
|
||||
</AlertDialogDescription>
|
||||
</AlertDialogHeader>
|
||||
{shown.length > 0 && (
|
||||
<ul className="max-h-40 overflow-y-auto rounded-md border bg-muted/40 px-3 py-2 text-sm">
|
||||
{shown.map((title) => (
|
||||
<li key={title} className="truncate py-0.5">
|
||||
{title}
|
||||
</li>
|
||||
))}
|
||||
{remaining > 0 && (
|
||||
<li className="py-0.5 text-muted-foreground">
|
||||
and {remaining} more
|
||||
</li>
|
||||
)}
|
||||
</ul>
|
||||
)}
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel onClick={() => resolve(false)}>
|
||||
Keep generating
|
||||
</AlertDialogCancel>
|
||||
<AlertDialogAction onClick={() => resolve(true)}>
|
||||
{unloads ? "Stop and unload" : "Stop and reload"}
|
||||
</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
);
|
||||
}
|
||||
|
|
@ -28,6 +28,7 @@ import {
|
|||
validateModel,
|
||||
} from "../api/chat-api";
|
||||
import { formatEta, formatRate } from "../utils/format-transfer";
|
||||
import { confirmStopRunningChatsIfNeeded } from "../utils/confirm-stop-running-chats";
|
||||
import {
|
||||
GPU_LAYERS_AUTO,
|
||||
isLocalModelPath,
|
||||
|
|
@ -463,7 +464,14 @@ export function useChatModelRuntime() {
|
|||
useChatRuntimeStore.getState().setModelLoading(true);
|
||||
void (async () => {
|
||||
try {
|
||||
// Unforced on purpose: a chat may stream on the PREVIOUS model and must not be killed by
|
||||
// cancelling this load. Nothing to report, since the route runs its stop-loading fast
|
||||
// path ahead of the active-chat refusal.
|
||||
await unloadModel({ model_path: model.id }).catch(() => {});
|
||||
// clearCheckpoint above assumed nothing was left loaded, but a forced switch keeps the
|
||||
// previous model resident until /load's teardown, and the stop-loading fast path leaves
|
||||
// it there. Take the answer from the backend, which reports none once it was evicted.
|
||||
await syncInferenceStatusToStore().catch(() => {});
|
||||
} finally {
|
||||
cancelUnloadPendingRef.current = false;
|
||||
if (!loadingModelRef.current) {
|
||||
|
|
@ -505,10 +513,11 @@ export function useChatModelRuntime() {
|
|||
// as a duplicate), don't start a second concurrent load and don't swallow the
|
||||
// request: surface it so the user waits or cancels. Centralized here so every
|
||||
// entry point is covered, not just the staged Load button.
|
||||
const inFlightLoad =
|
||||
loadingModelRef.current ??
|
||||
useChatRuntimeStore.getState().loadingModelPick;
|
||||
if (inFlightLoad) {
|
||||
const bailIfLoadInFlight = (): boolean => {
|
||||
const inFlightLoad =
|
||||
loadingModelRef.current ??
|
||||
useChatRuntimeStore.getState().loadingModelPick;
|
||||
if (!inFlightLoad) return false;
|
||||
if (typeof selection !== "string" && selection.previousConfig) {
|
||||
applyPerModelConfigToRuntime(selection.previousConfig);
|
||||
}
|
||||
|
|
@ -516,7 +525,7 @@ export function useChatModelRuntime() {
|
|||
inFlightLoad.id === modelId &&
|
||||
(inFlightLoad.ggufVariant ?? null) === (ggufVariant ?? null) &&
|
||||
(inFlightLoad.nativePathToken ?? null) === (nativePathToken ?? null);
|
||||
if (loadingSamePick) return;
|
||||
if (loadingSamePick) return true;
|
||||
const message =
|
||||
"Another model is already loading. Wait for it to finish or cancel it first.";
|
||||
setModelsError(message);
|
||||
|
|
@ -524,8 +533,61 @@ export function useChatModelRuntime() {
|
|||
toast.info("Another model is already loading", {
|
||||
description: "Wait for it to finish or cancel it first.",
|
||||
});
|
||||
return true;
|
||||
};
|
||||
if (bailIfLoadInFlight()) return;
|
||||
|
||||
// Picking an external provider leaves the local model resident and stops the status poll
|
||||
// mirroring it, so params.checkpoint cannot tell whether this pick is that same model.
|
||||
// Ask the backend before prompting: /load answers already_loaded ahead of its cancel
|
||||
// hook, so the dialog would promise to stop chats this pick never interrupts. A staged
|
||||
// config always carries forceReload, so Apply still reloads and prompts.
|
||||
const selectedCheckpoint =
|
||||
useChatRuntimeStore.getState().params.checkpoint;
|
||||
if (!forceReload && isExternalModelId(selectedCheckpoint)) {
|
||||
const residentStatus = await getInferenceStatus().catch(() => null);
|
||||
if (
|
||||
residentStatus &&
|
||||
resolveInferenceCheckpointId(residentStatus) === modelId &&
|
||||
(residentStatus.gguf_variant ?? null) === (ggufVariant ?? null)
|
||||
) {
|
||||
// Same window as the confirm below: a rival load may have started during that GET,
|
||||
// and it owns the resident model now.
|
||||
if (bailIfLoadInFlight()) return;
|
||||
// Roll back the config pre-applied for the load that is not happening BEFORE hydrating,
|
||||
// so the resident model's status wins over the staged snapshot.
|
||||
if (typeof selection !== "string" && selection.previousConfig) {
|
||||
applyPerModelConfigToRuntime(selection.previousConfig);
|
||||
}
|
||||
const previousGgufVariant =
|
||||
useChatRuntimeStore.getState().activeGgufVariant;
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.setCheckpoint(modelId, residentStatus.gguf_variant);
|
||||
applyActiveModelStatusToStore(residentStatus, {
|
||||
previousCheckpoint: selectedCheckpoint,
|
||||
previousGgufVariant,
|
||||
});
|
||||
syncModelCapabilities(modelId, residentStatus);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Every chat decodes on the llama-server this load replaces, so ask first, then allow the
|
||||
// cancel; the 409 gate stays armed for callers that never confirmed.
|
||||
const stopDecision = await confirmStopRunningChatsIfNeeded(
|
||||
forceReload ? "Applying these settings" : "Loading a different model",
|
||||
);
|
||||
if (!stopDecision.proceed) {
|
||||
if (typeof selection !== "string" && selection.previousConfig) {
|
||||
applyPerModelConfigToRuntime(selection.previousConfig);
|
||||
}
|
||||
return;
|
||||
}
|
||||
// Re-check: the confirm above awaits a GET, so a pick in that window would start a rival
|
||||
// load over the same refs. Nothing awaits before the reservation below.
|
||||
if (bailIfLoadInFlight()) return;
|
||||
const forceCancelActive = stopDecision.forceCancelActive;
|
||||
|
||||
const explicitIsLora =
|
||||
typeof selection === "string" ? undefined : selection.isLora;
|
||||
|
|
@ -777,6 +839,10 @@ export function useChatModelRuntime() {
|
|||
upgrade: validation.transformers_upgrade,
|
||||
// No installable release: custom-code models may fall back to the trust_remote_code gate below.
|
||||
trustRemoteCodeFallback: validation.requires_trust_remote_code,
|
||||
// The install refuses while chats generate and takes no force flag of its own, so
|
||||
// without this the "Stop and reload" the user just confirmed dies here: Retry hits
|
||||
// the same 409, and this path leaves chats running.
|
||||
forceCancelActive,
|
||||
});
|
||||
// The install unloads the previous model before the swap (even when
|
||||
// the swap then fails), so any exit after this point must roll back.
|
||||
|
|
@ -820,7 +886,14 @@ export function useChatModelRuntime() {
|
|||
: undefined;
|
||||
|
||||
if (currentCheckpoint) {
|
||||
await unloadModel({ model_path: currentCheckpoint });
|
||||
// With chats generating, skip this preliminary unload: it cancels them ahead of /load's
|
||||
// preflight, so a rejected target truncates replies for a model that never loads
|
||||
// (/load evicts past those checks itself). Idle, unload first and free VRAM early.
|
||||
if (!forceCancelActive) {
|
||||
await unloadModel({ model_path: currentCheckpoint });
|
||||
}
|
||||
// Set either way: /load can still leave no model resident, and an unneeded rollback
|
||||
// hits already_loaded before the gate.
|
||||
previousWasUnloaded = true;
|
||||
}
|
||||
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
|
||||
|
|
@ -935,6 +1008,7 @@ export function useChatModelRuntime() {
|
|||
n_cpu_moe: loadNCpuMoe,
|
||||
tensor_split: loadSplitRatio ?? undefined,
|
||||
gpu_ids: loadSelectedGpuIds ?? undefined,
|
||||
force_cancel_active: forceCancelActive,
|
||||
});
|
||||
|
||||
// If cancelled while loading, don't update UI to show
|
||||
|
|
@ -1173,6 +1247,8 @@ export function useChatModelRuntime() {
|
|||
n_cpu_moe: stateBeforeUnload.loadedNCpuMoe ?? 0,
|
||||
tensor_split: stateBeforeUnload.loadedSplitRatio ?? undefined,
|
||||
gpu_ids: stateBeforeUnload.loadedGpuIds ?? undefined,
|
||||
// The failed swap already unloaded the server those runs used.
|
||||
force_cancel_active: true,
|
||||
});
|
||||
const rollbackSpeculativeType = normalizeSpeculativeType(
|
||||
rollbackResponse.speculative_type,
|
||||
|
|
@ -1571,13 +1647,15 @@ export function useChatModelRuntime() {
|
|||
if (!params.checkpoint) {
|
||||
return false;
|
||||
}
|
||||
const runtime = useChatRuntimeStore.getState();
|
||||
if (runtime.modelLoading || runtime.loadingModelPick) {
|
||||
const bailIfLoading = (): boolean => {
|
||||
const runtime = useChatRuntimeStore.getState();
|
||||
if (!runtime.modelLoading && !runtime.loadingModelPick) return false;
|
||||
toast.info("A model is loading", {
|
||||
description: "Wait for it to finish or cancel it first.",
|
||||
});
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
if (bailIfLoading()) return false;
|
||||
setModelsError(null);
|
||||
if (isExternalModelId(params.checkpoint)) {
|
||||
clearCheckpoint();
|
||||
|
|
@ -1585,8 +1663,21 @@ export function useChatModelRuntime() {
|
|||
return true;
|
||||
}
|
||||
try {
|
||||
// Ejecting tears down llama-server, so every chat stops. Same prompt, but it
|
||||
// leaves no model loaded, so it must not be worded as a reload.
|
||||
const stopDecision = await confirmStopRunningChatsIfNeeded(
|
||||
"Unloading the model",
|
||||
"unload",
|
||||
);
|
||||
if (!stopDecision.proceed) return false;
|
||||
// Same window as selectModel: a load may have started during the confirm.
|
||||
if (bailIfLoading()) return false;
|
||||
|
||||
async function performUnload(): Promise<void> {
|
||||
await unloadModel({ model_path: params.checkpoint });
|
||||
await unloadModel({
|
||||
model_path: params.checkpoint,
|
||||
force_cancel_active: stopDecision.forceCancelActive,
|
||||
});
|
||||
clearCheckpoint();
|
||||
await refresh();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import {
|
|||
updateStoredChatThread,
|
||||
} from "../utils/chat-history-storage";
|
||||
import { clearComposerDraft } from "../utils/composer-draft";
|
||||
import { stopChatThread } from "../utils/stop-chat-thread";
|
||||
import {
|
||||
markChatThreadsDeleted,
|
||||
removeChatThreadTombstones,
|
||||
|
|
@ -25,6 +26,8 @@ import {
|
|||
export interface SidebarItem {
|
||||
type: "single" | "compare";
|
||||
id: string;
|
||||
/** The pane threads behind this row id; `runningByThreadId` is keyed per pane thread. */
|
||||
threadIds?: string[];
|
||||
title: string;
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
|
|
@ -56,11 +59,13 @@ export function groupThreads(
|
|||
const existing = pairItems.get(t.pairId);
|
||||
if (existing) {
|
||||
existing.updatedAt = Math.max(existing.updatedAt, lastActivityAt(t));
|
||||
existing.threadIds?.push(t.id);
|
||||
continue;
|
||||
}
|
||||
const item: SidebarItem = {
|
||||
type: "compare",
|
||||
id: t.pairId,
|
||||
threadIds: [t.id],
|
||||
title: t.title,
|
||||
createdAt: t.createdAt,
|
||||
updatedAt: lastActivityAt(t),
|
||||
|
|
@ -160,10 +165,9 @@ export function useChatSidebarItems(options?: {
|
|||
}
|
||||
|
||||
function cancelIfRunning(threadId: string): void {
|
||||
const { runningByThreadId, cancelByThreadId } =
|
||||
useChatRuntimeStore.getState();
|
||||
if (!runningByThreadId[threadId]) return;
|
||||
cancelByThreadId[threadId]?.();
|
||||
// Reaches a background thread, which cancelByThreadId cannot: a deleted chat must stop,
|
||||
// or the run keeps writing to a conversation that is gone.
|
||||
stopChatThread(threadId);
|
||||
}
|
||||
|
||||
export async function renameChatItem(
|
||||
|
|
|
|||
|
|
@ -55,8 +55,12 @@ export {
|
|||
export {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
toolThreadScope,
|
||||
useToolOutputFor,
|
||||
useUnresolvedToolPaneScope,
|
||||
useToolPaneScope,
|
||||
} from "./tool-output-scope";
|
||||
export { useToolAwaitingApproval } from "./tool-approval";
|
||||
export { PermissionModeDropdown } from "./permission-mode-select";
|
||||
export { useChatSearchStore } from "./stores/chat-search-store";
|
||||
export { usePinnedChatsStore } from "./stores/pinned-chats-store";
|
||||
|
|
@ -80,6 +84,7 @@ export {
|
|||
export { ApiProviderLogo } from "./api-provider-logo";
|
||||
export { useExternalProvidersStore } from "./stores/external-providers-store";
|
||||
export { ChatSearchDialog } from "./components/chat-search-dialog";
|
||||
export { StopRunningChatsDialog } from "./components/stop-running-chats-dialog";
|
||||
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
|
||||
export type { ProjectRecord } from "./types";
|
||||
export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
|
||||
|
|
|
|||
|
|
@ -85,7 +85,11 @@ import { requestPromptQueueStop } from "./utils/prompt-queue-boundary";
|
|||
import { isAssistantLocalThreadId } from "./utils/thread-ids";
|
||||
|
||||
const pendingHistoryAppendByMessageId = new Map<string, Promise<void>>();
|
||||
const pendingRunStartReadyByMessageId = new Map<string, Promise<void>>();
|
||||
// Resolves to the thread id assigned when this message's chat was first persisted.
|
||||
const pendingRunStartReadyByMessageId = new Map<
|
||||
string,
|
||||
Promise<string | undefined>
|
||||
>();
|
||||
|
||||
type TitleResponse = {
|
||||
choices?: Array<{
|
||||
|
|
@ -699,6 +703,10 @@ function createStudioDbAdapter(
|
|||
|
||||
async initialize(threadId: string) {
|
||||
await ensureThreadRecord({ threadId, modelType, pairId, projectId });
|
||||
// A run already streaming on this thread filed its handles under "__default" because
|
||||
// the id did not exist yet. Re-key them now, or the sidebar row and Stop look up an
|
||||
// id nothing is registered against.
|
||||
useChatRuntimeStore.getState().adoptDefaultThreadRun(threadId);
|
||||
return { remoteId: threadId, externalId: undefined };
|
||||
},
|
||||
|
||||
|
|
@ -835,8 +843,8 @@ function trackHistoryAppend(
|
|||
|
||||
function trackRunStartReady(
|
||||
messageId: string,
|
||||
ready: Promise<void>,
|
||||
): Promise<void> {
|
||||
ready: Promise<string | undefined>,
|
||||
): Promise<string | undefined> {
|
||||
pendingRunStartReadyByMessageId.set(messageId, ready);
|
||||
const cleanup = () => {
|
||||
setTimeout(() => {
|
||||
|
|
@ -851,7 +859,7 @@ function trackRunStartReady(
|
|||
|
||||
async function waitForRunStartHistoryAppend(
|
||||
messages: Parameters<ChatModelAdapter["run"]>[0]["messages"],
|
||||
): Promise<void> {
|
||||
): Promise<string | undefined> {
|
||||
// Deep Research reserves an assistant placeholder before invoking the model
|
||||
// adapter, so the user message is not necessarily the final entry here.
|
||||
const userMessage = [...messages]
|
||||
|
|
@ -862,15 +870,16 @@ async function waitForRunStartHistoryAppend(
|
|||
}
|
||||
const runStartReady = pendingRunStartReadyByMessageId.get(userMessage.id);
|
||||
const historyAppendReady = pendingHistoryAppendByMessageId.get(userMessage.id);
|
||||
const pending = [runStartReady, historyAppendReady].filter(
|
||||
(ready): ready is Promise<void> => ready !== undefined,
|
||||
);
|
||||
if (pending.length === 0) {
|
||||
return;
|
||||
if (runStartReady === undefined && historyAppendReady === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
let didBecomeReady = false;
|
||||
let adoptedThreadId: string | undefined;
|
||||
try {
|
||||
await Promise.all(pending);
|
||||
[adoptedThreadId] = await Promise.all([
|
||||
runStartReady ?? Promise.resolve(undefined),
|
||||
historyAppendReady?.then(() => undefined),
|
||||
]);
|
||||
didBecomeReady = true;
|
||||
} finally {
|
||||
if (
|
||||
|
|
@ -881,14 +890,22 @@ async function waitForRunStartHistoryAppend(
|
|||
pendingRunStartReadyByMessageId.delete(userMessage.id);
|
||||
}
|
||||
}
|
||||
return adoptedThreadId;
|
||||
}
|
||||
|
||||
function createPersistedRunAdapter(adapter: ChatModelAdapter): ChatModelAdapter {
|
||||
return {
|
||||
...adapter,
|
||||
async *run(options) {
|
||||
await waitForRunStartHistoryAppend(options.messages);
|
||||
const result = adapter.run(options);
|
||||
const adoptedThreadId = await waitForRunStartHistoryAppend(options.messages);
|
||||
// The thread has an id by the time that resolves, but assistant-ui bound unstable_threadId
|
||||
// before the await. Hand the run its real id so a first turn never files its handles
|
||||
// under the unresolved key that concurrent runs share.
|
||||
const result = adapter.run(
|
||||
!options.unstable_threadId && adoptedThreadId
|
||||
? { ...options, unstable_threadId: adoptedThreadId }
|
||||
: options,
|
||||
);
|
||||
if (!result) {
|
||||
return;
|
||||
}
|
||||
|
|
@ -1153,7 +1170,13 @@ function useStudioRuntimeAdapters(
|
|||
: typeof store.ggufContextLength === "number" &&
|
||||
store.ggufContextLength > 0;
|
||||
if (savedUsage && withinLocalLimit && modelMatches) {
|
||||
store.setContextUsage(savedUsage);
|
||||
// Key by the thread this loader read, not whichever is active when the await resolves:
|
||||
// a switch inside it would file this thread's usage under the incoming one. Same rule
|
||||
// the adapter's end-of-run write follows.
|
||||
store.setThreadContextUsage(remoteId, savedUsage);
|
||||
if (store.activeThreadId === remoteId) {
|
||||
store.setContextUsage(savedUsage);
|
||||
}
|
||||
}
|
||||
|
||||
// If any message has a stored parentId, reconstruct the tree so
|
||||
|
|
@ -1179,7 +1202,10 @@ function useStudioRuntimeAdapters(
|
|||
|
||||
append({ parentId, message }: ExportedMessageRepositoryItem) {
|
||||
const initializeThread = aui.threadListItem().initialize();
|
||||
trackRunStartReady(message.id, initializeThread.then(() => undefined));
|
||||
trackRunStartReady(
|
||||
message.id,
|
||||
initializeThread.then(({ remoteId }) => remoteId),
|
||||
);
|
||||
const write = (async () => {
|
||||
const { remoteId } = await initializeThread;
|
||||
if (isChatThreadDeleted(remoteId)) {
|
||||
|
|
@ -1308,17 +1334,6 @@ function createRuntimeHook(modelType: ModelType, pairId?: string) {
|
|||
};
|
||||
}
|
||||
|
||||
function stopChatRun(threadId: string | null | undefined) {
|
||||
if (!threadId) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
useChatRuntimeStore.getState().cancelByThreadId[threadId]?.();
|
||||
} catch {
|
||||
// The run may have ended while navigation was mounting.
|
||||
}
|
||||
}
|
||||
|
||||
function ThreadAutoSwitch({
|
||||
threadId,
|
||||
syncActiveThreadId = true,
|
||||
|
|
@ -1333,8 +1348,9 @@ function ThreadAutoSwitch({
|
|||
useEffect(() => {
|
||||
if (!isLoading && mainThreadId !== threadId) {
|
||||
if (syncActiveThreadId) {
|
||||
requestPromptQueueStop();
|
||||
stopChatRun(mainThreadId);
|
||||
// Stop queueing prompts to the outgoing thread but leave its run alone: its runtime
|
||||
// stays mounted and keeps streaming. Only an explicit Stop cancels one.
|
||||
requestPromptQueueStop({ cancelActiveRun: false });
|
||||
}
|
||||
const switchResult = aui.threads().switchToThread(threadId) as unknown;
|
||||
if (
|
||||
|
|
@ -1365,16 +1381,14 @@ function ThreadNewChatSwitch({
|
|||
}: { nonce: string }): ReactElement | null {
|
||||
const aui = useAui();
|
||||
const isLoading = useAuiState(({ threads }) => threads.isLoading);
|
||||
const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId);
|
||||
const mainThreadIdRef = useRef(mainThreadId);
|
||||
mainThreadIdRef.current = mainThreadId;
|
||||
|
||||
// The outgoing thread is not read here: New Chat leaves it running.
|
||||
useEffect(() => {
|
||||
if (isLoading) {
|
||||
return;
|
||||
}
|
||||
requestPromptQueueStop();
|
||||
stopChatRun(mainThreadIdRef.current);
|
||||
// New Chat leaves the previous conversation generating: its runtime stays mounted and
|
||||
// the sidebar spins. Stopping it is its own Stop button's job.
|
||||
requestPromptQueueStop({ cancelActiveRun: false });
|
||||
// Switch to a fresh local thread without persisting it yet; persistence
|
||||
// still happens on first message append.
|
||||
void aui.threads().switchToNewThread();
|
||||
|
|
|
|||
|
|
@ -762,6 +762,30 @@ export function isDownloadableHubRepo(x: {
|
|||
);
|
||||
}
|
||||
|
||||
type ContextUsageSnapshot = {
|
||||
promptTokens: number;
|
||||
completionTokens: number;
|
||||
totalTokens: number;
|
||||
cachedTokens: number;
|
||||
// Anthropic-only; optional so pre-cache-stats persisted entries load.
|
||||
cacheWriteTokens?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* One live run behind `runningByThreadId[id]`, with the `local` flag it started with so the
|
||||
* model-swap gate can tell llama-server runs from external ones when runs share a key.
|
||||
*/
|
||||
type ThreadRunOwner = {
|
||||
owner: () => void;
|
||||
local: boolean;
|
||||
};
|
||||
|
||||
type ToolStatusEntry = {
|
||||
status: string;
|
||||
startedAt: number;
|
||||
owner?: () => void;
|
||||
};
|
||||
|
||||
type ChatRuntimeStore = {
|
||||
settingsHydrated: boolean;
|
||||
params: InferenceParams;
|
||||
|
|
@ -771,7 +795,25 @@ type ChatRuntimeStore = {
|
|||
models: ChatModelSummary[];
|
||||
loras: ChatLoraSummary[];
|
||||
runningByThreadId: Record<string, boolean>;
|
||||
/**
|
||||
* The subset of `runningByThreadId` decoding on the local llama-server. Swapping the local
|
||||
* model neither interrupts an external-provider chat nor needs its consent, which is why
|
||||
* the backend keeps those out of `active_generations` too.
|
||||
*/
|
||||
localRunByThreadId: Record<string, boolean>;
|
||||
/**
|
||||
* Which runs set `runningByThreadId[id]`; see `setThreadRunning`'s `owner`. A list, not one
|
||||
* entry: runs without a resolved thread id share the "__default" key, so one entry would let
|
||||
* a newer run's clear delete an older run's flag while it still generates.
|
||||
*/
|
||||
runOwnerByThreadId: Record<string, ThreadRunOwner[]>;
|
||||
cancelByThreadId: Record<string, () => void>;
|
||||
/**
|
||||
* Backend cancels for the threads generating in the background. `cancelByThreadId` only holds
|
||||
* the visible thread's `cancelRun()`, so the adapter parks a closure here that POSTs that
|
||||
* run's own cancel_id. A list for the same reason as `runOwnerByThreadId`: "__default" is shared.
|
||||
*/
|
||||
serverCancelByThreadId: Record<string, (() => void)[]>;
|
||||
autoTitle: boolean;
|
||||
hfToken: string;
|
||||
modelsError: string | null;
|
||||
|
|
@ -892,7 +934,16 @@ type ChatRuntimeStore = {
|
|||
* consulted when `providerSupportsBuiltinWebFetch` is true.
|
||||
*/
|
||||
webFetchToolsEnabled: boolean;
|
||||
toolStatus: string | null;
|
||||
/**
|
||||
* Live tool status per conversation ("Running Python: ...") with its start time. Keyed by
|
||||
* thread, or one chat's tool call shows above every other composer; the timestamp keeps the
|
||||
* counter running across a thread switch.
|
||||
*/
|
||||
/**
|
||||
* Per-run entries, newest last. Unresolved threads share "__default", so one scalar per key
|
||||
* meant a finishing run's clear removed a sibling's status while its tool was still running.
|
||||
*/
|
||||
toolStatusByThreadId: Record<string, ToolStatusEntry[]>;
|
||||
/** Live stdout/stderr from running tools, keyed by toolCallId. Transient:
|
||||
* appended by tool_output, cleared on tool_end or run end. */
|
||||
toolLiveOutput: Record<string, string>;
|
||||
|
|
@ -967,9 +1018,12 @@ type ChatRuntimeStore = {
|
|||
/** Active model is a block-diffusion model (DiffusionGemma): drives the
|
||||
* denoising-canvas artifact auto-render. */
|
||||
loadedIsDiffusion: boolean;
|
||||
/** Live denoising frame for the in-progress diffusion message. Transient: set
|
||||
* per step, cleared when the run ends, never persisted into the transcript. */
|
||||
activeDiffusionCanvas: DiffusionCanvasFrame | null;
|
||||
/**
|
||||
* Live denoising frame per conversation ("__default" until the id exists). Transient: set per
|
||||
* step, cleared when the run ends, never persisted. Keyed, not global: two denoising chats
|
||||
* overwrote each other's frame, so the visible preview flickered or vanished.
|
||||
*/
|
||||
activeDiffusionCanvasByThreadId: Record<string, DiffusionCanvasFrame>;
|
||||
customContextLength: number | null;
|
||||
/** The pinned context the loaded model used (null = Auto), so dirty-tracking
|
||||
* and a later fit Apply can tell an explicit pin apart from Auto. */
|
||||
|
|
@ -992,14 +1046,13 @@ type ChatRuntimeStore = {
|
|||
pendingAudioBase64: string | null;
|
||||
pendingAudioName: string | null;
|
||||
pendingImageEditReference: PendingImageEditReference | null;
|
||||
contextUsage: {
|
||||
promptTokens: number;
|
||||
completionTokens: number;
|
||||
totalTokens: number;
|
||||
cachedTokens: number;
|
||||
// Anthropic-only; optional so pre-cache-stats persisted entries load.
|
||||
cacheWriteTokens?: number;
|
||||
} | null;
|
||||
contextUsage: ContextUsageSnapshot | null;
|
||||
/**
|
||||
* Per-thread copy of the above, so the bar survives a switch away and back. `contextUsage` is
|
||||
* the VISIBLE conversation's usage and a background run may not write it, so without this a
|
||||
* run finishing off-screen leaves nothing to restore.
|
||||
*/
|
||||
contextUsageByThreadId: Record<string, ContextUsageSnapshot>;
|
||||
modelLoading: boolean;
|
||||
loadingModelPick: LoadingModelPick | null;
|
||||
activeNativePathToken: string | null;
|
||||
|
|
@ -1018,9 +1071,35 @@ type ChatRuntimeStore = {
|
|||
setActivePresetSource: (source: ChatPresetSource) => void;
|
||||
setModels: (models: ChatModelSummary[]) => void;
|
||||
setLoras: (loras: ChatLoraSummary[]) => void;
|
||||
setThreadRunning: (threadId: string, running: boolean) => void;
|
||||
/**
|
||||
* `local` defaults to true, so an unqualified caller still counts for the model-swap gate.
|
||||
* `owner` narrows the clear to the run that set the flag: unresolved thread ids share the
|
||||
* "__default" key, so a blind delete would drop a sibling's live entry. Owners accumulate,
|
||||
* so the flag survives until the last one clears.
|
||||
*/
|
||||
setThreadRunning: (
|
||||
threadId: string,
|
||||
running: boolean,
|
||||
options?: { local?: boolean; owner?: () => void },
|
||||
) => void;
|
||||
/**
|
||||
* Re-key a first turn's run handles once its thread is persisted.
|
||||
*
|
||||
* A run that starts before its id exists files everything under "__default". Nothing moved it
|
||||
* afterwards, so once the user navigated away the sidebar found no run and showed no spinner;
|
||||
* stopChatThread had no handle either and the generation carried on holding a slot.
|
||||
*/
|
||||
adoptDefaultThreadRun: (threadId: string) => void;
|
||||
/**
|
||||
* Which key this run's handles live under now. `adoptDefaultThreadRun` re-keys them mid-run,
|
||||
* so a run that started under "__default" must look its owner up instead of reusing the key
|
||||
* it captured, or its writes and its final clear miss the entries.
|
||||
*/
|
||||
runKeyForOwner: (fallbackKey: string, owner: () => void) => string;
|
||||
registerThreadCancel: (threadId: string, cancel: () => void) => void;
|
||||
clearThreadCancel: (threadId: string) => void;
|
||||
registerThreadServerCancel: (threadId: string, cancel: () => void) => void;
|
||||
clearThreadServerCancel: (threadId: string, cancel?: () => void) => void;
|
||||
setAutoTitle: (enabled: boolean) => void;
|
||||
setHfToken: (token: string) => void;
|
||||
setModelsError: (error: string | null) => void;
|
||||
|
|
@ -1074,7 +1153,15 @@ type ChatRuntimeStore = {
|
|||
setRagAutoInjectMinScore: (score: number) => void;
|
||||
setRagOcrScanned: (enabled: boolean) => void;
|
||||
setRagCaptionFigures: (enabled: boolean) => void;
|
||||
setToolStatus: (status: string | null) => void;
|
||||
/**
|
||||
* `owner` is the run's identity token, as for `setThreadRunning`: unresolved threads share
|
||||
* "__default", so without it one run's cleanup clears a concurrent run's status.
|
||||
*/
|
||||
setToolStatus: (
|
||||
threadId: string,
|
||||
status: string | null,
|
||||
owner?: () => void,
|
||||
) => void;
|
||||
appendToolLiveOutput: (toolCallId: string, text: string) => void;
|
||||
/** Clear one tool's live output, or all when no id is given. */
|
||||
clearToolLiveOutput: (toolCallId?: string) => void;
|
||||
|
|
@ -1083,7 +1170,13 @@ type ChatRuntimeStore = {
|
|||
/** Drop a stale preserved full output (a new run is reusing the id). */
|
||||
clearToolFullOutput: (toolCallId: string) => void;
|
||||
setGeneratingStatus: (status: string | null) => void;
|
||||
setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void;
|
||||
setActiveDiffusionCanvas: (
|
||||
threadId: string | null,
|
||||
canvas: DiffusionCanvasFrame,
|
||||
) => void;
|
||||
/** Drop only `threadId`'s canvas: a run ending in a background chat must not wipe the
|
||||
* frame another chat is still painting. */
|
||||
clearActiveDiffusionCanvasForThread: (threadId: string | null) => void;
|
||||
setAutoHealToolCalls: (enabled: boolean) => void;
|
||||
setNudgeToolCalls: (enabled: boolean) => void;
|
||||
setMaxToolCallsPerMessage: (value: number) => void;
|
||||
|
|
@ -1103,6 +1196,11 @@ type ChatRuntimeStore = {
|
|||
) => void;
|
||||
clearPendingImageEditReference: () => void;
|
||||
setContextUsage: (usage: ChatRuntimeStore["contextUsage"]) => void;
|
||||
/** A finished run's usage, kept per thread so switching back re-applies it. */
|
||||
setThreadContextUsage: (
|
||||
threadId: string,
|
||||
usage: ContextUsageSnapshot,
|
||||
) => void;
|
||||
};
|
||||
|
||||
type PersistedChatSettings = Awaited<
|
||||
|
|
@ -1318,7 +1416,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
models: [],
|
||||
loras: [],
|
||||
runningByThreadId: {},
|
||||
localRunByThreadId: {},
|
||||
runOwnerByThreadId: {},
|
||||
cancelByThreadId: {},
|
||||
serverCancelByThreadId: {},
|
||||
autoTitle: false,
|
||||
hfToken: useHfTokenStore.getState().token,
|
||||
modelsError: null,
|
||||
|
|
@ -1382,11 +1483,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
),
|
||||
ragOcrScanned: loadBool(CHAT_RAG_OCR_KEY, DEFAULT_RAG_OCR),
|
||||
ragCaptionFigures: loadBool(CHAT_RAG_CAPTION_KEY, DEFAULT_RAG_CAPTION),
|
||||
toolStatus: null,
|
||||
toolStatusByThreadId: {},
|
||||
toolLiveOutput: {},
|
||||
toolFullOutput: {},
|
||||
generatingStatus: null,
|
||||
activeDiffusionCanvas: null,
|
||||
activeDiffusionCanvasByThreadId: {},
|
||||
autoHealToolCalls: true,
|
||||
nudgeToolCalls: true,
|
||||
maxToolCallsPerMessage: 25,
|
||||
|
|
@ -1433,6 +1534,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
pendingAudioName: null,
|
||||
pendingImageEditReference: null,
|
||||
contextUsage: null,
|
||||
contextUsageByThreadId: {},
|
||||
modelLoading: false,
|
||||
loadingModelPick: null,
|
||||
activeNativePathToken: null,
|
||||
|
|
@ -1505,7 +1607,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
const checkpointChanged = state.params.checkpoint !== params.checkpoint;
|
||||
return {
|
||||
params,
|
||||
...(checkpointChanged ? { contextUsage: null } : {}),
|
||||
...(checkpointChanged
|
||||
? { contextUsage: null, contextUsageByThreadId: {} }
|
||||
: {}),
|
||||
};
|
||||
}),
|
||||
setCustomPresets: (customPresets) =>
|
||||
|
|
@ -1528,16 +1632,94 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
}),
|
||||
setModels: (models) => set({ models }),
|
||||
setLoras: (loras) => set({ loras }),
|
||||
setThreadRunning: (threadId, running) =>
|
||||
setThreadRunning: (threadId, running, options) =>
|
||||
set((state) => {
|
||||
const next = { ...state.runningByThreadId };
|
||||
const nextLocal = { ...state.localRunByThreadId };
|
||||
const nextOwner = { ...state.runOwnerByThreadId };
|
||||
const owners = state.runOwnerByThreadId[threadId] ?? [];
|
||||
const local = options?.local !== false;
|
||||
if (running) {
|
||||
next[threadId] = true;
|
||||
if (options?.owner) {
|
||||
nextOwner[threadId] = [...owners, { owner: options.owner, local }];
|
||||
}
|
||||
// Any local owner keeps the key counted by the model-swap gate, so an external run
|
||||
// joining a shared key must not clear a sibling's flag.
|
||||
if (local) {
|
||||
nextLocal[threadId] = true;
|
||||
} else if (!owners.some((o) => o.local)) {
|
||||
delete nextLocal[threadId];
|
||||
}
|
||||
} else {
|
||||
delete next[threadId];
|
||||
const remaining = options?.owner
|
||||
? owners.filter((o) => o.owner !== options.owner)
|
||||
: [];
|
||||
// An owner missing from the list was already cleared, or the key belongs to siblings
|
||||
// only: either way this run must change nothing.
|
||||
if (options?.owner && remaining.length === owners.length) return state;
|
||||
// An ownerless clear predates per-run tracking, so it must not speak for runs that
|
||||
// own the key: leave them to clear themselves.
|
||||
if (!options?.owner && owners.length > 0) return state;
|
||||
if (remaining.length > 0) {
|
||||
nextOwner[threadId] = remaining;
|
||||
if (remaining.some((o) => o.local)) {
|
||||
nextLocal[threadId] = true;
|
||||
} else {
|
||||
delete nextLocal[threadId];
|
||||
}
|
||||
} else {
|
||||
delete next[threadId];
|
||||
delete nextLocal[threadId];
|
||||
delete nextOwner[threadId];
|
||||
}
|
||||
}
|
||||
return { runningByThreadId: next };
|
||||
return {
|
||||
runningByThreadId: next,
|
||||
localRunByThreadId: nextLocal,
|
||||
runOwnerByThreadId: nextOwner,
|
||||
};
|
||||
}),
|
||||
adoptDefaultThreadRun: (threadId) =>
|
||||
set((state) => {
|
||||
const key = "__default";
|
||||
if (!threadId || threadId === key) return state;
|
||||
// Two first turns can share "__default", and nothing links a run there to the thread being
|
||||
// persisted. Moving the arrays wholesale handed this thread the sibling's owner and stop
|
||||
// handle too, so stopping one aborted both. Adopt only when the key holds a single run.
|
||||
if ((state.runOwnerByThreadId[key]?.length ?? 0) > 1) return state;
|
||||
// Only the transient run maps move. Anything already filed under the real id wins,
|
||||
// since that is a later, better-identified run.
|
||||
const moved: Partial<ChatRuntimeStore> = {};
|
||||
const move = <T,>(
|
||||
map: Record<string, T>,
|
||||
name: keyof ChatRuntimeStore,
|
||||
) => {
|
||||
const entry = map[key];
|
||||
if (entry === undefined || map[threadId] !== undefined) return;
|
||||
const next = { ...map };
|
||||
delete next[key];
|
||||
next[threadId] = entry;
|
||||
(moved as Record<string, unknown>)[name as string] = next;
|
||||
};
|
||||
move(state.runningByThreadId, "runningByThreadId");
|
||||
move(state.localRunByThreadId, "localRunByThreadId");
|
||||
move(state.runOwnerByThreadId, "runOwnerByThreadId");
|
||||
move(state.cancelByThreadId, "cancelByThreadId");
|
||||
move(state.serverCancelByThreadId, "serverCancelByThreadId");
|
||||
move(state.toolStatusByThreadId, "toolStatusByThreadId");
|
||||
move(
|
||||
state.activeDiffusionCanvasByThreadId,
|
||||
"activeDiffusionCanvasByThreadId",
|
||||
);
|
||||
return Object.keys(moved).length > 0 ? moved : state;
|
||||
}),
|
||||
runKeyForOwner: (fallbackKey, owner) => {
|
||||
for (const [key, entries] of Object.entries(get().runOwnerByThreadId)) {
|
||||
if (entries.some((e) => e.owner === owner)) return key;
|
||||
}
|
||||
return fallbackKey;
|
||||
},
|
||||
registerThreadCancel: (threadId, cancel) =>
|
||||
set((state) => {
|
||||
const next = { ...state.cancelByThreadId };
|
||||
|
|
@ -1551,6 +1733,29 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
delete next[threadId];
|
||||
return { cancelByThreadId: next };
|
||||
}),
|
||||
registerThreadServerCancel: (threadId, cancel) =>
|
||||
set((state) => {
|
||||
const next = { ...state.serverCancelByThreadId };
|
||||
next[threadId] = [...(state.serverCancelByThreadId[threadId] ?? []), cancel];
|
||||
return { serverCancelByThreadId: next };
|
||||
}),
|
||||
// `cancel` narrows removal to the run that registered it: unresolved thread ids share the
|
||||
// "__default" key, so a blind delete would drop a live sibling.
|
||||
clearThreadServerCancel: (threadId, cancel) =>
|
||||
set((state) => {
|
||||
const current = state.serverCancelByThreadId[threadId];
|
||||
if (current === undefined) return state;
|
||||
const remaining =
|
||||
cancel === undefined ? [] : current.filter((c) => c !== cancel);
|
||||
if (remaining.length === current.length) return state;
|
||||
const next = { ...state.serverCancelByThreadId };
|
||||
if (remaining.length > 0) {
|
||||
next[threadId] = remaining;
|
||||
} else {
|
||||
delete next[threadId];
|
||||
}
|
||||
return { serverCancelByThreadId: next };
|
||||
}),
|
||||
setAutoTitle: (autoTitle) =>
|
||||
set((state) => {
|
||||
setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle);
|
||||
|
|
@ -1598,14 +1803,24 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
maxTokens: nextMaxTokens,
|
||||
},
|
||||
activeGgufVariant: ggufVariant ?? null,
|
||||
...(checkpointChanged ? { contextUsage: null } : {}),
|
||||
...(checkpointChanged
|
||||
? { contextUsage: null, contextUsageByThreadId: {} }
|
||||
: {}),
|
||||
// Switching to an external provider disables Deep Research, which only
|
||||
// applies to the local base model.
|
||||
...(isExternalModelId(modelId) ? { deepResearchEnabled: false } : {}),
|
||||
};
|
||||
}),
|
||||
// Re-apply the incoming thread's own usage rather than blanking the bar: a run that finished
|
||||
// in the background never wrote the visible value, and a still-mounted runtime skips the
|
||||
// history loader on the way back.
|
||||
setActiveThreadId: (activeThreadId) =>
|
||||
set({ activeThreadId, contextUsage: null }),
|
||||
set((state) => ({
|
||||
activeThreadId,
|
||||
contextUsage: activeThreadId
|
||||
? (state.contextUsageByThreadId[activeThreadId] ?? null)
|
||||
: null,
|
||||
})),
|
||||
setActiveProjectId: (activeProjectId) => set({ activeProjectId }),
|
||||
setIncognito: (incognito) => {
|
||||
if (incognito) saveBool(CHAT_DEEP_RESEARCH_ENABLED_KEY, false);
|
||||
|
|
@ -1636,6 +1851,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
ggufNativeContextLength: null,
|
||||
modelRequiresTrustRemoteCode: false,
|
||||
contextUsage: null,
|
||||
contextUsageByThreadId: {},
|
||||
supportsReasoning: false,
|
||||
reasoningAlwaysOn: false,
|
||||
reasoningEnabled: true,
|
||||
|
|
@ -1657,10 +1873,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
webFetchToolsEnabled: false,
|
||||
// Only the per-session enable pill resets; source/mode/top_k persist.
|
||||
ragEnabled: false,
|
||||
toolStatus: null,
|
||||
toolStatusByThreadId: {},
|
||||
toolLiveOutput: {},
|
||||
toolFullOutput: {},
|
||||
activeDiffusionCanvas: null,
|
||||
activeDiffusionCanvasByThreadId: {},
|
||||
kvCacheDtype: null,
|
||||
loadedKvCacheDtype: null,
|
||||
speculativeType: readPersistedSpeculativeType(),
|
||||
|
|
@ -1957,7 +2173,31 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
saveBool(CHAT_RAG_CAPTION_KEY, ragCaptionFigures);
|
||||
return { ragCaptionFigures };
|
||||
}),
|
||||
setToolStatus: (toolStatus) => set({ toolStatus }),
|
||||
setToolStatus: (threadId, status, owner) =>
|
||||
set((state) => {
|
||||
const next = { ...state.toolStatusByThreadId };
|
||||
const entries = state.toolStatusByThreadId[threadId] ?? [];
|
||||
const mine = entries.find((e) => e.owner === owner);
|
||||
if (!status) {
|
||||
// Drop only this run's entry: a sibling behind the same key may still be running a tool,
|
||||
// and its status has to survive this clear.
|
||||
if (mine === undefined) return state;
|
||||
const rest = entries.filter((e) => e !== mine);
|
||||
if (rest.length > 0) {
|
||||
next[threadId] = rest;
|
||||
} else {
|
||||
delete next[threadId];
|
||||
}
|
||||
} else {
|
||||
// Same text from the same run means the same call, so keep startedAt: only a new tool restarts it.
|
||||
if (mine?.status === status) return state;
|
||||
const entry = { status, startedAt: Date.now(), owner };
|
||||
next[threadId] = mine
|
||||
? entries.map((e) => (e === mine ? entry : e))
|
||||
: [...entries, entry];
|
||||
}
|
||||
return { toolStatusByThreadId: next };
|
||||
}),
|
||||
appendToolLiveOutput: (toolCallId, text) =>
|
||||
set((state) => ({
|
||||
toolLiveOutput: {
|
||||
|
|
@ -1995,8 +2235,21 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
delete next[toolCallId];
|
||||
return { toolLiveOutput: next };
|
||||
}),
|
||||
setActiveDiffusionCanvas: (activeDiffusionCanvas) =>
|
||||
set({ activeDiffusionCanvas }),
|
||||
setActiveDiffusionCanvas: (threadId, canvas) =>
|
||||
set((state) => ({
|
||||
activeDiffusionCanvasByThreadId: {
|
||||
...state.activeDiffusionCanvasByThreadId,
|
||||
[threadId || "__default"]: canvas,
|
||||
},
|
||||
})),
|
||||
clearActiveDiffusionCanvasForThread: (threadId) =>
|
||||
set((state) => {
|
||||
const key = threadId || "__default";
|
||||
if (state.activeDiffusionCanvasByThreadId[key] === undefined) return state;
|
||||
const next = { ...state.activeDiffusionCanvasByThreadId };
|
||||
delete next[key];
|
||||
return { activeDiffusionCanvasByThreadId: next };
|
||||
}),
|
||||
setGeneratingStatus: (generatingStatus) => set({ generatingStatus }),
|
||||
setAutoHealToolCalls: (autoHealToolCalls) =>
|
||||
set((state) => {
|
||||
|
|
@ -2062,7 +2315,27 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
set({ pendingImageEditReference }),
|
||||
clearPendingImageEditReference: () =>
|
||||
set({ pendingImageEditReference: null }),
|
||||
setContextUsage: (contextUsage) => set({ contextUsage }),
|
||||
// Write through to the visible thread's own entry, so a value restored by the history loader
|
||||
// survives a switch away and back: that loader runs once per mount and setActiveThreadId
|
||||
// reads the map, so without this the bar goes blank on return.
|
||||
setContextUsage: (contextUsage) =>
|
||||
set((state) => {
|
||||
if (!state.activeThreadId) return { contextUsage };
|
||||
const next = { ...state.contextUsageByThreadId };
|
||||
if (contextUsage) {
|
||||
next[state.activeThreadId] = contextUsage;
|
||||
} else {
|
||||
delete next[state.activeThreadId];
|
||||
}
|
||||
return { contextUsage, contextUsageByThreadId: next };
|
||||
}),
|
||||
setThreadContextUsage: (threadId, usage) =>
|
||||
set((state) => ({
|
||||
contextUsageByThreadId: {
|
||||
...state.contextUsageByThreadId,
|
||||
[threadId]: usage,
|
||||
},
|
||||
})),
|
||||
}));
|
||||
|
||||
// Mirror token edits made through the shared store (e.g. Unsloth's field).
|
||||
|
|
|
|||
|
|
@ -0,0 +1,69 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { create } from "zustand";
|
||||
|
||||
type Resolver = (confirmed: boolean) => void;
|
||||
|
||||
/** What confirming does to the model: reload it, or leave none loaded. */
|
||||
export type StopRunningChatsEffect = "reload" | "unload";
|
||||
|
||||
// One at a time: a new request declines any pending one so no promise leaks.
|
||||
let pendingResolver: Resolver | null = null;
|
||||
|
||||
interface StopRunningChatsDialogStore {
|
||||
open: boolean;
|
||||
/** How many conversations the pending action would stop. */
|
||||
count: number;
|
||||
/** Titles of those conversations, when known, for the dialog body. */
|
||||
titles: string[];
|
||||
/** What the user is about to do, e.g. "Loading a different model". */
|
||||
action: string;
|
||||
/** The set includes an embeddings/completions/audio request, which is not a chat. */
|
||||
hasNonChat: boolean;
|
||||
/** Ejecting leaves no model loaded, so it must not be described as a reload. */
|
||||
effect: StopRunningChatsEffect;
|
||||
requestConfirm: (args: {
|
||||
count: number;
|
||||
titles?: string[];
|
||||
action?: string;
|
||||
hasNonChat?: boolean;
|
||||
effect?: StopRunningChatsEffect;
|
||||
}) => Promise<boolean>;
|
||||
resolve: (confirmed: boolean) => void;
|
||||
}
|
||||
|
||||
export const useStopRunningChatsDialogStore =
|
||||
create<StopRunningChatsDialogStore>()((set) => ({
|
||||
open: false,
|
||||
count: 0,
|
||||
titles: [],
|
||||
action: "",
|
||||
hasNonChat: false,
|
||||
effect: "reload",
|
||||
requestConfirm: ({
|
||||
count,
|
||||
titles = [],
|
||||
action = "",
|
||||
hasNonChat = false,
|
||||
effect = "reload",
|
||||
}) =>
|
||||
new Promise<boolean>((resolve) => {
|
||||
pendingResolver?.(false);
|
||||
pendingResolver = resolve;
|
||||
set({ open: true, count, titles, action, hasNonChat, effect });
|
||||
}),
|
||||
resolve: (confirmed) => {
|
||||
const resolver = pendingResolver;
|
||||
pendingResolver = null;
|
||||
set({
|
||||
open: false,
|
||||
count: 0,
|
||||
titles: [],
|
||||
action: "",
|
||||
hasNonChat: false,
|
||||
effect: "reload",
|
||||
});
|
||||
resolver?.(confirmed);
|
||||
},
|
||||
}));
|
||||
18
studio/frontend/src/features/chat/tool-approval.ts
Normal file
18
studio/frontend/src/features/chat/tool-approval.ts
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"use client";
|
||||
|
||||
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
|
||||
|
||||
/**
|
||||
* True while this card's call is parked on the Allow / Deny prompt, so it can say it is
|
||||
* waiting rather than counting up "Running". Set when the backend gates the call.
|
||||
*/
|
||||
export function useToolAwaitingApproval(toolCallId?: string): boolean {
|
||||
return useChatRuntimeStore(
|
||||
(s) =>
|
||||
!!toolCallId &&
|
||||
Object.prototype.hasOwnProperty.call(s.toolConfirmations, toolCallId),
|
||||
);
|
||||
}
|
||||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
"use client";
|
||||
|
||||
import { useAuiState } from "@assistant-ui/react";
|
||||
import { createContext, useContext } from "react";
|
||||
import type { ModelType } from "./types";
|
||||
|
||||
|
|
@ -20,10 +21,58 @@ export function toolPaneScope(modelType?: ModelType, pairId?: string): string {
|
|||
return `${modelType ?? "base"}\u0000${pairId ?? ""}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Narrow a pane scope to one conversation: two threads in a pane can both be mid "call_0",
|
||||
* so without the thread in the key they share a store entry and swap outputs.
|
||||
*/
|
||||
export function toolThreadScope(paneScope: string, threadId?: string): string {
|
||||
return `${paneScope}\u0000${threadId ?? ""}`;
|
||||
}
|
||||
|
||||
export const ToolPaneScopeContext = createContext<string>(toolPaneScope());
|
||||
|
||||
/**
|
||||
* Store-key scope for the conversation this component renders in, taken from the surrounding
|
||||
* runtime so reader and writer agree without a prop.
|
||||
*
|
||||
* `remoteId`, not `id`: the adapter gets `unstable_threadId`, which assistant-ui sources from
|
||||
* `remoteId`, and an uninitialized thread has `id` but no `remoteId`. Reading `id` split the
|
||||
* keys apart for the first turn of every New Chat, so live tool output never reached the card.
|
||||
*/
|
||||
export function useToolPaneScope(): string {
|
||||
return useContext(ToolPaneScopeContext);
|
||||
const paneScope = useContext(ToolPaneScopeContext);
|
||||
const threadId = useAuiState(({ threadListItem }) => threadListItem.remoteId);
|
||||
return toolThreadScope(paneScope, threadId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Read a tool-output map for one call, tolerating a run that started before its thread had an id.
|
||||
*
|
||||
* The adapter captures its scope once at run start, so a first turn writes under the unresolved
|
||||
* scope for its whole life. The autosave can assign `remoteId` mid-run, which moves this
|
||||
* component's key but not the writer's, and the card went blank. Falling back to the pane-wide
|
||||
* scope keeps those entries reachable; only an unpersisted first turn can be filed there.
|
||||
*/
|
||||
/** The scope a run that started before its thread had an id writes under. */
|
||||
export function useUnresolvedToolPaneScope(): string {
|
||||
return toolThreadScope(useContext(ToolPaneScopeContext), undefined);
|
||||
}
|
||||
|
||||
export function useToolOutputFor(
|
||||
map: Record<string, string>,
|
||||
paneScope: string,
|
||||
toolCallId: string,
|
||||
): string {
|
||||
// Unconditional: hooks cannot sit behind the early return below.
|
||||
const unresolvedScope = useUnresolvedToolPaneScope();
|
||||
// Only a thread mid-run can be the one that just gained its id. Local ids repeat
|
||||
// ("call_0"), so an unconditional fallback showed a live first turn's stdout in every
|
||||
// older conversation whose own entry had been cleared.
|
||||
const isRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const own = map[toolOutputKey(paneScope, toolCallId)];
|
||||
if (own !== undefined) return own;
|
||||
if (!isRunning) return "";
|
||||
return map[toolOutputKey(unresolvedScope, toolCallId)] ?? "";
|
||||
}
|
||||
|
||||
/** Store key for the live/full tool output maps: pane scope + tool call id. */
|
||||
|
|
|
|||
|
|
@ -35,6 +35,11 @@ export interface ListLorasResponse {
|
|||
|
||||
export interface LoadModelRequest {
|
||||
model_path: string;
|
||||
/**
|
||||
* Stop any chats still generating instead of getting a 409: a load replaces the single
|
||||
* llama-server they all decode on. Set only after the user confirms.
|
||||
*/
|
||||
force_cancel_active?: boolean;
|
||||
nativePathLease?: string | null;
|
||||
hf_token: string | null;
|
||||
max_seq_length: number;
|
||||
|
|
@ -213,6 +218,9 @@ export interface LoadModelResponse {
|
|||
|
||||
export interface UnloadModelRequest {
|
||||
model_path: string;
|
||||
/** Stop any chats still generating instead of getting a 409: the unload takes down the
|
||||
* llama-server they all decode on. */
|
||||
force_cancel_active?: boolean;
|
||||
}
|
||||
|
||||
export interface InferenceStatusResponse {
|
||||
|
|
@ -309,6 +317,12 @@ export interface ApiMonitorEntry {
|
|||
completion_tokens?: number | null;
|
||||
total_tokens?: number | null;
|
||||
error?: string | null;
|
||||
// "lifecycle" is a model load/unload/download: event/reason instead of a prompt.
|
||||
kind?: "request" | "lifecycle";
|
||||
event?: "load" | "unload" | "download" | null;
|
||||
reason?: "manual" | "idle" | "api" | null;
|
||||
// 0-100 while a download row is running.
|
||||
progress?: number | null;
|
||||
}
|
||||
|
||||
export interface ApiMonitorResponse {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,100 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { getActiveGenerations } from "../api/chat-api";
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
import {
|
||||
type StopRunningChatsEffect,
|
||||
useStopRunningChatsDialogStore,
|
||||
} from "../stores/stop-running-chats-dialog-store";
|
||||
import { listStoredChatThreads } from "./chat-history-storage";
|
||||
|
||||
export interface StopRunningChatsDecision {
|
||||
/** False when the user chose to keep generating; the caller must not load. */
|
||||
proceed: boolean;
|
||||
/** Pass as `force_cancel_active`. True only after an explicit confirmation, so the backend's 409 still guards every other caller. */
|
||||
forceCancelActive: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Gate a model load / reload on the chats still generating: they share one llama-server,
|
||||
* so a reload ends all of them. Ask first, then let the backend cancel them once the load
|
||||
* is past preflight. External-provider chats are left out of both.
|
||||
*/
|
||||
export async function confirmStopRunningChatsIfNeeded(
|
||||
action = "Loading a different model",
|
||||
effect: StopRunningChatsEffect = "reload",
|
||||
): Promise<StopRunningChatsDecision> {
|
||||
// Local runs only: an external-provider chat is not stopped by the swap, so counting it
|
||||
// would block a safe load behind a dialog. The backend excludes them for the same reason.
|
||||
const { runningByThreadId, localRunByThreadId } =
|
||||
useChatRuntimeStore.getState();
|
||||
let running = Object.entries(runningByThreadId)
|
||||
.filter(([threadId, on]) => on && localRunByThreadId[threadId])
|
||||
.map(([threadId]) => threadId);
|
||||
let count = running.length;
|
||||
let hasNonChat = false;
|
||||
|
||||
// Always merge the backend snapshot: runningByThreadId is this tab's memory, empty after a
|
||||
// reload and blind to a second tab, while force_cancel_active cancels every backend run.
|
||||
// The union stays local-only, since external-provider runs are never in it.
|
||||
try {
|
||||
const active = await getActiveGenerations();
|
||||
const entries = active.active ?? [];
|
||||
const merged = new Set(running);
|
||||
for (const threadId of active.thread_ids ?? []) {
|
||||
merged.add(threadId);
|
||||
}
|
||||
running = [...merged];
|
||||
// Count conversations, not handles: one chat holds several at once while a tool
|
||||
// continuation registers its next leg before the previous unwinds, and active.count
|
||||
// counts those separately. A first turn started before its id was persisted has no
|
||||
// id to merge, so add those back or the prompt names fewer chats than will stop.
|
||||
const unnamed = entries.filter((entry) => !entry.thread_id).length;
|
||||
count = entries.length
|
||||
? running.length + unnamed
|
||||
: Math.max(active.count ?? 0, running.length);
|
||||
// Embeddings / completions / audio share the model but are not conversations, so the
|
||||
// prompt must not offer to stop chats that do not exist.
|
||||
hasNonChat = entries.some((entry) => (entry.kind ?? "chat") !== "chat");
|
||||
} catch {
|
||||
// Backend unreachable / older build: fall back to the local map only.
|
||||
}
|
||||
|
||||
if (count === 0) {
|
||||
return { proceed: true, forceCancelActive: false };
|
||||
}
|
||||
|
||||
let titles: string[] = [];
|
||||
try {
|
||||
const threads = await listStoredChatThreads();
|
||||
const byId = new Map(threads.map((t) => [t.id, t]));
|
||||
// A compare conversation runs two pane threads, and the sidebar and the route both treat
|
||||
// it as one chat. Counting the raw ids asked to stop two and listed its title twice. Fold
|
||||
// panes onto their pairId, keeping the backend's count when it is higher.
|
||||
const seen = new Set<string>();
|
||||
for (const id of running) {
|
||||
const thread = byId.get(id);
|
||||
const key = thread?.pairId ?? id;
|
||||
if (seen.has(key)) continue;
|
||||
seen.add(key);
|
||||
titles.push(thread?.title || "Untitled chat");
|
||||
}
|
||||
count = Math.max(seen.size, count - (running.length - seen.size));
|
||||
} catch {
|
||||
// Titles are decoration; the count alone is enough to make the choice.
|
||||
titles = [];
|
||||
}
|
||||
|
||||
const confirmed = await useStopRunningChatsDialogStore
|
||||
.getState()
|
||||
.requestConfirm({ count, titles, action, hasNonChat, effect });
|
||||
|
||||
if (!confirmed) {
|
||||
return { proceed: false, forceCancelActive: false };
|
||||
}
|
||||
|
||||
// Deliberately no local stop: the backend holds the cancel until the load clears preflight,
|
||||
// so stopping now would truncate every chat even for a rejected load.
|
||||
return { proceed: true, forceCancelActive: true };
|
||||
}
|
||||
|
|
@ -1,8 +1,17 @@
|
|||
export const PROMPT_QUEUE_STOP_EVENT = "unsloth:prompt-queue-stop";
|
||||
|
||||
export function requestPromptQueueStop() {
|
||||
export interface PromptQueueStopOptions {
|
||||
/** Also cancel the prompt the queue already dispatched. Navigation passes `false` to
|
||||
* leave it generating; an explicit stop passes `true` (the default). */
|
||||
cancelActiveRun?: boolean;
|
||||
}
|
||||
|
||||
export function requestPromptQueueStop(options: PromptQueueStopOptions = {}) {
|
||||
if (typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
window.dispatchEvent(new Event(PROMPT_QUEUE_STOP_EVENT));
|
||||
const { cancelActiveRun = true } = options;
|
||||
window.dispatchEvent(
|
||||
new CustomEvent(PROMPT_QUEUE_STOP_EVENT, { detail: { cancelActiveRun } }),
|
||||
);
|
||||
}
|
||||
|
|
|
|||
39
studio/frontend/src/features/chat/utils/stop-chat-thread.ts
Normal file
39
studio/frontend/src/features/chat/utils/stop-chat-thread.ts
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
|
||||
/**
|
||||
* Stop one conversation's generation, visible or not. Returns true if a stop was dispatched.
|
||||
*
|
||||
* `cancelByThreadId` is assistant-ui's `cancelRun()`, registered only for the thread on screen;
|
||||
* `serverCancelByThreadId` is registered for every run and POSTs that run's own `cancel_id`, so
|
||||
* it is the only handle a background conversation has. Both are per-run. Runs with an unresolved
|
||||
* thread id share the "__default" key, so stop every handle filed under it.
|
||||
*/
|
||||
export function stopChatThread(threadId: string | null | undefined): boolean {
|
||||
if (!threadId) return false;
|
||||
const { runningByThreadId, cancelByThreadId, serverCancelByThreadId } =
|
||||
useChatRuntimeStore.getState();
|
||||
if (!runningByThreadId[threadId]) return false;
|
||||
let stopped = false;
|
||||
try {
|
||||
const cancel = cancelByThreadId[threadId];
|
||||
if (cancel) {
|
||||
cancel();
|
||||
stopped = true;
|
||||
}
|
||||
} catch {
|
||||
// The run may have ended between the read above and this call.
|
||||
}
|
||||
// Also after cancelRun(): a proxy that swallows the fetch abort leaves the backend decoding.
|
||||
for (const serverCancel of serverCancelByThreadId[threadId] ?? []) {
|
||||
try {
|
||||
serverCancel();
|
||||
stopped = true;
|
||||
} catch {
|
||||
// Same as above.
|
||||
}
|
||||
}
|
||||
return stopped;
|
||||
}
|
||||
|
|
@ -2,7 +2,7 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { getInferenceStatus, loadModel } from "@/features/chat";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { createLoadingToastIcon, toast } from "@/lib/toast";
|
||||
import { toastError } from "@/shared/toast";
|
||||
import { useCallback, useEffect, useState } from "react";
|
||||
import { useShallow } from "zustand/react/shallow";
|
||||
|
|
@ -238,8 +238,15 @@ async function loadLocalModelSelection(
|
|||
): Promise<string | null> {
|
||||
const { target, ggufVariant } = selection;
|
||||
const modelLabel = ggufVariant ? `${target} (${ggufVariant})` : target;
|
||||
const toastId = toast.loading(`Loading ${modelLabel}...`, {
|
||||
let loadToastDismissed = false;
|
||||
const toastId = toast.message(`Loading ${modelLabel}...`, {
|
||||
description: "Starting the local inference server for this recipe.",
|
||||
duration: Number.POSITIVE_INFINITY,
|
||||
closeButton: true,
|
||||
icon: createLoadingToastIcon(),
|
||||
onDismiss: () => {
|
||||
loadToastDismissed = true;
|
||||
},
|
||||
});
|
||||
try {
|
||||
const isGguf = GGUF_MODEL_PATTERN.test(target) || Boolean(ggufVariant);
|
||||
|
|
@ -267,7 +274,16 @@ async function loadLocalModelSelection(
|
|||
// biome-ignore lint/style/useNamingConvention: api schema
|
||||
tensor_parallel: false,
|
||||
});
|
||||
toast.success(`Loaded ${modelLabel}`, { id: toastId, duration: 2000 });
|
||||
const successOptions = {
|
||||
description: undefined,
|
||||
duration: 2000,
|
||||
icon: undefined,
|
||||
};
|
||||
if (loadToastDismissed) {
|
||||
toast.success(`Loaded ${modelLabel}`, successOptions);
|
||||
} else {
|
||||
toast.success(`Loaded ${modelLabel}`, { ...successOptions, id: toastId });
|
||||
}
|
||||
return null;
|
||||
} catch (error) {
|
||||
toast.dismiss(toastId);
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ export type OpenAIAutoSwitchSettings = {
|
|||
idleUnloadActive: boolean;
|
||||
// Persist the KV cache to disk on idle unload and restore it on reload.
|
||||
autoUnloadKeepKv: boolean;
|
||||
// Fetch a GGUF named in an API request; stored independently of `enabled`, gated on it.
|
||||
autoDownloadModel: boolean;
|
||||
};
|
||||
|
||||
type ApiOpenAIAutoSwitchSettings = {
|
||||
|
|
@ -25,6 +27,8 @@ type ApiOpenAIAutoSwitchSettings = {
|
|||
idle_unload_active?: boolean;
|
||||
// biome-ignore lint/style/useNamingConvention: API schema
|
||||
auto_unload_keep_kv?: boolean;
|
||||
// biome-ignore lint/style/useNamingConvention: API schema
|
||||
auto_download_model?: boolean;
|
||||
};
|
||||
|
||||
let cachedSettings: OpenAIAutoSwitchSettings | null = null;
|
||||
|
|
@ -39,6 +43,7 @@ function fromApi(
|
|||
defaultEnabled: settings.default_enabled,
|
||||
idleUnloadActive: settings.idle_unload_active ?? false,
|
||||
autoUnloadKeepKv: settings.auto_unload_keep_kv ?? true,
|
||||
autoDownloadModel: settings.auto_download_model ?? false,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -73,6 +78,7 @@ export async function updateOpenAIAutoSwitchSettings(
|
|||
enabled: boolean,
|
||||
autoUnloadIdleSeconds?: number,
|
||||
autoUnloadKeepKv?: boolean,
|
||||
autoDownloadModel?: boolean,
|
||||
): Promise<OpenAIAutoSwitchSettings> {
|
||||
const res = await authFetch("/api/settings/openai-auto-switch", {
|
||||
method: "PUT",
|
||||
|
|
@ -88,6 +94,10 @@ export async function updateOpenAIAutoSwitchSettings(
|
|||
? {}
|
||||
: // biome-ignore lint/style/useNamingConvention: API schema
|
||||
{ auto_unload_keep_kv: autoUnloadKeepKv }),
|
||||
...(autoDownloadModel === undefined
|
||||
? {}
|
||||
: // biome-ignore lint/style/useNamingConvention: API schema
|
||||
{ auto_download_model: autoDownloadModel }),
|
||||
}),
|
||||
});
|
||||
if (!res.ok) {
|
||||
|
|
|
|||
42
studio/frontend/src/features/settings/api/openai-models.ts
Normal file
42
studio/frontend/src/features/settings/api/openai-models.ts
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { authFetch } from "@/features/auth";
|
||||
|
||||
export type OpenAIModel = {
|
||||
id: string;
|
||||
// Resident in memory now; the rest are downloaded and servable.
|
||||
loaded?: boolean;
|
||||
// On-disk GGUF quant. Ids stay bare for OpenAI compat, so append `:quant` to pin it.
|
||||
quant?: string;
|
||||
};
|
||||
|
||||
type ApiOpenAIModelList = {
|
||||
data?: { id?: unknown; loaded?: unknown; quant?: unknown }[];
|
||||
};
|
||||
|
||||
/**
|
||||
* The models this server can serve: `/v1/models` returns exactly the ids
|
||||
* `/v1/chat/completions` resolves against, and accepts the UI session JWT.
|
||||
*/
|
||||
export async function listOpenAIModels(): Promise<OpenAIModel[]> {
|
||||
const res = await authFetch("/v1/models");
|
||||
if (!res.ok) {
|
||||
throw new Error(`Failed to list models (${res.status})`);
|
||||
}
|
||||
const body = (await res.json()) as ApiOpenAIModelList;
|
||||
if (!Array.isArray(body?.data)) {
|
||||
return [];
|
||||
}
|
||||
return body.data.flatMap((entry) =>
|
||||
typeof entry?.id === "string" && entry.id
|
||||
? [
|
||||
{
|
||||
id: entry.id,
|
||||
loaded: entry.loaded === true,
|
||||
quant: typeof entry.quant === "string" ? entry.quant : undefined,
|
||||
},
|
||||
]
|
||||
: [],
|
||||
);
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue