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

Merge main into fix/studio-install-readiness

Browse source
This commit is contained in:
Daniel Han 2026-07-27 18:06:00 +00:00
commit 5c1530186a
115 changed files with 12537 additions and 841 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

8
studio/backend/auth/storage.py
View file

@ -44,7 +44,7 @@ def generate_bootstrap_password() -> str:
# Persisted from a previous run?
if _BOOTSTRAP_PW_PATH.is_file():
_bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
_bootstrap_password = _BOOTSTRAP_PW_PATH.read_text(encoding = "utf-8").strip()
if _bootstrap_password:
return _bootstrap_password
@ -57,7 +57,7 @@ def generate_bootstrap_password() -> str:
# Persist so the same passphrase survives restarts until password change.
ensure_dir(_BOOTSTRAP_PW_PATH.parent)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password, encoding = "utf-8")
try:
os.chmod(_BOOTSTRAP_PW_PATH, 0o600)
except OSError:
@ -76,7 +76,7 @@ def _load_bootstrap_password() -> Optional[str]:
global _bootstrap_password
_bootstrap_password = None
if _BOOTSTRAP_PW_PATH.is_file():
bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
bootstrap_password = _BOOTSTRAP_PW_PATH.read_text(encoding = "utf-8").strip()
if bootstrap_password:
_bootstrap_password = bootstrap_password
return _bootstrap_password
@ -99,7 +99,7 @@ def clear_bootstrap_password() -> None:
# stale plaintext can't be re-seeded by generate_bootstrap_password()
# if a later reset-password deletes auth.db and re-validates it.
try:
_BOOTSTRAP_PW_PATH.write_text("")
_BOOTSTRAP_PW_PATH.write_text("", encoding = "utf-8")
cleared = True
except OSError:
cleared = False

6
studio/backend/colab.py
View file

@ -90,7 +90,7 @@ def _store_colab_login_credentials(username: str, password: str) -> None:
path = _colab_login_credentials_path()
try:
path.parent.mkdir(parents = True, exist_ok = True)
path.write_text(f"{username}\n{password}\n")
path.write_text(f"{username}\n{password}\n", encoding = "utf-8")
try:
import os
os.chmod(path, 0o600)
@ -106,10 +106,10 @@ def _load_colab_login_credentials() -> "tuple[str, str] | None":
try:
if not path.is_file():
return None
lines = path.read_text().splitlines()
lines = path.read_text(encoding = "utf-8").splitlines()
if len(lines) >= 2 and lines[0] and lines[1]:
return lines[0], lines[1]
except OSError as e:
except (OSError, UnicodeDecodeError) as e:
logger.info(f"Could not load Colab login credentials ({e}).")
return None

4
studio/backend/core/export/export.py
View file

@ -241,7 +241,7 @@ def _offline_window_if(local_files_only):
def _is_wsl():
"""Detect if running under Windows Subsystem for Linux."""
try:
return "microsoft" in open("/proc/version").read().lower()
return "microsoft" in open("/proc/version", encoding = "utf-8").read().lower()
except Exception:
return False
@ -574,7 +574,7 @@ class ExportBackend:
)
metadata = {"base_model": base_model}
metadata_path = os.path.join(save_directory, "export_metadata.json")
with open(metadata_path, "w") as f:
with open(metadata_path, "w", encoding = "utf-8") as f:
json.dump(metadata, f, indent = 2)
logger.info(f"Wrote export metadata to {metadata_path}")
except Exception as e:

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:

2
studio/backend/core/inference/inference.py
View file

@ -567,7 +567,7 @@ class InferenceBackend:
_meta_path = Path(config.path) / "export_metadata.json"
try:
if _meta_path.exists():
_meta = json.loads(_meta_path.read_text())
_meta = json.loads(_meta_path.read_text(encoding = "utf-8"))
if _meta.get("base_model"):
processor_source = _meta["base_model"]
except Exception:

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

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

@ -246,7 +246,7 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]":
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
if "microsoft" not in fh.read().lower():
return []
except OSError:
except (OSError, UnicodeDecodeError):
return []
out: "list[str]" = []
for d in ("/opt/rocm/lib", "/opt/rocm/lib64"):
@ -569,11 +569,11 @@ def _load_swa_cache() -> dict:
if _SWA_CACHE is not None:
return _SWA_CACHE
try:
with open(_swa_cache_path()) as f:
with open(_swa_cache_path(), encoding = "utf-8") as f:
_SWA_CACHE = json.load(f)
if not isinstance(_SWA_CACHE, dict):
_SWA_CACHE = {}
except (FileNotFoundError, json.JSONDecodeError, OSError):
except (FileNotFoundError, json.JSONDecodeError, OSError, UnicodeDecodeError):
_SWA_CACHE = {}
return _SWA_CACHE
@ -583,10 +583,10 @@ def _save_swa_cache(cache: dict) -> None:
path = _swa_cache_path()
path.parent.mkdir(parents = True, exist_ok = True)
tmp = path.with_suffix(".json.tmp")
with open(tmp, "w") as f:
with open(tmp, "w", encoding = "utf-8") as f:
json.dump(cache, f, indent = 2, sort_keys = True)
tmp.replace(path)
except OSError:
except (OSError, UnicodeDecodeError):
pass
@ -620,7 +620,7 @@ def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
repo_type = "model",
cache_dir = active_hf_hub_cache(),
)
with open(cfg_path) as f:
with open(cfg_path, encoding = "utf-8") as f:
cfg = json.load(f)
except Exception:
return None
@ -3116,8 +3116,9 @@ class LlamaCppBackend:
prefer_rocr masks at the ROCr/HSA layer instead (clearing HIP). A HIP mask
filters only AFTER the HSA runtime enumerates every agent, and that
enumeration segfaults at startup on a GPU the build has no kernels for
(e.g. a gfx1103 iGPU under a gfx110X prebuilt), before llama-server logs a
line. ROCR drops the device at the driver layer, consuming physical ids.
(e.g. a gfx1036 iGPU under a gfx103X prebuilt: that bundle maps only
gfx1030/1031/1032/1034), before llama-server logs a line. ROCR drops the
device at the driver layer, consuming physical ids.
The CPU-only sentinel ("-1") has no portable ROCR spelling, so it keeps
the HIP mask. Windows keeps the HIP mask too: ROCR_VISIBLE_DEVICES is a
Linux ROCr variable (Windows HIP has no ROCr layer), so the ROCR pin
@ -3501,18 +3502,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 +3537,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 +3557,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 +3615,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: "
@ -3596,7 +3633,7 @@ class LlamaCppBackend:
except Exception:
pass
try:
with open("/proc/meminfo") as f:
with open("/proc/meminfo", encoding = "utf-8") as f:
for line in f:
if line.startswith("MemAvailable:"):
return int(line.split()[1]) // 1024 # kB -> MiB
@ -5138,7 +5175,7 @@ class LlamaCppBackend:
self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log"
self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1)
logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}")
except OSError as e:
except (OSError, UnicodeDecodeError) as e:
logger.debug(f"Could not open diffusion runner log file: {e}")
# The shim (and its visual server) die with this backend process, so a
@ -6355,7 +6392,7 @@ class LlamaCppBackend:
buffering = 1,
)
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
except OSError as e:
except (OSError, UnicodeDecodeError) as e:
# Best-effort; never block the load on logging.
logger.debug(f"Could not open llama-server log file: {e}")
self._llama_log_path = None
@ -6635,12 +6672,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")
@ -8236,7 +8284,7 @@ class LlamaCppBackend:
env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# Mask on AMD at the ROCr/HSA layer: HIP-only masking still
# enumerates every agent first, which segfaults on a deselected
# unsupported GPU (e.g. gfx1103 iGPU under a gfx110X prebuilt).
# unsupported GPU (e.g. gfx1036 iGPU under a gfx103X prebuilt).
self._emit_child_gpu_visibility(
env, ",".join(str(i) for i in gpu_indices), prefer_rocr = True
)
@ -8302,7 +8350,7 @@ class LlamaCppBackend:
buffering = 1,
)
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
except OSError as e:
except (OSError, UnicodeDecodeError) as e:
# Best-effort; never block the load on logging.
logger.debug(f"Could not open llama-server log file: {e}")
self._llama_log_path = None
@ -9477,7 +9525,7 @@ class LlamaCppBackend:
return
try:
path.parent.mkdir(parents = True, exist_ok = True)
path.write_text(f"{pid}:{cls._pid_start_identity(pid)}")
path.write_text(f"{pid}:{cls._pid_start_identity(pid)}", encoding = "utf-8")
except Exception as e:
logger.debug(f"Could not write llama-server pidfile: {e}")
@ -9611,7 +9659,7 @@ class LlamaCppBackend:
pid = -1
identity = ""
try:
pid_str, _, identity = path.read_text().strip().partition(":")
pid_str, _, identity = path.read_text(encoding = "utf-8").strip().partition(":")
pid = int(pid_str)
except Exception:
pid = -1

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

4
studio/backend/core/inference/sandbox_site/sitecustomize.py
View file

@ -111,7 +111,7 @@ def _load_sidecar(cwd):
"""Return the persisted ``source -> healed target`` map, or {} on any error
(missing/corrupt/foreign sidecar degrades to in-process-only behaviour)."""
try:
with open(_sidecar_path(cwd)) as fh:
with open(_sidecar_path(cwd), encoding = "utf-8") as fh:
data = json.load(fh)
except Exception: # noqa: BLE001 - a bad sidecar must never break user code
return {}
@ -131,7 +131,7 @@ def _record_sidecar(cwd, source, target):
return
data[source] = target
tmp = _sidecar_path(cwd) + ".tmp"
with open(tmp, "w") as fh:
with open(tmp, "w", encoding = "utf-8") as fh:
json.dump(data, fh)
os.replace(tmp, _sidecar_path(cwd))
except Exception: # noqa: BLE001 - persistence is best effort only

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

7
studio/backend/core/inference/worker.py
View file

@ -151,7 +151,7 @@ def _resolve_lora_4bit(mc, load_in_4bit: bool) -> bool:
import json
try:
with open(adapter_cfg_path) as f:
with open(adapter_cfg_path, encoding = "utf-8") as f:
adapter_cfg = json.load(f)
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora" and load_in_4bit:
@ -961,7 +961,10 @@ def run_inference_process(
if _local_adapter_cfg.is_file():
try:
_lora_base = (
_json.loads(_local_adapter_cfg.read_text()).get("base_model_name_or_path") or None
_json.loads(_local_adapter_cfg.read_text(encoding = "utf-8")).get(
"base_model_name_or_path"
)
or None
)
except Exception:
_lora_base = None

4
studio/backend/core/rag/embeddings.py
View file

@ -100,7 +100,7 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
path = Path(normalize_path(name)).expanduser() / "modules.json"
if not path.is_file():
return ()
data = json.loads(path.read_text())
data = json.loads(path.read_text(encoding = "utf-8"))
else:
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
@ -115,7 +115,7 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
)
except EntryNotFoundError:
return ()
data = json.loads(open(local).read())
data = json.loads(open(local, encoding = "utf-8").read())
subdirs = []
for module in data or ():
sub = str((module or {}).get("path", "")).strip().strip("/")

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,

8
studio/backend/hub/services/models/ollama.py
View file

@ -215,8 +215,8 @@ def _ollama_model_info_from_manifest(
return None
try:
manifest = json.loads(tag_file.read_text())
except (json.JSONDecodeError, OSError) as e:
manifest = json.loads(tag_file.read_text(encoding = "utf-8"))
except (json.JSONDecodeError, OSError, UnicodeDecodeError) as e:
logger.debug("Skipping unreadable/invalid Ollama manifest %s: %s", tag_file, e)
return None
@ -228,10 +228,10 @@ def _ollama_model_info_from_manifest(
config_blob = _ollama_blob_path(blobs_dir, config_digest)
if config_blob is not None and _safe_is_file(config_blob):
try:
cfg = json.loads(config_blob.read_text())
cfg = json.loads(config_blob.read_text(encoding = "utf-8"))
model_type = cfg.get("model_type", "")
file_type = cfg.get("file_type", "")
except (json.JSONDecodeError, OSError) as e:
except (json.JSONDecodeError, OSError, UnicodeDecodeError) as e:
logger.debug("Could not parse Ollama config blob %s: %s", config_blob, e)
layers = manifest.get("layers") or []

6
studio/backend/hub/utils/download_registry.py
View file

@ -462,8 +462,8 @@ def _read_marker_value(marker: Path) -> Optional[str]:
try:
if not marker.exists():
return None
value = marker.read_text().strip()
except OSError:
value = marker.read_text(encoding = "utf-8").strip()
except (OSError, UnicodeDecodeError):
return None
return value if value in VALID_TRANSPORTS else None
@ -473,7 +473,7 @@ def _write_marker_value(marker: Path, mode: str) -> None:
# tmp + rename so a SIGKILL mid-write can't leave a half-written marker.
# The tmp name is per-process so concurrent writers don't clobber tmps.
tmp = marker.with_name(f"{marker.name}.tmp-{os.getpid()}")
tmp.write_text(mode)
tmp.write_text(mode, encoding = "utf-8")
os.replace(tmp, marker)
except OSError:
# Best-effort: a missing marker next run purges the partial defensively,

4
studio/backend/hub/utils/paths.py
View file

@ -103,7 +103,7 @@ def _is_wsl() -> bool:
if sys.platform == "win32":
return False
try:
return "microsoft" in Path("/proc/version").read_text().lower()
return "microsoft" in Path("/proc/version").read_text(encoding = "utf-8").lower()
except Exception:
return False
@ -124,7 +124,7 @@ def _wsl_automount_root() -> str:
import configparser
parser = configparser.ConfigParser(inline_comment_prefixes = ("#", ";"))
parser.read("/etc/wsl.conf")
parser.read("/etc/wsl.conf", encoding = "utf-8")
root = parser.get("automount", "root", fallback = "").strip().strip("\"'")
except Exception:
return default

28
studio/backend/main.py
View file

@ -254,7 +254,11 @@ def _read_studio_install_id() -> str:
/api/health emits "" and the launcher accepts any healthy backend.
Carries no install-path info (matters when Unsloth runs -H 0.0.0.0)."""
try:
token = (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
token = (
(_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id")
.read_text(encoding = "utf-8")
.strip()
)
except (OSError, ValueError):
return ""
return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else ""
@ -358,7 +362,7 @@ def get_unsloth_version() -> str:
for line in version_file.read_text(encoding = "utf-8").splitlines():
if line.startswith("__version__ = "):
return line.split("=", 1)[1].strip().strip('"').strip("'")
except OSError:
except (OSError, UnicodeDecodeError):
pass
return "dev"
@ -1245,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
@ -1280,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

10
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py
View file

@ -20,7 +20,7 @@ class StateStore:
self._data: Dict[str, Any] = {}
if self.path.exists():
try:
with self.path.open() as f:
with self.path.open(encoding = "utf-8") as f:
self._data = json.load(f)
except Exception:
self._data = {}
@ -51,7 +51,7 @@ class StateStore:
def _flush(self) -> None:
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
with tmp.open("w") as f:
with tmp.open("w", encoding = "utf-8") as f:
json.dump(self._data, f, indent = 2, default = str)
os.replace(tmp, self.path)
@ -63,12 +63,14 @@ class JsonlWriter:
self.path = Path(path)
self.path.parent.mkdir(parents = True, exist_ok = True)
self._lock = threading.Lock()
self._fh = self.path.open("a", buffering = 1)
self._fh = self.path.open("a", buffering = 1, encoding = "utf-8")
self._count_seen_keys: set[str] = set()
# Preload seen keys for dedup across resumes
if self.path.exists() and self.path.stat().st_size > 0:
try:
with self.path.open() as f:
# No guess is safe for a file an older build wrote in the
# operator's locale, so read past whatever will not decode.
with self.path.open(encoding = "utf-8", errors = "replace") as f:
for line in f:
try:
obj = json.loads(line)

4
studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py
View file

@ -27,9 +27,9 @@ class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]):
orig_name = path_obj.name
if meta_path.exists():
try:
meta = json_mod.loads(meta_path.read_text())
meta = json_mod.loads(meta_path.read_text(encoding = "utf-8"))
orig_name = meta.get("original_filename", path_obj.name)
except (json_mod.JSONDecodeError, OSError):
except (json_mod.JSONDecodeError, OSError, UnicodeDecodeError):
pass
file_entries.append((path_obj, orig_name))

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

File diff suppressed because it is too large Load diff

10
studio/backend/routes/models.py
View file

@ -722,8 +722,8 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
stem_hash = hashlib.sha256(manifest_key.encode()).hexdigest()[:10]
try:
manifest = json.loads(tag_file.read_text())
except (json.JSONDecodeError, OSError) as e:
manifest = json.loads(tag_file.read_text(encoding = "utf-8"))
except (json.JSONDecodeError, OSError, UnicodeDecodeError) as e:
logger.debug(
"Skipping unreadable/invalid Ollama manifest %s: %s",
tag_file,
@ -738,10 +738,10 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
config_blob = blobs_dir / config_digest.replace(":", "-")
if config_blob.is_file():
try:
cfg = json.loads(config_blob.read_text())
cfg = json.loads(config_blob.read_text(encoding = "utf-8"))
model_type = cfg.get("model_type", "")
file_type = cfg.get("file_type", "")
except (json.JSONDecodeError, OSError) as e:
except (json.JSONDecodeError, OSError, UnicodeDecodeError) as e:
logger.debug(
"Could not parse Ollama config blob %s: %s",
config_blob,
@ -1042,7 +1042,7 @@ def _dir_has_downloaded_model(directory: Path, max_entries: int = 4000) -> bool:
if not m.is_file():
continue
try:
manifest = json.loads(m.read_text())
manifest = json.loads(m.read_text(encoding = "utf-8"))
except (json.JSONDecodeError, OSError, ValueError):
continue
for layer in manifest.get("layers") or []:

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

8
studio/backend/run.py
View file

@ -774,7 +774,7 @@ def _write_pid_file():
"""Write the current process PID to the studio PID file."""
try:
_PID_FILE.parent.mkdir(parents = True, exist_ok = True)
_PID_FILE.write_text(str(os.getpid()))
_PID_FILE.write_text(str(os.getpid()), encoding = "utf-8")
except OSError:
pass
@ -783,10 +783,10 @@ def _remove_pid_file():
"""Remove the PID file if it belongs to this process."""
try:
if _PID_FILE.is_file():
stored = _PID_FILE.read_text().strip()
stored = _PID_FILE.read_text(encoding = "utf-8").strip()
if stored == str(os.getpid()):
_PID_FILE.unlink(missing_ok = True)
except OSError:
except (OSError, UnicodeDecodeError):
pass
@ -934,7 +934,7 @@ def _iter_frontend_fallback_candidates() -> "list[Path]":
for finder in sp.glob("__editable___*_finder.py"):
try:
src = finder.read_text(encoding = "utf-8")
except OSError:
except (OSError, UnicodeDecodeError):
continue
# Tolerate single/multi-line dict literals; [^}]* rejects nested
# dicts, which the setuptools editable template never emits.

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"

8
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])
@ -1048,7 +1052,7 @@ def _rocm_torch_stub(monkeypatch):
def test_subset_pin_masks_via_rocr_on_rocm(monkeypatch):
# A GPU-subset pin must exclude the rest at the ROCr/HSA layer: HIP masking
# still enumerates every agent first, which segfaults the build on an
# unsupported deselected GPU (e.g. a gfx1103 iGPU under a gfx110X prebuilt).
# unsupported deselected GPU (e.g. a gfx1036 iGPU under a gfx103X prebuilt).
# ROCR drops it at the driver layer; only one mask is set (HIP cleared).
_rocm_torch_stub(monkeypatch)
env = {"HIP_VISIBLE_DEVICES": "9"} # stale/inherited HIP mask must not survive

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,

831
studio/backend/tests/test_install_resolve_prebuilt.py
View file

@ -6,7 +6,9 @@ by default; --published-repo overrides).
These back the in-app update for source-build (markerless) installs: the backend
asks the installer whether an official prebuilt exists for this host without
downloading. Network and host detection are stubbed; no GPU or internet needed.
downloading. Network and host detection are stubbed; no GPU or internet needed. The one
exception is the windows-rocm floor guard, which reads the fork's published manifest
because nothing in-tree mirrors it, and skips when that release is unreachable.
"""
from __future__ import annotations
@ -32,6 +34,18 @@ FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp
UPSTREAM = ilp.UPSTREAM_REPO # ggml-org/llama.cpp
@pytest.fixture(autouse = True)
def _no_ambient_hip_device_mask(monkeypatch):
"""These tests describe hosts through HostInfo, not through the environment.
A mask inherited from the shell (ML boxes commonly export CUDA_VISIBLE_DEVICES) means
the arch probe saw only part of the GPUs, which the Windows auto-Vulkan guard treats as
an unknown physical inventory. Clear all three so a host is described by its fields
alone; the tests that are about the mask set it explicitly."""
for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
monkeypatch.delenv(_env, raising = False)
def _host(**kw):
base = dict(
system = "Linux",
@ -407,7 +421,9 @@ def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
# Routing fork -> upstream also drops the fork release pin, which is in a
# different tag namespace and would make the upstream resolver miss.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False)
routed, repo, tag, _persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "b9596-mix-abc", force_cpu = False
)
assert repo == UPSTREAM
assert tag == ""
assert routed.has_intel_gpu is True
@ -416,7 +432,9 @@ def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
# A pin set WITH an explicit upstream repo is already on upstream -> kept.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
_routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False)
_routed, repo, tag, _persist = ilp._route_to_vulkan_prebuilt(
host, UPSTREAM, "b9596", force_cpu = False
)
assert repo == UPSTREAM
assert tag == "b9596"
@ -424,7 +442,9 @@ def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
# --cpu-fallback suppresses Vulkan routing even for an Intel host.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True)
routed, repo, tag, _persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "b9596-mix-abc", force_cpu = True
)
assert repo == FORK
assert tag == "b9596-mix-abc"
assert routed is host
@ -536,20 +556,20 @@ def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
has_physical_nvidia = True,
has_usable_nvidia = False,
)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
_routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted():
# An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
_routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_non_intel_unchanged():
host = _host(is_linux = True, is_x86_64 = True)
routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
assert routed is host
@ -797,3 +817,800 @@ def test_detect_host_cim_rescues_exploding_registry(monkeypatch):
)
assert host.has_intel_gpu is True
assert "powershell" in captured
def _windows_amd_host(**overrides):
defaults = dict(
system = "Windows",
machine = "amd64",
is_windows = True,
is_linux = False,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = True,
has_intel_gpu = False,
)
defaults.update(overrides)
return ilp.HostInfo(**defaults)
def test_route_to_vulkan_prebuilt_auto_fallback_for_legacy_amd_gfx():
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert persist == "vulkan"
assert routed.has_intel_gpu is True
assert routed.has_rocm is False
def test_route_to_vulkan_prebuilt_keeps_hip_when_one_gpu_is_supported():
host = _windows_amd_host(
rocm_gfx_target = "gfx1201",
rocm_gfx_targets = ["gfx1201", "gfx803"],
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_route_to_vulkan_prebuilt_auto_fallback_skips_hip_masked_hosts():
# A HIP mask can hide a HIP-capable dGPU, but the Vulkan runtime honours none of them,
# so auto-routing would let the installed backend grab the gfx1201 the user masked
# off.
host = _windows_amd_host(
rocm_gfx_target = "gfx803",
rocm_gfx_targets = ["gfx1201", "gfx803"],
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == FORK
assert persist is None
assert routed is host
def test_route_to_vulkan_prebuilt_auto_fallback_when_no_amd_gpu_reaches_floor():
# Every physical AMD device is below the floor, so no card can be exposed to HIP and
# the #7357 auto-Vulkan fallback still fires.
host = _windows_amd_host(
rocm_gfx_target = "gfx900",
rocm_gfx_targets = ["gfx803", "gfx900"],
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert persist == "vulkan"
assert routed.has_rocm is False
@pytest.mark.parametrize(
"mask_env", ["HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"]
)
def test_auto_vulkan_declines_when_a_hip_device_mask_filtered_the_probe(mask_env, monkeypatch):
# hipinfo is a HIP application, so under a mask rocm_gfx_targets is the VISIBLE set and
# a HIP-capable card can be hidden entirely. "No AMD GPU here reaches the floor" is then
# unprovable, and Vulkan honours none of these masks, so the auto fallback must decline
# rather than hand it the reserved card.
monkeypatch.setenv(mask_env, "1")
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
@pytest.mark.parametrize("mask_value", ["", " ", "-1"])
def test_auto_vulkan_declines_when_the_mask_hides_every_amd_gpu(mask_value, monkeypatch):
# An all-hiding mask is the strongest form of the same signal, not an exemption:
# detect_host() resolves no arch under it, but a forwarded --rocm-gfx still reconstructs
# one (setup infers it from the display-adapter name, which no HIP mask touches), so
# auto-routing would hand Vulkan every AMD GPU the user hid from HIP.
monkeypatch.setenv("HIP_VISIBLE_DEVICES", mask_value)
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
assert ilp._active_rocm_gfx_target(host) == "gfx803"
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_hip_device_mask_check_is_presence_not_value(monkeypatch):
# Presence is the whole test: any value means the HIP view is not the physical one, and
# no value can be read as "the probe saw everything".
assert ilp._hip_visible_device_mask_set() is False
for value in ("", " ", "-1", "0", "1", "0,1"):
monkeypatch.setenv("HIP_VISIBLE_DEVICES", value)
assert ilp._hip_visible_device_mask_set() is True, value
monkeypatch.delenv("HIP_VISIBLE_DEVICES")
monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "0")
assert ilp._hip_visible_device_mask_set() is True
monkeypatch.delenv("ROCR_VISIBLE_DEVICES")
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
assert ilp._hip_visible_device_mask_set() is True
def test_masked_probe_suppression_does_not_touch_non_amd_auto_paths(monkeypatch):
# The mask says nothing about an Intel iGPU, whose Vulkan auto path is unrelated.
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
host = _host(
system = "Windows",
is_windows = True,
has_intel_gpu = True,
has_rocm = False,
has_physical_nvidia = False,
has_usable_nvidia = False,
)
_routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "pin", force_cpu = False
)
assert repo == UPSTREAM
def test_route_to_vulkan_prebuilt_hip_masked_host_still_honours_explicit_optin(monkeypatch):
# The mask guard only suppresses the AUTOMATIC fallback; an explicit opt-in is the user
# taking responsibility for the Vulkan device mask themselves.
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
host = _windows_amd_host(
rocm_gfx_target = "gfx803",
rocm_gfx_targets = ["gfx1201", "gfx803"],
)
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "pin", force_cpu = False, llama_backend = "vulkan"
)
assert repo == UPSTREAM
assert persist == "vulkan"
def test_auto_vulkan_is_repository_specific_for_fork_only_gfx():
# gfx1034 is served only by the fork's gfx103X bundle: ggml-org's windows-hip radeon
# build does not target it and direct_upstream_release_plan() offers win-hip then CPU
# with no Vulkan branch, so the predicate must answer per repo.
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
assert ilp._should_auto_vulkan_for_amd_windows(host, UPSTREAM) is True
# An arch upstream really does build stays on HIP for both repos.
supported = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
assert ilp._should_auto_vulkan_for_amd_windows(supported, FORK) is False
assert ilp._should_auto_vulkan_for_amd_windows(supported, UPSTREAM) is False
# A family label is a bundle name, not an arch: upstream builds every member but
# gfx1034 / gfx1103, and the label cannot say which card this is, so it stays on HIP
# rather than moving the covered members onto Vulkan.
family = _windows_amd_host(rocm_gfx_target = "gfx110X", rocm_gfx_targets = ["gfx110X"])
assert ilp._should_auto_vulkan_for_amd_windows(family, UPSTREAM) is False
@pytest.mark.parametrize(
"repo", ["acme/llama.cpp-mirror", "GGML-ORG/llama.cpp", "unslothAI/llama.cpp"]
)
def test_fork_only_gfx_coverage_is_not_granted_to_other_repos(repo):
# Only the fork is planned from a manifest: resolve_simple_install_release_plans()
# compares == DEFAULT_PUBLISHED_REPO and sends everything else, mirrors and differently
# cased spellings alike, to direct_upstream_release_plan(). Granting a fork-only arch
# coverage there lands it on win-hip-radeon or CPU instead of Vulkan, so the predicate
# must gate on the fork rather than exempt one name.
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
assert ilp._should_auto_vulkan_for_amd_windows(host, repo) is True
supported = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
assert ilp._should_auto_vulkan_for_amd_windows(supported, repo) is False
@pytest.mark.parametrize("repo", [None, ""])
def test_empty_published_repo_gets_fork_coverage(repo):
# Negative control: the resolver defaults an empty repo to the fork, so the predicate
# must too, or the default install path loses its fork-only archs.
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
assert ilp._should_auto_vulkan_for_amd_windows(host, repo) is False
def test_upstream_windows_hip_targets_are_a_subset_of_the_combined_floor():
# The floor must stay a superset, else auto-Vulkan steals a host upstream builds for.
assert ilp.UPSTREAM_WINDOWS_HIP_GFX_TARGETS <= ilp.WINDOWS_HIP_PREBUILT_GFX_TARGETS
# The fork-only extras are exactly the archs that must route to Vulkan upstream.
assert ilp.WINDOWS_HIP_PREBUILT_GFX_TARGETS - ilp.UPSTREAM_WINDOWS_HIP_GFX_TARGETS == {
"gfx908",
"gfx90a",
"gfx1034",
"gfx1103",
}
def test_route_to_vulkan_prebuilt_unknown_gfx_does_not_auto_fallback():
host = _windows_amd_host(
has_rocm = True,
rocm_gfx_target = None,
rocm_gfx_targets = [],
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_route_to_vulkan_prebuilt_family_gfx_token_keeps_rocm():
host = _windows_amd_host(rocm_gfx_target = "gfx110X", rocm_gfx_targets = ["gfx110X"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_route_to_vulkan_prebuilt_gfx1103_keeps_rocm():
host = _windows_amd_host(rocm_gfx_target = "gfx1103", rocm_gfx_targets = ["gfx1103"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_route_to_vulkan_prebuilt_gfx1034_keeps_rocm():
# gfx1034 (RX 6500/6400-class) is covered by the fork's gfx103X bundle.
host = _windows_amd_host(rocm_gfx_target = "gfx1034", rocm_gfx_targets = ["gfx1034"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_route_to_vulkan_prebuilt_explicit_opt_in_on_mixed_amd(monkeypatch):
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
host = _windows_amd_host(
rocm_gfx_target = "gfx1201",
rocm_gfx_targets = ["gfx1201", "gfx803"],
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert persist == "vulkan"
assert routed.has_rocm is False
def test_direct_upstream_windows_amd_legacy_gfx_routes_to_vulkan():
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
rel = _upstream_release(
"b9925",
[
"llama-b9925-bin-win-hip-radeon-x64.zip",
"llama-b9925-bin-win-vulkan-x64.zip",
"llama-b9925-bin-win-cpu-x64.zip",
],
)
plan = ilp.direct_upstream_release_plan(rel, routed, repo, "latest")
assert persist == "vulkan"
assert plan.attempts[0].install_kind == "windows-vulkan"
def test_llama_backend_env_requests_vulkan(monkeypatch):
assert ilp.llama_backend_from_env() is None
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
assert ilp.llama_backend_from_env() == "vulkan"
assert ilp.force_vulkan_requested() is True
def test_llama_cpp_backend_env_does_not_trigger_vulkan(monkeypatch):
# UNSLOTH_LLAMA_CPP_BACKEND is a separate setup variable (auto/cpu) whose other values
# setup warns about and ignores, so reading it here would opt in behind that warning.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "vulkan")
assert ilp.llama_backend_from_env() is None
assert ilp.force_vulkan_requested() is False
def test_route_to_vulkan_prebuilt_hidden_physical_nvidia_amd_not_rerouted():
# Vulkan ignores CUDA_VISIBLE_DEVICES, so a CUDA-masked NVIDIA card next to a legacy
# AMD gfx must not auto-route: Vulkan could grab the reserved NVIDIA GPU.
host = _windows_amd_host(
rocm_gfx_target = "gfx803",
rocm_gfx_targets = ["gfx803"],
has_physical_nvidia = True,
has_usable_nvidia = False,
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_route_to_vulkan_prebuilt_explicit_opt_in_overrides_hidden_nvidia(monkeypatch):
# The physical-NVIDIA guard only gates the AMD auto path; an explicit opt-in wins.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
host = _windows_amd_host(
rocm_gfx_target = "gfx803",
rocm_gfx_targets = ["gfx803"],
has_physical_nvidia = True,
has_usable_nvidia = False,
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert persist == "vulkan"
# The gfx archs the fork's llama-prebuilt-manifest.json maps to a windows-rocm bundle.
# Static because parametrisation happens at import time and the routing tests below must
# stay offline; the guard further down re-derives it from the published manifest and fails
# on drift, so this is a checked mirror, not a second source of truth.
_FORK_WINDOWS_ROCM_GFX = (
"gfx908",
"gfx90a",
"gfx1030",
"gfx1031",
"gfx1032",
"gfx1034",
"gfx1100",
"gfx1101",
"gfx1102",
"gfx1103",
"gfx1150",
"gfx1151",
"gfx1200",
"gfx1201",
)
def _published_fork_windows_rocm_artifacts():
"""The fork's windows-rocm artifact records, read the way an install reads them.
_download_host_resolved_release is the path a default fork install takes first: it
resolves the latest release off the download host and hands llama-prebuilt-manifest.json
to parse_published_release_bundle, so these are the very records
published_rocm_choice_for_host later matches a host gfx against. No api.github.com call,
hence no shared rate-limit bucket to exhaust.
The manifest ships only as a release asset and nothing in-tree mirrors it, so this is
the one honest source. Only OSError and the release-side PrebuiltFallback become a skip,
so an offline run stays quiet while a manifest that fetches but no longer parses still
fails loudly."""
try:
resolved = ilp._download_host_resolved_release(FORK)
except OSError as exc:
pytest.skip(f"{FORK} release manifest unreachable: {exc}")
except ilp.PrebuiltFallback as exc:
pytest.skip(f"{FORK} latest release was rejected before its manifest parsed: {exc}")
if resolved is None:
pytest.skip(f"{FORK} published no resolvable latest release")
tag = resolved.bundle.release_tag
artifacts = [
artifact
for artifact in resolved.bundle.artifacts
if artifact.install_kind == "windows-rocm"
]
assert artifacts, f"{FORK}@{tag} manifest listed no windows-rocm artifacts"
return tag, artifacts
def test_windows_hip_gfx_floor_covers_every_fork_windows_rocm_bundle():
# Derived from the published manifest, not a second literal: a gfx the fork builds but
# the floor omits bypasses the fork manifest, downgrading a hash-approved windows-rocm
# bundle to an unhashed upstream Vulkan build. A newly published arch must redden here.
tag, artifacts = _published_fork_windows_rocm_artifacts()
# published_rocm_choice_for_host serves a bundle on a concrete mapped_targets entry or on
# the umbrella gfx_target itself, so both spellings must clear a floor. A gfx_target
# absent from its own mapped_targets is the family label (gfx110X); one present in it is
# a standalone bundle (gfx908) already counted as concrete.
concrete = {target.lower() for artifact in artifacts for target in artifact.mapped_targets}
labels = {
artifact.gfx_target.lower()
for artifact in artifacts
if artifact.gfx_target and artifact.gfx_target.lower() not in concrete
}
unfloored = sorted(concrete - ilp.WINDOWS_HIP_PREBUILT_GFX_TARGETS)
assert (
not unfloored
), f"auto-Vulkan would steal windows-rocm archs published in {FORK}@{tag}: {unfloored}"
unlabelled = sorted(labels - ilp.WINDOWS_ROCM_FAMILY_GFX_LABELS)
assert not unlabelled, (
f"update markers forward family labels {FORK}@{tag} publishes but "
f"WINDOWS_ROCM_FAMILY_GFX_LABELS omits: {unlabelled}"
)
# Keep the import-time tuple the offline routing tests parametrise on an exact mirror.
assert set(_FORK_WINDOWS_ROCM_GFX) == concrete, (
f"_FORK_WINDOWS_ROCM_GFX drifted from {FORK}@{tag}: "
f"gained {sorted(concrete - set(_FORK_WINDOWS_ROCM_GFX))}, "
f"lost {sorted(set(_FORK_WINDOWS_ROCM_GFX) - concrete)}"
)
@pytest.mark.parametrize("gfx", _FORK_WINDOWS_ROCM_GFX)
def test_route_to_vulkan_prebuilt_keeps_every_fork_windows_rocm_arch(gfx, monkeypatch):
# No ambient opt-in: this asserts the AUTO path leaves covered archs alone.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
host = _windows_amd_host(rocm_gfx_target = gfx, rocm_gfx_targets = [gfx])
routed, repo, tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert (repo, tag) == (FORK, "pin")
assert persist is None
def test_forwarded_gfx_does_not_undo_visible_device_auto_vulkan(monkeypatch):
# Mixed-AMD Windows host: GPU 0 = gfx1100 (HIP prebuilt exists), GPU 1 = gfx1010 (none).
# Under CUDA_VISIBLE_DEVICES=1 setup.ps1 still resolves GPU 0 and forwards gfx1100, but
# detect_host() resolved the visible gfx1010, so folding the forward in must not
# reinstate gfx1100 and install a HIP bundle the visible GPU cannot run.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx1010", rocm_gfx_targets = ["gfx1100", "gfx1010"])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx1100")
assert ilp._active_rocm_gfx_target(host) == "gfx1010"
assert host.rocm_gfx_targets == ["gfx1100", "gfx1010"]
# gfx1100 is masked off, not absent, and Vulkan does not honour the HIP mask, so the
# automatic fallback stays off and the HIP / fork path is kept.
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == FORK
assert persist is None
def test_forwarded_gfx_absent_from_probe_keeps_the_physical_hip_card(monkeypatch):
# Mixed-AMD Windows host: GPU 0 = gfx1100 (HIP prebuilt exists), GPU 1 = gfx803 (below
# the floor). CUDA_VISIBLE_DEVICES=1 reserves the gfx1100, so detect_host() picks gfx803
# as active but still reports both cards, and setup forwards a third arch the probe never
# saw (a stale env var, or name inference reading the other card). That forward selects
# the HIP target but must not delete the probe's inventory, or the floor check concludes
# no AMD GPU here reaches HIP and auto-routes to Vulkan, which ignores the HIP mask and
# enumerates the reserved gfx1100.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx1100", "gfx803"])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx900")
assert ilp._active_rocm_gfx_target(host) == "gfx900"
assert host.rocm_gfx_targets == ["gfx1100", "gfx803", "gfx900"]
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == FORK
assert persist is None
def test_forwarded_gfx_absent_from_probe_keeps_a_single_probed_hip_card(monkeypatch):
# Same rule on a single-GPU box: a stale below-floor forward over a probe-confirmed
# gfx1100 must not auto-route that machine to Vulkan.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
assert ilp._active_rocm_gfx_target(host) == "gfx803"
assert host.rocm_gfx_targets == ["gfx1100", "gfx803"]
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is False
def test_forwarded_gfx_absent_from_probe_still_allows_explicit_vulkan(monkeypatch):
# The physical-inventory rule gates the AUTO path only; naming the backend wins.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert persist == "vulkan"
def test_forwarded_gfx_on_unprobed_host_still_auto_vulkans(monkeypatch):
# Negative control: a driver-only AMD host runs no successful probe (no hipinfo, amd-smi
# suppressed), so --rocm-gfx is the ONLY source of the arch and there is no inventory to
# preserve. This is the #7357 path the feature exists for; it must still reach Vulkan.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
assert host.rocm_gfx_targets == ["gfx803"]
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is True
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert persist == "vulkan"
def test_forwarded_gfx_still_fills_an_unprobed_arch(monkeypatch):
# Negative control: on an amd-smi-only host detect_host() reports no arch, so the
# forward is the only source and must still apply.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx1151")
assert ilp._active_rocm_gfx_target(host) == "gfx1151"
assert ilp._should_auto_vulkan_for_amd_windows(host) is False
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == FORK
assert persist is None
def test_llama_backend_hip_opts_out_of_auto_vulkan(monkeypatch):
# hip names a backend, so it keeps the fork path even on an auto-fallback arch.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "hip")
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
assert ilp.force_vulkan_requested() is False
def test_explicit_backend_beats_legacy_force_vulkan(monkeypatch):
# A stale UNSLOTH_FORCE_VULKAN must not overrule UNSLOTH_LLAMA_BACKEND=rocm (== hip).
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "rocm")
assert ilp.resolved_llama_backend() == "hip"
assert ilp.force_vulkan_requested() is False
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == FORK
assert persist is None
def test_unknown_llama_backend_value_falls_through_to_legacy_flag(monkeypatch):
# An unrecognised value is ignored, not an error, so the legacy flag still works.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "banana")
assert ilp.resolved_llama_backend() is None
assert ilp.force_vulkan_requested() is False
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
assert ilp.force_vulkan_requested() is True
def test_llama_backend_flag_beats_conflicting_env(monkeypatch):
# --llama-backend is the caller's explicit request and outranks the env.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "hip")
assert ilp.force_vulkan_requested("vulkan") is True
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "pin", force_cpu = False, llama_backend = "vulkan"
)
assert repo == UPSTREAM
assert persist == "vulkan"
def _windows_arm64_host(**overrides):
defaults = dict(
system = "Windows",
machine = "ARM64",
is_windows = True,
is_linux = False,
is_macos = False,
is_x86_64 = False,
is_arm64 = True,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = False,
has_intel_gpu = False,
)
defaults.update(overrides)
return ilp.HostInfo(**defaults)
@pytest.mark.parametrize(
"env, flag",
[
({"UNSLOTH_LLAMA_BACKEND": "vulkan"}, None),
({"UNSLOTH_FORCE_VULKAN": "1"}, None),
({}, "vulkan"),
],
)
def test_vulkan_opt_in_ignored_on_windows_arm64(monkeypatch, env, flag):
# Upstream builds win-vulkan for x64 only (arm64 gets CPU + opencl-adreno), so rewriting
# the host would only swap the published arm64 bundle for the upstream CPU one.
for name, value in env.items():
monkeypatch.setenv(name, value)
host = _windows_arm64_host()
routed, repo, tag, persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "pin", force_cpu = False, llama_backend = flag
)
assert routed is host
assert (repo, tag) == (FORK, "pin")
assert persist is None
def test_vulkan_opt_in_still_routes_on_windows_x64(monkeypatch):
# Negative control for the arm64 guard: x64 keeps its Vulkan routing.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert persist == "vulkan"
def _choice(install_kind, name = "asset.zip"):
return ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = name,
url = f"https://example/{name}",
source_label = "upstream",
install_kind = install_kind,
)
@pytest.mark.parametrize("kind", ["windows-vulkan", "linux-vulkan"])
def test_persisted_llama_backend_keeps_vulkan_for_a_vulkan_bundle(kind):
assert ilp.persisted_llama_backend("vulkan", _choice(kind)) == "vulkan"
@pytest.mark.parametrize("kind", ["windows-arm64", "windows-cpu", "linux-cpu", "windows-rocm"])
def test_persisted_llama_backend_drops_vulkan_for_a_non_vulkan_bundle(kind):
# _plan_llama_phase re-asserts the marker's backend on every later update, so a Vulkan
# request that fell through to CPU must not leave a marker claiming Vulkan.
assert ilp.persisted_llama_backend("vulkan", _choice(kind)) is None
def test_persisted_llama_backend_passes_none_through():
assert ilp.persisted_llama_backend(None, _choice("windows-vulkan")) is None
def test_marker_records_no_backend_when_vulkan_fell_back_to_cpu(tmp_path):
# End to end over write_prebuilt_metadata: describe the CPU attempt that actually won,
# so the next update re-detects instead of re-asserting Vulkan forever.
checksums = ilp.ApprovedReleaseChecksums(
repo = UPSTREAM,
release_tag = "b9925",
upstream_tag = "b9925",
source_repo = UPSTREAM,
source_repo_url = f"https://github.com/{UPSTREAM}",
)
cpu = _choice("windows-arm64", "llama-b9925-bin-win-cpu-arm64.zip")
ilp.write_prebuilt_metadata(
tmp_path,
requested_tag = "latest",
llama_tag = "b9925",
release_tag = "b9925",
choice = cpu,
approved_checksums = checksums,
prebuilt_fallback_used = False,
llama_backend = "vulkan",
)
marker = json.loads((tmp_path / "UNSLOTH_PREBUILT_INFO.json").read_text())
assert marker["asset"] == "llama-b9925-bin-win-cpu-arm64.zip"
assert marker["llama_backend"] is None
vulkan = _choice("windows-vulkan", "llama-b9925-bin-win-vulkan-x64.zip")
ilp.write_prebuilt_metadata(
tmp_path,
requested_tag = "latest",
llama_tag = "b9925",
release_tag = "b9925",
choice = vulkan,
approved_checksums = checksums,
prebuilt_fallback_used = False,
llama_backend = "vulkan",
)
marker = json.loads((tmp_path / "UNSLOTH_PREBUILT_INFO.json").read_text())
assert marker["llama_backend"] == "vulkan"
# UNSLOTH_LLAMA_CPP_BACKEND (setup.sh/setup.ps1, "auto"|"cpu") and
# UNSLOTH_LLAMA_BACKEND (this module, a backend name) are different variables at
# different layers, and both accept "cpu". setup translates its own =cpu into
# --force-cpu to pin the CPU-only bundle on a GPU host, which is what keeps Intel
# iGPU Vulkan crashes away (#7213). Vulkan is opt-in here, so no trigger it adds
# may outrank that flag on any host.
_SIM_PLATFORMS = {
# WSL presents as Linux to this resolver, so it rides the Linux row.
"Linux": dict(
system = "Linux",
is_windows = False,
is_linux = True,
is_macos = False,
machine = "x86_64",
is_x86_64 = True,
is_arm64 = False,
),
"Windows": dict(
system = "Windows",
is_windows = True,
is_linux = False,
is_macos = False,
machine = "amd64",
is_x86_64 = True,
is_arm64 = False,
),
"macOS": dict(
system = "Darwin",
is_windows = False,
is_linux = False,
is_macos = True,
machine = "arm64",
is_x86_64 = False,
is_arm64 = True,
),
}
_SIM_GPUS = {
"nvidia": dict(
has_physical_nvidia = True,
has_usable_nvidia = True,
has_rocm = False,
has_intel_gpu = False,
nvidia_smi = "/usr/bin/nvidia-smi",
driver_cuda_version = "12.4",
compute_caps = ["8.9"],
),
"amd": dict(
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = True,
has_intel_gpu = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
rocm_gfx_target = "gfx803",
rocm_gfx_targets = ["gfx803"],
),
"intel": dict(
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = False,
has_intel_gpu = True,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
),
"cpu_only": dict(
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = False,
has_intel_gpu = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
),
}
def _sim_host(platform_name, gpu_name):
base = dict(visible_cuda_devices = None)
base.update(_SIM_PLATFORMS[platform_name])
base.update(_SIM_GPUS[gpu_name])
return ilp.HostInfo(**base)
@pytest.mark.parametrize("platform_name", sorted(_SIM_PLATFORMS))
@pytest.mark.parametrize("gpu_name", sorted(_SIM_GPUS))
@pytest.mark.parametrize("backend_env", [None, "vulkan", "hip", "rocm", "cpu"])
def test_forced_cpu_outranks_every_vulkan_trigger(
monkeypatch, platform_name, gpu_name, backend_env
):
"""A deliberate CPU install stays CPU on every host, whatever asks for Vulkan."""
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
if backend_env is None:
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
else:
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", backend_env)
# The legacy switch too, so a stale one cannot smuggle Vulkan past --force-cpu.
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
repo, tag = "unslothai/llama.cpp-prebuilt", "latest"
_, out_repo, _, persist = ilp._route_to_vulkan_prebuilt(
_sim_host(platform_name, gpu_name),
repo,
tag,
force_cpu = True,
llama_backend = "vulkan",
)
assert out_repo == repo, (platform_name, gpu_name, backend_env)
assert persist is None, (platform_name, gpu_name, backend_env)
def test_the_forced_cpu_guard_is_not_vacuous():
"""The same host DOES take Vulkan once the CPU pin is gone, or the check above
would pass on a resolver that had stopped routing to Vulkan entirely."""
repo, tag = "unslothai/llama.cpp-prebuilt", "latest"
_, out_repo, _, persist = ilp._route_to_vulkan_prebuilt(
_sim_host("Linux", "amd"),
repo,
tag,
force_cpu = False,
llama_backend = "vulkan",
)
assert out_repo != repo or persist == "vulkan"

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

4
studio/backend/tests/test_llama_cpp_update.py
View file

@ -473,6 +473,7 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
def _on_start(cmd):
captured["cmd"] = cmd
_write_install(
install_dir,
"b9518",
@ -480,6 +481,7 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
)
captured: dict = {}
popen_kwargs: dict = {}
_patch_installer_popen(
monkeypatch,
@ -497,6 +499,8 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
time.sleep(0.05)
assert job["state"] == "success", job
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
assert popen_kwargs["env"]["UNSLOTH_LLAMA_BACKEND"] == "vulkan"
assert "--llama-backend" in captured["cmd"] and "vulkan" in captured["cmd"]
@pytest.mark.parametrize(

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)

8
studio/backend/tests/test_research_runs_hardening.py
View file

@ -861,7 +861,9 @@ def test_stream_completion_timeout_is_absolute_despite_keepalives(monkeypatch):
def test_wall_clock_timeout_supports_python_without_asyncio_timeout(monkeypatch):
monkeypatch.delattr(research_runs.asyncio, "timeout")
# raising=False: on Python 3.10 asyncio.timeout does not exist to begin with,
# which is the very case these tests cover.
monkeypatch.delattr(research_runs.asyncio, "timeout", raising = False)
async def run():
async with research_runs._wall_clock_timeout(0.01):
@ -872,7 +874,9 @@ def test_wall_clock_timeout_supports_python_without_asyncio_timeout(monkeypatch)
def test_wall_clock_timeout_does_not_swallow_shutdown_cancellation(monkeypatch):
monkeypatch.delattr(research_runs.asyncio, "timeout")
# raising=False: on Python 3.10 asyncio.timeout does not exist to begin with,
# which is the very case these tests cover.
monkeypatch.delattr(research_runs.asyncio, "timeout", raising = False)
async def run(cleanup_started: asyncio.Event):
async with research_runs._wall_clock_timeout(0.01):

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

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

@ -776,7 +776,7 @@ def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]:
files = glob.glob("/sys/class/drm/card*/device/gpu_busy_percent")
if not files:
return None
values = [int(open(f).read().strip()) for f in files]
values = [int(open(f, encoding = "utf-8").read().strip()) for f in files]
return round(sum(values) / len(values), 1)
except Exception:
return None
@ -790,7 +790,7 @@ def _rocm_linux_sysfs_temp_c() -> Optional[float]:
files = glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input")
if not files:
return None
temps = [int(open(f).read().strip()) / 1000.0 for f in files]
temps = [int(open(f, encoding = "utf-8").read().strip()) / 1000.0 for f in files]
return round(max(temps), 1)
except Exception:
return None
@ -807,7 +807,9 @@ def _rocm_linux_sysfs_power_w() -> Optional[float]:
):
files = glob.glob(pattern)
if files:
watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files)
watts = sum(
int(open(f, encoding = "utf-8").read().strip()) / 1_000_000.0 for f in files
)
return round(watts, 1)
return None
except Exception:
@ -852,8 +854,8 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
total_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_total")
if not used_files or not total_files:
return None, None
used_bytes = sum(int(open(f).read().strip()) for f in used_files)
total_bytes = sum(int(open(f).read().strip()) for f in total_files)
used_bytes = sum(int(open(f, encoding = "utf-8").read().strip()) for f in used_files)
total_bytes = sum(int(open(f, encoding = "utf-8").read().strip()) for f in total_files)
if total_bytes == 0:
return None, None
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
@ -893,7 +895,7 @@ def _rocm_kfd_gpu_pci_ids() -> list[str]:
continue
props: dict[str, int] = {}
try:
with open(os.path.join(node_dir, "properties")) as f:
with open(os.path.join(node_dir, "properties"), encoding = "utf-8") as f:
for line in f:
parts = line.split()
if len(parts) == 2:
@ -901,7 +903,7 @@ def _rocm_kfd_gpu_pci_ids() -> list[str]:
props[parts[0]] = int(parts[1])
except ValueError:
continue
except OSError:
except (OSError, UnicodeDecodeError):
return [] # unreadable node could be a GPU: fail closed, don't shift
if props.get("simd_count", 0) <= 0:
continue # CPU node, not a GPU
@ -979,9 +981,9 @@ def _rocm_linux_sysfs_vram_by_pci_gb() -> dict[str, tuple[float, float]]:
if not bdf:
continue
try:
with open(os.path.join(dev_dir, "mem_info_vram_used")) as f:
with open(os.path.join(dev_dir, "mem_info_vram_used"), encoding = "utf-8") as f:
used_bytes = int(f.read().strip())
with open(os.path.join(dev_dir, "mem_info_vram_total")) as f:
with open(os.path.join(dev_dir, "mem_info_vram_total"), encoding = "utf-8") as f:
total_bytes = int(f.read().strip())
except (OSError, ValueError):
continue
@ -1704,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",
}
@ -2598,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),
@ -2725,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",
}

18
studio/backend/utils/llama_cpp_update.py
View file

@ -403,6 +403,7 @@ def _run_llama_phase(
pin_release_tag: Optional[str],
set_progress,
force_cpu: bool = False,
llama_backend: Optional[str] = None,
) -> dict:
"""The llama phase of a chained update: put the backend into a maintenance
state, run the installer for the latest prebuilt, then refresh caches so the
@ -454,14 +455,15 @@ def _run_llama_phase(
# updates. A natural fallback (or a legacy marker without the flag) heals to GPU (#6097).
if force_cpu:
cmd.append("--force-cpu")
if llama_backend == "vulkan":
cmd.extend(["--llama-backend", "vulkan"])
logger.info("llama update: installing", cmd = " ".join(cmd))
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
# Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm
# box would otherwise re-route and silently replace the Vulkan build.
# Re-assert it via the same env flag setup uses (mirrors
# _rocm_install_args).
if asset and "vulkan" in asset.lower():
# Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm box would
# otherwise re-route and silently replace it. Re-assert via setup's env/CLI flags.
if llama_backend == "vulkan" or (asset and "vulkan" in asset.lower()):
env["UNSLOTH_FORCE_VULKAN"] = "1"
env["UNSLOTH_LLAMA_BACKEND"] = "vulkan"
_flow.stream_installer(
cmd,
env,
@ -578,6 +580,9 @@ def _plan_llama_phase() -> dict:
from_tag = marker.get("tag") or marker.get("release_tag")
asset = marker.get("asset")
force_cpu = bool(marker.get("force_cpu"))
llama_backend = marker.get("llama_backend")
if llama_backend == "vulkan" or (asset and "vulkan" in str(asset).lower()):
llama_backend = "vulkan"
# Install exactly the release the banner offered: the installer's own
# "latest" is commit-date ordered and can lag the published_at pick
# above, reinstalling the current build in a loop (the #6219 class).
@ -621,6 +626,7 @@ def _plan_llama_phase() -> dict:
asset = (res or {}).get("asset")
# Source builds carry no forced-CPU marker, so nothing to preserve here.
force_cpu = False
llama_backend = None
# No pin: source-build detection resolves via --resolve-prebuilt latest,
# the same resolver the unpinned apply uses, so the two already agree.
pin_release_tag = None
@ -643,6 +649,7 @@ def _plan_llama_phase() -> dict:
"pin_release_tag": pin_release_tag,
"from_tag": from_tag,
"force_cpu": force_cpu,
"llama_backend": llama_backend,
}
}
@ -695,6 +702,7 @@ def start_update() -> dict:
llama_spec["pin_release_tag"],
set_progress,
force_cpu = llama_spec.get("force_cpu", False),
llama_backend = llama_spec.get("llama_backend"),
)
)
if llama_spec

8
studio/backend/utils/models/checkpoints.py
View file

@ -129,7 +129,7 @@ def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]:
if not trainer_state.exists():
return None
try:
with open(trainer_state) as f:
with open(trainer_state, encoding = "utf-8") as f:
state = json.load(f)
log_history = state.get("log_history", [])
if log_history:
@ -174,18 +174,18 @@ def scan_checkpoints(
metadata: dict = {}
try:
if adapter_config.exists():
cfg = json.loads(adapter_config.read_text())
cfg = json.loads(adapter_config.read_text(encoding = "utf-8"))
metadata["base_model"] = cfg.get("base_model_name_or_path")
metadata["peft_type"] = cfg.get("peft_type")
metadata["lora_rank"] = cfg.get("r")
elif config_file.exists():
cfg = json.loads(config_file.read_text())
cfg = json.loads(config_file.read_text(encoding = "utf-8"))
metadata["base_model"] = cfg.get("_name_or_path")
# Detect BNB quantization from config.json
if config_file.exists():
if "cfg" not in dir():
cfg = json.loads(config_file.read_text())
cfg = json.loads(config_file.read_text(encoding = "utf-8"))
quant_cfg = cfg.get("quantization_config")
if (
isinstance(quant_cfg, dict)

26
studio/backend/utils/models/model_config.py
View file

@ -631,7 +631,7 @@ def _raw_config_has_vision_config(
cache_dir = active_hf_hub_cache(),
)
)
config = json.loads(config_path.read_text())
config = json.loads(config_path.read_text(encoding = "utf-8"))
architectures = config.get("architectures") or []
model_type = config.get("model_type")
explicit_vision = (
@ -1083,7 +1083,7 @@ def _detect_audio_from_tokenizer(
]:
tok_file = snapshot / tok_path
if tok_file.exists():
tok_config = json.loads(tok_file.read_text())
tok_config = json.loads(tok_file.read_text(encoding = "utf-8"))
read_any = True
result = _check_token_patterns(tok_config)
if result:
@ -2283,7 +2283,7 @@ def scan_exported_models(
export_meta = run_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
meta = json.loads(export_meta.read_text(encoding = "utf-8"))
base_model = meta.get("base_model")
except Exception:
pass
@ -2312,7 +2312,7 @@ def scan_exported_models(
if adapter_config.exists():
export_type = "lora"
try:
cfg = json.loads(adapter_config.read_text())
cfg = json.loads(adapter_config.read_text(encoding = "utf-8"))
base_model = cfg.get("base_model_name_or_path")
except Exception:
pass
@ -2321,7 +2321,7 @@ def scan_exported_models(
export_meta = checkpoint_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
meta = json.loads(export_meta.read_text(encoding = "utf-8"))
base_model = meta.get("base_model")
except Exception:
pass
@ -2334,7 +2334,7 @@ def scan_exported_models(
export_meta = meta_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
meta = json.loads(export_meta.read_text(encoding = "utf-8"))
base_model = meta.get("base_model")
if base_model:
break
@ -2354,7 +2354,7 @@ def scan_exported_models(
outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json"
try:
if outputs_adapter_cfg.exists():
cfg = json.loads(outputs_adapter_cfg.read_text())
cfg = json.loads(outputs_adapter_cfg.read_text(encoding = "utf-8"))
base_model = cfg.get("base_model_name_or_path")
except Exception:
pass
@ -2380,7 +2380,7 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]:
adapter_config_path = checkpoint_path_obj / "adapter_config.json"
if adapter_config_path.exists():
with open(adapter_config_path, "r") as f:
with open(adapter_config_path, "r", encoding = "utf-8") as f:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
@ -2389,7 +2389,7 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]:
config_path = checkpoint_path_obj / "config.json"
if config_path.exists():
with open(config_path, "r") as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = json.load(f)
for key in ("model_name", "_name_or_path"):
base_model = config.get(key)
@ -2445,7 +2445,7 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
# adapter_config.json first
adapter_config_path = lora_path_obj / "adapter_config.json"
if adapter_config_path.exists():
with open(adapter_config_path, "r") as f:
with open(adapter_config_path, "r", encoding = "utf-8") as f:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
@ -2535,7 +2535,7 @@ def get_base_model_from_lora_identifier(
last_exc = exc
continue
try:
with open(cfg_path, "r") as f:
with open(cfg_path, "r", encoding = "utf-8") as f:
base_model = json.load(f).get("base_model_name_or_path")
except Exception as exc:
logger.warning("Could not parse adapter_config.json for '%s': %s", identifier, exc)
@ -2781,7 +2781,7 @@ class ModelConfig:
meta_path = gguf_dir / "export_metadata.json"
if meta_path.exists():
try:
meta = json.loads(meta_path.read_text())
meta = json.loads(meta_path.read_text(encoding = "utf-8"))
base = meta.get("base_model")
if base and is_vision_model(base, hf_token = hf_token):
base_is_vision = True
@ -2912,7 +2912,7 @@ class ModelConfig:
token = hf_token,
cache_dir = active_hf_hub_cache(),
)
with open(config_path, "r") as f:
with open(config_path, "r", encoding = "utf-8") as f:
adapter_config = json.load(f)
base_model = adapter_config.get("base_model_name_or_path")
if base_model:

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

2
studio/backend/utils/paths/path_utils.py
View file

@ -34,7 +34,7 @@ def _is_wsl() -> bool:
if sys.platform == "win32":
return False
try:
with open("/proc/version", "r") as f:
with open("/proc/version", "r", encoding = "utf-8") as f:
return "microsoft" in f.read().lower()
except Exception:
return False

4
studio/backend/utils/paths/storage_roots.py
View file

@ -126,7 +126,7 @@ def _xdg_user_dir(key: str) -> Path | None:
config = Path.home() / ".config" / "user-dirs.dirs"
try:
lines = config.read_text(encoding = "utf-8").splitlines()
except OSError:
except (OSError, UnicodeDecodeError):
return None
prefix = f"{key}="
for line in lines:
@ -212,7 +212,7 @@ def lmstudio_model_dirs() -> list[Path]:
settings_path = Path.home() / ".lmstudio" / "settings.json"
if settings_path.is_file():
try:
with open(settings_path) as f:
with open(settings_path, encoding = "utf-8") as f:
settings = json.load(f)
downloads = settings.get("downloadsFolder", "")
if downloads:

4
studio/backend/utils/security/consent.py
View file

@ -142,7 +142,7 @@ def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -
for name in _REMOTE_CODE_CONFIG_FILES:
p = root / name
if p.is_file():
configs.append(json.loads(p.read_text()))
configs.append(json.loads(p.read_text(encoding = "utf-8")))
return configs
from huggingface_hub import hf_hub_download
@ -164,7 +164,7 @@ def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -
# Transient/auth failure is not "absent" -> fail closed to "unknown" so
# the caller scans (a tokenizer/processor-only auto_map must not slip by).
return None
configs.append(json.loads(Path(p).read_text()))
configs.append(json.loads(Path(p).read_text(encoding = "utf-8")))
# Every config was read or a genuine 404 -> an empty list is a definitive
# "no auto_map", not "unknown".
return configs

6
studio/backend/utils/security/file_security.py
View file

@ -199,7 +199,9 @@ def _indexed_shard_paths(
inconclusive = True # transient: an index that might exist could not be read
continue
try:
weight_map = (json.loads(open(index_path).read()) or {}).get("weight_map") or {}
weight_map = (json.loads(open(index_path, encoding = "utf-8").read()) or {}).get(
"weight_map"
) or {}
for shard in weight_map.values():
shard_norm = _normalize_repo_path(str(shard))
# weight_map paths are relative to the index file's directory.
@ -326,7 +328,7 @@ def _st_load_roots(snapshot: Path) -> list:
roots = [snapshot]
try:
import json
modules = json.loads((snapshot / "modules.json").read_text())
modules = json.loads((snapshot / "modules.json").read_text(encoding = "utf-8"))
except (OSError, ValueError):
return roots # no / invalid modules.json -> snapshot root is the only load root
for module in modules or ():

4
studio/backend/utils/security/remote_code_approvals.py
View file

@ -69,7 +69,7 @@ def approval_target_key(targets) -> str:
def _load() -> dict:
"""Parsed store, or an empty skeleton on any error (fail-safe = re-prompt)."""
try:
with open(_store_path()) as f:
with open(_store_path(), encoding = "utf-8") as f:
data = json.load(f)
# Validate the shape, not just the version: a hand-edited ``subjects`` that is not a
# dict (e.g. ``[]``) would otherwise crash lookup/record instead of failing safe.
@ -92,7 +92,7 @@ def _save(data: dict) -> None:
storage_roots.ensure_dir(path.parent)
tmp = path.parent / f".{path.name}.tmp-{os.getpid()}"
try:
with open(tmp, "w") as f:
with open(tmp, "w", encoding = "utf-8") as f:
json.dump(data, f, indent = 2)
try:
os.chmod(tmp, 0o600)

28
studio/backend/utils/security/remote_code_scan.py
View file

@ -21,6 +21,8 @@ canonical scanner loads in-repo so the fallback never silently takes over.
from __future__ import annotations
import hashlib
import io
import tokenize
import importlib.util
import pathlib
import re
@ -392,6 +394,18 @@ def scan_remote_code_files(files: dict[str, str]) -> ScanResult:
return result
def _read_python_source(path) -> str:
"""Decode a .py the way Python will execute it: a PEP 263 cookie
(`# coding: cp1252`) wins, so forcing utf-8 would scan something other than
what runs."""
data = path.read_bytes()
try:
encoding = tokenize.detect_encoding(io.BytesIO(data).readline)[0]
except (SyntaxError, ValueError):
encoding = "utf-8"
return data.decode(encoding, errors = "replace")
def remote_code_fingerprint(files: dict[str, str]) -> str:
"""Stable sha256 over the (sorted) file contents, for pinning consent."""
h = hashlib.sha256()
@ -430,7 +444,7 @@ def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> d
# for an RCE gate (HIGH stays approvable; only CRITICAL hard-blocks).
for p in root.rglob("*.py"):
if p.is_file():
files[str(p.relative_to(root))] = p.read_text(errors = "replace")
files[str(p.relative_to(root))] = _read_python_source(p)
# A local config can still point auto_map at an EXTERNAL Hub repo
# (owner/name--module.Class) that executes on load, so fetch it. Every config
# that can declare auto_map is checked, so a custom processor's external code
@ -440,7 +454,7 @@ def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> d
p = root / name
if p.is_file():
try:
ext_refs |= _auto_map_refs(json.loads(p.read_text()))
ext_refs |= _auto_map_refs(json.loads(p.read_text(encoding = "utf-8")))
except Exception:
pass
if not _add_external_refs(files, ext_refs, hf_token, model_name):
@ -469,7 +483,7 @@ def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> d
f"{model_name}: config {cfg_name} could not be fetched ({exc})"
) from exc
try:
refs |= _auto_map_refs(json.loads(Path(cfg_path).read_text()))
refs |= _auto_map_refs(json.loads(Path(cfg_path).read_text(encoding = "utf-8")))
except Exception:
pass
own_refs = {fn for repo, fn in refs if repo is None}
@ -519,7 +533,7 @@ def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> d
raise RemoteCodeUnscannable(
f"{model_name}: present file {fn} could not be fetched ({exc})"
) from exc
files[fn] = Path(fp).read_text(errors = "replace")
files[fn] = _read_python_source(Path(fp))
# Code referenced from another repo executes too: scan it or fail closed.
if not _add_external_refs(files, refs, hf_token, model_name):
raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable")
@ -602,7 +616,7 @@ def external_auto_map_repos(model_name: str, hf_token: Optional[str] = None) ->
if not p.is_file():
continue
try:
refs = _auto_map_refs(json.loads(p.read_text()))
refs = _auto_map_refs(json.loads(p.read_text(encoding = "utf-8")))
except Exception:
continue
repos.update(repo for repo, _fn in refs if repo)
@ -624,7 +638,7 @@ def external_auto_map_repos(model_name: str, hf_token: Optional[str] = None) ->
except Exception:
continue
try:
refs = _auto_map_refs(json.loads(Path(cfg_path).read_text()))
refs = _auto_map_refs(json.loads(Path(cfg_path).read_text(encoding = "utf-8")))
except Exception:
continue
repos.update(repo for repo, _fn in refs if repo)
@ -701,5 +715,5 @@ def _add_external_refs(files: dict, refs, hf_token, model_name: str) -> bool:
exc,
)
return False
files[f"{repo}--{fn}"] = Path(fp).read_text(errors = "replace")
files[f"{repo}--{fn}"] = _read_python_source(Path(fp))
return True

35
studio/backend/utils/transformers_version.py
View file

@ -420,7 +420,7 @@ def _resolve_base_model(model_name: str) -> str:
adapter_cfg_path = local_path / "adapter_config.json"
if _safe_is_file(adapter_cfg_path):
try:
with open(adapter_cfg_path) as f:
with open(adapter_cfg_path, encoding = "utf-8") as f:
cfg = json.load(f)
base = cfg.get("base_model_name_or_path")
if base:
@ -437,7 +437,7 @@ def _resolve_base_model(model_name: str) -> str:
config_json_path = local_path / "config.json"
if _safe_is_file(config_json_path):
try:
with open(config_json_path) as f:
with open(config_json_path, encoding = "utf-8") as f:
cfg = json.load(f)
# Unsloth writes model_name, HF writes _name_or_path; skip a self-reference.
for _key in ("model_name", "_name_or_path"):
@ -534,14 +534,19 @@ def _adapter_base_from_hf_cache(model_name: str) -> str | None:
try:
if ref_main.is_file():
candidates.append(
repo_dir / "snapshots" / ref_main.read_text().strip() / "adapter_config.json"
repo_dir
/ "snapshots"
/ ref_main.read_text(encoding = "utf-8").strip()
/ "adapter_config.json"
)
candidates += sorted(
repo_dir.glob("snapshots/*/adapter_config.json"), key = _mtime, reverse = True
)
for cfg_path in candidates:
if cfg_path.is_file():
base = json.loads(cfg_path.read_text()).get("base_model_name_or_path")
base = json.loads(cfg_path.read_text(encoding = "utf-8")).get(
"base_model_name_or_path"
)
return base or None
except Exception as exc:
logger.debug("HF cache adapter_config.json lookup failed for '%s': %s", model_name, exc)
@ -611,7 +616,7 @@ def _check_tokenizer_config_needs_v5(model_name: str, hf_token: str | None = Non
local_tc = local_path / "tokenizer_config.json"
if _safe_is_file(local_tc):
try:
with open(local_tc) as f:
with open(local_tc, encoding = "utf-8") as f:
data = json.load(f)
tokenizer_class = data.get("tokenizer_class", "")
result = tokenizer_class in _TRANSFORMERS_5_TOKENIZER_CLASSES
@ -688,7 +693,12 @@ def _config_json_from_hf_cache(model_name: str) -> dict | None:
ref_main = repo_dir / "refs" / "main"
try:
if ref_main.is_file():
candidates.append(repo_dir / "snapshots" / ref_main.read_text().strip() / "config.json")
candidates.append(
repo_dir
/ "snapshots"
/ ref_main.read_text(encoding = "utf-8").strip()
/ "config.json"
)
# No refs/main (e.g. commit-pinned downloads): newest snapshot by mtime, not a stale
# lexicographically-first SHA, matching what the Hub cache would actually load.
candidates += sorted(
@ -696,7 +706,7 @@ def _config_json_from_hf_cache(model_name: str) -> dict | None:
)
for cfg_path in candidates:
if cfg_path.is_file():
with open(cfg_path) as f:
with open(cfg_path, encoding = "utf-8") as f:
return json.load(f)
except Exception as exc:
logger.debug("HF cache config.json lookup failed for '%s': %s", model_name, exc)
@ -721,7 +731,7 @@ def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | No
local_cfg = Path(model_name) / "config.json"
if _safe_is_file(local_cfg):
try:
with open(local_cfg) as f:
with open(local_cfg, encoding = "utf-8") as f:
cfg = json.load(f)
_config_json_cache[cache_key] = cfg
return cfg
@ -1755,7 +1765,7 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool:
metadata = di / "METADATA"
if not metadata.is_file():
continue
for line in metadata.read_text(errors = "replace").splitlines():
for line in metadata.read_text(errors = "replace", encoding = "utf-8").splitlines():
if line.startswith("Version:"):
installed_ver = line.split(":", 1)[1].strip()
if installed_ver != pkg_version:
@ -2391,7 +2401,10 @@ def _llmcompressor_shadow_is_valid() -> bool:
"""True if the shadow dir exists with a marker matching the current pin fingerprint."""
marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER
try:
return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT
return (
marker.is_file()
and marker.read_text(encoding = "utf-8").strip() == _LLMC_SHADOW_FINGERPRINT
)
except Exception:
return False
@ -2460,7 +2473,7 @@ def _ensure_venv_llmcompressor_exists() -> bool:
if result.returncode == 0:
try:
(Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER).write_text(
_LLMC_SHADOW_FINGERPRINT
_LLMC_SHADOW_FINGERPRINT, encoding = "utf-8"
)
except Exception:
pass

4
studio/backend/utils/utils.py
View file

@ -108,13 +108,13 @@ def hf_cache_snapshot_dir(model_name: str) -> Optional[Path]:
ref = repo_dir / "refs" / "main"
if not ref.is_file():
continue
commit = ref.read_text().strip()
commit = ref.read_text(encoding = "utf-8").strip()
if not commit:
continue
snapshot = repo_dir / "snapshots" / commit
if snapshot.is_dir():
return snapshot
except OSError:
except (OSError, UnicodeDecodeError):
continue
return None

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

464
studio/install_llama_prebuilt.py
View file

@ -65,6 +65,54 @@ EXIT_ERROR = 1
EXIT_BUSY = 3
EXIT_NO_SPACE = 4
# Every gfx a Windows AMD host can be served: ggml-org release.yml windows-hip GPU_TARGETS
# plus the fork's windows-rocm bundles. Must stay a superset of the manifest's windows-rocm
# mapped_targets, else auto-Vulkan steals a host the fork already builds for. Below this
# floor (e.g. gfx803 / RX 480) HIP has no prebuilt and Vulkan is the practical Windows
# llama-server backend (#7357).
WINDOWS_HIP_PREBUILT_GFX_TARGETS = frozenset(
{
"gfx908",
"gfx90a",
"gfx1030",
"gfx1031",
"gfx1032",
"gfx1034",
"gfx1100",
"gfx1101",
"gfx1102",
"gfx1103",
"gfx1150",
"gfx1151",
"gfx1200",
"gfx1201",
}
)
# Family labels forwarded by update markers / --rocm-gfx (gfx110X.zip assets).
WINDOWS_ROCM_FAMILY_GFX_LABELS = frozenset({"gfx103x", "gfx110x", "gfx120x"})
# Exactly ggml-org release.yml's windows-hip "radeon" gpu_targets. The set above adds the
# fork-only bundles (gfx1034, gfx1103, gfx908, gfx90a), served only against the fork.
UPSTREAM_WINDOWS_HIP_GFX_TARGETS = frozenset(
{
"gfx1030",
"gfx1031",
"gfx1032",
"gfx1100",
"gfx1101",
"gfx1102",
"gfx1150",
"gfx1151",
"gfx1200",
"gfx1201",
}
)
# install_kinds that really are a Vulkan bundle. A Vulkan request can still end on a CPU
# bundle (no Vulkan archive on Windows arm64; x64 falls through when it is missing or fails
# validation), so check against this to keep the marker honest (#7357).
VULKAN_INSTALL_KINDS = frozenset({"linux-vulkan", "windows-vulkan"})
# DiskPart-prompt suppression. RunAsInvoker does NOT stop amd-smi's runtime
# elevation (its manifest is asInvoker), so this is just harmless belt-and-
# suspenders for manifest-elevating tools. The real guard is _amd_smi_allowed():
@ -282,6 +330,7 @@ class HostInfo:
has_rocm: bool = False
has_intel_gpu: bool = False
rocm_gfx_target: str | None = None
rocm_gfx_targets: list[str] = field(default_factory = list)
# (major, minor) from platform.mac_ver(); None off macOS or if unparseable.
# Skips a macos prebuilt whose minimum-OS exceeds this host.
macos_version: tuple[int, int] | None = None
@ -2159,47 +2208,37 @@ def run_capture(
return result
def _pick_rocm_gfx_target(out: str) -> str | None:
"""Choose the gfx target rocminfo / hipinfo report for the active GPU.
def _list_rocm_gfx_targets(out: str) -> list[str]:
"""List gfx targets rocminfo / hipinfo report, one entry per physical GPU.
A bare first-match picked the wrong device on mixed APU + dGPU hosts
(e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100). Respect
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES so the
asset matches what HIP actually runs on. Falls back to the first GPU when
no env var is set.
rocminfo / hipinfo print the same gfx token multiple times per GPU (Name,
ISA, marketing-name). We first try to split the output on per-GPU section
headers (rocminfo: "Agent N" blocks, hipinfo: "device#N" entries) and take
exactly one gfx token per section. This gives the correct per-GPU list even
on same-arch multi-GPU hosts (e.g. two RX 7900 XTX cards) where global
dict.fromkeys dedup would collapse both cards to a single entry and make
HIP_VISIBLE_DEVICES=1 point out of range.
Falls back to insertion-order dedup when the output has no recognisable
section markers (flat gfx-string inputs, unit-test stubs, etc.).
Empty / "-1" env values mean no AMD GPU is visible to HIP: return None.
Both repeat the same gfx token per GPU (Name, ISA, marketing-name), so split on per-GPU
section headers to keep two entries on a dual same-arch host; flat strings and test stubs
fall back to insertion-order dedup.
"""
# Try to build a per-GPU token list by splitting on section boundaries.
# rocminfo sections are introduced by "Agent N" lines (optionally between
# rows of asterisks). hipinfo sections start with "device#N".
_sections = re.split(
r"(?mi)^\s*\*+\s*$\s*agent\s+\d+\s*$|\bdevice\s*#\s*\d+\b",
out,
)
if len(_sections) > 1:
# Section-based: one gfx token per GPU section preserves physical order.
_tokens: list[str] = []
for _sec in _sections[1:]:
_m = re.search(r"gfx[1-9][0-9a-z]{2,3}", _sec.lower())
if _m:
_tokens.append(_m.group(0))
else:
# Fallback: insertion-order dedup (handles flat strings / unknown formats).
_raw = re.findall(r"gfx[1-9][0-9a-z]{2,3}", out.lower())
_tokens = list(dict.fromkeys(_raw))
return _tokens
def _pick_rocm_gfx_target(out: str) -> str | None:
"""Choose the gfx target rocminfo / hipinfo report for the active GPU.
A bare first-match picked the wrong device on mixed APU + dGPU hosts (Strix Halo gfx1151
+ RX 7900 gfx1100), so honour HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES /
CUDA_VISIBLE_DEVICES; no env var means the first GPU, empty / "-1" means none (None).
"""
_tokens = _list_rocm_gfx_targets(out)
if not _tokens:
return None
@ -2407,6 +2446,7 @@ def detect_host() -> HostInfo:
has_rocm = False
rocm_gfx_target: str | None = None
rocm_gfx_targets: list[str] = []
if is_linux and not has_usable_nvidia:
# WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg
# only when HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and
@ -2444,6 +2484,7 @@ def detect_host() -> HostInfo:
if _result.returncode == 0 and _result.stdout.strip():
if _check(_result.stdout):
has_rocm = True
rocm_gfx_targets = _list_rocm_gfx_targets(_result.stdout)
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
elif is_windows and not has_usable_nvidia:
@ -2489,6 +2530,7 @@ def detect_host() -> HostInfo:
if _check(_result.stdout):
has_rocm = True
# hipinfo reports "gcnArchName: gfx1100" -- extract if present
rocm_gfx_targets = _list_rocm_gfx_targets(_result.stdout)
rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout)
break
# Note: amdhip64.dll presence alone is NOT treated as GPU evidence
@ -2504,11 +2546,11 @@ def detect_host() -> HostInfo:
if is_linux:
for _vendor_file in glob.glob("/sys/class/drm/card*/device/vendor"):
try:
with open(_vendor_file) as _vf:
with open(_vendor_file, encoding = "utf-8") as _vf:
if _vf.read().strip().lower() == "0x8086":
has_intel_gpu = True
break
except OSError:
except (OSError, UnicodeDecodeError):
continue
elif is_windows:
# Registry first (in-process; see windows_intel_gpu_in_registry).
@ -2551,6 +2593,7 @@ def detect_host() -> HostInfo:
has_rocm = has_rocm,
has_intel_gpu = has_intel_gpu,
rocm_gfx_target = rocm_gfx_target,
rocm_gfx_targets = rocm_gfx_targets,
macos_version = macos_version,
)
@ -2573,12 +2616,12 @@ def _apply_host_overrides(
force_cpu: bool = False,
) -> HostInfo:
"""Fold setup.sh/setup.ps1's forwarded detection into the host profile.
A forwarded gfx (--rocm-gfx or UNSLOTH_ROCM_GFX_ARCH) is authoritative and
implies ROCm: the installer's own hipinfo/amd-smi probe can miss the arch on
amd-smi-only hosts or when setup inferred it from the GPU name, leaving
rocm_gfx_target None and no per-gfx ROCm prebuilt selected. force_cpu is the
opposite explicit signal (arm64 Linux GPU host whose source build failed):
drop GPU attributes so the CPU prebuilt for this OS/arch is selected."""
A forwarded gfx (--rocm-gfx or UNSLOTH_ROCM_GFX_ARCH) implies ROCm and fills the gap
where our own hipinfo/amd-smi probe misses the arch (amd-smi-only hosts, or setup
inferring it from the GPU name), leaving no per-gfx ROCm prebuilt selected; it stays
authoritative except for the two advisory shapes narrowed below. force_cpu is the
opposite explicit signal (arm64 Linux GPU host whose source build failed): drop GPU
attributes so the CPU prebuilt for this OS/arch is selected."""
if force_cpu:
return dataclasses_replace(
host,
@ -2586,11 +2629,43 @@ def _apply_host_overrides(
has_physical_nvidia = False,
has_rocm = False,
rocm_gfx_target = None,
rocm_gfx_targets = [],
has_intel_gpu = False,
)
gfx = _normalize_forwarded_gfx(override_rocm_gfx)
if gfx:
return dataclasses_replace(host, has_rocm = True, rocm_gfx_target = gfx)
# setup.ps1's pick is not fully visible-device aware (neither branch reads
# CUDA_VISIBLE_DEVICES; amd-smi matches a bare integer only, so "1,0" falls back to
# GPU 0), while _pick_rocm_gfx_target() honours all three vars with HIP's semantics.
# So keep a probed active arch when the forward is only advisory, else
# _should_auto_vulkan_for_amd_windows() reads a HIP-supported GPU the user masked
# off and installs an unusable HIP bundle instead of Vulkan. Advisory means:
# * another GPU the probe saw ON THIS HOST, i.e. setup picked a different card;
# * a family label (gfx110X), a bundle name the update path emits from the marker
# asset and never a real arch -- it would upgrade an in-generation-but-unbuilt
# GPU (gfx1033) into a bundle it must not be served.
# Anything else the probe never reported is an operator override for a host whose
# arch the probe gets wrong or stale, exactly what --rocm-gfx documents, so it stays
# authoritative. UNSLOTH_ROCM_GFX_ARCH also still wins.
_manual = _normalize_forwarded_gfx(os.environ.get("UNSLOTH_ROCM_GFX_ARCH"))
_physical = _host_rocm_gfx_targets(host)
_active = _active_rocm_gfx_target(host)
_advisory = gfx in _physical or gfx in WINDOWS_ROCM_FAMILY_GFX_LABELS
if gfx != _manual and _active and gfx != _active and _advisory:
return dataclasses_replace(host, has_rocm = True)
return dataclasses_replace(
host,
has_rocm = True,
rocm_gfx_target = gfx,
# ADD the forwarded arch to the probe's per-GPU list, never replace it: that
# list is the PHYSICAL inventory _should_auto_vulkan_for_amd_windows() reads,
# and a forward says which GPU HIP should target, not which cards exist.
# Dropping a probe-confirmed GPU would let a stale below-floor forward
# auto-route a box with a HIP-capable card to Vulkan, which then enumerates
# that card regardless of HIP_VISIBLE_DEVICES. An empty probe still yields
# [gfx], so the driver-only host the forward exists for keeps auto-Vulkan.
rocm_gfx_targets = list(dict.fromkeys([*_physical, gfx])),
)
if override_has_rocm and not host.has_rocm:
return dataclasses_replace(host, has_rocm = True)
return host
@ -3050,7 +3125,7 @@ def _detect_host_rocm_version() -> tuple[int, int] | None:
os.path.join(rocm_root, "lib", "rocm_version"),
):
try:
with open(path) as fh:
with open(path, encoding = "utf-8") as fh:
parts = fh.read().strip().split("-")[0].split(".")
# Explicit length guard avoids relying on the broad except
# below to swallow IndexError when the version file contains
@ -4264,7 +4339,7 @@ def free_local_port() -> int:
def read_log_excerpt(log_path: Path, *, max_lines: int = 60) -> str:
try:
content = log_path.read_text(encoding = "utf-8", errors = "replace")
except FileNotFoundError:
except (FileNotFoundError, UnicodeDecodeError):
return ""
return "\n".join(content.splitlines()[-max_lines:])
@ -4680,7 +4755,7 @@ def _wsl_system_rocm_lib_dirs() -> list[str]:
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
if "microsoft" not in fh.read().lower():
return []
except OSError:
except (OSError, UnicodeDecodeError):
return []
out: list[str] = []
for d in ("/opt/rocm/lib", "/opt/rocm/lib64"):
@ -5464,6 +5539,18 @@ def _fork_manifest_release_plans(
raise PrebuiltFallback("no installable published llama.cpp releases were found")
def persisted_llama_backend(llama_backend: str | None, choice: AssetChoice) -> str | None:
"""The backend to record for an install that actually landed ``choice``.
A Vulkan request can end on a non-Vulkan bundle (no upstream Vulkan archive for Windows
arm64; x64 falls through to win-cpu-x64 when it is missing or fails validation), and
recording "vulkan" there would make the updater re-assert a backend that was never
installed. Mirrors force_cpu: persist the real outcome only."""
if llama_backend == "vulkan" and choice.install_kind not in VULKAN_INSTALL_KINDS:
return None
return llama_backend
def write_prebuilt_metadata(
install_dir: Path,
*,
@ -5474,6 +5561,7 @@ def write_prebuilt_metadata(
approved_checksums: ApprovedReleaseChecksums,
prebuilt_fallback_used: bool,
force_cpu: bool = False,
llama_backend: str | None = None,
) -> None:
source_asset_name, source_sha256 = selected_source_archive_metadata(
approved_checksums,
@ -5502,6 +5590,9 @@ def write_prebuilt_metadata(
# so a forced CPU install is not re-routed to a GPU bundle (#7213). An automatic
# --cpu-fallback (e.g. arm64 GPU-build recovery) stays False so it can heal to GPU.
"force_cpu": force_cpu,
# Deliberate or auto-selected Vulkan backend (#7357); the updater re-asserts it so
# AMD hosts are not swapped back to HIP. Dropped if the winning attempt was not Vulkan.
"llama_backend": persisted_llama_backend(llama_backend, choice),
"asset_sha256": choice.expected_sha256,
"source": choice.source_label,
# Binary-side repo/tag for non-fork sources (e.g. the ggml-org upstream
@ -5526,7 +5617,9 @@ def write_prebuilt_metadata(
"prebuilt_fallback_used": prebuilt_fallback_used,
"installed_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
}
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(metadata, indent = 2) + "\n")
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(
json.dumps(metadata, indent = 2) + "\n", encoding = "utf-8"
)
def sync_marker_force_cpu(install_dir: Path, persist_force_cpu: bool) -> None:
@ -5537,16 +5630,33 @@ def sync_marker_force_cpu(install_dir: Path, persist_force_cpu: bool) -> None:
GPU/Vulkan bundle that revives the crash (#7213)."""
marker_path = install_dir / "UNSLOTH_PREBUILT_INFO.json"
try:
marker = json.loads(marker_path.read_text())
marker = json.loads(marker_path.read_text(encoding = "utf-8"))
except (OSError, ValueError):
return
if not isinstance(marker, dict) or bool(marker.get("force_cpu")) == persist_force_cpu:
return
marker["force_cpu"] = persist_force_cpu
marker_path.write_text(json.dumps(marker, indent = 2) + "\n")
marker_path.write_text(json.dumps(marker, indent = 2) + "\n", encoding = "utf-8")
log(f"existing install reused; recorded force_cpu={persist_force_cpu} from this run")
def sync_marker_llama_backend(install_dir: Path, llama_backend: str | None) -> None:
"""Sync the persisted llama.cpp backend when the bundle is reused unchanged."""
marker_path = install_dir / "UNSLOTH_PREBUILT_INFO.json"
try:
marker = json.loads(marker_path.read_text())
except (OSError, ValueError):
return
if not isinstance(marker, dict) or marker.get("llama_backend") == llama_backend:
return
if llama_backend is None:
marker.pop("llama_backend", None)
else:
marker["llama_backend"] = llama_backend
marker_path.write_text(json.dumps(marker, indent = 2) + "\n")
log(f"existing install reused; recorded llama_backend={llama_backend!r} from this run")
def expected_install_fingerprint(
*,
llama_tag: str,
@ -5839,6 +5949,7 @@ def validate_prebuilt_choice(
prebuilt_fallback_used: bool,
quantized_path: Path,
force_cpu: bool = False,
llama_backend: str | None = None,
) -> tuple[Path, Path]:
source_repo, source_ref, source_archive, exact_source = preferred_source_archive(
approved_checksums, llama_tag
@ -5880,6 +5991,7 @@ def validate_prebuilt_choice(
approved_checksums = approved_checksums,
prebuilt_fallback_used = prebuilt_fallback_used,
force_cpu = force_cpu,
llama_backend = llama_backend,
)
# Hashless external prebuilts are not in the approved-sha256
# manifest and rely on the functional smoke test as their only integrity gate,
@ -5967,6 +6079,7 @@ def validate_prebuilt_attempts(
initial_fallback_used: bool = False,
existing_install_dir: Path | None = None,
force_cpu: bool = False,
llama_backend: str | None = None,
) -> tuple[AssetChoice, Path, bool]:
attempt_list = list(attempts)
if not attempt_list:
@ -6028,6 +6141,7 @@ def validate_prebuilt_attempts(
prebuilt_fallback_used = tried_fallback,
quantized_path = quantized_path,
force_cpu = force_cpu,
llama_backend = llama_backend,
)
except Exception as exc:
remove_tree(staging_dir)
@ -6053,12 +6167,43 @@ def validate_prebuilt_attempts(
raise PrebuiltFallback("no prebuilt bundle passed validation")
def force_vulkan_requested() -> bool:
"""Whether UNSLOTH_FORCE_VULKAN opts this host into the Vulkan llama.cpp
prebuilt instead of its detected CUDA/ROCm backend (e.g. so an AMD user can
run the Vulkan build for inference). Scoped to the llama.cpp backend; the
torch/training stack installs separately and still sees the real GPU.
def _normalized_llama_backend(value: str | None) -> str | None:
if not value:
return None
backend = value.strip().lower()
if backend in {"vulkan", "hip", "rocm", "cpu"}:
return "hip" if backend == "rocm" else backend
return None
def llama_backend_from_env() -> str | None:
"""Read an explicit llama.cpp backend preference from the environment.
Only ``UNSLOTH_LLAMA_BACKEND`` is honored. ``UNSLOTH_LLAMA_CPP_BACKEND`` is a separate
setup variable meaning ``auto``/``cpu`` (not a backend name) that setup warns about and
otherwise ignores, so reading it here would force Vulkan behind that warning.
"""
return _normalized_llama_backend(os.environ.get("UNSLOTH_LLAMA_BACKEND"))
def resolved_llama_backend(llama_backend: str | None = None) -> str | None:
"""The explicit backend for this run: --llama-backend, else the env var. None when
neither is set or the value is not a backend name we know."""
return _normalized_llama_backend(llama_backend) or llama_backend_from_env()
def force_vulkan_requested(llama_backend: str | None = None) -> bool:
"""Whether this run should install the upstream Vulkan llama.cpp prebuilt.
Triggered by ``UNSLOTH_LLAMA_BACKEND=vulkan``, legacy ``UNSLOTH_FORCE_VULKAN``, or
``--llama-backend vulkan``. Scoped to the llama.cpp backend; the torch/training stack
installs separately and still sees the real GPU.
"""
backend = resolved_llama_backend(llama_backend)
if backend is not None:
# Authoritative, so =hip is a real opt-out a stale UNSLOTH_FORCE_VULKAN cannot
# overrule.
return backend == "vulkan"
return os.environ.get("UNSLOTH_FORCE_VULKAN", "").strip().lower() in (
"1",
"true",
@ -6066,6 +6211,117 @@ def force_vulkan_requested() -> bool:
)
def _host_rocm_gfx_targets(host: HostInfo) -> list[str]:
if host.rocm_gfx_targets:
return [target.lower() for target in host.rocm_gfx_targets]
if host.rocm_gfx_target:
return [host.rocm_gfx_target.lower()]
return []
def _active_rocm_gfx_target(host: HostInfo) -> str | None:
"""The gfx HIP will run on (visible-device aware), not every physical GPU."""
if host.rocm_gfx_target:
return host.rocm_gfx_target.lower().strip()
return None
def _hip_visible_device_mask_set() -> bool:
"""Whether a HIP visible-device mask is in force for this process.
The Windows arch probe is hipinfo, itself a HIP application, so under a mask it
enumerates the VISIBLE devices, not the physical ones. Presence is the whole test: a
partial mask leaves the inventory unknowable, and an all-hiding "" / "-1" is the
strongest form of that, not an exemption, since --rocm-gfx can still supply an arch
(setup infers it from the display adapter, which no HIP mask touches) and would
auto-route a host on which the user hid every AMD GPU. Reads the same three vars as
_pick_rocm_gfx_target, so the two cannot disagree about the host."""
return any(
os.environ.get(_env) is not None
for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES")
)
def _windows_hip_gfx_targets(published_repo: str | None) -> frozenset[str]:
"""gfx targets the Windows HIP bundle of ``published_repo`` is actually built for.
The combined floor above includes the FORK's windows-rocm bundles, and only the fork is
planned from a manifest: resolve_simple_install_release_plans() sends every other
--published-repo (ggml-org, but equally any mirror of upstream-standard assets) to
direct_upstream_release_plan(), whose AMD branch offers win-hip-radeon then CPU and never
Vulkan. Answering "supported" there for a fork-only arch would silently land it on
HIP/CPU instead of the Vulkan bundle that would actually run, so gate on the fork rather
than exempting one repo, mirroring that dispatch exactly, spelling included: an empty
value defaults to the fork, and a differently cased repo really does take the upstream
path and must be answered with upstream coverage."""
if (published_repo or DEFAULT_PUBLISHED_REPO) == DEFAULT_PUBLISHED_REPO:
return WINDOWS_HIP_PREBUILT_GFX_TARGETS
return UPSTREAM_WINDOWS_HIP_GFX_TARGETS
def _gfx_is_windows_hip_supported(gfx: str, published_repo: str | None = None) -> bool:
token = gfx.lower().strip()
if token in WINDOWS_ROCM_FAMILY_GFX_LABELS:
# A bundle name, not an arch, so the concrete GPU is unknown here. The fork builds
# every member; upstream builds all but gfx1034 / gfx1103. Answering "unsupported"
# to cover that pair would move gfx1030..1032 / gfx1100..1102 off a working HIP
# build onto Vulkan for a member the label cannot identify, so HIP serves it either
# way. A concrete arch below still answers per repo, which is where gfx1034 and
# gfx1103 do reach Vulkan.
return True
return token in _windows_hip_gfx_targets(published_repo)
def _host_has_windows_hip_prebuilt_gfx(host: HostInfo, published_repo: str | None = None) -> bool:
active = _active_rocm_gfx_target(host)
if not active:
return False
return _gfx_is_windows_hip_supported(active, published_repo)
def _should_auto_vulkan_for_amd_windows(host: HostInfo, published_repo: str | None = None) -> bool:
"""True when NO AMD GPU on the host reaches the Windows HIP prebuilt floor."""
active = _active_rocm_gfx_target(host)
if not active:
# ROCm confirmed but gfx unknown (--has-rocm only): keep the HIP / fork / source path.
return False
if not (
host.is_windows
and host.has_rocm
# PHYSICAL, not merely usable: Vulkan ignores CUDA_VISIBLE_DEVICES and would
# enumerate a card hidden by it. Same gate as the Intel auto path below.
and not host.has_physical_nvidia
):
return False
# Judge every PHYSICAL AMD gfx, not just the active one. The visible-device vars choose
# `active`, but the Vulkan runtime honours none of them (it enumerates through
# GGML_VK_VISIBLE_DEVICES and Vulkan ordinals), so masking down to a below-floor card
# must not route the install to Vulkan: the installed backend would happily enumerate
# the HIP-capable card the user deliberately hid, possibly one reserved for another
# workload. Auto-fall back only when no AMD device on the box can be exposed to HIP; an
# explicit vulkan opt-in is unaffected.
#
# Under a mask the probe cannot supply that inventory at all (hipinfo sees only visible
# devices), so "no AMD GPU here reaches the floor" is unprovable and guessing wrong is
# the same reserved-card handover. An all-hiding "" / "-1" is included: a forwarded
# --rocm-gfx still reconstructs an arch there, and auto-routing would then hand Vulkan
# every AMD GPU the user hid. A mask is only ever set deliberately, and the driver-only
# single-GPU host this fallback exists for does not set one.
if _hip_visible_device_mask_set():
return False
targets = list(dict.fromkeys([*_host_rocm_gfx_targets(host), active]))
return not any(_gfx_is_windows_hip_supported(target, published_repo) for target in targets)
def _has_no_vulkan_prebuilt(host: HostInfo) -> bool:
"""Platforms that ship no Vulkan prebuilt at all, so routing there is pointless.
Upstream builds win-vulkan for x64 only; Windows arm64 gets CPU plus opencl-adreno, so
rewriting it to Vulkan-only would just swap the published bundle for the upstream CPU
one. macOS is handled separately (Metal)."""
return host.is_windows and host.is_arm64
def _vulkan_only_host(host: HostInfo) -> HostInfo:
"""Rewrite ``host`` so the asset selectors take their Vulkan branch.
@ -6079,52 +6335,79 @@ def _vulkan_only_host(host: HostInfo) -> HostInfo:
has_usable_nvidia = False,
has_physical_nvidia = False,
has_rocm = False,
rocm_gfx_target = None,
rocm_gfx_targets = [],
has_intel_gpu = True,
)
def _route_to_vulkan_prebuilt(
host: HostInfo, published_repo: str, published_release_tag: str, *, force_cpu: bool
) -> tuple[HostInfo, str, str]:
host: HostInfo,
published_repo: str,
published_release_tag: str,
*,
force_cpu: bool,
llama_backend: str | None = None,
) -> tuple[HostInfo, str, str, str | None]:
"""Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes
from UPSTREAM_REPO. Two triggers route here, both suppressed when a CPU flag
(--cpu-fallback or --force-cpu, folded into force_cpu) wins:
* UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend;
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose
of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
Applied by BOTH the install path and the --resolve-prebuilt probe so the
"is a prebuilt available" answer matches what actually gets installed.
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes from
UPSTREAM_REPO. Three triggers route here, all suppressed when a CPU flag (--cpu-fallback
or --force-cpu, folded into force_cpu) wins:
* ``UNSLOTH_LLAMA_BACKEND=vulkan`` / ``UNSLOTH_FORCE_VULKAN`` / ``--llama-backend
vulkan`` forces Vulkan over the detected CUDA/ROCm backend;
* Windows AMD with no HIP-prebuilt gfx arch auto-falls back to Vulkan (#7357);
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm, the purpose of the
has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
Applied by BOTH the install path and the --resolve-prebuilt probe so the "is a prebuilt
available" answer matches what actually gets installed.
Returns the (possibly rewritten) host, repo, and release tag.
Returns the (possibly rewritten) host, repo, release tag, and a backend to persist in
the install marker when updates must re-assert Vulkan.
"""
forced = force_vulkan_requested()
# Gate auto-routing on no PHYSICAL NVIDIA, not merely no usable one: a mixed
# NVIDIA+Intel host that hides NVIDIA with CUDA_VISIBLE_DEVICES=""/-1 keeps
# has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores
# CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the
# reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides.
forced = force_vulkan_requested(llama_backend)
# Auto-fall back only when the run named no backend: an explicit hip/cpu is the opt-out.
explicit_backend = resolved_llama_backend(llama_backend)
auto_no_hip = explicit_backend is None and _should_auto_vulkan_for_amd_windows(
host, published_repo
)
# No PHYSICAL NVIDIA, not merely no usable one: Vulkan ignores CUDA_VISIBLE_DEVICES, so
# auto-routing a host that hides its NVIDIA card would let it grab the reserved GPU.
auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
if force_cpu or not (forced or auto_intel):
return host, published_repo, published_release_tag
if force_cpu or not (forced or auto_intel or auto_no_hip):
return host, published_repo, published_release_tag, None
if host.is_macos:
if forced:
log(
"UNSLOTH_FORCE_VULKAN is set but ignored on macOS "
"UNSLOTH_LLAMA_BACKEND=vulkan is set but ignored on macOS "
"(Metal is used; there is no Vulkan prebuilt)"
)
return host, published_repo, published_release_tag
if forced:
return host, published_repo, published_release_tag, None
if _has_no_vulkan_prebuilt(host):
if forced:
log(
"Vulkan llama.cpp backend requested but ignored on Windows arm64 "
"(upstream ships no Vulkan arm64 prebuilt); keeping the published bundle"
)
return host, published_repo, published_release_tag, None
if auto_no_hip:
active = _active_rocm_gfx_target(host) or "unknown"
log(
"UNSLOTH_FORCE_VULKAN is set; installing the upstream Vulkan "
"llama.cpp prebuilt instead of the detected GPU backend"
"Active AMD GPU arch is not supported by the Windows HIP prebuilt "
f"({active}); installing the upstream Vulkan llama.cpp prebuilt instead"
)
# Forcing may override a detected NVIDIA/ROCm host, so normalize it to
# Vulkan-only; an auto-detected Intel host already is.
host = _vulkan_only_host(host)
persist_backend = "vulkan"
elif forced:
log(
"Vulkan llama.cpp backend requested; installing the upstream Vulkan "
"prebuilt instead of the detected GPU backend"
)
host = _vulkan_only_host(host)
persist_backend = "vulkan"
else:
log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt")
persist_backend = None
# Swapping the fork for upstream invalidates a fork release pin: the two use
# different tag namespaces (fork b9596-mix-<sha> vs upstream b9596), so a
# pinned fork tag would make the upstream resolver query a nonexistent
@ -6133,7 +6416,7 @@ def _route_to_vulkan_prebuilt(
# (repo unchanged here) is preserved.
if published_repo != UPSTREAM_REPO:
published_release_tag = ""
return host, UPSTREAM_REPO, published_release_tag
return host, UPSTREAM_REPO, published_release_tag, persist_backend
def diffusion_visual_server_backfill_needed(
@ -6297,6 +6580,7 @@ def install_prebuilt(
override_rocm_gfx: str | None = None,
force_cpu: bool = False,
persist_force_cpu: bool = False,
llama_backend: str | None = None,
instruction_cleanup_root: Path | None = None,
) -> None:
# force_cpu drops GPU detection (mechanism, both --cpu-fallback and --force-cpu);
@ -6308,8 +6592,12 @@ def install_prebuilt(
override_rocm_gfx = override_rocm_gfx,
force_cpu = force_cpu,
)
host, published_repo, published_release_tag = _route_to_vulkan_prebuilt(
host, published_repo, published_release_tag, force_cpu = force_cpu
host, published_repo, published_release_tag, persist_llama_backend = _route_to_vulkan_prebuilt(
host,
published_repo,
published_release_tag,
force_cpu = force_cpu,
llama_backend = llama_backend,
)
choice: AssetChoice | None = None
cleanup_root = install_dir if instruction_cleanup_root is None else instruction_cleanup_root
@ -6354,6 +6642,10 @@ def install_prebuilt(
# Reused bundle is unchanged, but a fresh --force-cpu still must be
# recorded so the updater re-asserts it (#7213).
sync_marker_force_cpu(install_dir, persist_force_cpu)
sync_marker_llama_backend(
install_dir,
persisted_llama_backend(persist_llama_backend, current.attempts[0]),
)
return
with scratch_dir("unsloth-llama-prebuilt-") as work_dir:
probe_path = work_dir / "stories260K.gguf"
@ -6374,6 +6666,10 @@ def install_prebuilt(
f"{plan.release_tag} upstream_tag={plan.llama_tag}; skipping reinstall"
)
sync_marker_force_cpu(install_dir, persist_force_cpu)
sync_marker_llama_backend(
install_dir,
persisted_llama_backend(persist_llama_backend, choice),
)
return
log(
"selected "
@ -6396,6 +6692,7 @@ def install_prebuilt(
existing_install_dir = install_dir,
# Persist only the deliberate choice, not a transient fallback.
force_cpu = persist_force_cpu,
llama_backend = persist_llama_backend,
)
except ExistingInstallSatisfied:
return
@ -6517,6 +6814,16 @@ def parse_args() -> argparse.Namespace:
"bundle that would revive the Intel iGPU crash (#7213)."
),
)
parser.add_argument(
"--llama-backend",
choices = ("vulkan",),
help = (
"Force the llama.cpp prebuilt backend. vulkan installs the upstream Vulkan "
"bundle and records the choice so Studio updates keep it; ignored on hosts "
"with no Vulkan prebuilt (macOS, Windows arm64). "
"Same effect as UNSLOTH_LLAMA_BACKEND=vulkan / UNSLOTH_FORCE_VULKAN=1."
),
)
resolve_group = parser.add_mutually_exclusive_group()
resolve_group.add_argument(
"--resolve-llama-tag",
@ -6681,8 +6988,12 @@ def main() -> int:
)
# Same Vulkan routing the install path applies, so the probe's answer
# matches what would install (an Intel/forced-Vulkan host -> upstream).
host, repo, release_tag = _route_to_vulkan_prebuilt(
host, args.published_repo, args.published_release_tag or "", force_cpu = _cpu_mechanism
host, repo, release_tag, _persist_llama_backend = _route_to_vulkan_prebuilt(
host,
args.published_repo,
args.published_release_tag or "",
force_cpu = _cpu_mechanism,
llama_backend = args.llama_backend,
)
try:
_requested, plans = resolve_simple_install_release_plans(
@ -6724,6 +7035,7 @@ def main() -> int:
# updater re-asserts it. --cpu-fallback stays transient and heals to GPU.
force_cpu = args.cpu_fallback or args.force_cpu,
persist_force_cpu = args.force_cpu,
llama_backend = args.llama_backend,
instruction_cleanup_root = install_arg.absolute(),
)
return EXIT_SUCCESS

10
studio/install_node_prebuilt.py
View file

@ -535,7 +535,9 @@ def install_lock(lock_path: Path) -> Iterator[None]:
break
except FileExistsError:
try:
raw = lock_path.read_text().strip()
# errors="replace" so an undecodable lock reaches the int()
# below and is treated as a stale PID, not retried forever.
raw = lock_path.read_text(encoding = "utf-8", errors = "replace").strip()
except FileNotFoundError:
continue
stale = False
@ -660,7 +662,7 @@ def write_metadata(install_dir: Path, *, version: str, asset: str, sha256: str)
"asset": asset,
"sha256": sha256,
}
metadata_path(install_dir).write_text(json.dumps(payload, indent = 2) + "\n")
metadata_path(install_dir).write_text(json.dumps(payload, indent = 2) + "\n", encoding = "utf-8")
def load_metadata(install_dir: Path) -> dict | None:
@ -668,8 +670,8 @@ def load_metadata(install_dir: Path) -> dict | None:
if not path.exists():
return None
try:
data = json.loads(path.read_text())
except (json.JSONDecodeError, OSError):
data = json.loads(path.read_text(encoding = "utf-8"))
except (json.JSONDecodeError, OSError, UnicodeDecodeError):
return None
return data if isinstance(data, dict) else None

178
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(
@ -503,7 +537,7 @@ def _detect_rocm_version() -> tuple[int, int] | None:
os.path.join(rocm_root, "lib", "rocm_version"),
):
try:
with open(path) as fh:
with open(path, encoding = "utf-8") as fh:
parts = fh.read().strip().split("-")[0].split(".")
# Explicit length guard: don't rely on the broad except below to
# swallow IndexError on a single-component version (e.g. "6\n").
@ -776,7 +810,7 @@ def _linux_amd_gfx_from_cpuinfo() -> "str | None":
"""Infer gfx arch from /proc/cpuinfo on integrated AMD APUs (Strix Halo/Point)."""
try:
text = Path("/proc/cpuinfo").read_text(encoding = "utf-8", errors = "replace")
except OSError:
except (OSError, UnicodeDecodeError):
return None
if re.search(r"Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo", text, re.IGNORECASE):
return "gfx1151"
@ -828,7 +862,7 @@ def _is_wsl() -> bool:
try:
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
return "microsoft" in fh.read().lower()
except OSError:
except (OSError, UnicodeDecodeError):
return False
@ -852,11 +886,11 @@ def _linux_amd_display_device_present() -> bool:
try:
for dev in Path("/sys/bus/pci/devices").iterdir():
try:
if (dev / "vendor").read_text().strip() != "0x1002":
if (dev / "vendor").read_text(encoding = "utf-8").strip() != "0x1002":
continue
if (dev / "class").read_text().strip().startswith("0x03"):
if (dev / "class").read_text(encoding = "utf-8").strip().startswith("0x03"):
return True
except OSError:
except (OSError, UnicodeDecodeError):
continue
except OSError:
pass
@ -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"
@ -1067,9 +1129,9 @@ def _has_rocm_gpu() -> bool:
for entry in os.listdir(kfd_nodes):
gpu_id_path = os.path.join(kfd_nodes, entry, "gpu_id")
try:
with open(gpu_id_path) as fh:
with open(gpu_id_path, encoding = "utf-8") as fh:
gpu_id = fh.read().strip()
except OSError:
except (OSError, UnicodeDecodeError):
continue
if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node
continue
@ -1079,9 +1141,9 @@ def _has_rocm_gpu() -> bool:
# false positive (e.g. NVIDIA open-driver KFD nodes lacking it).
props_path = os.path.join(kfd_nodes, entry, "properties")
try:
with open(props_path) as fh:
with open(props_path, encoding = "utf-8") as fh:
props = fh.read()
except OSError:
except (OSError, UnicodeDecodeError):
continue # can't confirm vendor -- skip
if not re.search(r"\bvendor_id\s+4098\b", props):
continue
@ -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

6
studio/prebuilt_core.py
View file

@ -1078,7 +1078,9 @@ def install_lock(lock_path: Path) -> Iterator[None]:
except FileExistsError:
stale = False
try:
raw = lock_path.read_text().strip()
# errors="replace" so an undecodable lock reaches the int()
# below and is treated as a corrupt PID, not retried forever.
raw = lock_path.read_text(encoding = "utf-8", errors = "replace").strip()
except FileNotFoundError:
# Lock vanished between our open and read -- retry
continue
@ -2070,7 +2072,7 @@ def load_prebuilt_metadata(ops: ModuleOps, install_dir: Path) -> dict[str, Any]
return None
try:
payload = json.loads(path.read_text(encoding = "utf-8"))
except (json.JSONDecodeError, OSError):
except (json.JSONDecodeError, OSError, UnicodeDecodeError):
return None
return payload if isinstance(payload, dict) else None

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

32
tests/sh/test_torch_constraint.sh
View file

@ -94,28 +94,34 @@ echo "=== Structural: TORCH_CONSTRAINT in install.sh ==="
_SH_CONTENT=$(cat "$INSTALL_SH")
_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
# Each hardware branch assigns its own triple, so counting every occurrence made
# adding a branch (gfx906 in #7354) a test edit. The default is the one assigned at
# top level; a branch's is always indented, so anchor on that instead of counting.
_count=$(grep -c '^TORCH_CONSTRAINT="torch>=2.4,<2.11.0"$' "$INSTALL_SH" || true)
assert_eq "default TORCH_CONSTRAINT assignment exists" "1" "$_count"
_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "tightened TORCH_CONSTRAINT assignment exists" "1" "$_count"
_has=$([ "$_count" -ge 1 ] && echo "yes" || echo "no")
assert_eq "tightened TORCH_CONSTRAINT assignment exists" "yes" "$_has"
_count=$(grep -c '"\$TORCH_CONSTRAINT"' "$INSTALL_SH" || true)
_has_var=$([ "$_count" -ge 1 ] && echo "yes" || echo "no")
assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var"
# Hardcoded torch>=2.4,<2.11.0 should only appear once (the default assignment)
_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded"
# What the old "appears exactly once" count was really guarding: an install line that
# spells the pin out ignores whatever the branch above it chose.
_literal=$(grep -cE 'uv pip install .*"torch>=' "$INSTALL_SH" || true)
assert_eq "no pip install hardcodes a torch pin" "0" "$_literal"
# Companions must be bounded to torch's window everywhere: the <2.11 bound appears
# twice (default assignments + the pinned custom-leaf block), never bare. torchaudio
# 2.11 dropped its exact torch pin, so a bare companion next to a <2.11-capped torch
# resolves a mismatched 2.11 build.
_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"' "$INSTALL_SH" || true)
assert_eq "torchvision bounded (<0.26) at default + custom-leaf" "2" "$_count"
_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "torchaudio bounded (<2.11) at default + custom-leaf" "2" "$_count"
# Companions must be bounded to torch's window everywhere, never bare: torchaudio 2.11
# dropped its exact torch pin, so a bare companion next to a <2.11-capped torch resolves
# a mismatched 2.11 build. Every assignment, not a fixed number of them.
_total=$(grep -cE '^[[:space:]]*TORCHVISION_CONSTRAINT="' "$INSTALL_SH" || true)
_bounded=$(grep -cE '^[[:space:]]*TORCHVISION_CONSTRAINT="torchvision>=[0-9][0-9.]*,<[0-9][0-9.]*"$' "$INSTALL_SH" || true)
assert_eq "every torchvision constraint is upper-bounded" "$_total" "$_bounded"
_total=$(grep -cE '^[[:space:]]*TORCHAUDIO_CONSTRAINT="' "$INSTALL_SH" || true)
_bounded=$(grep -cE '^[[:space:]]*TORCHAUDIO_CONSTRAINT="torchaudio>=[0-9][0-9.]*,<[0-9][0-9.]*"$' "$INSTALL_SH" || true)
assert_eq "every torchaudio constraint is upper-bounded" "$_total" "$_bounded"
_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision"$' "$INSTALL_SH" || true)
assert_eq "no bare torchvision companion remains" "0" "$_count"
_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio"$' "$INSTALL_SH" || true)

4
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -1253,6 +1253,7 @@ def test_install_prebuilt_falls_back_to_older_release_plan(
initial_fallback_used = False,
existing_install_dir = None,
force_cpu = False,
llama_backend = None,
):
call_log.append((llama_tag, initial_fallback_used))
if llama_tag == "b9002":
@ -2457,6 +2458,7 @@ def test_install_prebuilt_skips_when_older_release_fallback_matches_existing_ins
initial_fallback_used = False,
existing_install_dir = None,
force_cpu = False,
llama_backend = None,
):
call_log.append(llama_tag)
raise PrebuiltFallback("validation failed for latest release")
@ -2605,6 +2607,7 @@ def test_install_prebuilt_skips_same_release_fallback_attempt_when_installed(
prebuilt_fallback_used,
quantized_path,
force_cpu = False,
llama_backend = None,
):
attempted_names.append(choice.name)
if choice.name == first_choice.name:
@ -2732,6 +2735,7 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p
initial_fallback_used = False,
existing_install_dir = None,
force_cpu = False,
llama_backend = None,
):
attempted.append((llama_tag, release_tag, attempts[0].source_label))
if llama_tag == "b9002":

442
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)
@ -4495,11 +4823,93 @@ class TestApplyHostOverrides:
assert out.has_rocm is True
assert out.rocm_gfx_target == "gfx1200"
def test_forwarded_gfx_is_authoritative(self):
# setup already applied visible-device selection; its value wins.
host = rocm_host(rocm_gfx_target = "gfx1100")
def test_forwarded_gfx_does_not_clobber_probed_arch(self, monkeypatch):
# setup.ps1's pick is not fully visible-device aware (ignores CUDA_VISIBLE_DEVICES,
# amd-smi branch drops comma masks), so when it resolved the host's OTHER physical
# GPU it must not replace the arch detect_host() picked for the visible one.
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = rocm_host(rocm_gfx_target = "gfx1010", rocm_gfx_targets = ["gfx1100", "gfx1010"])
out = _apply_host_overrides(host, override_rocm_gfx = "gfx1100")
assert out.rocm_gfx_target == "gfx1010"
assert out.rocm_gfx_targets == ["gfx1100", "gfx1010"]
assert out.has_rocm is True
def test_forwarded_gfx_absent_from_host_stays_authoritative(self, monkeypatch):
# An arch no probe here ever reported is not a setup mispick: it is an explicit
# --rocm-gfx for a host whose probe is wrong or stale, so it must still win.
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = rocm_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151")
assert out.rocm_gfx_target == "gfx1151"
# ... but it says which arch HIP targets, not which cards exist, so the probed
# gfx1100 is still in the box and stays in the per-GPU list.
assert out.rocm_gfx_targets == ["gfx1100", "gfx1151"]
assert out.has_rocm is True
def test_forwarded_family_label_never_overrides_a_probed_arch(self, monkeypatch):
# The update path re-derives --rocm-gfx from the marker's family-named asset, so a
# family label is a bundle name, not a real arch, and must stay advisory: gfx1033 is
# in-generation but unbuilt, so gfx103X winning would serve a bundle it cannot run.
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = rocm_host(rocm_gfx_target = "gfx1033", rocm_gfx_targets = ["gfx1033"])
out = _apply_host_overrides(host, override_rocm_gfx = "gfx103X")
assert out.rocm_gfx_target == "gfx1033"
assert out.has_rocm is True
def test_forwarded_family_label_still_fills_an_unprobed_arch(self, monkeypatch):
# Negative control: with no probed arch the forward is the only source, so it
# applies.
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
out = _apply_host_overrides(cpu_host(), override_rocm_gfx = "gfx110X")
assert out.rocm_gfx_target == "gfx110x"
assert out.has_rocm is True
def test_forwarded_gfx_matching_active_keeps_physical_gfx_list(self, monkeypatch):
# When the forward agrees with the probe the per-GPU list must survive: collapsing
# it would hide the host's other AMD cards from the Windows auto-Vulkan floor
# check.
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = rocm_host(rocm_gfx_target = "gfx1010", rocm_gfx_targets = ["gfx1100", "gfx1010"])
out = _apply_host_overrides(host, override_rocm_gfx = "gfx1010")
assert out.rocm_gfx_target == "gfx1010"
assert out.rocm_gfx_targets == ["gfx1100", "gfx1010"]
def test_forwarded_gfx_never_drops_a_probed_physical_gpu(self, monkeypatch):
# The per-GPU list is the PHYSICAL inventory the Windows auto-Vulkan floor check
# reads, so a forwarded arch the probe never saw must be ADDED, not replace it:
# dropping the probe-confirmed gfx1100 would tell that check no AMD GPU on the box
# reaches the HIP floor when one plainly does.
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = rocm_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx1100", "gfx803"])
out = _apply_host_overrides(host, override_rocm_gfx = "gfx900")
assert out.rocm_gfx_target == "gfx900"
assert out.rocm_gfx_targets == ["gfx1100", "gfx803", "gfx900"]
def test_forwarded_gfx_not_duplicated_when_already_probed(self, monkeypatch):
# UNSLOTH_ROCM_GFX_ARCH makes the forward win over the probe's visible-device
# pick, so this reaches the same branch; the list must stay deduplicated.
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx803")
host = rocm_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100", "gfx803"])
out = _apply_host_overrides(host, override_rocm_gfx = "gfx803")
assert out.rocm_gfx_target == "gfx803"
assert out.rocm_gfx_targets == ["gfx1100", "gfx803"]
def test_forwarded_gfx_on_an_unprobed_host_lists_only_itself(self, monkeypatch):
# Negative control for the two above: nothing probed means no inventory to
# preserve, so the driver-only host keeps a single-entry list and auto-Vulkan.
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
out = _apply_host_overrides(cpu_host(), override_rocm_gfx = "gfx803")
assert out.rocm_gfx_target == "gfx803"
assert out.rocm_gfx_targets == ["gfx803"]
def test_manual_env_override_still_wins_over_probe(self, monkeypatch):
# UNSLOTH_ROCM_GFX_ARCH is the manual escape hatch for hosts whose arch the probes
# get wrong, so it stays authoritative.
monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx1151")
host = rocm_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151")
assert out.rocm_gfx_target == "gfx1151"
assert out.rocm_gfx_targets == ["gfx1100", "gfx1151"]
def test_has_rocm_only_keeps_probe_gfx(self):
out = _apply_host_overrides(cpu_host(), override_has_rocm = True)

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

17
tests/studio/test_ui_font_scale_contract.py
View file

@ -39,6 +39,17 @@ def _frontend_sources():
yield path
def _rel(path):
"""Source-relative path with forward slashes on every OS.
The allowlists above are written with "/", so a plain str(relative_to(SRC))
silently stops matching on Windows and every allowlisted file reports as an
offender. Keeping the separator normalised here also keeps failure messages
identical across platforms.
"""
return path.relative_to(SRC).as_posix()
def test_preference_writes_a_scale_not_the_root_font_size():
assert 'setVar("--ui-font-scale"' in STORE
assert 'el.setAttribute("data-ui-font-size"' in STORE
@ -136,7 +147,7 @@ def test_no_raw_pixel_text_utilities():
for path in _frontend_sources():
text = path.read_text(encoding = "utf-8")
for m in re.finditer(r"(?<![\w-])(?:text|leading)-\[[0-9.]+px\]", text):
offenders.append(f"{path.relative_to(SRC)}: {m.group(0)}")
offenders.append(f"{_rel(path)}: {m.group(0)}")
assert offenders == [], (
"Raw px text utilities ignore the UI font size preference; use the "
f"text-ui-* / leading-ui-* tokens in index.css instead: {offenders[:10]}"
@ -157,7 +168,7 @@ def test_css_font_sizes_reference_the_scale():
continue
if "1px" in decl:
continue # library layout tricks (KaTeX-style), not text
offenders.append(f"{path.relative_to(SRC)}: {decl.strip()[:80]}")
offenders.append(f"{_rel(path)}: {decl.strip()[:80]}")
assert offenders == [], (
"CSS typography must multiply by var(--ui-font-scale, 1) or be "
f"allowlisted here with a reason: {offenders[:10]}"
@ -167,7 +178,7 @@ def test_css_font_sizes_reference_the_scale():
def test_inline_font_size_styles_reference_the_scale():
offenders = []
for path in _frontend_sources():
rel = str(path.relative_to(SRC))
rel = _rel(path)
if rel in FONTSIZE_STYLE_ALLOWLIST:
continue
text = path.read_text(encoding = "utf-8")

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"

407
tests/test_runtime_text_encoding.py Normal file
View file

@ -0,0 +1,407 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Guard: shipping code must name an encoding on every text read and write.
`Path.read_text()`, `Path.write_text()`, `Path.open()` and builtin `open()` fall back
to `locale.getencoding()`: UTF-8 on the Linux and macOS runners, cp1252 on a stock
Windows install. Every file this repo reads at runtime is UTF-8 (HF `config.json` /
`tokenizer_config.json` / `adapter_config.json`, Ollama manifests, GGUF export
metadata), so on Windows those reads crash or, worse, succeed with mojibake: a
DeepSeek or Qwen tokenizer_config.json carries U+FF5C and U+2581 in its chat
template, and at utils/models/model_config.py that read sits inside a broad
`except Exception: logger.debug(...)`, so the token-pattern check silently
returned the wrong answer.
Unlike the import-time rule in test_source_read_encoding.py this is scope agnostic:
runtime reads live inside functions, and shipping code has no legitimate reason to
let the operator's locale decide. No reachability analysis to get wrong, so no
allowlist and no false positives.
Binary handles are skipped (no encoding to name, and passing one is a ValueError),
and a non-constant mode counts as unknown rather than text: demanding `encoding =`
on a call that may resolve to "rb" would leave no compliant way to write it.
Known limitation, deliberately not closed: `configparser.ConfigParser.read()` also
defaults to the locale encoding, but cannot be matched by name without resolving the
receiver, since `f.read(n)`, `resp.read(limit)` and `handle.read(chunk)` are spelled
identically. Flagging it would be a false positive with no compliant fix, the exact
failure mode this guard avoids. The one live `ConfigParser.read` (/etc/wsl.conf,
hub/utils/paths.py) is pinned by hand; a future one has to be caught in review.
"""
# `str | None` below is evaluated at import on Python 3.9 (requires-python >= 3.9).
from __future__ import annotations
import ast
import subprocess
from pathlib import Path
REPO = Path(__file__).resolve().parent.parent
# Everything that ships. `studio/` covers the installers too: install_python_stack.py
# reads /sys/class/kfd, the same detection path as utils/hardware/hardware.py. Test
# trees fall under the narrower import-time rule in test_source_read_encoding.py.
ROOTS = (REPO / "unsloth", REPO / "studio", REPO / "unsloth_cli")
# The frontend tree is TypeScript; node_modules is vendored third-party code.
SKIP_DIRS = {"build", "dist", "frontend", "node_modules", "src-tauri", ".venv", "site-packages"}
GUARDED_METHODS = {"read_text", "write_text"}
# Path classes, so an unbound `Path.open(p)` shifts every argument one right.
PATH_CLASSES = {"Path", "PosixPath", "PurePath", "WindowsPath"}
# Values that re-select the platform default when passed as the encoding.
PLATFORM_DEFAULT_ENCODINGS = (None, "locale")
# Calls that return the platform default, so naming one pins nothing.
PLATFORM_DEFAULT_CALLS = {"getdefaultencoding", "getencoding", "getpreferredencoding"}
# Modules whose `open` IS the builtin: same signature, same platform default.
BUILTIN_OPEN_MODULES = {"builtins", "io"}
# Take an encoding in "t" mode but default to "rb". Value is its positional slot.
COMPRESSED_OPENERS = {"bz2": 3, "gzip": 3, "lzma": None}
# Distinct from None so that "no mode argument at all" still means text.
UNKNOWN_MODE = object()
def _mode(call: ast.Call, positional_index: int):
"""The call's mode, or UNKNOWN_MODE when it is not a literal."""
# A splat hides the mode, so it is unknown rather than absent: falling through to
# "r" would flag a call that may resolve to binary, with no compliant way to fix it.
if any(isinstance(a, ast.Starred) for a in call.args):
return UNKNOWN_MODE
if any(kw.arg is None for kw in call.keywords):
return UNKNOWN_MODE
if len(call.args) > positional_index:
node = call.args[positional_index]
return node.value if isinstance(node, ast.Constant) else UNKNOWN_MODE
for kw in call.keywords:
if kw.arg == "mode":
return kw.value.value if isinstance(kw.value, ast.Constant) else UNKNOWN_MODE
return "r"
def _names_encoding(call: ast.Call) -> bool:
"""True only for an encoding that actually pins one.
`encoding = None` and `encoding = "locale"` re-select the platform default, so the
keyword being present is not enough. A `**kwargs` splat may carry an encoding we
cannot see, so it counts as named rather than as an unsatisfiable demand.
"""
for kw in call.keywords:
if kw.arg is None:
return True
if kw.arg != "encoding":
continue
if isinstance(kw.value, ast.Constant) and kw.value.value in PLATFORM_DEFAULT_ENCODINGS:
return False
if isinstance(kw.value, ast.Call) and _callee_name(kw.value.func) in PLATFORM_DEFAULT_CALLS:
return False # locale.getencoding() is the default, spelled out
return True
return False
def _is_text(call: ast.Call, positional_index: int) -> bool:
mode = _mode(call, positional_index)
return mode is not UNKNOWN_MODE and "b" not in str(mode)
def _imports_at_each_call(tree: ast.Module) -> dict:
"""The imports visible at every call, keyed by node id.
A function's own imports stay in that function: hoisting them would let one local
`from PIL.Image import open` turn off the builtin check for the whole file.
"""
visible_at = {}
def walk(node, visible):
if isinstance(node, ast.Call):
visible_at[id(node)] = visible
for child in ast.iter_child_nodes(node):
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
walk(child, {**visible, **_imported_names(child)})
else:
walk(child, visible)
walk(tree, _imported_names(tree))
return visible_at
def _foreign_names(tree: ast.Module) -> set:
"""Names bound to an object another library built.
`z = zipfile.ZipFile(p)` then `z.open(name)` is a binary member stream taking no
encoding, so demanding one leaves no correct edit.
"""
modules = _imported_names(tree)
names = set()
for node in ast.walk(tree):
if not isinstance(node, ast.Assign) or not isinstance(node.value, ast.Call):
continue
if _foreign_receiver(node.value, modules):
names.update(t.id for t in node.targets if isinstance(t, ast.Name))
return names
def _imported_names(tree) -> dict:
"""Names this module's imports bind, mapped to where they came from.
The name alone settles nothing: `import tarfile as tf` hides an opener that takes
no encoding, and `from PIL.Image import open` puts another behind the most familiar
name there is. Resolving the origin covers both, with no module list to maintain.
"""
bound = {}
stack = list(ast.iter_child_nodes(tree))
while stack:
node = stack.pop()
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
continue # that function's business, not this scope's
if isinstance(node, ast.Import):
for a in node.names:
bound[(a.asname or a.name).split(".")[0]] = a.name
elif isinstance(node, ast.ImportFrom):
for a in node.names:
bound[a.asname or a.name] = f"{node.module}.{a.name}" if node.module else a.name
else:
stack.extend(ast.iter_child_nodes(node))
return bound
def _callee_name(func):
"""The bare name a callee ends in, whether or not it is qualified."""
return func.id if isinstance(func, ast.Name) else getattr(func, "attr", None)
def _origin_root(name, modules) -> str:
"""The top-level module a bound name came from, or the name itself."""
return modules.get(name, name).split(".")[0]
def _compressed_key(name, modules):
"""The COMPRESSED_OPENERS entry this receiver resolves to, if any."""
for candidate in (name, _origin_root(name, modules)):
if candidate in COMPRESSED_OPENERS:
return candidate
return None
def _open_alias(name, modules):
"""What a bare callable resolves to: "builtin", a COMPRESSED_OPENERS key, or None."""
origin = modules.get(name)
if origin is None:
return "builtin" if name == "open" else None
parts = origin.split(".")
if parts[-1] != "open":
return None
if parts[0] in BUILTIN_OPEN_MODULES or origin == "open":
return "builtin"
return parts[0] if parts[0] in COMPRESSED_OPENERS else None
def _is_path_class(name, modules) -> bool:
"""True for a pathlib class, including under an alias."""
if name is None:
return False
return (modules.get(name) or name).split(".")[-1] in PATH_CLASSES
def _is_path_attr(node) -> bool:
"""True for a qualified path class, as in `pathlib.Path`."""
return isinstance(node, ast.Attribute) and node.attr in PATH_CLASSES
def _foreign_receiver(node, modules) -> bool:
"""True when the thing before `.open` is an object another library built.
`zipfile.ZipFile(p).open(name)` returns a binary member stream taking no encoding,
so it needs the same exemption as the bare `zipfile.open` spelling.
"""
if not isinstance(node, ast.Call):
return False
func = node.func
if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name):
root, name = func.value.id, func.attr
elif isinstance(func, ast.Name):
root = name = func.id
else:
return False
return root in modules and not _is_path_class(name, modules)
def _offender(
call: ast.Call,
modules = None,
foreign = (),
) -> str | None:
"""The call's name if it does text I/O without pinning an encoding."""
modules = {} if modules is None else modules
func = call.func
if isinstance(func, ast.Attribute):
receiver = func.value.id if isinstance(func.value, ast.Name) else None
# `Path.read_text(p)` is `p.read_text()` unbound: the instance takes slot 0,
# so every argument shifts one place right.
shift = 1 if _is_path_class(receiver, modules) or _is_path_attr(func.value) else 0
if func.attr in GUARDED_METHODS:
if func.attr == "read_text" and not shift and call.args:
first = call.args[0]
# Bound read_text takes encoding first, so None or "locale" there is a
# platform-default read. Any other positional means the receiver is
# importlib.metadata's Distribution: a filename, and no encoding at all.
if isinstance(first, ast.Constant) and first.value in PLATFORM_DEFAULT_ENCODINGS:
return "read_text()"
return None
return None if _names_encoding(call) else f"{func.attr}()"
if func.attr == "open":
if receiver is not None and _origin_root(receiver, modules) in BUILTIN_OPEN_MODULES:
return (
None if not _is_text(call, 1) or _names_encoding(call) else f"{receiver}.open()"
)
compressed = _compressed_key(receiver, modules) if receiver else None
if compressed is not None:
# "rb" by default, so only an explicit text mode is in scope.
mode = _mode(call, 1)
if mode is UNKNOWN_MODE or "t" not in str(mode):
return None
return None if _names_encoding(call) else f"{compressed}.open()"
# Any other imported receiver is somebody else's opener: tarfile takes a
# compression mode, Image a binary file. Neither has an encoding to name.
if receiver is not None and receiver in modules and receiver not in PATH_CLASSES:
return None
if _foreign_receiver(func.value, modules) or receiver in foreign:
return None
if not _is_text(call, shift):
return None
return None if _names_encoding(call) else "Path.open()"
return None
if isinstance(func, ast.Name):
alias = _open_alias(func.id, modules)
if alias == "builtin":
if not _is_text(call, 1):
return None
return None if _names_encoding(call) else "open()"
if alias is not None:
mode = _mode(call, 1)
if mode is UNKNOWN_MODE or "t" not in str(mode):
return None
return None if _names_encoding(call) else f"{alias}.open()"
return None
def _is_test_path(path: Path) -> bool:
parts = path.relative_to(REPO).parts
if SKIP_DIRS.intersection(parts):
return True
if "tests" in parts or "test" in parts:
return True
return path.name.startswith("test_") or path.name.endswith("_test.py")
def _offenders_in(src: str, label: str = "<snippet>"):
tree = ast.parse(src, filename = label)
visible_at = _imports_at_each_call(tree)
foreign = _foreign_names(tree)
found = []
for node in ast.walk(tree):
if isinstance(node, ast.Call):
name = _offender(node, visible_at.get(id(node), {}), foreign)
if name is not None:
found.append((node.lineno, name))
return found
def _tracked_sources():
"""Shipping *.py that git is actually tracking.
A walk also picks up whatever is lying in the checkout (a built `build/lib` copy,
a nested worktree, a vendored dep). None of those are ours to police, and a stale
artifact would fail this for everybody who has one.
"""
listed = subprocess.run(
["git", "-C", str(REPO), "ls-files", "-z", "--", "*.py"],
capture_output = True,
timeout = 60,
)
if listed.returncode != 0:
return None # not a checkout, so fall back to walking
names = listed.stdout.decode("utf-8", errors = "replace").split("\0")
return [REPO / n for n in names if n]
def _walked_sources():
return [p for root in ROOTS if root.is_dir() for p in sorted(root.rglob("*.py"))]
def test_shipping_code_names_an_encoding():
offenders = []
sources = _tracked_sources()
if sources is None:
sources = _walked_sources()
roots = {r.resolve() for r in ROOTS}
for path in sorted(sources):
if not roots.intersection(path.resolve().parents) or _is_test_path(path):
continue
try:
tree = ast.parse(path.read_text(encoding = "utf-8"), filename = str(path))
except SyntaxError:
continue
rel = path.relative_to(REPO).as_posix()
visible_at = _imports_at_each_call(tree)
foreign = _foreign_names(tree)
for node in ast.walk(tree):
if isinstance(node, ast.Call):
name = _offender(node, visible_at.get(id(node), {}), foreign)
if name is not None:
offenders.append(f"{rel}:{node.lineno}: {name}")
assert offenders == [], (
f"{len(offenders)} text read/write call sites in shipping code let the "
"operator's locale decide the encoding, so they crash or silently "
'produce mojibake on Windows. Pass encoding = "utf-8": ' + repr(offenders)
)
# The assertion above passes vacuously once the trees are clean, so it cannot tell a
# working detector from one that always returns None. These pin the detector itself.
def test_detects_the_plain_cases():
assert _offenders_in("from pathlib import Path\np = Path('x')\ns = p.read_text()\n")
assert _offenders_in("p.write_text('hi')\n")
assert _offenders_in("f = open('x')\n")
assert _offenders_in("f = open('x', 'w')\n")
assert _offenders_in("f = p.open()\n")
# Inside a function body too: shipping reads are not import-time.
assert _offenders_in("def load(p):\n return p.read_text()\n")
def test_rejects_encoding_that_reselects_the_platform_default():
assert _offenders_in("s = p.read_text(encoding = None)\n")
assert _offenders_in("s = p.read_text(encoding = 'locale')\n")
def test_accepts_a_pinned_encoding():
assert not _offenders_in("s = p.read_text(encoding = 'utf-8')\n")
assert not _offenders_in("f = open('x', 'w', encoding = 'utf-8')\n")
assert not _offenders_in("f = p.open(encoding = 'utf-8')\n")
assert not _offenders_in("s = p.read_text(encoding = 'utf-8', errors = 'replace')\n")
def test_skips_binary_handles():
# Binary has no encoding to name; passing one is a ValueError.
assert not _offenders_in("f = open('x', 'rb')\n")
assert not _offenders_in("f = open('x', mode = 'wb')\n")
assert not _offenders_in("f = p.open('rb')\n")
def test_skips_unknown_modes():
# A call that may resolve to "rb" has no compliant way to name an encoding.
assert not _offenders_in("mode = 'rb' if binary else 'r'\nf = open(path, mode)\n")
assert not _offenders_in("f = open(path, mode = chosen)\n")
def test_skips_foreign_openers_and_readers():
assert not _offenders_in("import fitz\nd = fitz.open(stream = b, filetype = 'pdf')\n")
assert not _offenders_in("import tarfile\nt = tarfile.open(p, 'r:gz')\n")
# importlib.metadata Distribution.read_text takes a positional filename.
assert not _offenders_in("s = dist.read_text('direct_url.json')\n")
def test_test_trees_are_out_of_scope():
assert _is_test_path(REPO / "tests" / "test_x.py")
assert _is_test_path(REPO / "studio" / "backend" / "tests" / "helpers.py")
assert not _is_test_path(REPO / "studio" / "backend" / "routes" / "inference.py")

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

Some files were not shown because too many files have changed in this diff Show more