ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge remote-tracking branch 'origin/main' into studio-wheelmap-torch212

Browse source
This commit is contained in:
Daniel Han 2026-07-27 13:19:04 +00:00
commit ef5059614d
76 changed files with 10566 additions and 561 deletions
Download patch file Download diff file Expand all files Collapse all files

6
.github/workflows/consolidated-tests-ci.yml vendored
View file

@ -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"

3
.gitignore vendored
View file

@ -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/

128
install.sh
View file

@ -257,6 +257,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
@ -3296,10 +3341,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 +3382,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 +3659,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 +3701,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 +3903,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 +3960,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

17
studio/backend/auth/authentication.py
View file

@ -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

71
studio/backend/core/inference/_vulkan_probe.py
View file

@ -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

124
studio/backend/core/inference/api_monitor.py
View file

@ -52,6 +52,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 +92,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
@ -127,6 +138,73 @@ 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.
"""
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 +290,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 +314,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 +337,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 +350,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:

274
studio/backend/core/inference/llama_admission.py
View file

@ -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,11 +214,23 @@ class _Waiter:
class LlamaAdmissionLease:
def __init__(self, queue: Optional["LlamaAdmissionQueue"]):
__slots__ = ("_queue", "_slot", "_released", "_release_lock")
def __init__(
self,
queue: Optional["LlamaAdmissionQueue"],
slot: Optional[int] = None,
):
self._queue = queue
self._slot = slot
self._released = False
self._release_lock = threading.Lock()
@property
def slot(self) -> Optional[int]:
"""Pool slot this lease holds, or None when admission is disabled."""
return self._slot
def release(self) -> None:
queue = None
with self._release_lock:
@ -154,7 +239,7 @@ class LlamaAdmissionLease:
self._released = True
queue = self._queue
if queue is not None:
queue.release()
queue.release(self._slot)
async def __aenter__(self) -> "LlamaAdmissionLease":
return self
@ -164,6 +249,8 @@ class LlamaAdmissionLease:
class LlamaAdmissionReservation:
__slots__ = ("_queue", "_lease", "_waiter", "snapshot")
def __init__(
self,
*,
@ -195,6 +282,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,35 +323,80 @@ 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")
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()
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) -> 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.
return bool(self._free) and self._held < self._capacity
def _take_slot_locked(self) -> Optional[int]:
if not self._can_admit_locked():
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))
if not config.enabled:
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()
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,13 +409,20 @@ 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 cancel(self, waiter: _Waiter) -> None:
@ -291,7 +437,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 +455,30 @@ class LlamaAdmissionQueue:
def is_idle(self) -> bool:
with self._lock:
self._prune_waiters_locked()
return self._active == 0 and not self._waiters
return self._in_use == 0 and not self._waiters
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():
waiter = self._waiters.popleft()
if waiter.cancelled or waiter.future.done():
continue
self._active += 1
lease = LlamaAdmissionLease(self)
slot = self._take_slot_locked()
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 +493,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)),
)

97
studio/backend/core/inference/llama_cpp.py
View file

@ -3501,18 +3501,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:
@ -3537,10 +3536,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(),
@ -3554,21 +3556,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(
@ -3577,7 +3614,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: "
@ -6635,12 +6671,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")

18
studio/backend/core/inference/llama_keepwarm.py
View file

@ -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)

26
studio/backend/core/inference/llama_server_args.py
View file

@ -118,6 +118,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
@ -193,9 +194,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
@ -306,6 +306,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]]:

167
studio/backend/core/inference/local_model_resolver.py
View file

@ -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

24
studio/backend/core/inference/model_ids.py
View file

@ -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
View 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()

11
studio/backend/core/inference/tool_loop_controller.py
View file

@ -212,7 +212,16 @@ def status_for_tool(tool_name: str, arguments: Mapping[str, Any]) -> str:
if tool_name == "web_search":
url = str(arguments.get("url") or "").strip()
if url:
parsed = urlparse(url)
# Bare hosts are fetched as https, so normalize first or the badge
# stays generic for exactly the URLs the fetch layer accepts.
from core.inference.tools import _normalize_url_scheme
try:
parsed = urlparse(_normalize_url_scheme(url))
except ValueError:
# Runs in prepare_call, outside the fetch's exception handler:
# raising here kills the turn instead of returning "Blocked:".
return "Reading page..."
if parsed.scheme in ("http", "https") and parsed.hostname:
host = parsed.hostname
if host.startswith("www."):

67
studio/backend/core/inference/tools.py
View file

@ -6332,7 +6332,8 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str
try:
infos = socket.getaddrinfo(hostname, port, type = socket.SOCK_STREAM)
except OSError as e:
except (OSError, UnicodeError) as e:
# IDNA encoding rejects a hostname with UnicodeError, not OSError.
return False, f"Failed to resolve host: {e}", ""
if not infos:
@ -6562,6 +6563,56 @@ def _read_capped_body(resp, max_bytes, timeout, deadline, cancel_event):
return None, b"".join(chunks)
_DOTTED_HOST_RE = re.compile(r"[A-Za-z0-9-]+(\.[A-Za-z0-9-]+)+")
# ASCII-only because str.isdigit() is True for digits int() refuses ("²"), and
# capped at 5 digits so the range check never converts an unbounded integer.
_PORT_RE = re.compile(r"[0-9]{1,5}")
def _normalize_url_scheme(url: str) -> str:
"""Prepend ``https://`` to bare hosts (``google.com``, ``example.com:8443``).
``urlparse`` reads the host of a ``host:port`` input as the scheme, so those
are recognised by a dotted host-like scheme with an empty netloc. Rewrites a
dotted host with an optional in-range port, and the ``//host`` form. Real
schemes (``file:``, ``javascript:``, including ``file:80``), root-relative
paths (``/login``) and bad ports are returned untouched so the caller
rejects them. A dotted scheme is indistinguishable from ``host:port``, so
``com.acme.app:443/cb`` is rewritten too; an empty port (``example.com:``)
is kept as-is, matching ``https://example.com:``.
The host is matched against the raw authority, never against what
``urlparse`` returned, because urlsplit strips tabs/newlines (3.10) and
leading C0/space (3.12). Anything it would strip fails the match, so the
decision and the rewritten string cannot disagree across versions."""
from urllib.parse import urlparse
url = url.strip()
try:
parsed = urlparse(url)
except ValueError:
# Unmatched IPv6 brackets, or an NFKC-decomposing netloc: not a bare host.
return url
if parsed.scheme:
if parsed.netloc or not _DOTTED_HOST_RE.fullmatch(parsed.scheme):
return url
rest = url
elif url.startswith("//"):
rest = url[2:]
elif url.startswith("/"):
return url
else:
rest = url
authority = re.split(r"[/?#]", rest, maxsplit = 1)[0]
host, _, port = authority.partition(":")
if not _DOTTED_HOST_RE.fullmatch(host):
return url
if port and not (_PORT_RE.fullmatch(port) and 1 <= int(port) <= 65535):
return url
return "https://" + rest
def _fetch_url_raw(
url: str,
timeout: int = 30,
@ -6575,6 +6626,8 @@ def _fetch_url_raw(
``error`` is a user-facing message string when the fetch failed (the
existing "Blocked:" / "Failed to fetch URL:" wording), else ``None``.
Blocks private/loopback/link-local targets and caps the download size.
No input reaches the caller as an exception: the URL is model-supplied, so
every malformed form resolves to one of these strings.
``deadline`` is an optional ``time.monotonic`` cutoff for the whole fetch
(redirect hops and body read included) and ``cancel_event`` aborts it when
@ -6583,11 +6636,15 @@ def _fetch_url_raw(
from urllib.parse import urlparse
from .web_access_policy import check_url_access
parsed = urlparse(url)
# Before the policy gate: it requires an http(s) scheme, so a bare host
# would be refused there and never reach the fetch.
url = _normalize_url_scheme(url)
allowed, reason, canonical_host = check_url_access(url, website_policy)
if not allowed:
return reason, "", ""
# check_url_access already parsed this and read .port, so this cannot raise.
parsed = urlparse(url)
port = parsed.port or (443 if parsed.scheme == "https" else 80)
ok, reason, pinned_ip = _resolve_with_budget(
canonical_host,
@ -6648,13 +6705,15 @@ def _fetch_url_raw(
if not location:
return "Failed to fetch URL: redirect missing Location header.", "", ""
current_url = urljoin(current_url, location)
rp = urlparse(current_url)
# Server-controlled, so never scheme-upgraded; the gate below
# reads .port first, so the parse after it cannot raise.
allowed, policy_reason, redirect_host = check_url_access(
current_url,
website_policy,
)
if not allowed:
return policy_reason, "", ""
rp = urlparse(current_url)
rp_port = rp.port or (443 if rp.scheme == "https" else 80)
ok2, reason2, pinned_ip = _resolve_with_budget(
redirect_host,
@ -6872,6 +6931,8 @@ def _fetch_page_text(
deadline = None if timeout is None else time.monotonic() + timeout
from .web_access_policy import check_url_access
# Before the policy gate (needs a scheme) and the README routing (reads host/path).
url = _normalize_url_scheme(url)
allowed, reason, _hostname = check_url_access(url, website_policy)
if not allowed:
return reason

26
studio/backend/hub/services/download_lifecycle.py
View file

@ -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}")

16
studio/backend/hub/services/models/downloads.py
View file

@ -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,

20
studio/backend/main.py
View file

@ -1249,16 +1249,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 +1286,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

1106
studio/backend/routes/inference.py
View file

File diff suppressed because it is too large Load diff

14
studio/backend/routes/settings.py
View file

@ -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,
)

20
studio/backend/tests/conftest.py
View file

@ -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.

973
studio/backend/tests/test_anthropic_admission.py Normal file
View file

@ -0,0 +1,973 @@
# 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; 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")
}
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_exits_the_passthrough_tracker(monkeypatch):
"""A disconnect before the body starts must still exit the cancel tracker.
The wrapper replaces the response's own pre-start hook, so it has to chain to
it. Asserting through _CANCEL_REGISTRY 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, "passthrough should have registered a tracker"
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())

11
studio/backend/tests/test_anthropic_messages.py
View file

@ -1523,6 +1523,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()

262
studio/backend/tests/test_anthropic_passthrough_respawn.py Normal file
View file

@ -0,0 +1,262 @@
# 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 = []
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, "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, "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, "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

97
studio/backend/tests/test_api_monitor.py
View file

@ -258,3 +258,100 @@ 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("...")
# ── 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
studio/backend/tests/test_gpu_memory_mode.py
View file

@ -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])

6
studio/backend/tests/test_gpu_selection.py
View file

@ -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,

537
studio/backend/tests/test_llama_admission.py
View file

@ -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,234 @@ 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())

40
studio/backend/tests/test_llama_server_args.py
View file

@ -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
@ -448,6 +449,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 ─────────────────────────────────────────────

17
studio/backend/tests/test_model_ids.py
View file

@ -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
View file

File diff suppressed because it is too large Load diff

549
studio/backend/tests/test_openai_auto_switch.py
View file

@ -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)
@ -1877,7 +1935,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 +3008,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 +3355,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 +3380,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 +4056,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 +4142,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

162
studio/backend/tests/test_openai_catalog.py
View file

@ -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

11
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -1611,7 +1611,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(
@ -6473,7 +6473,14 @@ class TestApiMonitorSafetensorsUsage:
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)

89
studio/backend/tests/test_system_vulkan_gpu_info.py
View file

@ -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

25
studio/backend/tests/test_tool_loop_controller.py
View file

@ -7,6 +7,8 @@ import json
import sys
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
@ -93,6 +95,29 @@ def test_status_and_provenance_match_local_event_conventions():
}
@pytest.mark.parametrize(
"url, expected",
[
# bare hosts are fetched, so the badge must name them
("google.com", "Reading: google.com"),
("www.google.com/x", "Reading: google.com"),
("//google.com", "Reading: google.com"),
("example.com:8443/path", "Reading: example.com"),
("github.com/unslothai/unsloth", "Reading: github.com"),
# still generic for what the fetch layer refuses
("/login", "Reading page..."),
("javascript:alert(1)", "Reading page..."),
# urlparse raises on these, outside the fetch's handler: degrade, not raise
("https://[::1", "Reading page..."),
("https://::1]", "Reading page..."),
("//exam／ple.com", "Reading page..."),
("//example.com＠", "Reading page..."),
],
)
def test_status_names_the_host_for_schemeless_urls(url, expected):
assert status_for_tool("web_search", {"url": url}) == expected
def test_prepare_execute_builds_visible_events_and_model_tool_message():
controller = ToolLoopController(tools = [_tool("web_search")])
decision = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))

30
studio/backend/tests/test_training_worker_flash_attn.py
View file

@ -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)

170
studio/backend/tests/test_web_fetch_scheme_normalization.py Normal file
View file

@ -0,0 +1,170 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Bare hosts ("google.com") must be fetched as https, not refused."""
from __future__ import annotations
import sys
from pathlib import Path
import pytest
_BACKEND = Path(__file__).resolve().parent.parent
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
from core.inference import tools # noqa: E402
@pytest.fixture
def resolved(monkeypatch):
seen: dict = {}
def fake_resolve(hostname, port, deadline, cancel_event):
seen["hostname"] = hostname
seen["port"] = port
return False, "stopped", None
monkeypatch.setattr(tools, "_resolve_with_budget", fake_resolve)
return seen
@pytest.mark.parametrize(
"url, hostname, port",
[
("google.com", "google.com", 443),
("www.google.com/x", "www.google.com", 443),
("//google.com", "google.com", 443),
("https://google.com", "google.com", 443),
("http://google.com", "google.com", 80),
("example.com:8443/path", "example.com", 8443),
("example.com:8443", "example.com", 8443),
("sub.example.co.uk:8080", "sub.example.co.uk", 8080),
],
)
def test_schemeless_urls_are_fetched_as_https(resolved, url, hostname, port):
err, _, _ = tools._fetch_url_raw(url)
assert resolved["hostname"] == hostname
assert resolved["port"] == port
assert "only http/https" not in (err or "")
@pytest.mark.parametrize(
"url",
[
"ftp://x.com",
"file:///etc/passwd",
"javascript:alert(1)",
"mailto:a@b.c",
# scheme:digits must not masquerade as host:port
"file:80",
"javascript:443/path",
"mailto:25",
# out-of-range ports are not host:port either
"example.com:99999",
"example.com:0",
# ports must match ASCII [0-9]: str.isdigit() is True for digits int() refuses
"example.com:²",
"example.com:²/x",
"example.com:①",
"example.com:1²",
"//example.com:²",
# non-ASCII decimal digits int() accepts are ports urlparse then refuses
"example.com:٤٤٣",
# root-relative paths have no host to fetch
"/login",
"/github.com/owner/repo",
],
)
def test_non_http_schemes_still_blocked(url):
err, _, _ = tools._fetch_url_raw(url)
assert err and "only http/https" in err
def test_absurdly_long_port_does_not_raise():
err, _, _ = tools._fetch_url_raw("example.com:" + "9" * 4400)
assert err and "only http/https" in err
def test_out_of_range_port_returns_error_instead_of_raising():
# check_url_access owns the wording; what matters is a string, not a raise.
err, _, _ = tools._fetch_url_raw("https://example.com:99999")
assert err and err.startswith("Blocked:")
def test_redirect_to_out_of_range_port_is_blocked(monkeypatch):
# A redirect target reads .port too, so it needs the same guard.
import urllib.request
from urllib.error import HTTPError
monkeypatch.setattr(
tools,
"_resolve_with_budget",
lambda host, port, deadline, cancel: (True, "", "93.184.216.34"),
)
class _Redirecting:
def open(self, req, **kw):
hdrs = {"Location": "https://example.org:99999/next"}
raise HTTPError(req.full_url, 302, "Found", hdrs, None)
monkeypatch.setattr(urllib.request, "build_opener", lambda *handlers: _Redirecting())
err, _, _ = tools._fetch_url_raw("https://example.com")
assert err and err.startswith("Blocked:")
@pytest.mark.parametrize(
"url",
[
# urlparse raises on these; a model-supplied URL must still return a string
"//exam／ple.com", # NFKC-decomposes into "/"
"//example.com＠", # NFKC-decomposes into "@"
"//example.com：", # NFKC-decomposes into ":"
"https://[::1", # unmatched IPv6 bracket
"https://::1]",
],
)
def test_malformed_url_is_blocked_instead_of_raising(url):
err, _, _ = tools._fetch_url_raw(url)
assert err and err.startswith("Blocked:")
def test_idna_failure_is_reported_instead_of_raising(monkeypatch):
# getaddrinfo raises UnicodeError, not OSError, when IDNA encoding fails.
import socket
def boom(*a, **k):
raise UnicodeError("encoding with 'idna' codec failed")
monkeypatch.setattr(socket, "getaddrinfo", boom)
err, _, _ = tools._fetch_url_raw("https://münich.example")
assert err and err.startswith("Failed to resolve host:")
@pytest.mark.parametrize(
"url, hostname",
[
(" google.com", "google.com"),
("google.com\n", "google.com"),
("\t example.com:8443 ", "example.com"),
],
)
def test_surrounding_whitespace_is_stripped(resolved, url, hostname):
# _web_search strips, but direct callers of the fetch layer do not.
tools._fetch_url_raw(url)
assert resolved["hostname"] == hostname
@pytest.mark.parametrize("url", ["127.0.0.1", "169.254.169.254", "10.0.0.1", "192.168.1.1"])
def test_normalization_does_not_bypass_ssrf_guard(url):
err, _, _ = tools._fetch_url_raw(url, timeout = 3)
assert err and "non-public address" in err
def test_schemeless_github_repo_still_routes_to_readme_api():
# Must run before _github_repo_readme_api_url, else a bare repo URL scrapes HTML.
normalized = tools._normalize_url_scheme("github.com/unslothai/unsloth")
assert tools._github_repo_readme_api_url(normalized) == (
"https://api.github.com/repos/unslothai/unsloth/readme"
)

31
studio/backend/utils/hardware/hardware.py
View file

@ -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",
}

40
studio/backend/utils/openai_auto_switch_settings.py
View file

@ -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()
),
)

18
studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx
View file

@ -23,6 +23,18 @@ const RE_BLOCK_SEP = /\n---\n/;
const RE_TITLE = /Title:\s*(.+)/;
const RE_URL = /URL:\s*(.+)/;
const RE_SNIPPET = /Snippet:\s*(.+)/s;
// Mirrors _normalize_url_scheme: a dotted host, optionally followed by a port
// that may be empty ("example.com:" fetches on the default port) but otherwise
// has to be in range, so the card names a host only when the backend fetches it.
const RE_BARE_HOST = /^[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)+(?::(\d{0,5}))?(?:[/?#]|$)/;
function isBareHostFetchedAsHttps(value: string): boolean {
const match = RE_BARE_HOST.exec(value);
if (!match) return false;
const port = match[1];
if (!port) return true;
return Number(port) >= 1 && Number(port) <= 65535;
}
/**
* Reject non-http(s) URLs. Web-search/fetch output is provider-controlled,
@ -72,8 +84,12 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({
const isUrlFetch = !!url;
const displayDomain = (() => {
if (!url) return "";
// new URL() throws on the bare hosts the backend fetches, so mirror that
// grammar or the card names no host for exactly the URLs it does fetch.
const bare = url.startsWith("//") ? url.slice(2) : url;
const candidate = isBareHostFetchedAsHttps(bare) ? `https://${bare}` : url;
try {
const parsed = new URL(url);
const parsed = new URL(candidate);
if (parsed.protocol !== "http:" && parsed.protocol !== "https:") return "";
return parsed.hostname.replace(/^www\./, "");
} catch {

6
studio/frontend/src/features/chat/types/api.ts
View file

@ -291,6 +291,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 {

10
studio/frontend/src/features/settings/api/openai-auto-switch.ts
View file

@ -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
View 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,
},
]
: [],
);
}

213
studio/frontend/src/features/settings/components/api-monitor-console.tsx
View file

@ -7,6 +7,7 @@ import {
ActivityIcon,
ChevronDownIcon,
CircleIcon,
PowerOffIcon,
RefreshCwIcon,
} from "lucide-react";
import {
@ -17,11 +18,19 @@ import {
useRef,
useState,
} from "react";
import { getApiMonitor, getApiMonitorEntry } from "../../chat/api/chat-api";
import {
getApiMonitor,
getApiMonitorEntry,
getInferenceStatus,
unloadModel,
} from "../../chat/api/chat-api";
import { resolveInferenceCheckpointId } from "../../chat/lib/apply-inference-status-to-store";
import { useChatRuntimeStore } from "../../chat/stores/chat-runtime-store";
import type { ApiMonitorEntry, ApiMonitorResponse } from "../../chat/types/api";
const API_INFERENCE_PREFIX_RE = /^\/api\/inference/;
const V1_PREFIX_RE = /^\/v1\//;
const PAGE_SIZE = 5;
function formatTime(value: number): string {
return new Date(value * 1000).toLocaleTimeString([], {
@ -87,6 +96,65 @@ function UsageBar({ value }: { value?: number | null }): ReactElement | null {
);
}
function isLifecycle(entry: ApiMonitorEntry): boolean {
return entry.kind === "lifecycle";
}
function lifecycleLabel(entry: ApiMonitorEntry): string {
if (entry.event === "unload") {
return entry.reason === "idle" ? "Model unloaded (idle)" : "Model unloaded";
}
if (entry.event === "download") {
if (entry.status === "running") {
const pct = entry.progress;
return typeof pct === "number"
? `Downloading model (${Math.round(pct)}%)`
: "Downloading model";
}
if (entry.status === "completed") return "Model downloaded";
// A cancel is deliberate, so saying it failed misreads the user's own action.
return entry.status === "cancelled"
? "Model download cancelled"
: "Model download failed";
}
if (entry.status === "running") {
return "Loading model";
}
if (entry.status === "completed") {
return "Model loaded";
}
return "Model load failed";
}
// Load/unload rows: label, model and time. No prompt or detail, so nothing to expand.
function LifecycleEntry({ entry }: { entry: ApiMonitorEntry }): ReactElement {
return (
<article className="min-w-0 rounded-lg border border-border/70 bg-muted/25">
<div className="flex w-full min-w-0 items-start justify-between gap-3 p-3">
<div className="min-w-0">
<div className="flex min-w-0 items-center gap-2">
<ActivityIcon
className={cn("size-3.5 shrink-0", statusTone(entry.status))}
/>
<span className="truncate text-xs font-medium">
{lifecycleLabel(entry)}
</span>
</div>
<div className="mt-1 truncate text-ui-11 text-muted-foreground">
{entry.model}
</div>
</div>
<div className="shrink-0 text-right text-ui-11 text-muted-foreground">
<div>{formatTime(entry.started_at)}</div>
{entry.event === "load" || entry.event === "download" ? (
<div>{formatDuration(entry.duration_ms)}</div>
) : null}
</div>
</div>
</article>
);
}
function MonitorEntry({
entry,
detail,
@ -191,6 +259,7 @@ export function ApiMonitorConsole(): ReactElement {
const [data, setData] = useState<ApiMonitorResponse | null>(null);
const [error, setError] = useState<string | null>(null);
const [refreshing, setRefreshing] = useState(false);
const [unloading, setUnloading] = useState(false);
const [expandedIds, setExpandedIds] = useState<Set<string>>(() => new Set());
const [details, setDetails] = useState<Record<string, ApiMonitorEntry>>({});
const [loadingDetails, setLoadingDetails] = useState<Set<string>>(
@ -211,6 +280,28 @@ export function ApiMonitorConsole(): ReactElement {
}
}, []);
// /unload matches on the internal id, which the monitor omits, so read it from status.
const unloadActiveModel = useCallback(async (): Promise<void> => {
setUnloading(true);
try {
const status = await getInferenceStatus();
const checkpoint = resolveInferenceCheckpointId(status);
if (!checkpoint) {
setError(null);
return;
}
await unloadModel({ model_path: checkpoint });
// Same as the chat eject flow: the store still holds the freed checkpoint.
useChatRuntimeStore.getState().clearCheckpoint();
setError(null);
await loadMonitor();
} catch (err: unknown) {
setError(err instanceof Error ? err.message : "Failed to unload the model");
} finally {
setUnloading(false);
}
}, [loadMonitor]);
useEffect(() => {
let cancelled = false;
let timer: number | undefined;
@ -253,6 +344,48 @@ export function ApiMonitorConsole(): ReactElement {
const statusLabel = data?.status ?? "idle";
const hasActive = (data?.active_requests ?? 0) > 0;
const entries = useMemo(() => data?.entries ?? [], [data]);
// Page 1 tracks the live list; paging back freezes the id order so history holds still.
const [page, setPage] = useState(0);
const [frozenIds, setFrozenIds] = useState<string[] | null>(null);
const byId = useMemo(
() => new Map(entries.map((entry) => [entry.id, entry])),
[entries],
);
const ordered = useMemo(() => {
if (frozenIds === null) {
return entries;
}
return frozenIds.flatMap((id) => {
const entry = byId.get(id);
return entry ? [entry] : [];
});
}, [byId, entries, frozenIds]);
const pageCount = Math.max(1, Math.ceil(ordered.length / PAGE_SIZE));
const pageIndex = Math.min(page, pageCount - 1);
const visible = ordered.slice(
pageIndex * PAGE_SIZE,
pageIndex * PAGE_SIZE + PAGE_SIZE,
);
const newerCount =
frozenIds === null
? 0
: entries.filter((entry) => !frozenIds.includes(entry.id)).length;
const goToPage = useCallback(
(next: number): void => {
if (next <= 0) {
setFrozenIds(null);
setPage(0);
return;
}
// Freeze on the way off page 1 so the history under the cursor holds still.
setFrozenIds((prev) => prev ?? entries.map((entry) => entry.id));
setPage(next);
},
[entries],
);
const loadDetail = useCallback(
(id: string): void => {
if (loadingDetailsRef.current.has(id)) {
@ -305,8 +438,9 @@ export function ApiMonitorConsole(): ReactElement {
);
useEffect(() => {
for (const entry of entries) {
if (!expandedIds.has(entry.id)) {
// Only rows on screen: an expanded row on another page would keep polling.
for (const entry of visible) {
if (isLifecycle(entry) || !expandedIds.has(entry.id)) {
continue;
}
const cached = detailsRef.current[entry.id];
@ -314,7 +448,7 @@ export function ApiMonitorConsole(): ReactElement {
loadDetail(entry.id);
}
}
}, [entries, expandedIds, loadDetail]);
}, [visible, expandedIds, loadDetail]);
return (
<section className="flex min-w-0 flex-col rounded-lg border border-border/70 bg-background">
@ -339,6 +473,22 @@ export function ApiMonitorConsole(): ReactElement {
<div className="rounded-full border border-border px-2.5 py-1 text-xs capitalize text-muted-foreground">
{statusLabel}
</div>
{/* Always rendered, disabled when idle: the only manual release must stay visible. */}
<Button
type="button"
variant="ghost"
size="sm"
onClick={() => void unloadActiveModel()}
disabled={unloading || !data?.active_model}
title={
data?.active_model
? "Unload the model and free its VRAM"
: "No model is loaded"
}
>
<PowerOffIcon className="size-3.5" />
{unloading ? "Unloading" : "Unload"}
</Button>
<Button
type="button"
variant="ghost"
@ -375,19 +525,54 @@ export function ApiMonitorConsole(): ReactElement {
</div>
) : (
<div className="grid gap-3">
{entries.map((entry) => (
<MonitorEntry
key={entry.id}
entry={entry}
detail={details[entry.id]}
expanded={expandedIds.has(entry.id)}
loading={loadingDetails.has(entry.id)}
onToggle={() => toggleEntry(entry)}
/>
))}
{visible.map((entry) =>
isLifecycle(entry) ? (
<LifecycleEntry key={entry.id} entry={entry} />
) : (
<MonitorEntry
key={entry.id}
entry={entry}
detail={details[entry.id]}
expanded={expandedIds.has(entry.id)}
loading={loadingDetails.has(entry.id)}
onToggle={() => toggleEntry(entry)}
/>
),
)}
</div>
)}
</div>
{/* Also while frozen: retention can shrink that list below one page, and hiding the
pager would strand the console on a stale snapshot. */}
{ordered.length > PAGE_SIZE || frozenIds !== null ? (
<div className="flex items-center justify-between gap-2 border-t border-border/60 px-4 py-2 text-xs text-muted-foreground">
<span>
Page {pageIndex + 1} of {pageCount}
{newerCount > 0 ? ` (${newerCount.toLocaleString()} new)` : ""}
</span>
<div className="flex items-center gap-1">
<Button
variant="ghost"
size="sm"
className="h-7 px-2 text-xs"
onClick={() => goToPage(pageIndex - 1)}
disabled={pageIndex === 0 && frozenIds === null}
>
Newer
</Button>
<Button
variant="ghost"
size="sm"
className="h-7 px-2 text-xs"
onClick={() => goToPage(pageIndex + 1)}
disabled={pageIndex >= pageCount - 1}
>
Older
</Button>
</div>
</div>
) : null}
</section>
);
}

19
studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
View file

@ -65,6 +65,7 @@ export function ModelAutoSwitchSection() {
idleSeconds: number | undefined,
syncDraft = true,
keepKv?: boolean,
autoDownload?: boolean,
) => {
setIsSaving(true);
setError(null);
@ -73,6 +74,7 @@ export function ModelAutoSwitchSection() {
enabled,
idleSeconds,
keepKv,
autoDownload,
);
setSettings(saved);
if (syncDraft) {
@ -117,6 +119,11 @@ export function ModelAutoSwitchSection() {
void persist(settings.enabled, undefined, false, keepKv);
};
const handleAutoDownloadToggle = (autoDownload: boolean) => {
if (!settings) return;
void persist(settings.enabled, undefined, false, undefined, autoDownload);
};
return (
<SettingsSection title={t("settings.general.modelAutoSwitch.sectionTitle")}>
<SettingsRow
@ -129,6 +136,18 @@ export function ModelAutoSwitchSection() {
onCheckedChange={handleToggle}
/>
</SettingsRow>
<SettingsRow
label={t("settings.general.modelAutoSwitch.autoDownload")}
description={t(
"settings.general.modelAutoSwitch.autoDownloadDescription",
)}
>
<Switch
checked={settings?.autoDownloadModel ?? false}
disabled={!settings?.enabled || isSaving}
onCheckedChange={handleAutoDownloadToggle}
/>
</SettingsRow>
<SettingsRow
label={t("settings.general.modelAutoSwitch.idleUnload")}
description={t(

329
studio/frontend/src/features/settings/components/usage-examples.tsx
View file

@ -29,11 +29,8 @@ import { HugeiconsIcon } from "@hugeicons/react";
import { useEffect, useMemo, useRef, useState } from "react";
import { Streamdown } from "streamdown";
import { loadCodingAgents } from "../api/coding-agents";
import {
type OpenAIAutoSwitchSettings,
loadOpenAIAutoSwitchSettings,
updateOpenAIAutoSwitchSettings,
} from "../api/openai-auto-switch";
import { loadOpenAIAutoSwitchSettings } from "../api/openai-auto-switch";
import { type OpenAIModel, listOpenAIModels } from "../api/openai-models";
import { buildAgentCommand, isLoopbackHost, normalizeHost } from "./agent-command";
type ExampleType =
@ -89,14 +86,7 @@ const JAVASCRIPT_TYPES = new Set<ExampleType>([
"javascriptAdvanced",
]);
const PROMPT = "Can Unsloth Studio do API calling?";
// Auto-switch demo: a second call naming a different downloaded GGUF so the
// example shows that the model field selects which model serves.
// A placeholder the user replaces with one of their downloaded GGUFs. A fixed
// repo is usually not one they have, so the resolver would fall through and the
// demo would keep serving the current model instead of switching.
const SWITCH_MODEL = "your-other-downloaded-GGUF";
const SWITCH_PROMPT = "Now answer as a different model.";
const PROMPT = "What is Unsloth Studio?";
// web_search + python + terminal are the reliable built-in tools.
const TOOLS = ["web_search", "python", "terminal"];
const ADV = {
@ -196,19 +186,13 @@ function winBody(model: string, variant: Variant): string {
return JSON.stringify(body, null, 2);
}
// A leading comment (valid in both bash and PowerShell) noting the model field
// selects the served model when auto-switch is on.
const SWITCH_NOTE =
'# "Switch model by request" is on: set "model" to any downloaded GGUF to switch.\n';
function curlUnix(
base: string,
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
return `${autoSwitch ? SWITCH_NOTE : ""}curl ${base}/v1/chat/completions \\
return `curl ${base}/v1/chat/completions \\
-H "Authorization: Bearer ${key}" \\
-H "Content-Type: application/json" \\
-d '${shSingle(curlBodyPretty(model, variant))}'`;
@ -219,9 +203,8 @@ function curlWindows(
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
return `${autoSwitch ? SWITCH_NOTE : ""}$body = '${psSingle(winBody(model, variant))}'
return `$body = '${psSingle(winBody(model, variant))}'
Set-Content -Path body.json -Value $body -Encoding ascii
curl.exe ${base}/v1/chat/completions \`
-H "Authorization: Bearer ${key}" \`
@ -229,29 +212,11 @@ curl.exe ${base}/v1/chat/completions \`
-d "@body.json"`;
}
// A second OpenAI call naming a different downloaded GGUF: with auto-switch on,
// Unsloth loads it before serving, so the model field selects the served model.
function pythonSwitchDemo(): string {
return `
# "Switch model by request" is on: replace the model below with another GGUF you
# have downloaded and Unsloth loads it before serving. Unknown names keep serving
# the current model.
response = client.chat.completions.create(
model=${j(SWITCH_MODEL)},
messages=[{"role": "user", "content": ${j(SWITCH_PROMPT)}}],
stream=True,
)
for chunk in response:
print(chunk.choices[0].delta.content or "", end="")`;
}
function pythonSnippet(
base: string,
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
const named =
variant === "advanced"
@ -296,7 +261,7 @@ response = client.chat.completions.create(
messages=[{"role": "user", "content": ${j(PROMPT)}}],${named}${extraBody}
stream=True,
)
${loop}${autoSwitch ? pythonSwitchDemo() : ""}`;
${loop}`;
}
function javascriptSnippet(
@ -304,7 +269,6 @@ function javascriptSnippet(
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
const options: string[] = [];
if (variant === "advanced") {
@ -343,23 +307,6 @@ const response = await client.chat.completions.create({
for await (const chunk of response) {
process.stdout.write(chunk.choices?.[0]?.delta?.content || "");
}${autoSwitch ? javascriptSwitchDemo() : ""}`;
}
function javascriptSwitchDemo(): string {
return `
// "Switch model by request" is on: replace the model below with another GGUF you
// have downloaded and Unsloth loads it before serving. Unknown names keep serving
// the current model.
const switchResponse = await client.chat.completions.create({
model: ${j(SWITCH_MODEL)},
messages: [{ role: "user", content: ${j(SWITCH_PROMPT)} }],
stream: true,
});
for await (const chunk of switchResponse) {
process.stdout.write(chunk.choices?.[0]?.delta?.content || "");
}`;
}
@ -368,31 +315,28 @@ function buildSnippets(
key: string,
model: string,
os: Os,
autoSwitch: boolean,
): Record<ExampleType, string> {
const curl = os === "windows" ? curlWindows : curlUnix;
return {
curl: curl(base, key, model, "plain", autoSwitch),
python: pythonSnippet(base, key, model, "plain", autoSwitch),
javascript: javascriptSnippet(base, key, model, "plain", autoSwitch),
curlTools: curl(base, key, model, "tools", autoSwitch),
pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch),
javascriptTools: javascriptSnippet(base, key, model, "tools", autoSwitch),
curlAdvanced: curl(base, key, model, "advanced", autoSwitch),
pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch),
javascriptAdvanced: javascriptSnippet(
base,
key,
model,
"advanced",
autoSwitch,
),
curl: curl(base, key, model, "plain"),
python: pythonSnippet(base, key, model, "plain"),
javascript: javascriptSnippet(base, key, model, "plain"),
curlTools: curl(base, key, model, "tools"),
pythonTools: pythonSnippet(base, key, model, "tools"),
javascriptTools: javascriptSnippet(base, key, model, "tools"),
curlAdvanced: curl(base, key, model, "advanced"),
pythonAdvanced: pythonSnippet(base, key, model, "advanced"),
javascriptAdvanced: javascriptSnippet(base, key, model, "advanced"),
};
}
const KEY_PLACEHOLDER = "sk-unsloth-YOUR_KEY";
const MODEL_FALLBACK = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL";
const USE_TUNNEL_KEY = "unsloth_api_use_tunnel";
// Slow retry while /v1 has nothing to name: a download or load moves no store state.
const CATALOG_RETRY_MS = 15000;
// Slower beat once something is servable: an idle unload frees a model without
// touching the store, so residency is never settled for good.
const CATALOG_IDLE_MS = 60000;
function readUseTunnelPref(): boolean {
if (typeof window === "undefined") return true;
@ -412,18 +356,111 @@ function writeUseTunnelPref(value: boolean): void {
}
}
function useLoadedModelName(): string {
// A checkpoint can be an on-disk load path, which /v1 never advertises. Mirrors _looks_like_path.
function looksLikePath(id: string): boolean {
return (
id.startsWith("/") ||
id.startsWith("~") ||
id.startsWith(".") ||
id.includes("\\") ||
id.toLowerCase().endsWith(".gguf") ||
(id.match(/\//g)?.length ?? 0) >= 2
);
}
// Same model, ignoring any ":quant" a caller pinned.
function sameBaseModelId(a: string, b: string): boolean {
const base = (id: string) => id.trim().toLowerCase().split(":")[0];
return a.trim().toLowerCase() === b.trim().toLowerCase() || base(a) === base(b);
}
// The model the examples name: always an id /v1 resolves against, null when there is none.
function useExampleModelName(): string | null {
const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint);
const ggufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
return useMemo(() => {
if (!checkpoint || checkpoint.startsWith("external::")) {
return MODEL_FALLBACK;
}
if (ggufVariant && !checkpoint.includes(":")) {
return `${checkpoint}:${ggufVariant}`;
}
return checkpoint;
// null until /v1/models answers: "not asked yet" must not read as "holds nothing".
const [catalog, setCatalog] = useState<OpenAIModel[] | null>(null);
// A downloaded but unloaded model is only runnable when switching is on.
const [autoSwitch, setAutoSwitch] = useState(false);
// Idle-unload on its own (UNSLOTH_MODEL_IDLE_TTL, switching off) reloads exactly
// what it freed: the stored checkpoint only, never an arbitrary catalog entry.
const [idleReload, setIdleReload] = useState(false);
const usableCheckpoint =
!!checkpoint && !checkpoint.startsWith("external::") && !looksLikePath(checkpoint);
// Always: a stored checkpoint can stop being servable without the store changing.
// biome-ignore lint/correctness/useExhaustiveDependencies: a load or unload must refetch the servable ids
useEffect(() => {
let cancelled = false;
let timeoutId: number | null = null;
const update = () => {
// null on failure, never [] or false: a transient error is no evidence that the
// server holds nothing, and those negatives blanked every example while the
// model was still servable. Keep the last answer and retry.
void Promise.all([
listOpenAIModels().catch(() => null),
loadOpenAIAutoSwitchSettings()
.then((s) => [s.enabled, s.idleUnloadActive] as const)
.catch(() => null),
])
.then(([models, settings]) => {
if (cancelled) return true;
if (models !== null) setCatalog(models);
if (settings !== null) {
setAutoSwitch(settings[0]);
setIdleReload(settings[1]);
}
// Resident only slows the polling; it never stops it.
return models !== null && models.some((m) => m.loaded);
})
.then((resolved) => {
if (cancelled) return;
timeoutId = window.setTimeout(
update,
resolved ? CATALOG_IDLE_MS : CATALOG_RETRY_MS,
);
});
};
update();
return () => {
cancelled = true;
if (timeoutId !== null) window.clearTimeout(timeoutId);
};
}, [checkpoint, ggufVariant]);
return useMemo(() => {
// Name something held here, with its quant to pin the file on disk.
const fromCatalog = (): string | null => {
const pick =
catalog?.find((m) => m.loaded) ?? (autoSwitch ? catalog?.[0] : undefined);
if (!pick) {
return null;
}
return pick.quant && !pick.id.includes(":")
? `${pick.id}:${pick.quant}`
: pick.id;
};
// The store keeps a checkpoint across an idle unload and across the model being
// deleted, so it only names a runnable model while the catalog still lists it:
// resident, or downloaded with switching able to reload it. A null catalog means
// /v1/models has not answered, which is not evidence against it.
const entry = catalog?.find((m) => sameBaseModelId(m.id, checkpoint ?? ""));
const backed =
catalog === null || (!!entry && (entry.loaded || autoSwitch || idleReload));
if (usableCheckpoint && checkpoint && backed) {
if (checkpoint.includes(":")) {
return checkpoint;
}
// Pin the quant the catalog advertises, not the stored one: membership proves the
// repo, and the saved quant can name a file deleted while another quant remains.
// Fall back to the store only before /v1/models answers.
const quant = catalog === null ? ggufVariant : entry?.quant;
return quant ? `${checkpoint}:${quant}` : checkpoint;
}
return fromCatalog();
}, [autoSwitch, catalog, checkpoint, ggufVariant, idleReload, usableCheckpoint]);
}
// Backend PATH detection is only safe in the desktop app, where the UI owns
@ -493,11 +530,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
const base =
useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin);
const localAgentDetection = canUseLocalAgentDetection(base);
// null while loading; the same setting the General tab exposes (shared cache).
const [autoSwitch, setAutoSwitch] = useState<OpenAIAutoSwitchSettings | null>(
null,
);
const [savingAutoSwitch, setSavingAutoSwitch] = useState(false);
useEffect(() => {
void fetchDeviceType({ force: true });
@ -575,27 +607,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
}
}, [agent, detectedAgents, activeGgufVariant, activeNativePathToken, ggufContextLength]);
useEffect(() => {
let cancelled = false;
void loadOpenAIAutoSwitchSettings()
.then((s) => {
if (!cancelled) setAutoSwitch(s);
})
.catch(() => {
// Best-effort: leave the toggle off if the setting can't be read.
});
return () => {
cancelled = true;
};
}, []);
const model = useLoadedModelName();
const model = useExampleModelName();
const key = apiKey || KEY_PLACEHOLDER;
const autoSwitchOn = autoSwitch?.enabled ?? false;
// Null model: nothing is servable, so there is no snippet worth copying.
const snippets = useMemo(
() => buildSnippets(base, key, model, os, autoSwitchOn),
[base, key, model, os, autoSwitchOn],
() => (model ? buildSnippets(base, key, model, os) : null),
[base, key, model, os],
);
// Agent command must target the server the panel shows, not the :8888 default.
const agentCommand = useMemo(
@ -613,6 +631,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
: "python";
const handleCopy = async () => {
if (!snippets) return;
if (await copyToClipboard(snippets[lang])) {
setCopied(true);
setTimeout(() => setCopied(false), 1800);
@ -624,20 +643,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
writeUseTunnelPref(next);
};
// Same setting as the General tab; persist optimistically and revert on failure
// so the examples reflect the live model-switch behavior.
const handleToggleAutoSwitch = (next: boolean) => {
const idle = autoSwitch?.autoUnloadIdleSeconds ?? 0;
setAutoSwitch((prev) => (prev ? { ...prev, enabled: next } : prev));
setSavingAutoSwitch(true);
void updateOpenAIAutoSwitchSettings(next, idle)
.then(setAutoSwitch)
.catch(() => {
setAutoSwitch((prev) => (prev ? { ...prev, enabled: !next } : prev));
})
.finally(() => setSavingAutoSwitch(false));
};
const handleCopyUrl = async () => {
if (cloudflareUrl && (await copyToClipboard(cloudflareUrl))) {
setCopiedUrl(true);
@ -658,41 +663,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
{t("settings.apiKeys.usageExamples")}
</h2>
<div className="min-w-0 max-w-full overflow-hidden rounded-lg border border-border bg-muted/20">
{/* Same setting as the General tab; surfaced here so the request `model`
actually switches the served model, which the examples below show. */}
<div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
<div className="flex shrink-0 items-center gap-1.5">
<Switch
size="sm"
checked={autoSwitchOn}
disabled={autoSwitch === null || savingAutoSwitch}
onCheckedChange={handleToggleAutoSwitch}
aria-label={t("settings.general.modelAutoSwitch.enable")}
/>
<span className="text-ui-11 font-medium text-foreground">
{t("settings.general.modelAutoSwitch.enable")}
</span>
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="flex items-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring"
aria-label={t(
"settings.general.modelAutoSwitch.enableDescription",
)}
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3.5"
/>
</button>
</TooltipTrigger>
<TooltipContent className="max-w-[260px] text-ui-11 leading-snug">
{t("settings.general.modelAutoSwitch.enableDescription")}
</TooltipContent>
</Tooltip>
</div>
</div>
{/* No model-auto-switch row: ModelAutoSwitchSection renders that setting just below. */}
{cloudflareUrl ? (
<div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
<div className="flex shrink-0 items-center gap-1.5">
@ -802,25 +773,33 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
</button>
</div>
) : null}
<div className="relative min-w-0">
<button
type="button"
onClick={handleCopy}
className="absolute right-2 top-2 z-10 flex items-center gap-1 rounded border border-border bg-background/80 px-1.5 py-1 text-ui-11 text-muted-foreground backdrop-blur transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring"
aria-label={t("settings.apiKeys.copySnippet")}
>
<HugeiconsIcon
icon={copied ? Tick02Icon : Copy01Icon}
className={cn("size-3.5", copied && "text-emerald-600")}
{snippets ? (
<div className="relative min-w-0">
<button
type="button"
onClick={handleCopy}
className="absolute right-2 top-2 z-10 flex items-center gap-1 rounded border border-border bg-background/80 px-1.5 py-1 text-ui-11 text-muted-foreground backdrop-blur transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring"
aria-label={t("settings.apiKeys.copySnippet")}
>
<HugeiconsIcon
icon={copied ? Tick02Icon : Copy01Icon}
className={cn("size-3.5", copied && "text-emerald-600")}
/>
{copied
? t("settings.apiKeys.copied")
: t("settings.apiKeys.copy")}
</button>
<HighlightedCode
key={snippets[lang]}
code={snippets[lang]}
language={shikiLang}
/>
{copied ? t("settings.apiKeys.copied") : t("settings.apiKeys.copy")}
</button>
<HighlightedCode
key={snippets[lang]}
code={snippets[lang]}
language={shikiLang}
/>
</div>
</div>
) : (
<div className="min-w-0 px-3 py-2.5 text-ui-11 leading-snug text-muted-foreground">
{t("settings.apiKeys.usageNoModel")}
</div>
)}
<div className="flex min-w-0 flex-col gap-1.5 border-t border-border px-3 py-2.5">
<span className="text-ui-11 font-semibold text-foreground">
{t("settings.apiKeys.codingAgents")}

4
studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
View file

@ -168,12 +168,12 @@ export function ApiKeysTab() {
)}
</section>
<ModelAutoSwitchSection />
<ApiMonitorConsole />
<UsageExamples apiKey={revealed} />
<ModelAutoSwitchSection />
<Dialog open={revokeTarget !== null} onOpenChange={(o) => !o && setRevokeTarget(null)}>
<DialogContent className="max-w-md">
<DialogHeader>

30
studio/frontend/src/hooks/use-gpu-info.ts
View file

@ -100,13 +100,29 @@ function toGpuInfo(
}
function toGpuDevices(data: SystemInfoResponse | null): SystemGpuDevice[] {
// Unpinnable configurations must hide every pick surface: XPU indices are
// torch-xpu ordinals no applicator speaks, and Vulkan-only builds pin ggml's
// own ordinals -- /load and /validate 400 picks on both, so the backend
// reports gpu.gguf_gpu_ids_supported and every gate keyed on physicalIndex
// (picker, persisted-pick reconcile) follows it. The device flavor lives on
// the TOP-LEVEL device_backend field; absent support info defaults to
// pinnable (older backend).
// GGUF loads run through llama-server, so on a Vulkan build the pickable set
// is the inference inventory, not the torch view: it can see cards torch
// cannot, and its indices are the ggml ordinals `--device Vulkan<i>` pins.
// The XPU ban does not apply there, it is about torch-xpu ordinals that no
// applicator speaks; a Vulkan pick does not use them.
const inference = data?.inference_gpu;
if (inference?.backend === "vulkan" && (inference.devices ?? []).length) {
const picksAccepted = inference.gguf_gpu_ids_supported !== false;
return (inference.devices ?? [])
.filter((d) => typeof d.index === "number")
.map((d) => ({
index: d.index as number,
name: d.name ?? `GPU ${d.index}`,
memoryTotalGb: d.memory_total_gb ?? 0,
memoryFreeGb: d.vram_free_gb ?? 0,
physicalIndex: picksAccepted && d.index_kind === "vulkan",
}));
}
// Otherwise the torch view is the pickable set. Unpinnable configurations
// must hide every pick surface: XPU indices are torch-xpu ordinals no
// applicator speaks, so /load and /validate 400 them, and the backend reports
// gpu.gguf_gpu_ids_supported. Absent support info defaults to pinnable
// (older backend).
const pinnableBackend =
data?.device_backend !== "xpu" &&
data?.gpu?.gguf_gpu_ids_supported !== false;

17
studio/frontend/src/i18n/locales/en.ts
View file

@ -284,20 +284,21 @@ export const en = {
sectionTitle: "Model auto-switch (OpenAI API)",
enable: "Switch model by request",
enableDescription:
"When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.",
"Load a downloaded GGUF named in an API request before serving. Off by default.",
autoDownload: "Download missing models",
autoDownloadDescription:
"Fetch a GGUF named in an API request that is not downloaded yet. Anyone with an API key can then use disk and bandwidth.",
idleUnload: "Idle auto-unload",
idleUnloadDescription:
"Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded. Minimum 60 seconds.",
idleNeedsEnable:
"Turn on Switch model by request so an unloaded model reloads on next use.",
idleActiveViaEnv:
"Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.",
"Free VRAM after this many idle seconds. 0 keeps it loaded, minimum 60.",
idleNeedsEnable: "Turn on Switch model by request first.",
idleActiveViaEnv: "Active via UNSLOTH_MODEL_IDLE_TTL.",
loadError: "Failed to load model auto-switch settings.",
saveError: "Failed to save model auto-switch settings.",
idleError: "Enter 0 to keep the model loaded, or at least 60 seconds.",
keepKv: "Keep chat context across idle unload",
keepKvDescription:
"Save the model's KV cache to disk before an idle unload and restore it on reload, so resumed chats skip re-reading their history. Chat context is written to disk (up to 10 GB) until it is restored or cleaned up.",
"Save the KV cache before an idle unload so resumed chats skip re-reading history. Up to 10 GB on disk.",
},
previewSharing: {
sectionTitle: "Preview sharing",
@ -818,6 +819,8 @@ export const en = {
copyAccessToken: "Copy access token",
copyNow: "Copy now - this won't be shown again.",
usageExamples: "Usage examples",
usageNoModel:
"Load or download a model to see runnable examples. This server has no model to name yet.",
usageTools: "Tools",
exampleCurlTools: "curl + tools",
examplePythonTools: "Python + tools",

158
studio/install_python_stack.py
View file

@ -94,6 +94,40 @@ def _strix_needs_amd_arch_index(ver: tuple[int, int]) -> bool:
return key is not None and key < _ROCM_ARCH_INDEX_FLOOR
# MI50 / Radeon VII (gfx906, Vega 20): rocm6.4+/7.x wheels bundle ROCm libraries
# whose Tensile kernels dropped gfx906 (rocBLAS "TensileLibrary.dat ... not read
# for gfx906", ROCm/TheRock#1844), failing 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). Uses the _default (<2.11) pkg specs -- the
# rocm7.2 floor of 2.11 cannot be satisfied there. Mirrors install.sh.
_GFX906_LEGACY_TAG = "rocm6.3"
def _gfx906_needs_legacy_index(ver: tuple[int, int]) -> bool:
"""True when the generic tag picked for the host ROCm version is newer than
rocm6.3, i.e. its wheels lack gfx906 kernels and must be rerouted."""
key = next((k for k in sorted(_ROCM_TORCH_INDEX, reverse = True) if ver >= k), None)
return key is not None and key > (6, 3)
def _runtime_target_is_gfx906() -> bool:
"""True when the runtime GPU target is gfx906 (MI50 / Radeon VII).
An explicit UNSLOTH_ROCM_GFX_ARCH wins (mirrors _infer_linux_amd_gfx_arch /
the display path), so a host whose rocminfo/amd-smi emit no gfx token can
still opt in. Otherwise report gfx906 only when it is the SOLE distinct arch:
_detect_amd_gfx_codes() de-duplicates arches, which loses per-device ordinals
on a mixed host, so a non-gfx906 selection is never mis-identified as gfx906
(and downgraded to rocm6.3). Mixed gfx906+dGPU hosts opt in with the env var.
"""
# Normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) so the
# feature-flag suffix does not defeat the exact comparison (mirrors device_type.py).
override = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower().split(":")[0]
if override:
return override == "gfx906"
return set(_detect_amd_gfx_codes()) == {"gfx906"}
# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug).
# Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare.
_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset(
@ -939,6 +973,34 @@ def _detect_bnb_rocm_dll_ver() -> str | None:
return max(all_vers, key = lambda v: int(v)) if all_vers else None
# Set right before the base unsloth install (which resolves its unconditional
# bitsandbytes dependency); read by _ensure_rocm_torch to drop a freshly pulled
# generic wheel on gfx906 while leaving a pre-existing source build untouched.
_GFX906_BNB_ABSENT_BEFORE_BASE = False
def _bitsandbytes_installed() -> bool:
"""True if bitsandbytes is importable in the target venv. Runs a fresh
subprocess so a package installed earlier this run is seen; only checks the
spec (does NOT import bitsandbytes)."""
try:
return (
subprocess.run(
[
sys.executable,
"-c",
"import importlib.util, sys; "
"sys.exit(0 if importlib.util.find_spec('bitsandbytes') else 1)",
],
capture_output = True,
timeout = 60,
).returncode
== 0
)
except Exception:
return False
_BNB_ROCM_SITECUSTOMIZE_BEGIN = "# BEGIN Unsloth BNB_ROCM_VERSION"
_BNB_ROCM_SITECUSTOMIZE_END = "# END Unsloth BNB_ROCM_VERSION"
_BNB_ROCM_VERSION_SOURCE_ENV = "UNSLOTH_BNB_ROCM_VERSION_SOURCE"
@ -1890,13 +1952,25 @@ def _ensure_rocm_torch() -> None:
)
rocm_torch_ready = True
# An explicit UNSLOTH_ROCM_GFX_ARCH=gfx906 pins the runtime target to the
# MI50 / Radeon VII path; it must win over the Strix probe-order detection
# below (a mixed Strix + MI50 host could otherwise route to gfx1151), so the
# Strix override is skipped when it is set.
_gfx906_arch_override = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower().split(
":"
)[0] == "gfx906"
# Strix Halo / Point (gfx1151 / gfx1150) need torch from AMD's per-gfx index
# (2.11+rocm7.13); any generic pytorch.org rocm index lacks the fixes (ROCm 7.1
# segfaults in _grouped_mm). See _strix_needs_amd_arch_index for the floor gate.
_strix_override_url: "str | None" = None
_strix_override_pkgs: "tuple[str, str, str] | None" = None
# An explicit ROCm pin is authoritative: never auto-reroute it.
if _strix_needs_amd_arch_index(ver) and _explicit_rocm_torch_index_url() is None:
if (
_strix_needs_amd_arch_index(ver)
and _explicit_rocm_torch_index_url() is None
and not _gfx906_arch_override
):
gfx_codes = _detect_amd_gfx_codes()
_strix_gfx = {"gfx1151", "gfx1150", "gfx1152"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
@ -1933,6 +2007,34 @@ def _ensure_rocm_torch() -> None:
f" skipping AMD per-gfx index override.\n"
)
# gfx906 (MI50 / Radeon VII): is this the runtime GPU target? Used below to skip
# the generic bitsandbytes wheel (no gfx906 kernels). This must hold even under
# an explicit torch-index pin: a gfx906 host that pins rocm6.3 (without also
# setting UNSLOTH_ROCM_GFX_ARCH) would otherwise reinstall the prebuilt bnb wheel
# over the user's source-built gfx906 bnb. So a pin suppresses only the torch
# reroute (_gfx906_override below), NOT the gfx906 detection for the bnb skip.
_runtime_is_gfx906 = _gfx906_arch_override or _runtime_target_is_gfx906()
# Reroute torch to the last gfx906-capable wheel family (rocm6.3) only when the
# host ROCm version would otherwise pick a newer, kernel-less index -- and never
# over an explicit pin or an active Strix reroute (the pin/Strix path installs
# its own index; only the bnb skip must still apply on those paths).
_gfx906_override = (
_runtime_is_gfx906
and _gfx906_needs_legacy_index(ver)
and _explicit_rocm_torch_index_url() is None
and _strix_override_url is None
)
if _gfx906_override:
print(
f"\n gfx906 (MI50 / Radeon VII / Vega 20) is the runtime target with ROCm "
f"{ver[0]}.{ver[1]}.\n"
f" Routing torch install to the {_GFX906_LEGACY_TAG} index: the last wheel\n"
f" family that runs on gfx906 (newer rocm wheels ship without gfx906 BLAS\n"
f" kernels and fail at first use). gfx906 is a community-maintained legacy\n"
f" path: 16-bit LoRA and full finetuning work; bitsandbytes 4-bit QLoRA\n"
f" requires a source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd).\n"
)
# The Strix override must fire even when has_hip_torch is True: an existing
# torch.version.hip == "7.1" is exactly the broken combo it repairs.
if _strix_override_url is not None and _strix_override_pkgs is not None:
@ -1954,6 +2056,29 @@ def _ensure_rocm_torch() -> None:
constrain = False,
)
rocm_torch_ready = True
# gfx906 fires even when has_hip_torch is True: a +rocm7.x build IS the broken
# combo it repairs. A torch already on rocm6.3 wheels is left alone (the tag
# check below is False, and rocm_torch_ready is already True from has_hip_torch,
# so the generic fallback is skipped).
elif _gfx906_override and _GFX906_LEGACY_TAG not in _installed_torch_ver:
index_url = f"{_PYTORCH_WHL_BASE}/{_GFX906_LEGACY_TAG}"
_torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["_default"]
print(
f" gfx906 legacy override -- installing torch from "
f"{_strip_index_url_credentials(index_url)}"
)
pip_install(
f"ROCm torch (gfx906, {_GFX906_LEGACY_TAG})",
"--force-reinstall",
"--no-cache-dir",
_torch_pkg,
_vision_pkg,
_audio_pkg,
"--index-url",
index_url,
constrain = False,
)
rocm_torch_ready = True
elif not rocm_torch_ready:
# Reinstall when torch is not ROCm yet, OR a ROCm build's family differs from a pin.
# Gate on rocm_torch_ready (not has_hip_torch alone) so a successful inferred-gfx
@ -2002,11 +2127,33 @@ def _ensure_rocm_torch() -> None:
)
rocm_torch_ready = True
# gfx906 has no prebuilt bitsandbytes: the continuous-release/PyPI wheels ship
# no gfx906 kernels, and force-reinstalling them would clobber a user's
# source-built bnb (the only 4-bit path on this arch) on every `studio update`.
# Skip the auto-install and leave whatever bnb is present.
if rocm_torch_ready and _runtime_is_gfx906:
print(
_dim(
" gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels). "
"Build bitsandbytes from source for 4-bit QLoRA -- "
"see docs.unsloth.ai/get-started/install-and-update/amd."
)
)
# The base install resolves unsloth's unconditional bitsandbytes dep to a
# generic CUDA wheel with no gfx906 kernels ("invalid device function" at
# 4-bit use). Drop it if this run pulled it in; a pre-existing source build
# (present before the base install) is left untouched.
if _GFX906_BNB_ABSENT_BEFORE_BASE and _bitsandbytes_installed():
print(_dim(" gfx906: removing generic bitsandbytes pulled in as a dependency"))
subprocess.run(
[sys.executable, "-m", "pip", "uninstall", "-y", "bitsandbytes"],
capture_output = True,
)
# Install bitsandbytes only when torch links against ROCm. Prefers the
# continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix), falling back
# to PyPI when the pre-release wheel won't install. Use pip for the
# pre-release wheel because uv rejects its filename/metadata version mismatch.
if rocm_torch_ready:
elif rocm_torch_ready:
_bnb_url = _bnb_rocm_prerelease_url()
_bnb_installed = False
if _bnb_url is not None:
@ -2767,6 +2914,13 @@ def install_python_stack() -> int:
"mlx-vlm",
)
# gfx906: the base install below resolves unsloth's unconditional bitsandbytes
# dep to a generic CUDA wheel (no gfx906 kernels). Record bnb's presence now so
# _ensure_rocm_torch can drop a freshly pulled wheel while keeping a source build.
global _GFX906_BNB_ABSENT_BEFORE_BASE
if not skip_base:
_GFX906_BNB_ABSENT_BEFORE_BASE = not _bitsandbytes_installed()
# 3. Core packages: unsloth-zoo + unsloth (or custom package name)
if skip_base:
pass

2
studio/src-tauri/tauri.conf.json
View file

@ -36,7 +36,7 @@
},
"plugins": {
"updater": {
"pubkey": "dW50cnVzdGVkIGNvbW1lbnQ6IG1pbmlzaWduIHB1YmxpYyBrZXk6IDE0NjlBNkEwQTc0QjY0Q0UKUldUT1pFdW5vS1pwRkl3TXRrOGlYdDNMeTh1bEhSbURrM3IyT2lTK1BiekpTc0Z2SXZFeVNibDIK",
"pubkey": "dW50cnVzdGVkIGNvbW1lbnQ6IG1pbmlzaWduIHB1YmxpYyBrZXk6IDFBQzA4RjczODM0RjE1QjcKUldTM0ZVK0RjNC9BR2t4R0RVaFR5cTkyUlRVQ1FwaGV0Nk04eWNwWXBhZnlzalJydllmZm1QTS8K",
"endpoints": [
"https://github.com/unslothai/unsloth/releases/download/desktop-latest/latest.json"
],

20
tests/python/test_tokenizers_and_torch_constraint.py
View file

@ -88,11 +88,21 @@ class TestStructuralTorchConstraint:
"""$TORCH_CONSTRAINT must appear in a uv pip install line."""
assert '"$TORCH_CONSTRAINT"' in self._sh
def test_hardcoded_torch_constraint_only_once(self):
"""The hard-coded torch>=2.4,<2.11.0 string should appear exactly once
in install.sh (the default assignment), not in pip install lines."""
count = self._sh.count('"torch>=2.4,<2.11.0"')
assert count == 1, f"Expected 1, found {count}"
def test_hardcoded_torch_constraint_only_on_assignments(self):
"""The hard-coded torch>=2.4,<2.11.0 string must only appear on
TORCH_CONSTRAINT= assignment lines, never on a pip/uv install line
(those must reference $TORCH_CONSTRAINT). Two assignments are expected:
the default, and the gfx906 (MI50) reroute that restores the default
<2.11 window after the rocm7.2 floor bump raised it to 2.11."""
hits = [ln for ln in self._sh.splitlines() if '"torch>=2.4,<2.11.0"' in ln]
assert hits, "default constraint literal missing from install.sh"
for ln in hits:
assert (
"TORCH_CONSTRAINT=" in ln
), f"torch>=2.4,<2.11.0 hardcoded off a TORCH_CONSTRAINT= assignment: {ln.strip()!r}"
assert (
"pip install" not in ln
), f"torch>=2.4,<2.11.0 hardcoded on a pip install line: {ln.strip()!r}"
def test_tightening_guarded_by_skip_torch(self):
"""The block must check SKIP_TORCH=false."""

67
tests/saving/run_offline_gguf_integration.py Normal file
View file

@ -0,0 +1,67 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Download real Gemma weights and run offline integration tests for #7481.
Sets ``UNSLOTH_INTEGRATION_IMPORT=1`` for the pytest subprocess so the
real-cache suite is not silently skipped. Requires a host that can import
unsloth (typically GPU).
Example:
python tests/saving/run_offline_gguf_integration.py
python tests/saving/run_offline_gguf_integration.py --download-only
"""
from __future__ import annotations
import os
import subprocess
import sys
import tempfile
from pathlib import Path
REPO = "unsloth/gemma-3-270m-it-bnb-4bit"
CACHE_ROOT = Path(
os.environ.get("HF_HOME") or os.path.join(tempfile.gettempdir(), "hf_offline_test_cache")
)
def download():
from huggingface_hub import snapshot_download
os.environ.setdefault("HF_HOME", str(CACHE_ROOT))
path = snapshot_download(REPO, cache_dir = str(CACHE_ROOT / "hub"))
print("cached at", path)
def run_tests():
os.environ.setdefault("HF_HOME", str(CACHE_ROOT))
# Real-cache suite is gated on this; without it every integration test skips
# and the runner reports success after only the fake-cache unit file ran.
env = os.environ.copy()
env["UNSLOTH_INTEGRATION_IMPORT"] = "1"
cmd = [
sys.executable,
"-m",
"pytest",
"tests/saving/test_offline_gguf_vlm_tokenizer_7481.py",
"tests/saving/test_offline_gguf_real_cache_integration.py",
"-q",
]
raise SystemExit(subprocess.call(cmd, cwd = str(Path(__file__).resolve().parents[2]), env = env))
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--download-only", action = "store_true")
args = parser.parse_args()
download()
if not args.download_only:
run_tests()
if __name__ == "__main__":
main()

122
tests/saving/test_offline_gguf_real_cache_integration.py Normal file
View file

@ -0,0 +1,122 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Integration tests for #7481 using real cached Gemma weights.
Requires a one-time online download into ``$HF_HOME`` (defaults to a
``hf_offline_test_cache`` directory under the platform temp dir):
HF_HOME=<cache> python -c \\
"from huggingface_hub import snapshot_download; snapshot_download('unsloth/gemma-3-270m-it-bnb-4bit', cache_dir='<cache>/hub')"
Every test here drives unsloth's own resolver. Resolving through
``hf_hub_download`` directly would pass with the fix reverted, since that is
plain huggingface_hub behaviour rather than anything this change touches.
Importing unsloth pulls the whole package graph, which CPU-only hosts cannot
do, so the suite is gated behind ``UNSLOTH_INTEGRATION_IMPORT=1``.
"""
from __future__ import annotations
import os
import socket
import tempfile
from pathlib import Path
import pytest
REPO = "unsloth/gemma-3-270m-it-bnb-4bit"
CACHE_ROOT = Path(
os.environ.get("HF_HOME") or os.path.join(tempfile.gettempdir(), "hf_offline_test_cache")
)
pytestmark = [
pytest.mark.integration,
pytest.mark.skipif(
os.environ.get("UNSLOTH_INTEGRATION_IMPORT") != "1",
reason = "full unsloth import needs a GPU host; set UNSLOTH_INTEGRATION_IMPORT=1 to enable",
),
]
def _require_cached_repo():
from huggingface_hub import scan_cache_dir
cache_dir = CACHE_ROOT / "hub"
if not cache_dir.exists():
pytest.skip(f"cache missing at {cache_dir}; run snapshot_download for {REPO}")
repos = [r.repo_id for r in scan_cache_dir(str(cache_dir)).repos]
if REPO not in repos:
pytest.skip(f"{REPO} not in {cache_dir}")
def _block_network(monkeypatch):
def _guard(*args, **kwargs):
raise OSError("network blocked for offline integration test")
# Patch the method, not the class: replacing socket.socket itself breaks any
# isinstance(x, socket.socket) in the stack under test.
monkeypatch.setattr(socket.socket, "connect", _guard)
monkeypatch.setattr(socket, "create_connection", _guard)
monkeypatch.setattr(socket, "getaddrinfo", _guard)
def _offline_env(monkeypatch):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
monkeypatch.setenv("HF_HOME", str(CACHE_ROOT))
def test_real_cached_snapshot_resolves_offline(monkeypatch):
_require_cached_repo()
_offline_env(monkeypatch)
_block_network(monkeypatch)
from unsloth.models.loader_utils import _resolve_hub_repo_local_dir
snap = Path(
_resolve_hub_repo_local_dir(
REPO,
cache_dir = str(CACHE_ROOT / "hub"),
local_files_only = True,
)
)
assert (snap / "tokenizer.json").is_file()
assert (snap / "tokenizer.model").is_file()
def test_real_cached_tokenizer_loads_from_snapshot_not_repo_id(monkeypatch):
"""The #7481 fix: the loader hands transformers a snapshot dir, not a repo id."""
_require_cached_repo()
_offline_env(monkeypatch)
_block_network(monkeypatch)
from unsloth.models.loader_utils import _load_pretrained_tokenizer_fast
tok = _load_pretrained_tokenizer_fast(
REPO,
local_files_only = True,
cache_dir = str(CACHE_ROOT / "hub"),
)
assert tok.vocab_size > 0
# A repo id here means the Hub metadata probe was reached, which is the bug.
assert tok.name_or_path != REPO
assert Path(tok.name_or_path).is_dir()
def test_real_cached_unsloth_helpers_offline(monkeypatch):
_require_cached_repo()
_offline_env(monkeypatch)
_block_network(monkeypatch)
from unsloth.models.loader_utils import _load_pretrained_tokenizer_fast
from unsloth.save import _has_tokenizer_model
tok = _load_pretrained_tokenizer_fast(
REPO,
local_files_only = True,
cache_dir = str(CACHE_ROOT / "hub"),
)
assert tok.vocab_size > 0
assert _has_tokenizer_model(tok) is True

336
tests/saving/test_offline_gguf_vlm_tokenizer_7481.py Normal file
View file

@ -0,0 +1,336 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Offline GGUF export must not probe the Hub for VLM tokenizer metadata (issue #7481).
Regression for ``PreTrainedTokenizerFast.from_pretrained`` on a repo id calling
``is_base_mistral()`` -> ``model_info()`` even with ``TRANSFORMERS_OFFLINE=1``.
Pure CPU, no network, no GPU.
"""
import json
import os
from types import SimpleNamespace
from unittest.mock import patch
from unsloth.models import loader_utils as L
_REPO = "llmfan46/gemma-4-E4B-it-ultra-uncensored-heretic"
_COMMIT = "5964fe4c7339c5974e879baba8982a09616f68ca"
def _write_gemma4_cache(
root,
repo_id = _REPO,
commit = _COMMIT,
):
"""Minimal cached snapshot matching the reporter's layout."""
org, name = repo_id.split("/")
repo_root = root / f"models--{org}--{name}"
snap = repo_root / "snapshots" / commit
snap.mkdir(parents = True)
refs = repo_root / "refs"
refs.mkdir(parents = True, exist_ok = True)
(refs / "main").write_text(commit, encoding = "utf-8")
(snap / "tokenizer_config.json").write_text(
json.dumps({"tokenizer_class": "GemmaTokenizer", "model_max_length": 8192}),
encoding = "utf-8",
)
(snap / "tokenizer.json").write_text(
json.dumps(
{
"version": "1.0",
"truncation": None,
"padding": None,
"added_tokens": [],
"normalizer": None,
"pre_tokenizer": None,
"post_processor": None,
"decoder": None,
"model": {"type": "BPE", "vocab": {"<pad>": 0}, "merges": []},
}
),
encoding = "utf-8",
)
(snap / "processor_config.json").write_text("{}", encoding = "utf-8")
(snap / "config.json").write_text(
json.dumps({"model_type": "gemma4"}),
encoding = "utf-8",
)
return snap
def _offline_env(monkeypatch, cache_root):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
monkeypatch.setenv("HF_HUB_CACHE", str(cache_root))
def test_resolve_hub_repo_cached_file_finds_tokenizer_model(tmp_path, monkeypatch):
snap = _write_gemma4_cache(tmp_path)
(snap / "tokenizer.model").write_bytes(b"sp-model")
_offline_env(monkeypatch, tmp_path)
got = L._resolve_hub_repo_cached_file(
_REPO,
"tokenizer.model",
local_files_only = True,
cache_dir = str(tmp_path),
)
assert got == str(snap / "tokenizer.model")
def test_resolve_hub_repo_local_dir_from_cached_snapshot(tmp_path, monkeypatch):
snap = _write_gemma4_cache(tmp_path)
_offline_env(monkeypatch, tmp_path)
got = L._resolve_hub_repo_local_dir(_REPO, local_files_only = True, cache_dir = str(tmp_path))
assert got == str(snap)
def test_hub_repo_or_local_path_prefers_snapshot_over_repo_id(tmp_path, monkeypatch):
snap = _write_gemma4_cache(tmp_path)
_offline_env(monkeypatch, tmp_path)
got = L._hub_repo_or_local_path(_REPO, local_files_only = True, cache_dir = str(tmp_path))
assert got == str(snap)
assert got != _REPO
def test_hub_repo_or_local_path_keeps_repo_id_online(tmp_path, monkeypatch):
snap = _write_gemma4_cache(tmp_path)
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
got = L._hub_repo_or_local_path(_REPO, local_files_only = False, cache_dir = str(tmp_path))
assert got == _REPO
assert got != str(snap)
def test_has_tokenizer_model_offline_does_not_cache_negative(tmp_path, monkeypatch):
from unsloth.save import _TOKENIZER_MODEL_CACHE, _has_tokenizer_model
snap = _write_gemma4_cache(tmp_path)
_offline_env(monkeypatch, tmp_path)
_TOKENIZER_MODEL_CACHE.clear()
tok = SimpleNamespace(name_or_path = _REPO)
assert _has_tokenizer_model(tok, token = None) is False
assert _REPO not in _TOKENIZER_MODEL_CACHE
(snap / "tokenizer.model").write_bytes(b"sp-model")
assert _has_tokenizer_model(tok, token = None) is True
def test_preserve_sentencepiece_offline_copies_cached_model(tmp_path, monkeypatch):
from unsloth.save import _TOKENIZER_MODEL_CACHE, _preserve_sentencepiece_tokenizer_assets
snap = _write_gemma4_cache(tmp_path)
(snap / "tokenizer.model").write_bytes(b"cached-sp-model")
_offline_env(monkeypatch, tmp_path)
_TOKENIZER_MODEL_CACHE.clear()
save_dir = tmp_path / "export"
save_dir.mkdir()
(save_dir / "tokenizer_config.json").write_text("{}", encoding = "utf-8")
tok = SimpleNamespace(name_or_path = _REPO)
_preserve_sentencepiece_tokenizer_assets(tok, str(save_dir))
assert (save_dir / "tokenizer.model").read_bytes() == b"cached-sp-model"
def test_load_pretrained_tokenizer_fast_passes_snapshot_not_repo_id(tmp_path, monkeypatch):
snap = _write_gemma4_cache(tmp_path)
_offline_env(monkeypatch, tmp_path)
seen_paths = []
class _FakeFast:
@classmethod
def from_pretrained(cls, path, **kwargs):
seen_paths.append(path)
assert kwargs.get("local_files_only") is True
return SimpleNamespace(name_or_path = path)
monkeypatch.setattr(
"transformers.PreTrainedTokenizerFast",
_FakeFast,
raising = False,
)
with patch("huggingface_hub.HfApi.model_info") as model_info:
model_info.side_effect = AssertionError("model_info must not run offline")
tok = L._load_pretrained_tokenizer_fast(_REPO, cache_dir = str(tmp_path))
assert seen_paths == [str(snap)]
assert tok.name_or_path == str(snap)
def test_has_tokenizer_model_offline_skips_model_info(tmp_path, monkeypatch):
from unsloth.save import _TOKENIZER_MODEL_CACHE, _has_tokenizer_model
_write_gemma4_cache(tmp_path)
_offline_env(monkeypatch, tmp_path)
_TOKENIZER_MODEL_CACHE.clear()
tok = SimpleNamespace(name_or_path = _REPO)
# A raising side_effect proves nothing: _has_tokenizer_model wraps the call
# in `except Exception: return False`, so it passes with the fix reverted.
with patch("huggingface_hub.HfApi.model_info") as model_info:
assert _has_tokenizer_model(tok, token = None) is False
assert model_info.call_count == 0
def test_has_tokenizer_model_probes_cache_before_model_info(tmp_path, monkeypatch):
from unsloth.save import _TOKENIZER_MODEL_CACHE, _has_tokenizer_model
snap = _write_gemma4_cache(tmp_path)
(snap / "tokenizer.model").write_bytes(b"sp-model")
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
_TOKENIZER_MODEL_CACHE.clear()
tok = SimpleNamespace(name_or_path = _REPO)
with patch("huggingface_hub.HfApi.model_info") as model_info:
model_info.side_effect = AssertionError("model_info must not run when cache hit")
assert _has_tokenizer_model(tok, token = None) is True
def test_offline_aware_load_persists_local_only_for_saving(tmp_path, monkeypatch):
"""An explicit ``local_files_only = True`` load must still be local-only at save time.
``transformers`` takes ``local_files_only`` as an explicit ``from_pretrained``
parameter, so it never reaches ``tokenizer.init_kwargs``, and
``_offline_aware_load`` restores the offline env vars once the load returns.
Without the stamp the request is invisible by the time we save.
"""
from unsloth.save import _TOKENIZER_MODEL_CACHE, _has_tokenizer_model
# Snapshot has tokenizer metadata but deliberately no tokenizer.model, so the
# cache probe misses and only the local-only stamp can stop the Hub request.
_write_gemma4_cache(tmp_path)
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
_TOKENIZER_MODEL_CACHE.clear()
@L._offline_aware_load
def _load(model_name, **kwargs):
assert os.environ.get("HF_HUB_OFFLINE") == "1"
# A processor keeps the Hub repo id and carries no local_files_only.
return object(), SimpleNamespace(
tokenizer = SimpleNamespace(name_or_path = model_name, init_kwargs = {}),
)
_model, processor = _load(_REPO, local_files_only = True)
assert os.environ.get("HF_HUB_OFFLINE") is None
assert processor.tokenizer.init_kwargs.get("local_files_only") is None
assert L._tokenizer_wants_local_only(processor.tokenizer) is True
with patch("huggingface_hub.HfApi.model_info") as model_info:
model_info.return_value = SimpleNamespace(
siblings = [SimpleNamespace(rfilename = "tokenizer.model")],
)
assert _has_tokenizer_model(processor, token = None) is False
assert model_info.call_count == 0
def test_preserve_sentencepiece_after_local_only_load_never_downloads(tmp_path, monkeypatch):
"""The save path inherits the load's local-only mode: no metadata probe, no download."""
import huggingface_hub
from unsloth.save import _TOKENIZER_MODEL_CACHE, _preserve_sentencepiece_tokenizer_assets
_write_gemma4_cache(tmp_path)
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
_TOKENIZER_MODEL_CACHE.clear()
@L._offline_aware_load
def _load(model_name, **kwargs):
return object(), SimpleNamespace(
tokenizer = SimpleNamespace(name_or_path = model_name, init_kwargs = {}),
)
_model, processor = _load(_REPO, local_files_only = True)
save_dir = tmp_path / "export"
save_dir.mkdir()
(save_dir / "tokenizer_config.json").write_text("{}", encoding = "utf-8")
real_download = huggingface_hub.hf_hub_download
seen_local_files_only = []
def _recording_download(*args, **kwargs):
seen_local_files_only.append(kwargs.get("local_files_only"))
return real_download(*args, **kwargs)
monkeypatch.setattr("huggingface_hub.hf_hub_download", _recording_download)
with patch("huggingface_hub.HfApi.model_info") as model_info:
model_info.return_value = SimpleNamespace(
siblings = [SimpleNamespace(rfilename = "tokenizer.model")],
)
_preserve_sentencepiece_tokenizer_assets(processor, str(save_dir), token = None)
assert model_info.call_count == 0
# Every hf_hub_download here must be a cache probe, never a Hub fetch.
assert seen_local_files_only and all(seen_local_files_only)
assert not (save_dir / "tokenizer.model").exists()
def test_has_tokenizer_model_local_files_only_skips_model_info(tmp_path, monkeypatch):
from unsloth.save import _TOKENIZER_MODEL_CACHE, _has_tokenizer_model
_write_gemma4_cache(tmp_path)
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
_TOKENIZER_MODEL_CACHE.clear()
tok = SimpleNamespace(
name_or_path = _REPO,
init_kwargs = {"local_files_only": True},
)
with patch("huggingface_hub.HfApi.model_info") as model_info:
assert _has_tokenizer_model(tok, token = None) is False
assert model_info.call_count == 0
def test_custom_cache_dir_survives_to_saving(tmp_path, monkeypatch):
"""A local-only load with a caller-supplied cache_dir that no env var points
at. Saving derives its cache from HF_HUB_CACHE / HF_HOME, so without the
stamp it probes the wrong place, and the local-only marker then stops it
falling back to the Hub, silently dropping tokenizer.model."""
from unsloth.save import _TOKENIZER_MODEL_CACHE, _has_tokenizer_model
custom_cache = tmp_path / "caller_cache"
custom_cache.mkdir()
snap = _write_gemma4_cache(custom_cache)
(snap / "tokenizer.model").write_bytes(b"sp-model")
# The environment points somewhere else entirely.
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "unrelated"))
_TOKENIZER_MODEL_CACHE.clear()
@L._offline_aware_load
def _load(**kwargs):
return SimpleNamespace(name_or_path = _REPO)
tok = _load(local_files_only = True, cache_dir = str(custom_cache))
assert L._tokenizer_cache_dir(tok) == str(custom_cache)
with patch("huggingface_hub.HfApi.model_info") as model_info:
assert _has_tokenizer_model(tok, token = None) is True
assert model_info.call_count == 0

354
tests/studio/install/test_rocm_support.py
View file

@ -826,8 +826,10 @@ class TestEnsureRocmTorch:
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""An explicit gfx wheel-index pin is authoritative: install from it verbatim
with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4
would otherwise pick the rocm6.4 wheel / trigger the Strix re-route)."""
with torch 2.11, and the pin must not be second-guessed (host ROCm 6.4 would
otherwise pick the rocm6.4 wheel / trigger the Strix re-route). The gfx probe
may run for the bnb-skip flag, but returning a Strix arch must not reroute the
pinned torch index."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall
@ -836,10 +838,7 @@ class TestEnsureRocmTorch:
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
# Would raise if the Strix block ran (it is skipped on an explicit pin).
with patch.object(
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
):
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151"]):
_ensure_rocm_torch()
assert mock_pip.call_count == 1
torch_call = str(mock_pip.call_args_list[0])
@ -931,7 +930,8 @@ class TestEnsureRocmTorch:
def test_gfx_pin_over_installed_pre211_rocm_reinstalls(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""A gfx* pin (2.11 line) over an installed pre-2.11 +rocm6.4 build reinstalls."""
"""A gfx* pin (2.11 line) over an installed pre-2.11 +rocm6.4 build reinstalls.
The gfx probe may run for the bnb-skip flag but must not alter the pinned index."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
@ -940,9 +940,7 @@ class TestEnsureRocmTorch:
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
with patch.object(
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
):
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151"]):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1151" in torch_call
@ -1027,9 +1025,9 @@ class TestEnsureRocmTorch:
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
with patch.object(
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
):
# The gfx probe may run for the bnb-skip flag; returning a Strix
# arch must not reroute the pinned torch index.
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151"]):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1151" in torch_call
@ -1158,6 +1156,336 @@ class TestEnsureRocmTorch:
mock_pip.assert_not_called()
# TEST: gfx906 (MI50 / Radeon VII) legacy reroute -- generic wheels after rocm6.3
# lack gfx906 code objects, so torch must come from the rocm6.3 index.
class TestGfx906LegacyReroute:
"""gfx906 hosts on ROCm >= 6.4 must be rerouted to the rocm6.3 torch index;
hosts already on gfx906-capable wheels are left alone."""
@staticmethod
def _gfx906_reroute_block(source: str) -> str:
"""The MI50/gfx906 reroute block, bounded on the ';;' that closes its
rocm[0-9]* case arm -- robust to comment growth (no magic char offset)."""
start = source.find("MI50 / Radeon VII (gfx906")
assert start >= 0, "gfx906 reroute block not found in install.sh"
end = source.find("\n ;;", start)
assert end >= 0, "end of gfx906 case arm not found"
return source[start:end]
def test_gfx906_needs_legacy_index_floor(self):
f = stack_mod._gfx906_needs_legacy_index
# rocm6.0-6.3 tags still ship gfx906 kernels: no reroute.
assert f((6, 3)) is False
assert f((6, 0)) is False
assert f((5, 0)) is False # below any known tag
# Anything that picks a tag newer than rocm6.3 must reroute.
assert f((6, 4)) is True
assert f((7, 2)) is True
assert f((7, 14)) is True
def test_runtime_target_is_gfx906_selection(self, monkeypatch):
"""Env override wins; else gfx906 only when it is the SOLE distinct arch."""
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
# Sole gfx906 (one or several identical MI50s de-dup to {'gfx906'}).
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]):
assert stack_mod._runtime_target_is_gfx906() is True
# Mixed host: gfx906 is NOT the sole arch -> not auto-selected (Codex #3:
# de-dup loses ordinals, so never downgrade a non-gfx906 selection).
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906", "gfx1100"]):
assert stack_mod._runtime_target_is_gfx906() is False
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = []):
assert stack_mod._runtime_target_is_gfx906() is False
# Explicit override wins even when probes see nothing (Codex #2).
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906")
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = []):
assert stack_mod._runtime_target_is_gfx906() is True
# ...and a non-gfx906 override is honored on a gfx906-present host.
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx1100")
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]):
assert stack_mod._runtime_target_is_gfx906() is False
# A copied HIP gcnArchName (gfx906:sramecc-:xnack-) normalizes to gfx906
# (Codex #4: the feature-flag suffix must not defeat the exact comparison).
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906:sramecc-:xnack-")
with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = []):
assert stack_mod._runtime_target_is_gfx906() is True
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"])
def test_gfx906_on_rocm72_routes_to_rocm63(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""CPU torch on a ROCm 7.2 MI50 host installs from rocm6.3, not rocm7.2."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "rocm6.3" in torch_call
assert "rocm7.2" not in torch_call
# The _default (<2.11) window: the rocm7.2 2.11 floor cannot be satisfied
# on the rocm6.3 index (torch <= 2.9.x there).
assert "torch>=2.4,<2.11.0" in torch_call
# gfx906 has no prebuilt bnb -- the generic wheel must not be installed.
assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"])
def test_gfx906_repairs_existing_rocm72_torch(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""An installed +rocm7.2 torch IS the broken combo: reinstall from rocm6.3
even though has_hip_torch is True."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "rocm6.3" in torch_call
assert "torch>=2.4,<2.11.0" in torch_call
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"])
def test_gfx906_already_on_rocm63_left_alone(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""torch already on rocm6.3 wheels must not be reinstalled (no update loop),
and the generic bnb wheel must not clobber a source build."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"6.3.42131|2.7.0+rocm6.3\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
mock_pip.assert_not_called()
# gfx906: prebuilt bnb is skipped entirely (no torch reinstall, no bnb).
mock_pip_try.assert_not_called()
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1100", "gfx906"])
def test_mixed_host_gfx906_not_sole_arch_skips_reroute(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""Mixed host (gfx906 + dGPU) with no explicit override: gfx906 is not the
sole arch, so the generic index is kept (Codex #3: never downgrade a
de-dup-ambiguous mixed host)."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "rocm7.2" in torch_call
assert "rocm6.3" not in torch_call
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = [])
def test_gfx906_env_override_forces_reroute(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""UNSLOTH_ROCM_GFX_ARCH=gfx906 reroutes even when probes emit no gfx token
(Codex #2: runtime-only ROCm hosts where rocminfo/amd-smi are absent)."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906")
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "rocm6.3" in torch_call
assert "rocm7.2" not in torch_call
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"])
def test_gfx906_bnb_skipped_even_when_index_pinned(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""A gfx906 user who pins the ROCm index AND sets the arch override still
skips the generic bnb wheel: the pin suppresses the torch reroute, not the
gfx906 runtime flag used for the bnb skip."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906")
monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", "https://download.pytorch.org/whl/rocm6.3")
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall from the pinned index
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
# torch is (re)installed from the pinned rocm6.3 index...
assert any("rocm6.3" in str(c) for c in mock_pip.call_args_list)
# ...but the prebuilt bnb wheel is never installed.
assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151", "gfx906"])
def test_gfx906_override_wins_over_strix_probe(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""Mixed Strix + MI50 host: UNSLOTH_ROCM_GFX_ARCH=gfx906 suppresses the Strix
override (which probe order would otherwise pick) and routes to rocm6.3."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906")
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "rocm6.3" in torch_call
assert "gfx1151" not in torch_call
# gfx906 target -> generic bnb wheel skipped.
assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
@patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"])
def test_gfx906_bnb_skipped_on_pinned_index_without_env_override(
self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch
):
"""Codex #2: a real gfx906 host that pins the ROCm index but does NOT set
UNSLOTH_ROCM_GFX_ARCH must still skip the prebuilt bnb wheel -- the pin
suppresses only the torch reroute, not the probe-driven gfx906 detection
used for the bnb skip (otherwise `studio update` clobbers source-built bnb)."""
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", "https://download.pytorch.org/whl/rocm6.3")
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall from the pinned index
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
# torch is (re)installed from the pinned rocm6.3 index...
assert any("rocm6.3" in str(c) for c in mock_pip.call_args_list)
# ...but the prebuilt bnb wheel is never installed (probe saw sole gfx906).
assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list)
def test_install_sh_gfx906_env_suppresses_strix(self):
"""install.sh must skip the Strix reroute when UNSLOTH_ROCM_GFX_ARCH=gfx906."""
source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
assert 'if [ "$_gfx906_env" != "gfx906" ]; then' in source
def test_install_sh_gfx906_normalizes_override_and_clears_radeon(self):
"""install.sh must (Codex #4) strip a gfx906:… feature suffix before the exact
comparison, and (Codex #3) clear the Radeon marketing flag for every gfx906
target -- not only when the >=6.4 reroute fires -- so a Radeon VII already on
rocm6.3 does not divert to the repo.radeon.com branch."""
source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
# Override normalization (both the reroute block and the bnb-skip helper):
# strip the gfx906:… feature suffix and trim whitespace (mirror py .strip()).
assert "_gfx906_env=${_gfx906_env%%:*}" in source
assert "_bnb_gfx_env=${_bnb_gfx_env%%:*}" in source
assert source.count("tr -d '[:space:]'") >= 2
# Radeon flag cleared as soon as gfx906 is the target, before the leaf gate.
block = self._gfx906_reroute_block(source)
clear_pos = block.find("_amd_gpu_radeon=false")
leaf_gate_pos = block.find("_rocm_leaf_below")
assert clear_pos >= 0 and leaf_gate_pos >= 0
# the unconditional clear must precede the >=6.4 leaf-gated reroute.
assert clear_pos < leaf_gate_pos
def test_install_sh_bnb_skip_probes_under_pin(self):
"""install.sh _is_gfx906_bnb_skip must probe gfx906 when the index is pinned
(Codex #1): a pin skips the reroute block that sets _gfx906_target, so the
helper falls back to _probe_amd_gfx_arch to catch a real gfx906 host."""
source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
start = source.find("_is_gfx906_bnb_skip() {")
assert start >= 0
body = source[start : start + 900]
assert "_torch_index_pinned" in body
assert "_probe_amd_gfx_arch" in body
def test_install_sh_has_gfx906_reroute(self):
"""install.sh must mirror the Python reroute: honor UNSLOTH_ROCM_GFX_ARCH,
gate on a gfx906 target, route to rocm6.3, with the same _default (<2.11)
trio, and skip the prebuilt bnb wheel."""
source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
block = self._gfx906_reroute_block(source)
assert "_gfx906_target=" in block
assert "UNSLOTH_ROCM_GFX_ARCH" in block
assert "/rocm6.3" in block
for spec in stack_mod._ROCM_TORCH_PKG_SPECS["_default"]:
assert spec in block
# The bnb skip helper must exist and be wired at the install sites.
assert "_is_gfx906_bnb_skip" in source
def test_device_type_defaults_compile_off_on_gfx906(self):
"""unsloth/device_type.py must default Dynamo/compile off on gfx906
(user-overridable via setdefault)."""
source = (PACKAGE_ROOT / "unsloth" / "device_type.py").read_text(encoding = "utf-8")
gate_start = source.find("gfx906")
assert gate_start >= 0
gate_body = source[gate_start : gate_start + 800]
assert 'setdefault("TORCHDYNAMO_DISABLE", "1")' in gate_body
assert 'setdefault("TORCH_COMPILE_DISABLE", "1")' in gate_body
assert 'setdefault("UNSLOTH_COMPILE_DISABLE", "1")' in gate_body
# TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692)

94
tests/studio/install/test_selection_logic.py
View file

@ -2935,6 +2935,100 @@ class TestPublishedRocmGfxSelection:
assert choice.name == "app-b9457-windows-x64-rocm-gfx120X.zip"
class TestPublishedRocmBundleCoverage:
"""Every arch the installer routes torch for should also have a llama.cpp
bundle, or be recorded here as a known gap. Routing an arch for torch while
no bundle covers it is silent: the GPU works for training and drops to a HIP
source build for inference, which is correct but much slower to install."""
# mapped_targets of each published ROCm bundle, mirroring
# unslothai/llama.cpp's llama-prebuilt-manifest.json.
PUBLISHED = {
"gfx103X": ["gfx1030", "gfx1031", "gfx1032", "gfx1034"],
"gfx110X": ["gfx1100", "gfx1101", "gfx1102", "gfx1103"],
"gfx120X": ["gfx1200", "gfx1201"],
"gfx1150": ["gfx1150"],
"gfx1151": ["gfx1151"],
"gfx908": ["gfx908"],
"gfx90a": ["gfx90a"],
}
# Arches _GFX_TO_AMD_INDEX_ARCH routes torch for that no bundle covers.
# gfx1033/1035/1036: RDNA 2 variants, never built.
# gfx1152: Krackan Point (Radeon 860M/840M). Torch goes to its own
# repo.amd.com/rocm/whl/gfx1152 leaf, but no llama.cpp bundle exists, so
# these hosts source-build. Publish a -gfx1152 bundle, or add gfx1152 to
# the gfx1150 bundle's mapped_targets if that build genuinely covers it,
# then drop it from this set.
KNOWN_GAPS = {"gfx1033", "gfx1035", "gfx1036", "gfx1152"}
def _release(self):
return make_release(
[
make_artifact(
f"app-b9457-linux-x64-rocm-{fam}.tar.gz",
install_kind = "linux-rocm",
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = None,
rank = 1000,
gfx_target = fam,
mapped_targets = targets,
)
for fam, targets in self.PUBLISHED.items()
],
upstream_tag = "b9457",
)
def _host(self, gfx):
return make_host(
machine = "x86_64",
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = True,
rocm_gfx_target = gfx,
)
def test_known_gaps_fall_back_to_source_build(self):
"""A gap arch must return None rather than be served a sibling bundle:
a wrong-ISA binary fails at the first BLAS call instead of installing
slowly, which is the worse of the two outcomes."""
release = self._release()
for gfx in sorted(self.KNOWN_GAPS):
assert (
INSTALL_LLAMA_PREBUILT.published_rocm_choice_for_host(
release, self._host(gfx), "linux-rocm"
)
is None
), f"{gfx} is in KNOWN_GAPS but a bundle now matches it; drop it from the set"
def test_every_torch_routed_arch_is_covered_or_a_known_gap(self):
"""The guard that would have caught gfx1152: adding an arch to
_GFX_TO_AMD_INDEX_ARCH without a bundle must be a deliberate entry in
KNOWN_GAPS, not an unnoticed drop to source builds."""
import re
# Read the table from source rather than importing the installer module,
# which pulls in a heavy dependency chain this suite does not need.
stack = (PACKAGE_ROOT / "studio" / "install_python_stack.py").read_text(encoding = "utf-8")
body = re.search(r"_GFX_TO_AMD_INDEX_ARCH.*?=\s*\{(.*?)\n\}", stack, re.S)
assert body, "_GFX_TO_AMD_INDEX_ARCH not found in install_python_stack.py"
routed = set(re.findall(r'"(gfx[0-9a-z]+)":', body.group(1)))
assert routed, "parsed no arches out of _GFX_TO_AMD_INDEX_ARCH"
covered = {t.lower() for targets in self.PUBLISHED.values() for t in targets}
uncovered = {a for a in routed if a.lower() not in covered}
assert uncovered == self.KNOWN_GAPS, (
f"llama.cpp bundle coverage drifted: {sorted(uncovered - self.KNOWN_GAPS)} "
f"newly uncovered, {sorted(self.KNOWN_GAPS - uncovered)} no longer a gap"
)
class TestPublishedMacosForkSelection:
"""macOS routes to the fork's llama-<tag>-bin-macos-<arch>.tar.gz, selected by install_kind."""

21
tests/studio/test_model_picker_contracts.py
View file

@ -578,3 +578,24 @@ def test_legacy_migration_is_idempotent_and_non_destructive():
# Layer 3: non-overwriting merge skips an existing (or default) key, so even a
# forced re-run cannot duplicate or clobber a user's config.
assert "if (isDefaultConfig(migrated) || Object.hasOwn(map, key)) {" in src
def test_vulkan_inference_devices_are_the_pickable_set():
"""GGUF loads run through llama-server, so on a Vulkan build the picker must
offer the inference inventory (ggml ordinals, the space `--device Vulkan<i>`
pins) rather than the torch view, which can miss cards llama-server drives.
The XPU ban must not apply there: it is about torch-xpu ordinals no
applicator speaks, and a Vulkan pick does not use them.
"""
src = " ".join(_read("hooks/use-gpu-info.ts").split())
# The Vulkan inventory is consulted first, and only when it has devices.
assert (
"const inference = data?.inference_gpu; "
'if (inference?.backend === "vulkan" && (inference.devices ?? []).length) {' in src
)
# Pinnable on the ggml ordinal space, gated on the backend's own support flag.
assert "const picksAccepted = inference.gguf_gpu_ids_supported !== false;" in src
assert 'physicalIndex: picksAccepted && d.index_kind === "vulkan",' in src
# The torch fallback keeps its physical-only gate and the XPU ban.
assert 'data?.device_backend !== "xpu" &&' in src
assert 'physicalIndex: pinnableBackend && d.index_kind === "physical",' in src

200
tests/studio/test_usage_examples_model_source_contract.py Normal file
View file

@ -0,0 +1,200 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Static contract for which model the API usage examples name, and for the
model-auto-switch control living in exactly one place on the API keys tab."""
import re
from pathlib import Path
REPO = Path(__file__).resolve().parents[2]
SETTINGS = REPO / "studio/frontend/src/features/settings"
USAGE_EXAMPLES_TSX = SETTINGS / "components/usage-examples.tsx"
OPENAI_MODELS_TS = SETTINGS / "api/openai-models.ts"
API_KEYS_TAB_TSX = SETTINGS / "tabs/api-keys-tab.tsx"
def test_examples_name_a_model_the_server_can_serve():
# A hardcoded repo id made copied curls 404; read the servable ids from /v1/models.
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
assert 'from "../api/openai-models"' in src
assert "function useExampleModelName(): string" in src
hook = src[src.find("function useExampleModelName") : src.find("// Backend PATH detection")]
assert "listOpenAIModels()" in hook
# Precedence: live checkpoint, then a loaded entry, then any entry if switching is on.
assert "catalog?.find((m) => m.loaded) ?? (autoSwitch ? catalog?.[0] : undefined)" in hook
# The snippet pins the quant so the request names the file on disk.
assert "`${pick.id}:${pick.quant}`" in hook
api = OPENAI_MODELS_TS.read_text(encoding = "utf-8")
assert 'authFetch("/v1/models")' in api
def test_examples_never_print_a_hardcoded_model_id():
# The bug this exists for: a `[]` catalog printed a snippet before /v1/models answered.
# It is tri-state now, and the panel asks for a model instead.
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
assert "MODEL_FALLBACK" not in src
# No repo-shaped literal anywhere: a snippet may only name what /v1 returns.
assert re.search(r'"unsloth/[^"]+"', src) is None
assert "function useExampleModelName(): string | null" in src
assert "useState<OpenAIModel[] | null>(null)" in src
# Nothing servable means nothing is built, so there is nothing to copy.
assert "(model ? buildSnippets(base, key, model, os) : null)" in src
assert "if (!snippets) return;" in src
assert "{snippets ? (" in src
assert 't("settings.apiKeys.usageNoModel")' in src
en = EN_TS.read_text(encoding = "utf-8")
assert "usageNoModel:" in en
def test_catalog_refresh_follows_the_loaded_model():
# A dep list missing these never re-ran, so a finished load left the first fetch's
# name. Nor may it be gated on having no checkpoint: the store keeps one across an
# idle unload, which changes nothing React can see.
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
hook = src[src.find("function useExampleModelName") : src.find("// Backend PATH detection")]
assert "}, [checkpoint, ggufVariant]);" in hook
assert "needsCatalog" not in hook
# A finishing download moves no store state, so the fetch retries on a timer too,
# and residency only slows that timer rather than stopping it.
assert "CATALOG_RETRY_MS" in hook and "CATALOG_IDLE_MS" in hook
assert "window.clearTimeout(timeoutId)" in hook
assert "const CATALOG_RETRY_MS = 15000;" in src
assert "const CATALOG_IDLE_MS = 60000;" in src
def test_a_stored_checkpoint_needs_catalog_evidence():
# The store keeps a checkpoint across an idle unload and across a deletion, so
# preferring it on the switch setting alone named a model /v1/models had proved
# absent, and the snippets 404d instead of falling back.
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
hook = src[src.find("function useExampleModelName") : src.find("// Backend PATH detection")]
assert 'const entry = catalog?.find((m) => sameBaseModelId(m.id, checkpoint ?? ""));' in hook
# Resident, or downloaded with something able to reload it. Never the setting alone.
assert "(!!entry && (entry.loaded || autoSwitch || idleReload))" in hook
assert "autoSwitch ||\n" not in hook
def test_standalone_idle_unload_still_names_the_stored_checkpoint():
# UNSLOTH_MODEL_IDLE_TTL without auto-switch reloads exactly what it freed, so the
# stored checkpoint stays runnable and the panel must keep showing it. The stash
# restores only that model, so it can never pick catalog[0].
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
hook = src[src.find("function useExampleModelName") : src.find("// Backend PATH detection")]
assert "const [idleReload, setIdleReload] = useState(false);" in hook
assert "setIdleReload(settings[1])" in hook
assert "s.idleUnloadActive" in hook
# fromCatalog stays gated on auto-switch alone.
assert "?? (autoSwitch ? catalog?.[0] : undefined)" in hook
assert "idleReload ? catalog" not in hook
def test_a_failed_refresh_does_not_erase_what_the_server_holds():
# Catching into [] and false made a transient error authoritative: the panel dropped
# a still-servable model and printed "No model". The catalog is deliberately
# tri-state, and a failure must stay the unknown state.
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
hook = src[src.find("function useExampleModelName") : src.find("// Backend PATH detection")]
assert "listOpenAIModels().catch(() => null)" in hook
assert ".catch(() => null)," in hook
assert "if (models !== null) setCatalog(models);" in hook
assert "if (settings !== null) {" in hook
# The old negatives must be gone entirely.
assert "catch(() => [] as OpenAIModel[])" not in hook
assert "catch(() => [false, false] as const)" not in hook
assert "catch(() => false)" not in hook
def test_the_pinned_quant_comes_from_the_catalog():
# Catalog membership proves the repo, not the saved quant: the stored one can name
# a file deleted while another quant remains, so pinning it 404d on a missing quant
# with a runnable one listed.
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
hook = src[src.find("function useExampleModelName") : src.find("// Backend PATH detection")]
assert "const quant = catalog === null ? ggufVariant : entry?.quant;" in hook
assert "`${checkpoint}:${ggufVariant}`" not in hook
def test_usage_examples_has_no_duplicate_auto_switch_control():
# ModelAutoSwitchSection renders this setting just below and shares no state with it.
src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8")
# Reading the setting is fine; writing it here is what would be a second control.
assert "updateOpenAIAutoSwitchSettings" not in src
assert "SWITCH_NOTE" not in src
assert "Switch model by request" not in src
assert "pythonSwitchDemo" not in src
assert "javascriptSwitchDemo" not in src
assert "modelAutoSwitch" not in src
tab = API_KEYS_TAB_TSX.read_text(encoding = "utf-8")
assert "<ModelAutoSwitchSection />" in tab
API_MONITOR_TSX = SETTINGS / "components/api-monitor-console.tsx"
def test_api_monitor_pages_five_at_a_time():
# The backend retains 50 terminal entries; the console used to dump them all at once.
src = API_MONITOR_TSX.read_text(encoding = "utf-8")
assert "const PAGE_SIZE = 5;" in src
assert "ordered.slice(" in src
# Paging back must freeze the id order, or live traffic reorders history under it.
assert "frozenIds" in src
assert "setFrozenIds((prev) => prev ?? entries.map((entry) => entry.id))" in src
def test_api_monitor_renders_lifecycle_rows():
src = API_MONITOR_TSX.read_text(encoding = "utf-8")
assert "function LifecycleEntry(" in src
assert 'entry.kind === "lifecycle"' in src
for label in ("Loading model", "Model loaded", "Model unloaded"):
assert label in src
# Lifecycle rows have no prompt/reply to fetch.
assert "isLifecycle(entry) || !expandedIds.has(entry.id)" in src
def test_auto_switch_section_sits_above_the_monitor():
tab = API_KEYS_TAB_TSX.read_text(encoding = "utf-8")
assert tab.index("<ModelAutoSwitchSection />") < tab.index("<ApiMonitorConsole />")
assert tab.index("<ApiMonitorConsole />") < tab.index("<UsageExamples")
AUTO_SWITCH_TSX = SETTINGS / "components/model-auto-switch-section.tsx"
EN_TS = REPO / "studio/frontend/src/i18n/locales/en.ts"
def test_api_monitor_renders_download_rows():
src = API_MONITOR_TSX.read_text(encoding = "utf-8")
assert 'entry.event === "download"' in src
for label in ("Downloading model", "Model downloaded", "Model download failed"):
assert label in src
def test_monitor_can_unload_the_loaded_model():
src = API_MONITOR_TSX.read_text(encoding = "utf-8")
assert "unloadActiveModel" in src
# Always rendered so the manual release stays discoverable; disabled, not hidden.
assert "disabled={unloading || !data?.active_model}" in src
assert "{data?.active_model ? (" not in src
# /unload matches on the internal id, omitted here (a host path), so read it from status.
assert "resolveInferenceCheckpointId(status)" in src
assert "unloadModel({ model_path: checkpoint })" in src
def test_auto_download_toggle_is_gated_on_auto_switch():
# Downloading what auto-switch cannot load fetches gigabytes nothing can serve.
src = AUTO_SWITCH_TSX.read_text(encoding = "utf-8")
assert "modelAutoSwitch.autoDownload" in src
assert "settings?.autoDownloadModel ?? false" in src
row = src[src.find("modelAutoSwitch.autoDownload") :]
assert "disabled={!settings?.enabled || isSaving}" in row[: row.find("</SettingsRow>")]
def test_auto_download_copy_warns_about_api_key_holders():
en = EN_TS.read_text(encoding = "utf-8")
start = en.find("autoDownloadDescription:")
assert start != -1
description = en[start : en.find("\n", en.find('",', start))]
assert "API key" in description

2
tests/test_bad_mappings_redirect.py
View file

@ -26,7 +26,7 @@ def _load_get_model_name():
namespace = dict(mapper_ns)
namespace["SUPPORTS_FOURBIT"] = True
namespace["_env_says_offline"] = lambda: True
namespace["_get_new_mapper"] = lambda: ({}, {}, {})
namespace["_get_new_mapper"] = lambda: ({}, {}, {}, {}, {})
wanted = {"__get_model_name", "_resolve_with_mappers", "get_model_name"}
for node in tree.body:

3
tests/test_get_model_name.py
View file

@ -6,7 +6,8 @@ from unsloth.models.mapper import FLOAT_TO_INT_MAPPER, MAP_TO_UNSLOTH_16bit
def _no_remote_mapper():
return {}, {}, {}
# int_to_float, float_to_int, map_to_16bit, fp8_block, fp8_row
return {}, {}, {}, {}, {}
class TestGetModelName(unittest.TestCase):

155
tests/test_new_mapper_fetched_fp8.py Normal file
View file

@ -0,0 +1,155 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Regression tests for what ``_get_new_mapper`` hands back to the upgrade probe.
``test_new_mapper_no_global_leak.py`` serves the repo's own ``mapper.py`` as both installed
and fetched source, so it cannot tell a fetched table from a fresh copy of the installed one.
Two gaps it misses:
1. The probe must answer for an fp8 repo only the FETCHED mapper knows, so an extra ``"8"``
entry is spliced into the fetched source only. Isolating the exec without returning the
fetched fp8 tables would silently drop the fp8 half of the upgrade check.
2. The probe must survive a fetched ``mapper.py`` with no fp8 tables (anything older, or a
future rename): reading them with ``[]`` raises ``KeyError`` into the bare ``except``,
taking the 4bit half, the probe's whole purpose, down with it.
``loader_utils`` imports torch, so ast-extract the resolvers and run them against a stubbed
``requests``, as in ``tests/test_bad_mappings_redirect.py``.
"""
import ast
import os
import sys
import types
_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models")
_WANTED = {"__get_model_name", "_resolve_with_mappers", "_get_new_mapper", "get_model_name"}
# An fp8 ("8") model, spliced into the FETCHED mapper only.
_NEW_KEY = "unsloth/Zeta-9B-Only-On-Main"
_NEW_OFFICIAL = "zeta-org/Zeta-9B-Only-On-Main-FP8"
_NEW_BLOCK = "unsloth/Zeta-9B-Only-On-Main-FP8-Block"
_NEW_ROW = "unsloth/Zeta-9B-Only-On-Main-FP8-Row"
_ANCHOR = ' "unsloth/Kimi-K2-Instruct-BF16" : ('
def _mapper_source():
with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f:
return f.read()
def _with_extra_fp8_model(source):
assert _ANCHOR in source, "anchor moved; update this test"
entry = (
f' "{_NEW_KEY}" : {{\n'
f' "16" : ("{_NEW_KEY}", "zeta-org/Zeta-9B-Only-On-Main"),\n'
f' "8" : ("{_NEW_OFFICIAL}", "{_NEW_BLOCK}", "{_NEW_ROW}"),\n'
f" }},\n"
)
return source.replace(_ANCHOR, entry + _ANCHOR, 1)
def _without_fp8_tables(source):
"""A mapper.py from before the fp8 tables existed."""
return source.replace("FLOAT_TO_FP8_BLOCK_MAPPER", "SOME_OTHER_BLOCK_TABLE").replace(
"FLOAT_TO_FP8_ROW_MAPPER", "SOME_OTHER_ROW_TABLE"
)
class _FakeResponse:
def __init__(self, text):
self.text = text
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def _install_fake_requests(monkeypatch, text):
module = types.ModuleType("requests")
module.get = lambda url, timeout = None: _FakeResponse(text)
monkeypatch.setitem(sys.modules, "requests", module)
def _install_fake_vllm_absent(monkeypatch, namespace):
"""vllm >= 0.12.0 returns early from __get_model_name, leaving the probe unreachable."""
monkeypatch.delitem(sys.modules, "vllm", raising = False)
fake = types.ModuleType("importlib")
fake.util = types.SimpleNamespace(find_spec = lambda name: None)
namespace["importlib"] = fake
def _load_resolver(installed_source):
"""Stand-in for loader_utils' module globals, built from `installed_source`."""
from unsloth_zoo.utils import Version
mapper_ns = {}
exec(compile(installed_source, "mapper.py", "exec"), mapper_ns)
namespace = {
"INT_TO_FLOAT_MAPPER": mapper_ns["INT_TO_FLOAT_MAPPER"],
"FLOAT_TO_INT_MAPPER": mapper_ns["FLOAT_TO_INT_MAPPER"],
"MAP_TO_UNSLOTH_16bit": mapper_ns["MAP_TO_UNSLOTH_16bit"],
"FLOAT_TO_FP8_BLOCK_MAPPER": mapper_ns["FLOAT_TO_FP8_BLOCK_MAPPER"],
"FLOAT_TO_FP8_ROW_MAPPER": mapper_ns["FLOAT_TO_FP8_ROW_MAPPER"],
"SUPPORTS_FOURBIT": True,
"transformers_version": Version("4.57.6"),
"Version": Version,
"os": os,
}
with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f:
tree = ast.parse(f.read())
for node in tree.body:
if isinstance(node, ast.Assign) and any(
getattr(t, "id", None) in ("BAD_MAPPINGS", "_OFFLINE_ENV_VALUES", "_OFFLINE_ENV_KEYS")
for t in node.targets
):
exec(compile(ast.Module([node], []), "<assign>", "exec"), namespace)
elif isinstance(node, ast.FunctionDef) and (
node.name in _WANTED or node.name == "_env_says_offline"
):
exec(compile(ast.Module([node], []), node.name, "exec"), namespace)
return namespace
def test_probe_answers_for_an_fp8_repo_only_the_fetched_mapper_knows(monkeypatch):
installed = _mapper_source()
namespace = _load_resolver(installed)
installed_block = namespace["FLOAT_TO_FP8_BLOCK_MAPPER"]
installed_row = namespace["FLOAT_TO_FP8_ROW_MAPPER"]
assert _NEW_OFFICIAL.lower() not in installed_block, "the installed table must not know it"
_install_fake_requests(monkeypatch, _with_extra_fp8_model(installed))
_install_fake_vllm_absent(monkeypatch, namespace)
try:
resolved = namespace["get_model_name"](
_NEW_OFFICIAL, load_in_4bit = False, load_in_fp8 = "block"
)
except NotImplementedError as error:
assert "not supported in your current Unsloth version" in str(error)
else:
raise AssertionError(
f"a fetched-only fp8 repo must raise the upgrade error, got {resolved!r}"
)
# Answering must not have adopted the fetched tables.
assert namespace["FLOAT_TO_FP8_BLOCK_MAPPER"] is installed_block
assert namespace["FLOAT_TO_FP8_ROW_MAPPER"] is installed_row
assert _NEW_OFFICIAL.lower() not in namespace["FLOAT_TO_FP8_BLOCK_MAPPER"]
def test_probe_survives_a_fetched_mapper_without_the_fp8_tables(monkeypatch):
installed = _mapper_source()
namespace = _load_resolver(installed)
_install_fake_requests(monkeypatch, _without_fp8_tables(installed))
int_to_float, float_to_int, map_to_16bit = namespace["_get_new_mapper"]()[:3]
assert (
int_to_float and float_to_int and map_to_16bit
), "a fetched mapper.py without the fp8 tables must not take the 4bit upgrade check down"

99
tests/test_new_mapper_no_global_leak.py Normal file
View file

@ -0,0 +1,99 @@
"""Regression test for ``_get_new_mapper`` leaking into ``loader_utils`` globals.
``get_model_name`` calls ``_get_new_mapper()`` whenever a name misses the local
tables, purely to answer "would a newer Unsloth support this?". It fetches
``mapper.py`` from GitHub main, prefixes the three mappers it wants with
``NEW_``, and ``exec``s the result into ``globals()``.
The slice starts at ``__INT_TO_FLOAT_MAPPER``, so it also carries
``FLOAT_TO_FP8_BLOCK_MAPPER``/``FLOAT_TO_FP8_ROW_MAPPER`` and the two
``_add_*`` helpers, and those names are NOT renamed. Exec'ing into
``globals()`` therefore rebinds the FP8 tables that ``loader_utils`` imported
from the installed ``mapper``, so every later ``get_model_name(...,
load_in_fp8 = ...)`` in the process resolves through GitHub main's table
instead of the installed one. The probe is supposed to read, not to swap the
installed mappings out from under the caller.
``loader_utils`` imports torch, so ast-extract ``_get_new_mapper`` and run it
against a stubbed ``requests`` rather than importing unsloth (which needs a GPU).
"""
import ast
import os
import sys
import types
_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models")
def _mapper_source():
with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f:
return f.read()
def _extract_get_new_mapper(namespace):
with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f:
tree = ast.parse(f.read())
for node in tree.body:
if isinstance(node, ast.FunctionDef) and node.name == "_get_new_mapper":
exec(compile(ast.Module([node], []), node.name, "exec"), namespace)
return namespace["_get_new_mapper"]
raise AssertionError("_get_new_mapper not found in loader_utils.py")
class _FakeResponse:
def __init__(self, text):
self.text = text
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def _install_fake_requests(monkeypatch, text):
module = types.ModuleType("requests")
module.get = lambda url, timeout = None: _FakeResponse(text)
monkeypatch.setitem(sys.modules, "requests", module)
def test_get_new_mapper_does_not_rebind_the_installed_fp8_tables(monkeypatch):
_install_fake_requests(monkeypatch, _mapper_source())
installed = {}
exec(compile(_mapper_source(), "mapper.py", "exec"), installed)
block = installed["FLOAT_TO_FP8_BLOCK_MAPPER"]
row = installed["FLOAT_TO_FP8_ROW_MAPPER"]
assert block and row, "the installed FP8 tables should not be empty"
# Stand in for loader_utils' module globals, which import the FP8 tables.
namespace = {"FLOAT_TO_FP8_BLOCK_MAPPER": block, "FLOAT_TO_FP8_ROW_MAPPER": row}
get_new_mapper = _extract_get_new_mapper(namespace)
int_to_float, float_to_int, map_to_16bit, fp8_block, fp8_row = get_new_mapper()
# _get_new_mapper swallows every exception and returns empty dicts, so assert
# it actually ran before trusting anything below.
assert int_to_float and float_to_int and map_to_16bit, "the fetch/exec path did not run"
# the probe has to hand the FETCHED fp8 tables back, or a newly added fp8 repo would
# miss both the installed tables and the probe and skip the upgrade message
assert fp8_block and fp8_row
assert fp8_block is not block and fp8_row is not row
assert namespace["FLOAT_TO_FP8_BLOCK_MAPPER"] is block
assert namespace["FLOAT_TO_FP8_ROW_MAPPER"] is row
def test_get_new_mapper_leaves_no_helpers_behind(monkeypatch):
_install_fake_requests(monkeypatch, _mapper_source())
namespace = {}
get_new_mapper = _extract_get_new_mapper(namespace)
before = set(namespace)
assert all(get_new_mapper()), "the fetch/exec path did not run"
leaked = set(namespace) - before
assert not leaked, f"_get_new_mapper leaked {sorted(leaked)} into its module globals"

17
unsloth/device_type.py
View file

@ -117,6 +117,23 @@ DEVICE_COUNT: int = get_device_count()
ALLOW_PREQUANTIZED_MODELS: bool = True
# HSA_STATUS_ERROR_EXCEPTION checks - sometimes AMD fails for BnB
ALLOW_BITSANDBYTES: bool = True
# gfx906 (MI50 / Radeon VII / Vega 20): Dynamo/Inductor codegen is broken on this
# legacy GCN arch (ROCm dropped it after 6.3) - compiled graphs crash or miscompile
# while the eager path trains fine. Default compile off; setdefault so a user
# override wins.
if DEVICE_TYPE == "hip":
try:
_gcn_arch = torch.cuda.get_device_properties(0).gcnArchName.split(":")[0].strip().lower()
except Exception:
_gcn_arch = ""
if _gcn_arch == "gfx906":
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("UNSLOTH_COMPILE_DISABLE", "1")
print(
"Unsloth: gfx906 (MI50 / Radeon VII) detected - torch.compile disabled "
"(community-maintained legacy GCN path)."
)
if DEVICE_TYPE == "hip":
try:
import bitsandbytes

8
unsloth/import_fixes.py
View file

@ -418,7 +418,13 @@ def fix_vllm_aimv2_issue():
spec = importlib.util.find_spec("vllm")
if spec is None:
return
vllm_version = importlib_version("vllm")
# A findable spec with unreadable dist metadata (broken/partial vllm install)
# must not crash unsloth import; every other vllm probe here guards this too.
try:
vllm_version = importlib_version("vllm")
except Exception as e:
logger.info(f"Unsloth: Skipping vLLM aimv2 fix -- vLLM version unreadable ({e})")
return
if Version(vllm_version) < Version("0.10.1"):
vllm_location = spec.origin
if vllm_location is None:

245
unsloth/models/loader_utils.py
View file

@ -191,19 +191,43 @@ def _get_new_mapper():
.replace("MAP_TO_UNSLOTH_16bit", "NEW_MAP_TO_UNSLOTH_16bit")
)
exec(new_mapper, globals())
# Exec into a throwaway namespace, never globals(). The slice also carries
# FLOAT_TO_FP8_BLOCK_MAPPER / FLOAT_TO_FP8_ROW_MAPPER, the _add_* helpers
# and the builder's loop variables, so exec'ing into globals() would swap
# the FP8 tables this module imported from the installed mapper for the
# ones on GitHub main. This is only a probe for "would a newer Unsloth
# support this name?", so it must not change what the installed version
# resolves; the fetched FP8 tables are returned for the probe to use
# instead of being written over the installed ones.
namespace = {}
exec(new_mapper, namespace)
return (
NEW_INT_TO_FLOAT_MAPPER,
NEW_FLOAT_TO_INT_MAPPER,
NEW_MAP_TO_UNSLOTH_16bit,
namespace["NEW_INT_TO_FLOAT_MAPPER"],
namespace["NEW_FLOAT_TO_INT_MAPPER"],
namespace["NEW_MAP_TO_UNSLOTH_16bit"],
# .get, not []: these two come from the fetched file under its own names (unlike
# the NEW_ names above, renamed here), so an older or renamed mapper.py would
# KeyError into the bare except and take the 4bit half of the probe down too.
# {} is safe: the probe runs only after the installed tables already missed.
namespace.get("FLOAT_TO_FP8_BLOCK_MAPPER", {}),
namespace.get("FLOAT_TO_FP8_ROW_MAPPER", {}),
)
except:
return {}, {}, {}
return {}, {}, {}, {}, {}
def _resolve_with_mappers(
model_name, load_in_4bit, load_in_fp8, int_to_float, float_to_int, map_to_unsloth_16bit
model_name,
load_in_4bit,
load_in_fp8,
int_to_float,
float_to_int,
map_to_unsloth_16bit,
fp8_block = None,
fp8_row = None,
):
# fp8_block/fp8_row default to the installed tables; the newer-mapper probe passes the
# fetched ones so it can answer for new FP8 repos without rebinding the installed ones.
return __get_model_name(
model_name = model_name,
load_in_4bit = load_in_4bit,
@ -211,8 +235,8 @@ def _resolve_with_mappers(
FLOAT_TO_INT_MAPPER = float_to_int,
MAP_TO_UNSLOTH_16bit = map_to_unsloth_16bit,
load_in_fp8 = load_in_fp8,
FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER,
FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER,
FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER if fp8_block is None else fp8_block,
FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER if fp8_row is None else fp8_row,
)
@ -252,9 +276,13 @@ def get_model_name(
and not _env_says_offline() # offline: skip the remote (raw GitHub) mapper refresh
):
# Try checking if a new Unsloth version allows it!
NEW_INT_TO_FLOAT_MAPPER, NEW_FLOAT_TO_INT_MAPPER, NEW_MAP_TO_UNSLOTH_16bit = (
_get_new_mapper()
)
(
NEW_INT_TO_FLOAT_MAPPER,
NEW_FLOAT_TO_INT_MAPPER,
NEW_MAP_TO_UNSLOTH_16bit,
NEW_FP8_BLOCK_MAPPER,
NEW_FP8_ROW_MAPPER,
) = _get_new_mapper()
upgraded_model_name = _resolve_with_mappers(
model_name = model_name,
load_in_4bit = load_in_4bit,
@ -262,6 +290,10 @@ def get_model_name(
int_to_float = NEW_INT_TO_FLOAT_MAPPER,
float_to_int = NEW_FLOAT_TO_INT_MAPPER,
map_to_unsloth_16bit = NEW_MAP_TO_UNSLOTH_16bit,
# the fp8 probe has to look at the FETCHED tables too, or a new fp8 repo would
# miss both here and in the installed tables and skip the upgrade message
fp8_block = NEW_FP8_BLOCK_MAPPER,
fp8_row = NEW_FP8_ROW_MAPPER,
)
if upgraded_model_name is not None:
raise NotImplementedError(
@ -843,6 +875,55 @@ def _get_effective_local_files_only(kwargs):
return _env_says_offline()
# Attribute stamped on a tokenizer/processor that was loaded local-only, so a later
# save still knows. transformers takes local_files_only as an explicit from_pretrained
# parameter and never copies it into tokenizer.init_kwargs, and _offline_aware_load
# restores the offline env vars when the load window closes, so without this stamp an
# explicit local_files_only = True load is invisible by the time we save (issue #7481).
_LOCAL_FILES_ONLY_ATTR = "_unsloth_local_files_only"
# The load's cache_dir travels with it too: saving derives one from HF_HUB_CACHE /
# HF_HOME, which does not see a caller-supplied cache.
_LOADED_CACHE_DIR_ATTR = "_unsloth_loaded_cache_dir"
def _mark_loaded_local_files_only(result, cache_dir = None):
"""Stamp a load's local-only mode and cache_dir onto the returned objects."""
for obj in result if isinstance(result, (tuple, list)) else (result,):
try:
# A processor keeps the tokenizer that _has_tokenizer_model unwraps to,
# so stamp both (a wrapped model can raise from its own __getattr__).
targets = (obj, getattr(obj, "tokenizer", None))
except Exception:
targets = (obj,)
for target in targets:
if target is None:
continue
# Objects that reject new attributes (__slots__) are skipped.
try:
setattr(target, _LOCAL_FILES_ONLY_ATTR, True)
if cache_dir:
setattr(target, _LOADED_CACHE_DIR_ATTR, str(cache_dir))
except Exception:
pass
return result
def _tokenizer_cache_dir(tokenizer):
"""The cache_dir the load used, when it was not the environment's."""
tokenizer = tokenizer.tokenizer if hasattr(tokenizer, "tokenizer") else tokenizer
return getattr(tokenizer, _LOADED_CACHE_DIR_ATTR, None)
def _tokenizer_wants_local_only(tokenizer):
"""True when Hub metadata probes should be skipped for this tokenizer."""
if _env_says_offline():
return True
if getattr(tokenizer, _LOCAL_FILES_ONLY_ATTR, False):
return True
init_kwargs = getattr(tokenizer, "init_kwargs", None) or {}
return bool(init_kwargs.get("local_files_only"))
def _is_offline_related_error(exc):
"""True if exc (or its cause/context chain) is a lost-connection error, not a
missing file. Plain FileNotFoundError propagates; LocalEntryNotFoundError is offline."""
@ -1067,7 +1148,9 @@ def _offline_aware_load(fn):
if _get_effective_local_files_only(kwargs):
kwargs["local_files_only"] = True
with _force_hf_offline():
return fn(*args, **kwargs)
# Stamp inside the window: the env vars are restored on exit, so the
# request has to travel on the objects themselves to reach saving.
return _mark_loaded_local_files_only(fn(*args, **kwargs), kwargs.get("cache_dir"))
_pb_were_disabled = _progress_bars_were_disabled() # restore before any retry
try:
return fn(*args, **kwargs)
@ -1126,6 +1209,144 @@ def _has_local_processor_files(path):
)
def _resolve_hub_repo_local_dir(
repo_id,
*,
token = None,
cache_dir = None,
# Default closed: a "resolve local dir" helper must not download. False here
# means five filenames each retried with backoff before it gives up.
local_files_only = True,
filenames = (
"tokenizer_config.json",
"config.json",
"tokenizer.json",
"preprocessor_config.json",
"processor_config.json",
),
):
"""Return a local snapshot directory for a Hub repo id when files are cached.
On transformers 4.57.2 through 5.5.4, ``PreTrainedTokenizerFast.from_pretrained``
on a repo id can still call ``model_info()`` when ``local_files_only=True`` and
no offline env var is set. Loading from the resolved snapshot dir avoids that
Hub probe. Upstream fixed this in transformers 5.6.0 (huggingface/transformers#43603);
this helper can be removed once the supported floor is past that version.
"""
if not isinstance(repo_id, str) or not repo_id:
return None
if os.path.isdir(repo_id):
return repo_id
if cache_dir is None:
cache_dir = os.environ.get("HF_HUB_CACHE")
from huggingface_hub import hf_hub_download
for filename in filenames:
try:
path = hf_hub_download(
repo_id = repo_id,
filename = filename,
token = token,
cache_dir = cache_dir,
local_files_only = local_files_only,
)
if path and os.path.isfile(path):
return os.path.dirname(path)
except Exception:
continue
return None
def _resolve_hub_repo_cached_file(
repo_id,
filename,
*,
token = None,
cache_dir = None,
local_files_only = True,
):
"""Return a cached file path under a Hub snapshot, or None if absent."""
local_dir = _resolve_hub_repo_local_dir(
repo_id,
token = token,
cache_dir = cache_dir,
local_files_only = local_files_only,
filenames = (filename,),
)
if local_dir is None:
return None
path = os.path.join(local_dir, filename)
return path if os.path.isfile(path) else None
def _hub_repo_or_local_path(
repo_id,
*,
token = None,
cache_dir = None,
local_files_only = False,
filenames = None,
):
"""Prefer a cached snapshot path over a Hub repo id when offline or ``local_files_only``."""
if isinstance(repo_id, str) and os.path.isdir(repo_id):
return repo_id
lfo = bool(local_files_only) or _env_says_offline()
if not lfo:
return repo_id
local_dir = _resolve_hub_repo_local_dir(
repo_id,
token = token,
cache_dir = cache_dir,
local_files_only = True,
filenames = filenames
or (
"tokenizer_config.json",
"config.json",
"tokenizer.json",
"preprocessor_config.json",
"processor_config.json",
),
)
return local_dir if local_dir is not None else repo_id
def _load_pretrained_tokenizer_fast(
tokenizer_name,
*,
padding_side = "left",
token = None,
trust_remote_code = False,
cache_dir = None,
local_files_only = False,
):
"""Load ``PreTrainedTokenizerFast`` without Hub metadata probes when cached/offline.
Needed on transformers 4.57.2-5.5.4; redundant once the floor is past 5.6.0.
"""
from transformers import PreTrainedTokenizerFast
lfo = bool(local_files_only) or _env_says_offline()
load_path = _hub_repo_or_local_path(
tokenizer_name,
token = token,
cache_dir = cache_dir,
local_files_only = lfo,
filenames = (
"tokenizer_config.json",
"tokenizer.json",
"tokenizer.model",
),
)
return PreTrainedTokenizerFast.from_pretrained(
load_path,
padding_side = padding_side,
token = token,
trust_remote_code = trust_remote_code,
cache_dir = cache_dir,
local_files_only = lfo,
)
def _resolve_checkpoint_tokenizer_name(
old_model_name,
kwargs,

30
unsloth/models/vision.py
View file

@ -631,7 +631,9 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
# Offline helpers live in loader_utils.py (shared canonical source).
from .loader_utils import (
_get_effective_local_files_only,
_hub_repo_or_local_path,
_is_offline_related_error,
_load_pretrained_tokenizer_fast,
_offline_aware_load,
)
@ -667,20 +669,27 @@ def _construct_vlm_processor_fallback(
tell an offline failure (retry from cache) from a genuine one."""
_fb_err = None
try:
from transformers import AutoImageProcessor, PreTrainedTokenizerFast, AutoConfig
from transformers import AutoImageProcessor, AutoConfig
from transformers.models.auto.processing_auto import PROCESSOR_MAPPING_NAMES
import json
load_path = _hub_repo_or_local_path(
tokenizer_name,
token = token,
cache_dir = cache_dir,
local_files_only = local_files_only,
)
# Load image processor
image_processor = AutoImageProcessor.from_pretrained(
tokenizer_name,
load_path,
token = token,
trust_remote_code = trust_remote_code,
cache_dir = cache_dir,
local_files_only = local_files_only,
)
# Load tokenizer via PreTrainedTokenizerFast (bypasses tokenizer_class check)
tok = PreTrainedTokenizerFast.from_pretrained(
# Load tokenizer via PreTrainedTokenizerFast (bypasses tokenizer_class check).
# Resolve the cached snapshot first so transformers does not call model_info (#7481).
tok = _load_pretrained_tokenizer_fast(
tokenizer_name,
padding_side = "left",
token = token,
@ -740,7 +749,7 @@ def _construct_vlm_processor_fallback(
# Try the top-level config.model_type which often has the processor mapping.
try:
config = AutoConfig.from_pretrained(
tokenizer_name,
load_path,
token = token,
trust_remote_code = trust_remote_code,
cache_dir = cache_dir,
@ -1653,9 +1662,15 @@ class FastBaseModel:
# Last resort: AutoTokenizer, then PreTrainedTokenizerFast (raise on network failure to retry).
def _last_resort_tokenizer(lfo):
from transformers import AutoTokenizer as _AutoTokenizer
load_path = _hub_repo_or_local_path(
tokenizer_name,
token = token,
cache_dir = kwargs.get("cache_dir"),
local_files_only = lfo,
)
try:
return _AutoTokenizer.from_pretrained(
tokenizer_name,
load_path,
padding_side = "left",
token = token,
trust_remote_code = trust_remote_code,
@ -1663,8 +1678,7 @@ class FastBaseModel:
local_files_only = lfo,
)
except Exception:
from transformers import PreTrainedTokenizerFast
return PreTrainedTokenizerFast.from_pretrained(
return _load_pretrained_tokenizer_fast(
tokenizer_name,
padding_side = "left",
token = token,

71
unsloth/save.py
View file

@ -52,7 +52,12 @@ import traceback
import psutil
import re
from transformers.models.llama.modeling_llama import logger
from .models.loader_utils import get_model_name
from .models.loader_utils import (
get_model_name,
_resolve_hub_repo_cached_file,
_tokenizer_cache_dir,
_tokenizer_wants_local_only,
)
from .models._utils import _convert_torchao_model
from .ollama_template_mappers import OLLAMA_TEMPLATES, MODEL_TO_OLLAMA_TEMPLATE_MAPPER
from transformers import ProcessorMixin, PreTrainedTokenizerBase
@ -446,6 +451,27 @@ def _has_tokenizer_model(tokenizer, token = None):
if source in _TOKENIZER_MODEL_CACHE:
return _TOKENIZER_MODEL_CACHE[source]
# Hub repo id: probe local cache before model_info (issue #7481).
cache_dir = _tokenizer_cache_dir(tokenizer) or os.environ.get("HF_HUB_CACHE")
if not cache_dir:
hf_home = os.environ.get("HF_HOME")
if hf_home:
cache_dir = os.path.join(hf_home, "hub")
cached_path = _resolve_hub_repo_cached_file(
source,
"tokenizer.model",
token = token,
local_files_only = True,
cache_dir = cache_dir,
)
if cached_path is not None:
_TOKENIZER_MODEL_CACHE[source] = True
return True
if _tokenizer_wants_local_only(tokenizer):
return False
try:
repo_info = HfApi(token = token).model_info(source, files_metadata = False)
except Exception:
@ -505,15 +531,33 @@ def _preserve_sentencepiece_tokenizer_assets(
if os.path.isfile(local_path):
downloaded_path = local_path
else:
from huggingface_hub import hf_hub_download
try:
downloaded_path = hf_hub_download(
repo_id = source,
filename = "tokenizer.model",
token = token,
)
except Exception:
downloaded_path = None
cache_dir = _tokenizer_cache_dir(tokenizer) or os.environ.get("HF_HUB_CACHE")
if not cache_dir:
hf_home = os.environ.get("HF_HOME")
if hf_home:
cache_dir = os.path.join(hf_home, "hub")
cached_path = _resolve_hub_repo_cached_file(
source,
"tokenizer.model",
token = token,
local_files_only = True,
cache_dir = cache_dir,
)
if cached_path is not None:
downloaded_path = cached_path
else:
from huggingface_hub import hf_hub_download
try:
downloaded_path = hf_hub_download(
repo_id = source,
filename = "tokenizer.model",
token = token,
local_files_only = _tokenizer_wants_local_only(tokenizer),
cache_dir = cache_dir,
)
except Exception:
downloaded_path = None
if not os.path.isfile(tokenizer_model) and downloaded_path is not None:
shutil.copy2(downloaded_path, tokenizer_model)
@ -3793,7 +3837,12 @@ def unsloth_convert_lora_to_ggml_and_save_locally(
return _unsloth_save_lora_gguf(self, tokenizer, save_directory, outtype = outtype)
from .models.loader_utils import get_model_name
from .models.loader_utils import (
get_model_name,
_resolve_hub_repo_cached_file,
_tokenizer_cache_dir,
_tokenizer_wants_local_only,
)
from unsloth_zoo.saving_utils import (
merge_and_overwrite_lora,
prepare_saving,

36
unsloth_cli/claude_subagent_mcp.py
View file

@ -29,6 +29,10 @@ from unsloth_cli.commands.start import (
_MAX_RESULT_CHARACTERS = 100_000
_CANCEL_POLL_SECONDS = 0.1
_CANCEL_GRACE_SECONDS = 2.0
# A local server that accepts the connection and then never answers leaves the
# child, and the parent waiting on it, blocked forever. Generous enough not to cut
# a long legitimate run short; 0 restores the unbounded wait.
_DEFAULT_TIMEOUT_SECONDS = 1800.0
def _required_env(name: str) -> str:
@ -38,6 +42,18 @@ def _required_env(name: str) -> str:
return value
def _timeout_seconds() -> float:
"""Wall-clock cap on one child run; 0 or unparsable means wait forever."""
raw = os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT")
if raw is None or not raw.strip():
return _DEFAULT_TIMEOUT_SECONDS
try:
parsed = float(raw.strip())
except ValueError:
return _DEFAULT_TIMEOUT_SECONDS
return parsed if parsed > 0 else 0.0
def _bounded(text: str) -> str:
if len(text) <= _MAX_RESULT_CHARACTERS:
return text
@ -153,6 +169,17 @@ def run_local_agent(
"--output-format",
"json",
"--no-session-persistence",
# Strip human-blocking tools so the child runs unattended. Only the read-only
# child's writers bite today, since a --print child is never offered the plan
# or prompt tools; those are listed anyway so a version that starts offering
# them cannot stall the subagent. Bash is denied read-only side because plan
# mode gates it through the same local model, which is not a write barrier.
"--disallowedTools",
(
"AskUserQuestion,EnterPlanMode,Edit,Write,NotebookEdit,Bash"
if read_only
else "AskUserQuestion,EnterPlanMode,ExitPlanMode"
),
"--append-system-prompt",
_SUBAGENT_PLAN_INSTRUCTIONS if read_only else _SUBAGENT_INSTRUCTIONS,
f"Task: {task}",
@ -185,6 +212,8 @@ def run_local_agent(
[executable, *command[1:]],
**popen_kwargs,
)
deadline = _timeout_seconds()
started_at = time.monotonic()
try:
while True:
try:
@ -194,6 +223,13 @@ def run_local_agent(
if cancel_event.is_set():
_stop_child(process)
raise RuntimeError("The local Claude agent was cancelled.")
waited = time.monotonic() - started_at
if deadline and waited > deadline:
_stop_child(process)
raise RuntimeError(
f"The local Claude agent produced nothing after {waited:.0f}s. "
"The local server is likely wedged; check that a model is loaded."
)
except BaseException:
if process.poll() is None:
_stop_child(process)

117
unsloth_cli/commands/start.py
View file

@ -4,6 +4,7 @@
"""`unsloth start` — launch a coding agent against a running Unsloth server."""
import atexit
import base64
import contextlib
import json
import os
@ -19,7 +20,7 @@ import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import NamedTuple, NoReturn, Optional
from typing import Literal, NamedTuple, NoReturn, Optional
from urllib.parse import urlencode, urlparse
import click
@ -182,6 +183,17 @@ _TENSOR_PARALLEL_OPTION = typer.Option(
rich_help_panel = _PANEL_MODEL,
help = "Split a GGUF across GPUs by tensor instead of by layer (multi-GPU only).",
)
_GPU_MEMORY_MODE_OPTION = typer.Option(
None,
"--gpu-memory-mode",
rich_help_panel = _PANEL_MODEL,
help = (
"GPU memory strategy for GGUF models loaded by this command. Auto lets "
"Unsloth manage placement. Manual with default layers and context delegates "
"placement and sizing to llama.cpp --fit. Omit when attaching to preserve "
"the running model's mode."
),
)
# Server knobs. Only used when `unsloth start` auto-starts the server (--serve);
# they have no effect when attaching to a server someone else already started.
@ -458,6 +470,7 @@ class LoadOptions(NamedTuple):
max_seq_length: int = 0
load_in_4bit: bool = True
tensor_parallel: bool = False
gpu_memory_mode: Optional[Literal["auto", "manual"]] = None
class ServerOptions(NamedTuple):
@ -992,6 +1005,8 @@ def _start_studio_server(
command += ["--no-load-in-4bit"]
if load.tensor_parallel:
command += ["--tensor-parallel"]
if load.gpu_memory_mode is not None:
command += ["--gpu-memory-mode", load.gpu_memory_mode]
log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log"
typer.echo("Starting Unsloth server")
@ -1462,7 +1477,11 @@ def _resolve_model(
# without reloading when the variant AND settings match, so a second session running
# the same command still attaches without evicting the first.
load_has_overrides = bool(
load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel
load.gguf_variant
or load.max_seq_length
or not load.load_in_4bit
or load.tensor_parallel
or load.gpu_memory_mode is not None
)
# /v1/models also lists cached-but-unloaded catalog entries (loaded == False);
# matching one would skip /api/inference/load and leave the agent pointed at a
@ -1516,6 +1535,10 @@ def _resolve_model(
payload["load_in_4bit"] = False
if load.tensor_parallel:
payload["tensor_parallel"] = True
if load.gpu_memory_mode is not None:
payload["gpu_memory_mode"] = load.gpu_memory_mode
if load.gpu_memory_mode == "manual":
payload["gpu_layers"] = -1
loaded = _load_model_with_progress(base, key, requested, load, payload)
if loaded.get("status") == "already_loaded":
typer.echo(f"Reusing loaded model: {_display_model_spec(requested, load.gguf_variant)}")
@ -2052,6 +2075,32 @@ def _opencode_subagent_inline_config(path: Path, permission: dict) -> dict:
return inline
def _b64_path(path: Path) -> str:
"""Path as base64, so it can cross a shell without being expanded."""
return base64.b64encode(str(path).encode("utf-8")).decode("ascii")
_CLAUDE_PLAN_GATE_SCRIPT = '''\
"""Deny the editing agent while the parent session is in plan mode."""
import json, sys
try:
mode = (json.load(sys.stdin) or {}).get("permission_mode")
except Exception:
sys.exit(0) # fail open: a hook error must never block the parent session
if mode == "plan":
print(json.dumps({"hookSpecificOutput": {
"hookEventName": "PreToolUse",
"permissionDecision": "deny",
"permissionDecisionReason": (
"Plan mode is active. Call the read-only Unsloth plan agent "
"(unsloth_plan_agent) instead of unsloth_agent."
),
}}))
sys.exit(0)
'''
def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path:
"""Write a session plugin that exposes the local Claude child through MCP."""
plugin = path / "unsloth-local-agent"
@ -2094,6 +2143,52 @@ def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path:
}
},
)
# Claude already refuses the editing tool in plan mode, since it advertises
# readOnlyHint false. This PreToolUse hook replaces that dead end with a reason
# naming the read-only tool to call instead. Skipped under the WSL bridge, where
# the gate is a Linux path but the hook would run beside the Windows claude.
gate = plugin / "hooks" / "plan_gate.py"
if command == "wsl.exe":
# A persisted plugin dir may still hold a gate from an earlier non-WSL run.
for stale in (gate, plugin / "hooks" / "hooks.json"):
stale.unlink(missing_ok = True)
else:
_write_private_text(gate, _CLAUDE_PLAN_GATE_SCRIPT)
_write_private_json(
plugin / "hooks" / "hooks.json",
{
"hooks": {
"PreToolUse": [
{
"matcher": _CLAUDE_SUBAGENT_TOOL,
"hooks": [
{
"type": "command",
# Run through runpy rather than handing the path to
# the interpreter: a missing gate is then an
# ordinary traceback (exit 1, fails open) instead
# of exit 2, which Claude treats as a blocking
# error and would deny the tool in every mode.
# The path is base64'd because this string goes
# through a shell: a temp root holding $(..) or a
# backtick expands under sh, %VAR% under cmd, and
# the gate then silently fails open. base64's
# alphabet has no metacharacter in either.
"command": (
f'"{sys.executable}" -c '
f'"import base64,runpy; runpy.run_path('
f"base64.b64decode('{_b64_path(gate)}').decode())\""
),
# A hook with no timeout stalls the parent for as
# long as it hangs; measured unbounded past 400s.
"timeout": 10,
}
],
}
]
}
},
)
skill = plugin / "skills" / "local-agent" / "SKILL.md"
skill.parent.mkdir(parents = True, exist_ok = True, mode = 0o700)
skill.write_text(
@ -2949,6 +3044,7 @@ def claude(
max_seq_length: int = _CONTEXT_OPTION,
load_in_4bit: bool = _LOAD_4BIT_OPTION,
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
gpu_memory_mode: Optional[Literal["auto", "manual"]] = _GPU_MEMORY_MODE_OPTION,
enable_tools: bool = _ENABLE_TOOLS_OPTION,
tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION,
tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION,
@ -2969,7 +3065,7 @@ def claude(
base, key, entry = _connect(
api_key,
model,
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel),
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel, gpu_memory_mode),
serve = serve,
launch = launch,
server_options = ServerOptions(
@ -3066,6 +3162,7 @@ def codex(
max_seq_length: int = _CONTEXT_OPTION,
load_in_4bit: bool = _LOAD_4BIT_OPTION,
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
gpu_memory_mode: Optional[Literal["auto", "manual"]] = _GPU_MEMORY_MODE_OPTION,
enable_tools: bool = _ENABLE_TOOLS_OPTION,
tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION,
tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION,
@ -3086,7 +3183,7 @@ def codex(
base, key, entry = _connect(
api_key,
model,
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel),
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel, gpu_memory_mode),
serve = serve,
launch = launch,
server_options = ServerOptions(
@ -3164,6 +3261,7 @@ def openclaw(
max_seq_length: int = _CONTEXT_OPTION,
load_in_4bit: bool = _LOAD_4BIT_OPTION,
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
gpu_memory_mode: Optional[Literal["auto", "manual"]] = _GPU_MEMORY_MODE_OPTION,
enable_tools: bool = _ENABLE_TOOLS_OPTION,
tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION,
tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION,
@ -3184,7 +3282,7 @@ def openclaw(
base, key, entry = _connect(
api_key,
model,
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel),
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel, gpu_memory_mode),
serve = serve,
launch = launch,
server_options = ServerOptions(
@ -3244,6 +3342,7 @@ def opencode(
max_seq_length: int = _CONTEXT_OPTION,
load_in_4bit: bool = _LOAD_4BIT_OPTION,
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
gpu_memory_mode: Optional[Literal["auto", "manual"]] = _GPU_MEMORY_MODE_OPTION,
enable_tools: bool = _ENABLE_TOOLS_OPTION,
tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION,
tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION,
@ -3264,7 +3363,7 @@ def opencode(
base, key, entry = _connect(
api_key,
model,
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel),
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel, gpu_memory_mode),
serve = serve,
launch = launch,
server_options = ServerOptions(
@ -3404,6 +3503,7 @@ def hermes(
max_seq_length: int = _CONTEXT_OPTION,
load_in_4bit: bool = _LOAD_4BIT_OPTION,
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
gpu_memory_mode: Optional[Literal["auto", "manual"]] = _GPU_MEMORY_MODE_OPTION,
enable_tools: bool = _ENABLE_TOOLS_OPTION,
tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION,
tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION,
@ -3426,7 +3526,7 @@ def hermes(
base, key, entry = _connect(
api_key,
model,
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel),
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel, gpu_memory_mode),
serve = serve,
launch = launch,
server_options = ServerOptions(
@ -3460,6 +3560,7 @@ def pi(
max_seq_length: int = _CONTEXT_OPTION,
load_in_4bit: bool = _LOAD_4BIT_OPTION,
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
gpu_memory_mode: Optional[Literal["auto", "manual"]] = _GPU_MEMORY_MODE_OPTION,
enable_tools: bool = _ENABLE_TOOLS_OPTION,
tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION,
tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION,
@ -3480,7 +3581,7 @@ def pi(
base, key, entry = _connect(
api_key,
model,
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel),
LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel, gpu_memory_mode),
serve = serve,
launch = launch,
server_options = ServerOptions(

19
unsloth_cli/commands/studio.py
View file

@ -19,7 +19,7 @@ import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Optional
from typing import List, Literal, Optional
import typer
from unsloth_cli.commands import _password_prompt
@ -1172,6 +1172,7 @@ def _load_model_via_http(
gguf_variant: Optional[str],
max_seq_length: int,
load_in_4bit: bool,
gpu_memory_mode: Literal["auto", "manual"] = "auto",
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
timeout: int = 600,
@ -1188,6 +1189,9 @@ def _load_model_via_http(
}
if gguf_variant:
payload["gguf_variant"] = gguf_variant
if gpu_memory_mode == "manual":
payload["gpu_memory_mode"] = "manual"
payload["gpu_layers"] = -1
if tensor_parallel:
payload["tensor_parallel"] = True
if llama_extra_args:
@ -1728,6 +1732,16 @@ def run(
rich_help_panel = _RUN_PANEL_MODEL,
help = "Runtime context length in tokens (0 = model default for GGUF; 2048 for hub models)",
),
gpu_memory_mode: Literal["auto", "manual"] = typer.Option(
"auto",
"--gpu-memory-mode",
rich_help_panel = _RUN_PANEL_MODEL,
help = (
"GPU memory strategy for GGUF models. Auto lets Unsloth select GPUs "
"and cap context to fit VRAM. Manual with default layers and context "
"delegates placement and sizing to llama.cpp --fit."
),
),
load_in_4bit: bool = typer.Option(
True, "--load-in-4bit/--no-load-in-4bit", rich_help_panel = _RUN_PANEL_MODEL
),
@ -2127,6 +2141,8 @@ def run(
"--host",
host,
]
if gpu_memory_mode != "auto":
args.extend(["--gpu-memory-mode", gpu_memory_mode])
if gguf_variant:
args.extend(["--gguf-variant", gguf_variant])
# Forward the explicit polarity; a future default flip on one
@ -2244,6 +2260,7 @@ def run(
gguf_variant = gguf_variant,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
gpu_memory_mode = gpu_memory_mode,
tensor_parallel = tensor_parallel,
llama_extra_args = extra_llama_args,
)

179
unsloth_cli/tests/test_claude_plan_gate.py Normal file
View file

@ -0,0 +1,179 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Deterministic plan-mode routing for the local Claude subagent.
SKILL.md asks the parent model to pick the read-only tool in plan mode, which a
small local model can forget. The generated plugin also ships a PreToolUse hook
that reads permission_mode directly, so the editing agent is denied by rule.
"""
from __future__ import annotations
import json
import subprocess
import sys
import pytest
from unsloth_cli.commands import start
def _plugin(tmp_path):
return start.write_claude_subagent_plugin(tmp_path, {"UNSLOTH_CLAUDE_SUBAGENT_MODEL": "m"})
def _run_gate(script, payload):
return subprocess.run(
[sys.executable, str(script)],
input = payload,
capture_output = True,
text = True,
timeout = 30,
)
def test_plugin_registers_a_pretooluse_hook_on_the_editing_tool(tmp_path):
plugin = _plugin(tmp_path)
hooks = json.loads((plugin / "hooks" / "hooks.json").read_text())["hooks"]["PreToolUse"]
[entry] = hooks
# Only the destructive tool is gated; the read-only agent stays reachable.
assert entry["matcher"] == start._CLAUDE_SUBAGENT_TOOL
assert start._CLAUDE_SUBAGENT_PLAN_TOOL not in json.dumps(hooks)
[hook] = entry["hooks"]
assert hook["type"] == "command"
assert sys.executable in hook["command"]
# The interpreter is quoted: unquoted, any space in the path splits the command.
assert f'"{sys.executable}"' in hook["command"]
# The gate path rides as base64, never as a literal the shell can expand.
encoded = start._b64_path(plugin / "hooks" / "plan_gate.py")
assert encoded in hook["command"]
assert str(plugin / "hooks" / "plan_gate.py") not in hook["command"]
# A hook with no timeout stalls the parent for as long as it hangs.
assert 0 < hook["timeout"] <= 30
def test_gate_script_is_written_and_compiles(tmp_path):
plugin = _plugin(tmp_path)
gate = plugin / "hooks" / "plan_gate.py"
compile(gate.read_text(), str(gate), "exec") # syntax-valid as shipped
def test_gate_denies_the_editing_tool_in_plan_mode(tmp_path):
gate = _plugin(tmp_path) / "hooks" / "plan_gate.py"
result = _run_gate(gate, json.dumps({"permission_mode": "plan"}))
assert result.returncode == 0
output = json.loads(result.stdout)["hookSpecificOutput"]
assert output["hookEventName"] == "PreToolUse"
assert output["permissionDecision"] == "deny"
# The reason is shown to the model, so it must name the tool to call instead.
assert "unsloth_plan_agent" in output["permissionDecisionReason"]
@pytest.mark.parametrize("mode", ["default", "acceptEdits", "bypassPermissions", "dontAsk", "auto"])
def test_gate_allows_every_non_plan_mode(tmp_path, mode):
gate = _plugin(tmp_path) / "hooks" / "plan_gate.py"
result = _run_gate(gate, json.dumps({"permission_mode": mode}))
assert result.returncode == 0
assert result.stdout.strip() == "" # no decision -> normal permission flow
@pytest.mark.parametrize("payload", ["", "not json", "[]", "null", "{}"])
def test_gate_fails_open_on_unusable_input(tmp_path, payload):
# A hook crash would block the parent session, so anything unparsable allows.
gate = _plugin(tmp_path) / "hooks" / "plan_gate.py"
result = _run_gate(gate, payload)
assert result.returncode == 0
assert result.stdout.strip() == ""
def test_plugin_still_writes_the_mcp_server_and_skill(tmp_path):
# The hook is additive; the existing wiring must be untouched.
plugin = _plugin(tmp_path)
assert (plugin / ".mcp.json").exists()
assert (plugin / "skills" / "local-agent" / "SKILL.md").exists()
assert (plugin / ".claude-plugin" / "plugin.json").exists()
def test_wsl_run_clears_a_gate_left_by_an_earlier_windows_run(tmp_path, monkeypatch):
# The plugin dir survives across runs when persisted, so a gate written by a
# Windows run would otherwise be shipped into the distro with an interpreter
# path it cannot execute.
plugin = _plugin(tmp_path)
gate = plugin / "hooks" / "plan_gate.py"
hooks = plugin / "hooks" / "hooks.json"
assert gate.exists() and hooks.exists()
monkeypatch.setattr(start, "_wsl_windows_executable", lambda _argv: True)
monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
_plugin(tmp_path)
assert not gate.exists()
assert not hooks.exists()
def test_hook_command_survives_a_missing_gate_and_a_path_with_spaces(tmp_path):
# Handing the path straight to the interpreter makes a missing gate exit 2,
# which Claude treats as a blocking error: the editing tool would then be
# denied in every mode, not just plan. Going through runpy makes it exit 1.
plugin = _plugin(tmp_path / "dir with space")
hook = json.loads((plugin / "hooks" / "hooks.json").read_text())
command = hook["hooks"]["PreToolUse"][0]["hooks"][0]["command"]
# Works normally through the real shell path Claude uses.
denied = subprocess.run(
command,
input = json.dumps({"permission_mode": "plan"}),
shell = True,
capture_output = True,
text = True,
timeout = 30,
)
assert denied.returncode == 0
assert json.loads(denied.stdout)["hookSpecificOutput"]["permissionDecision"] == "deny"
(plugin / "hooks" / "plan_gate.py").unlink()
gone = subprocess.run(
command,
input = json.dumps({"permission_mode": "default"}),
shell = True,
capture_output = True,
text = True,
timeout = 30,
)
assert gone.returncode != 2, "exit 2 blocks the tool in every mode"
assert gone.stdout.strip() == ""
@pytest.mark.parametrize("hostile", ["sub$(echo X)", "tick`echo X`", "var$HOME", "pct%TEMP%pct"])
def test_gate_survives_shell_metacharacters_in_its_path(tmp_path, hostile):
# The hook command is run by a shell. A temp root holding these expands under
# sh (or cmd, for %VAR%) before Python sees the path, so the gate is not found
# and exits 1, which fails open and silently drops the routing message.
plugin = _plugin(tmp_path / hostile)
command = json.loads((plugin / "hooks" / "hooks.json").read_text())["hooks"]["PreToolUse"][0][
"hooks"
][0]["command"]
denied = subprocess.run(
command,
input = json.dumps({"permission_mode": "plan"}),
shell = True,
capture_output = True,
text = True,
timeout = 60,
)
assert denied.returncode == 0, denied.stderr
decision = json.loads(denied.stdout)["hookSpecificOutput"]["permissionDecision"]
assert decision == "deny"

50
unsloth_cli/tests/test_claude_subagent_mcp.py
View file

@ -15,6 +15,16 @@ import pytest
import unsloth_cli.claude_subagent_mcp as bridge
def _stub_env(monkeypatch, tmp_path):
"""Minimum env + claude lookup for driving run_local_agent under a fake Popen."""
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888")
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test")
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M")
monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path))
monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude")
monkeypatch.setattr(bridge, "_claude_flags", lambda model: ["--settings", "{}"])
def test_protocol_lists_and_calls_local_agent():
initialized = bridge._response(
{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}},
@ -229,6 +239,8 @@ def test_local_child_uses_unsloth_without_overwriting_parent_auth(
assert command[:3] == ["/usr/local/bin/claude", "--model", "unsloth/model-GGUF:Q4_K_M"]
assert command[command.index("--permission-mode") + 1] == permission
assert "--no-session-persistence" in command
disallowed = command[command.index("--disallowedTools") + 1]
assert disallowed == "AskUserQuestion,EnterPlanMode,ExitPlanMode"
assert captured["cwd"] == str(tmp_path)
assert captured["stdin"] is bridge.subprocess.DEVNULL
assert captured["stdout"] is bridge.subprocess.PIPE
@ -275,6 +287,10 @@ def test_read_only_local_child_uses_plan_mode(monkeypatch, tmp_path):
assert bridge.run_local_agent("plan this", read_only = True) == "PLAN_OK"
command = captured["command"]
assert command[command.index("--permission-mode") + 1] == "plan"
disallowed = command[command.index("--disallowedTools") + 1]
assert disallowed == "AskUserQuestion,EnterPlanMode,Edit,Write,NotebookEdit,Bash"
# Bash matters: plan mode routes it through a classifier served by this same
# local model, so without the deny a "read-only" child can still write files.
prompt = command[command.index("--append-system-prompt") + 1]
assert "read-only local coding subagent" in prompt
@ -403,3 +419,37 @@ def test_stop_child_kills_survivors_after_leader_exit(monkeypatch, tmp_path):
def test_result_parser_accepts_diagnostics_before_json():
output = "connector warning\n" + json.dumps({"is_error": False, "result": "OK"})
assert bridge._result_text(output) == "OK"
def test_child_is_stopped_when_it_produces_nothing_before_the_deadline(monkeypatch, tmp_path):
# A local server that accepts and never answers used to block the child, and
# the parent waiting on it, indefinitely. Measured past 400s before this.
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT", "0.3")
_stub_env(monkeypatch, tmp_path)
stopped = []
class _Hanging:
returncode = None
def communicate(self, timeout = None):
raise subprocess.TimeoutExpired("claude", timeout)
def poll(self):
return None
monkeypatch.setattr(bridge, "_stop_child", lambda proc: stopped.append(proc))
monkeypatch.setattr(bridge.subprocess, "Popen", lambda *a, **k: _Hanging())
with pytest.raises(RuntimeError, match = "produced nothing"):
bridge.run_local_agent("hello")
assert stopped, "a timed-out child must be killed, not left running"
def test_timeout_can_be_disabled(monkeypatch):
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT", "0")
assert bridge._timeout_seconds() == 0.0
for bad in ("", " ", "abc"):
monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT", bad)
assert bridge._timeout_seconds() == bridge._DEFAULT_TIMEOUT_SECONDS
monkeypatch.delenv("UNSLOTH_CLAUDE_SUBAGENT_TIMEOUT")
assert bridge._timeout_seconds() == bridge._DEFAULT_TIMEOUT_SECONDS

61
unsloth_cli/tests/test_start.py
View file

@ -1993,6 +1993,7 @@ def test_start_studio_server_forwards_tool_flags_via_command_and_env(monkeypatch
start._start_studio_server("http://127.0.0.1:8888", "unsloth/M-GGUF", start.LoadOptions())
cmd, env = captured["command"], captured["kwargs"]["env"]
assert "--disable-tools" in cmd and "--enable-tools" not in cmd
assert "--gpu-memory-mode" not in cmd
assert env["UNSLOTH_DISABLE_TOOL_CALL_HEALING"] == "0"
assert env["UNSLOTH_TOOL_CALL_NUDGE"] == "1"
@ -2206,6 +2207,59 @@ def test_connect_load_knobs_reach_server_even_when_id_loaded(fake_studio):
]
@pytest.mark.parametrize(
"command_name", ["claude", "codex", "openclaw", "opencode", "hermes", "pi"]
)
def test_start_agents_expose_gpu_memory_mode_option(command_name):
import inspect
command = getattr(start, command_name)
opt = inspect.signature(command).parameters["gpu_memory_mode"].default
assert set(getattr(opt, "param_decls", None) or []) == {"--gpu-memory-mode"}
assert getattr(opt, "default", None) is None
assert getattr(opt, "rich_help_panel", None) == start._PANEL_MODEL
@pytest.mark.parametrize(
"mode,expected",
[
("auto", {"model_path": MODEL["id"], "gpu_memory_mode": "auto"}),
(
"manual",
{
"model_path": MODEL["id"],
"gpu_memory_mode": "manual",
"gpu_layers": -1,
},
),
],
)
def test_start_gpu_memory_mode_reaches_running_server(fake_studio, mode, expected):
result = CliRunner().invoke(
start.start_app,
[
"claude",
"--no-launch",
"--model",
MODEL["id"],
"--gpu-memory-mode",
mode,
],
)
assert result.exit_code == 0, result.output
loads = [call for call in fake_studio if call[1].endswith("/api/inference/load")]
assert loads == [("POST", f"{BASE}/api/inference/load", expected)]
def test_start_rejects_invalid_gpu_memory_mode(fake_studio):
result = CliRunner().invoke(
start.start_app,
["claude", "--no-launch", "--gpu-memory-mode", "invalid"],
)
assert result.exit_code != 0
assert "Invalid value for '--gpu-memory-mode'" in result.output
def test_connect_model_variant_suffix_loads_split_repo(fake_studio):
# When the model is not already loaded, the `:QUANT` suffix becomes the gguf_variant
# and the load uses the bare (valid) repo id, mirroring `unsloth run repo --gguf-variant`.
@ -2617,7 +2671,11 @@ def test_start_studio_server_builds_command_and_waits(monkeypatch, capsys):
"http://127.0.0.1:8888",
"unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL",
start.LoadOptions(
gguf_variant = "UD-Q4_K_XL", max_seq_length = 8192, load_in_4bit = True, tensor_parallel = True
gguf_variant = "UD-Q4_K_XL",
max_seq_length = 8192,
load_in_4bit = True,
tensor_parallel = True,
gpu_memory_mode = "manual",
),
)
cmd = captured["command"]
@ -2627,6 +2685,7 @@ def test_start_studio_server_builds_command_and_waits(monkeypatch, capsys):
assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL"
assert cmd[cmd.index("--context-length") + 1] == "8192"
assert "--tensor-parallel" in cmd
assert cmd[cmd.index("--gpu-memory-mode") + 1] == "manual"
assert "--start-api-key-marker" not in cmd
assert captured["kwargs"]["env"][start._START_API_KEY_MARKER_ENV] == "1"
assert start.os.environ[start._START_API_KEY_MARKER_ENV] == "parent"

84
unsloth_cli/tests/test_studio_run_parallel_flag.py
View file

@ -13,7 +13,9 @@ canonicaliser and the legacy `-m` / `-hfr` / `-f` shim.
from __future__ import annotations
import json
import sys
from io import BytesIO
from pathlib import Path
import pytest
@ -63,6 +65,19 @@ def test_context_length_alias_is_registered():
assert "--context-length" in flags
def test_gpu_memory_mode_option_is_registered_with_auto_default():
"""The GPU placement policy is a first-class model option."""
studio_mod = _load_run_command()
import inspect
sig = inspect.signature(studio_mod.run)
opt = sig.parameters["gpu_memory_mode"].default
flags = set(getattr(opt, "param_decls", None) or [])
assert flags == {"--gpu-memory-mode"}
assert getattr(opt, "default", None) == "auto"
assert getattr(opt, "rich_help_panel", None) == "Model"
def test_parallel_default_is_four():
"""Default must stay at 4 so plain `unsloth studio run` is unchanged."""
studio_mod = _load_run_command()
@ -377,6 +392,75 @@ def test_reexec_forwards_context_length_alias(monkeypatch):
assert "--context-length" not in argv, argv
def test_reexec_forwards_manual_gpu_memory_mode(monkeypatch):
"""An explicit manual policy must survive the Studio venv re-exec."""
result, captured = _invoke_run(
monkeypatch,
_BASE + ["--gpu-memory-mode", "manual"],
)
assert len(captured) == 1, result.output
argv = captured[0]["argv"]
assert _value_after(argv, "--gpu-memory-mode") == "manual", argv
def test_reexec_omits_default_gpu_memory_mode(monkeypatch):
"""The default stays compatible with older Studio venv launchers."""
result, captured = _invoke_run(monkeypatch, _BASE)
assert len(captured) == 1, result.output
assert "--gpu-memory-mode" not in captured[0]["argv"]
def test_run_rejects_invalid_gpu_memory_mode(monkeypatch):
result, captured = _invoke_run(
monkeypatch,
_BASE + ["--gpu-memory-mode", "invalid"],
)
assert result.exit_code != 0
assert captured == []
@pytest.mark.parametrize(
"mode,expected",
[
("auto", {"model_path": "owner/model-GGUF", "max_seq_length": 0, "load_in_4bit": True}),
(
"manual",
{
"model_path": "owner/model-GGUF",
"max_seq_length": 0,
"load_in_4bit": True,
"gpu_memory_mode": "manual",
"gpu_layers": -1,
},
),
],
)
def test_load_model_http_payload_for_gpu_memory_mode(monkeypatch, mode, expected):
"""Manual plus untouched layer and context settings matches the UI payload."""
studio_mod = _load_run_command()
captured = {}
def urlopen(request, timeout):
captured["request"] = request
captured["timeout"] = timeout
return BytesIO(b'{"model": "owner/model-GGUF"}')
monkeypatch.setattr(studio_mod.urllib.request, "urlopen", urlopen)
result = studio_mod._load_model_via_http(
port = 8888,
api_key = "sk-test",
model = "owner/model-GGUF",
gguf_variant = None,
max_seq_length = 0,
load_in_4bit = True,
gpu_memory_mode = mode,
)
assert result == {"model": "owner/model-GGUF"}
assert json.loads(captured["request"].data) == expected
assert captured["request"].get_header("Authorization") == "Bearer sk-test"
def test_reexec_mixed_parallel_with_passthrough(monkeypatch):
"""--parallel + llama-server pass-through flags must all reach the child."""
result, captured = _invoke_run(