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

Merge branch 'main' into studio/api-monitor-and-per-model-settings

Browse source 
Ports the lifecycle and download rows and the manual unload control from
main's monitor console onto the new monitor page, so nothing regresses
when the console is removed. Re-points the contract tests that read the
console file at the page and at the settings link that replaced it.
This commit is contained in:
Unsloth 2026-07-27 05:40:27 -07:00
commit a3f202a0de
101 changed files with 10279 additions and 702 deletions
Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -279,6 +279,8 @@ unsloth studio -H 0.0.0.0 -p 8888
```
The Cloudflare tunnel is **off by default**: `-H 0.0.0.0` exposes the raw port only, not a public internet URL. Pair the wildcard bind with `--cloudflare` (`unsloth studio -H 0.0.0.0 --cloudflare`) to also publish a public `https://*.trycloudflare.com` link, or prefer `--secure` (above), which keeps the raw port private. `--cloudflare` has no effect on a loopback bind.
On a wildcard bind Unsloth works out the address to share by asking `ifconfig.me` for the public IP, then asks `check-host.net` whether that port is reachable so it can tell you if a firewall is in the way. Both contact a third party. Set `UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK=1` to skip them; the banner then shows the LAN address and no reachability line.
The first time Unsloth is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Unsloth shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI.
For headless setups that cannot answer that prompt, set the initial admin password non-interactively with `--password` (only takes effect when no password is set yet; if one already exists it is a hard error, so rotate later with `unsloth studio reset-password`):

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

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

@ -164,6 +164,22 @@ async def get_current_subject_allow_password_change(
)
# The literal the examples ship with; pasting one unedited is likelier 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 unedited example placeholder is called out;
every real key still gets one indistinguishable message, so this reveals
nothing about which keys exist."""
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 +192,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

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

@ -56,6 +56,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
@ -90,6 +97,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
@ -136,6 +147,75 @@ class ApiMonitor:
self._trim_terminal_locked()
return entry.id
def record_lifecycle(
self,
*,
event: str,
model: str,
reason: Optional[str] = None,
running: bool = False,
) -> str:
"""Record a model load/unload alongside the request traffic that caused it.
``running=True`` opens the row (a load in progress) and the caller closes
it with the usual :meth:`finish` / :meth:`fail`; an unload is terminal on
arrival. Rows are shared, so every subject sees them, and share the same
retention budget as requests.
"""
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 (the
caller only has the load path up front, 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 (a load that was already
satisfied, so nothing was actually loaded)."""
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
@ -221,6 +301,19 @@ 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` this never touches an
entry that already finished, so a catch-all in a ``finally`` cannot stamp
an error onto a request that in fact 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
@ -233,15 +326,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,
@ -253,7 +349,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(
@ -266,16 +362,19 @@ 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, *, subject: Optional[str] = None) -> None:
@ -291,6 +390,10 @@ class ApiMonitor:
return
self._entries = deque(entry for entry in self._entries if entry.subject != subject)
@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)),
)

121
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
@ -3501,18 +3501,17 @@ class LlamaCppBackend:
return []
@staticmethod
def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
"""Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
def _run_vulkan_probe(binary: Optional[str] = None) -> list[dict]:
"""Run ``_vulkan_probe.py`` and parse its per-device lines.
Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance
in this process) and returns (device_index, free_mib, total_mib) sorted
by index. The index is ggml's compact Vulkan ordinal -- the one the
registry names ``Vulkan<index>`` and load_model pins with ``--device``,
NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set
``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the
list already reflects it. iGPUs leave a host-RAM margin (see
``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
their real total through. [] when no Vulkan build or device is reachable.
Returns raw (uncapped) rows sorted by index:
``{"index", "free_mib", "total_mib", "is_igpu", "name"}``. The index is
ggml's compact Vulkan ordinal -- the one the registry names
``Vulkan<index>`` and load_model pins with ``--device``, NOT the raw
``GGML_VK_VISIBLE_DEVICES`` space. A user-set ``GGML_VK_VISIBLE_DEVICES``
is honored by ggml (passed through), so the list already reflects it.
``name`` is ggml's device description; "" from an older 4-column probe.
[] when no Vulkan build or device is reachable.
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if not binary:
@ -3537,10 +3536,13 @@ class LlamaCppBackend:
)
probe_script = Path(__file__).with_name("_vulkan_probe.py")
try:
# UTF-8 to match the probe's stdout reconfigure: device names can be
# non-ASCII, and the platform-default decode (cp1252) could throw.
result = subprocess.run(
[sys.executable, str(probe_script), str(binary_dir)],
capture_output = True,
text = True,
encoding = "utf-8",
errors = "replace",
timeout = 15,
env = env,
**_windows_hidden_subprocess_kwargs(),
@ -3554,21 +3556,56 @@ class LlamaCppBackend:
logger.debug(f"vulkan GPU probe failed: {e}")
return []
gpus: list[tuple[int, int, int]] = []
rows: list[dict] = []
for line in result.stdout.strip().splitlines():
parts = line.split("\t")
if len(parts) != 4:
# 4 columns from an older probe (no name); 5 with the name column.
if len(parts) not in (4, 5):
continue
try:
idx = int(parts[0])
free_mib = int(parts[1]) // (1024 * 1024)
is_igpu = parts[2] == "1"
# iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
# fit stays on free*frac (the host reserve below is its
# headroom); a discrete card passes its real total through.
total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024)
rows.append(
{
"index": int(parts[0]),
"free_mib": int(parts[1]) // (1024 * 1024),
"is_igpu": parts[2] == "1",
"total_mib": int(parts[3]) // (1024 * 1024),
"name": parts[4].strip() if len(parts) == 5 else "",
}
)
except ValueError:
continue
rows.sort(key = lambda r: r["index"])
return rows
@staticmethod
def vulkan_device_inventory(binary: Optional[str] = None) -> list[dict]:
"""UI-facing Vulkan device list: the devices llama-server will actually
use, with real totals (an iGPU keeps its shared-RAM total here -- the
caller labels it, unlike the fit which zeroes it). Same rows as
``_run_vulkan_probe``; names fall back to ``Vulkan<i>``.
"""
rows = LlamaCppBackend._run_vulkan_probe(binary)
for row in rows:
if not row["name"]:
row["name"] = f"Vulkan{row['index']}"
return rows
@staticmethod
def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
"""Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
Fit-oriented view of ``_run_vulkan_probe``: returns (device_index,
free_mib, total_mib) sorted by index. iGPUs leave a host-RAM margin (see
``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
their real total through. [] when no Vulkan build or device is reachable.
"""
gpus: list[tuple[int, int, int]] = []
for row in LlamaCppBackend._run_vulkan_probe(binary):
idx, free_mib, is_igpu = row["index"], row["free_mib"], row["is_igpu"]
# iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
# fit stays on free*frac (the host reserve below is its
# headroom); a discrete card passes its real total through.
total_mib = 0 if is_igpu else row["total_mib"]
capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu)
if capped < free_mib:
logger.info(
@ -3577,7 +3614,6 @@ class LlamaCppBackend:
f"({free_mib}->{capped}MiB usable)"
)
gpus.append((idx, capped, total_mib))
gpus.sort(key = lambda g: g[0])
if gpus:
logger.info(
"Vulkan GPU memory detected: "
@ -3596,7 +3632,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 +5174,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 +6391,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 +6671,23 @@ class LlamaCppBackend:
# Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server;
# serve them with the diffusion runner (same OpenAI-compat interface).
if self._is_diffusion:
# Final defense: route and pre-teardown preflights reject before Phase 1.
if is_vulkan_backend and gpu_ids:
raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR)
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
# prior load (this path skips the command builder that clears it).
self._layer_preserves_tensor_intent = False
# On a Vulkan build gpu_ids are ggml Vulkan ordinals, but the diffusion
# runner selects its device by CUDA physical index (_diffusion_gpu_arg
# forwards gpu_ids[0] as a CUDA/DG_GPU token) with no mapping to them.
# The route rejects a CONFIRMED-diffusion pick up front; an uncached GGUF
# only classified as diffusion post-download still reaches here with a
# pin, so drop it and serve on the default device (like an unpinned load).
if gpu_ids and is_vulkan_backend:
logger.warning(
"Ignoring gpu_ids %s for diffusion GGUF on a Vulkan build: "
"the diffusion runner cannot map ggml Vulkan ordinals; "
"serving on the default device.",
gpu_ids,
)
gpu_ids = None
with self._lock:
if self._cancel_event.is_set():
logger.info("Load cancelled before diffusion server start")
@ -8302,7 +8349,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 +9524,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 +9658,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:
"""Record an idle auto-unload in the API monitor, using the advertised repo id
from the stash so the row never shows the on-disk load path. Best-effort."""
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)

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

@ -34,6 +34,16 @@ 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. The
# retained index covers what was already known; nothing covers the one that just
# landed until the next scan, and the request path must not call it 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 +113,28 @@ 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. A bare id means whichever quant a plain load would take, so put
# that first: everything downstream reads [0], and answering with the
# largest can evict a working model and then OOM starting it.
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 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,
and which quant to name, from a single scan."""
from pathlib import Path
path = getattr(info, "path", None)
@ -123,8 +144,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 +314,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, rather than waiting out the TTL.
Keeps the entries. Callers on the request path read this cache without
scanning, so emptying it would leave them with no evidence about any local
model until the rebuild lands, and a bare request for one of them would be
answered by whatever is resident. Only a completed download invalidates, and
that only ever adds models, 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 +358,78 @@ 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 here
would park the request path on the very scan it is trying to stay off. Reading
``_scan[0]`` is 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.
Callers that cannot afford the scan use this plus ``allow_scan=False``, so this
is the only thing that ever refreshes the index for them. It has to cover a
stale index and 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 those callers for the life of the process.
Never touches ``_lock``, which the scan holds throughout, and never blocks.
"""
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 callers on the request path: the scan walks several model dirs and HF
caches, takes seconds on a large install, and holds a lock every other
caller queues behind. A stale answer is fine there, since what is on 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 +445,45 @@ 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, ())``. Splits
the name like the resolver so the two agree. 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

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

@ -39,10 +39,30 @@ 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 straight out of the HF cache has a snapshot directory as its
identifier, whose basename is a commit hash. Recover the repo id so callers
show ``unsloth/gemma-4-31B-it-GGUF`` rather than ``c1ac76e99d55...``.
"""
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 +71,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)]

831
studio/backend/core/inference/openai_auto_download.py Normal file
View file

@ -0,0 +1,831 @@
# 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
downloaded in the background instead of erroring, and the request is told to
retry rather than being 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 throughout, and the retry that lands after the download is served by the
new model 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 it is a GGUF repo.
``gpt-4`` and ``anthropic/claude-3.5-sonnet`` alike fall through to the
resident model as before: a namespace is not evidence of intent, since LiteLLM
and OpenRouter address every provider that way.
- GGUF repos only, decided from the remote file list, not the repo name. GGUF
runs under llama.cpp, which never imports repo Python.
- Anything declaring ``auto_map`` is refused, so ``trust_remote_code`` can only
ever be 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 both run while the
# provisional slot is held, so an unresponsive Hub would pin the single flight and stall
# the request long past the metadata budget. The code probe fetches up to three small
# configs, so it gets more room than the single 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 already 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 a client 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 turns it into an
HTTPException with 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. The slot is kept until a retry surfaces it, since
# the advertised retry interval is far longer than the watcher's poll and the
# client would otherwise just 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: an unrecognized GGUF below a subdirectory keys on its path
("build/llama-13b", which is_valid_gguf_variant allows and the catalog advertises),
while "C:/models/x.gguf" leaves a drive letter that is no repo id at all.
"""
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. That keeps ``gpt-4`` and other foreign ids
falling through untouched, and avoids the bare-name ``unsloth/`` prefixing in
ModelConfig.from_identifier 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.
``vendor/model`` is how LiteLLM and OpenRouter address every provider, and
``name:latest`` is how Ollama tags one, so neither a namespace nor a colon
proves a request was meant for this server. 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, which here would be
the server owner's. False is what actually means anonymous.
"""
return hf_token or False
def _servable_key(repo_id: str, hf_token: Optional[str]) -> str:
"""Cache key, per credential.
The Hub answers 404 for a private repo the caller cannot see, so a verdict
reached without a token says nothing about a caller who has one. Keyed on a
digest so the token itself is never 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 stops waiting and takes *default*, which each call site chooses so that 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 of their own, and sharded quants sum
across their shards. The byte total comes from the download plan, which folds
the companions back into every quant, so the disk reserve is measured against
what the worker fetches rather than the main files alone.
"""
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 two extractors part ways: this one
# takes the last hyphenated segment ("7b" of llama-7b) while the plan and
# the worker key the whole stem. Advertising ours dispatches a variant the
# worker cannot resolve, so take theirs for the unrecognized case only.
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 the release on ``repo_id`` alone let a stale operation clear a newer
one for the same repo: variant A errors, an adopting request frees the slot,
a retry starts variant B, and A's watcher then matches on the repo and clears
B on its way out -- admitting a second repository download alongside B.
Identity ties every release to the operation that actually took the slot.
"""
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 worker still running still owns the slot: releasing it on the
# clock alone would admit a second multi-GB download alongside it.
# "unknown" cannot confirm it is alive, so stop holding it then,
# or a broken probe would wedge auto-download for good.
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
# started the warm, and a second one 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 rather than
# silently starting the same download again.
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: falling
through to the resident model is what such a label does anyway, and refusing
it 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 rather than quants, so a repo holding only those
# is not downloadable here either. Answering otherwise would hold an ordinary
# foreign label at model_download_busy for the length of 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
spending gigabytes on weights that cannot answer the request that asked for
them; 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 not a downloadable GGUF repo (LiteLLM/OpenRouter style) would be told
# to wait out a multi-hour download instead of falling through to the resident
# model. Only a label that could itself be downloaded is a second 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.
# Same rule as the metadata probe and the worker.
# None on timeout, which refuses: an unchecked repo is not a cleared one.
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 what /v1/models advertises all have to 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
an Ollama-style tag that names no quant at all (":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 it is: a repo of generically named GGUFs has
# real variants like "llama-13b" that are valid worker keys but do not look
# like quants, 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("/")

27
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,32 @@ 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 the request path
# from a cached scan with no watcher, and would otherwise report the model
# absent and let the request be served by whatever is resident. Models only,
# since datasets share this path and noting one as a local model would refuse
# a bare request naming that id 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 rather than on the first request that needs it, so the
# new model resolves without a scan on 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
supplied 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))

1104
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

@ -38,6 +38,7 @@ from picker.schemas import MAX_CHAT_TEMPLATE_BYTES, chat_template_byte_length
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,
@ -45,6 +46,7 @@ from utils.openai_auto_switch_settings import (
get_openai_auto_switch_enabled,
resolve_model_override_key,
get_stored_auto_unload_idle_seconds,
get_stored_openai_auto_download_enabled,
set_model_override,
set_openai_auto_switch,
)
@ -114,6 +116,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):
@ -125,6 +128,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
# A quant suffix, as modelOverrideKey builds it. Matched against the loader's own
@ -304,6 +309,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(),
)
@ -312,8 +318,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(
@ -333,6 +342,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,
)

48
studio/backend/run.py
View file

@ -111,13 +111,26 @@ from startup_banner import print_studio_access_banner, print_studio_stop_hint
logger = get_logger(__name__)
DISABLE_PUBLIC_CHECK_ENV = "UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK"
def public_check_disabled() -> bool:
"""True when the operator has turned off the third-party startup lookups.
On a wildcard bind Unsloth asks ifconfig.me for the public IP and check-host.net
whether the port is reachable. Both are useful for sharing a Studio but both tell
an outside service this machine is running one, which lab and privacy-sensitive
deployments do not want (#7307 Problem 8). Set the var to opt out.
"""
return os.environ.get(DISABLE_PUBLIC_CHECK_ENV, "").strip().lower() in {"1", "true", "yes"}
def _resolve_external_ip() -> str:
"""Resolve the machine's external IP address.
Tries, in order:
1. GCE metadata server (instant on Google Cloud VMs)
2. ifconfig.me (anywhere with internet)
2. ifconfig.me (anywhere with internet, skipped by UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK)
3. LAN IP via UDP socket trick (fallback)
"""
import urllib.request
@ -136,14 +149,15 @@ def _resolve_external_ip() -> str:
except Exception:
pass
# 2. Public IP service.
try:
with urllib.request.urlopen("https://ifconfig.me", timeout = 3) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
except Exception:
pass
# 2. Public IP service. Third-party, so skippable; the LAN address below still works.
if not public_check_disabled():
try:
with urllib.request.urlopen("https://ifconfig.me", timeout = 3) as resp:
ip = resp.read().decode().strip()
if ip:
return ip
except Exception:
pass
# 3. Fallback: LAN IP via UDP socket trick
try:
@ -304,7 +318,8 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
"""Probe check-host.net to confirm display_host:port is reachable from the
public internet. Synchronous so output lands between the banner URLs and the
stop hint. Bounded at ~15s; failures swallowed (verifier failing != Unsloth
failing). Only meaningful for a wildcard bind."""
failing). Only meaningful for a wildcard bind, and skipped entirely by
UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK."""
global _public_reachable
# Reset to "unknown" each run; set True/False only when the probe decides.
_public_reachable = None
@ -344,6 +359,11 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
# Not an IP literal; probe by hostname.
pass
# The probe hands display_host:port to a third party and asks it to connect.
if public_check_disabled():
logger.debug("Skipping the check-host.net probe (%s).", DISABLE_PUBLIC_CHECK_ENV)
return
try:
qs = urllib.parse.urlencode({"host": f"{display_host}:{port}", "max_nodes": 3})
req = urllib.request.Request(
@ -754,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
@ -763,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
@ -914,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.

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

@ -58,6 +58,28 @@ 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. That is right in a
server and wrong here: it walks the developer's actual caches, which on a large
install takes seconds, and the resulting I/O starves the loop under the
timing-sensitive streaming tests. Tests that exercise the warm 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 background warm still left the
# cold path walking those caches synchronously inside the admission wait, so on a
# large install the assertion became a 503 "still indexing". Tests that want the
# cold path set _scan back themselves (and stub the scan). _build_index is left
# alone so the tests that call it directly still exercise the real walk.
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

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

@ -314,3 +314,98 @@ def test_api_monitor_records_whether_the_caller_used_an_api_key():
by_id = {entry["id"]: entry for entry in monitor.snapshot(subject = "u")}
assert by_id[ui]["via_api_key"] is False
assert by_id[api]["via_api_key"] is True
# ── model lifecycle rows (load / unload) ────────────────────────────
def test_lifecycle_load_row_opens_running_then_closes():
monitor = ApiMonitor(max_entries = 5)
event_id = monitor.record_lifecycle(event = "load", model = "org/A-GGUF", running = True)
row = monitor.snapshot()[0]
assert row["kind"] == "lifecycle" and row["event"] == "load"
assert row["status"] == "running" and row["duration_ms"] is None
# A load in progress is not an in-flight API request.
assert monitor.active_count() == 0
monitor.relabel(event_id, "org/A-GGUF:Q4_K_M")
monitor.finish(event_id)
row = monitor.snapshot()[0]
assert row["status"] == "completed"
assert row["model"] == "org/A-GGUF:Q4_K_M"
assert row["duration_ms"] is not None
def test_lifecycle_unload_row_is_terminal_on_arrival():
monitor = ApiMonitor(max_entries = 5)
monitor.record_lifecycle(event = "unload", model = "org/A-GGUF", reason = "idle")
row = monitor.snapshot()[0]
assert row["status"] == "completed"
assert (row["event"], row["reason"]) == ("unload", "idle")
assert monitor.active_count() == 0
def test_lifecycle_rows_are_visible_to_every_subject():
# A load is server-wide, so it must not vanish for other API keys like a request does.
monitor = ApiMonitor(max_entries = 5)
monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
subject = "alice",
)
event_id = monitor.record_lifecycle(event = "unload", model = "org/A-GGUF")
bob = monitor.snapshot(subject = "bob")
assert [r["kind"] for r in bob] == ["lifecycle"]
assert monitor.get(event_id, subject = "bob") is not None
assert len(monitor.snapshot(subject = "alice")) == 2
def test_request_rows_stay_private_to_their_subject():
monitor = ApiMonitor(max_entries = 5)
rid = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
subject = "alice",
)
assert monitor.snapshot(subject = "bob") == []
assert monitor.get(rid, subject = "bob") is None
def test_discard_drops_a_row_that_never_happened():
# A load that found the model already resident must leave no trace.
monitor = ApiMonitor(max_entries = 5)
event_id = monitor.record_lifecycle(event = "load", model = "org/A-GGUF", running = True)
monitor.discard(event_id)
assert monitor.snapshot() == []
monitor.discard(event_id) # idempotent
def test_fail_open_never_touches_a_finished_row():
# Called from a finally, so it must not stamp an error onto a load that succeeded.
monitor = ApiMonitor(max_entries = 5)
event_id = monitor.record_lifecycle(event = "load", model = "org/A-GGUF", running = True)
monitor.finish(event_id)
monitor.fail_open(event_id, "Load did not complete")
row = monitor.snapshot()[0]
assert row["status"] == "completed" and row["error"] is None
still_open = monitor.record_lifecycle(event = "load", model = "org/B-GGUF", running = True)
monitor.fail_open(still_open, "Load did not complete")
assert monitor.snapshot()[0]["status"] == "error"
def test_lifecycle_rows_share_the_retention_budget():
monitor = ApiMonitor(max_entries = 2)
for i in range(4):
monitor.record_lifecycle(event = "unload", model = f"org/M{i}")
models = [r["model"] for r in monitor.snapshot()]
assert models == ["org/M3", "org/M2"]
def test_request_rows_report_kind_request():
monitor = ApiMonitor(max_entries = 2)
monitor.start(endpoint = "/v1/chat/completions", method = "POST", model = "m", prompt = "hi")
assert monitor.snapshot()[0]["kind"] == "request"

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

@ -427,14 +427,16 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
self.assertTrue(result["available"])
self.assertEqual(result["backend"], "vulkan")
self.assertEqual(result["index_kind"], "relative")
# ggml Vulkan ordinals are the space `--device Vulkan<i>` pins, so they
# are selectable, unlike a torch-xpu relative ordinal.
self.assertEqual(result["index_kind"], "vulkan")
self.assertEqual(result["parent_visible_gpu_ids"], [])
self.assertEqual(
result["devices"],
[
{
"index": 0,
"index_kind": "relative",
"index_kind": "vulkan",
"visible_ordinal": 0,
"name": "Vulkan0",
"memory_total_gb": 8.0,

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

@ -16,6 +16,7 @@ from core.inference.llama_admission import (
ADMISSION_CONTROL_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
ADMISSION_QUEUE_PER_SLOT_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
DEFAULT_ADMISSION_MAX_QUEUE,
@ -28,8 +29,23 @@ from core.inference.llama_admission import (
)
_ADMISSION_ENV = (
ADMISSION_CONTROL_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
ADMISSION_QUEUE_PER_SLOT_ENV,
*llama_admission._LEGACY_ENV.values(),
)
@pytest.fixture(autouse = True)
def _reset_queues():
def _reset_queues(monkeypatch):
# Clear ambient settings for every test, not just the ones that remember to:
# a canonical name set on the machine silently beats the legacy name a test
# is exercising, and the queue registry is process-global.
for name in _ADMISSION_ENV:
monkeypatch.delenv(name, raising = False)
reset_llama_admission_queues()
yield
reset_llama_admission_queues()
@ -41,15 +57,25 @@ def test_admission_config_defaults(monkeypatch):
ADMISSION_QUEUE_TIMEOUT_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
ADMISSION_QUEUE_PER_SLOT_ENV,
"UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL",
"UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT",
"UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL",
"UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE",
):
monkeypatch.delenv(name, raising = False)
config = llama_admission_config_from_env()
# Literals, not the module constants: comparing a default to itself would let
# any future value change through silently.
assert config.enabled is True
assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE
assert config.queue_timeout_s is None # wait forever
assert config.keepalive_interval_s == 5.0
assert config.max_queue is None # no absolute cap
assert config.queue_per_slot == 16
assert (DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, DEFAULT_ADMISSION_MAX_QUEUE) == (None, None)
assert DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S == 5.0
def test_admission_config_env_overrides(monkeypatch):
@ -66,6 +92,25 @@ def test_admission_config_env_overrides(monkeypatch):
assert config.max_queue is None
def test_admission_config_honors_legacy_openai_compat_env(monkeypatch):
# The queue is shared with /v1/messages now, but existing OPENAI_COMPAT
# settings must keep working.
monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", "7")
monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL", "off")
config = llama_admission_config_from_env()
assert config.max_queue == 7
assert config.enabled is False
def test_admission_config_prefers_neutral_env_over_legacy(monkeypatch):
monkeypatch.setenv("UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", "7")
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "3")
assert llama_admission_config_from_env().max_queue == 3
def test_admission_config_positive_queue_timeout_env(monkeypatch):
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600")
@ -106,6 +151,160 @@ def test_fifo_capacity_one_grants_next_waiter_on_release():
asyncio.run(_run())
def test_pool_hands_out_distinct_slots_and_reuses_them():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
leases = [queue.reserve(capacity = 3, config = config).lease_nowait() for _ in range(3)]
assert sorted(lease.slot for lease in leases) == [0, 1, 2] # one slot each
snapshot = queue.snapshot()
assert (snapshot.active, snapshot.free, snapshot.capacity) == (3, 0, 3)
# A freed slot returns to the pool and is handed to the next caller.
freed = leases[1].slot
leases[1].release()
assert queue.snapshot().free == 1
reused = queue.reserve(capacity = 3, config = config).lease_nowait()
assert reused.slot == freed
reused.release()
leases[0].release()
leases[2].release()
snapshot = queue.snapshot()
assert (snapshot.active, snapshot.free) == (0, 3)
asyncio.run(_run())
def test_pool_waiter_is_handed_a_real_slot():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
held = queue.reserve(capacity = 1, config = config).lease_nowait()
waiting = queue.reserve(capacity = 1, config = config)
assert waiting.lease_nowait() is None
assert queue.snapshot().free == 0
held.release()
granted = await waiting.wait(0.1)
assert granted is not None and granted.slot == 0 # the slot just freed
granted.release()
asyncio.run(_run())
def test_shrinking_capacity_retires_slots_beyond_the_new_pool():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
leases = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
assert queue.snapshot().capacity == 4
# llama-server reloaded with fewer --parallel slots; in-flight holders keep
# running and their slots retire instead of returning to the smaller pool.
shrunk = queue.reserve(capacity = 2, config = config)
assert shrunk.lease_nowait() is None # all 4 still held, nothing free
for lease in leases:
lease.release()
granted = await shrunk.wait(0.1)
assert granted is not None and granted.slot < 2
granted.release()
snapshot = queue.snapshot()
assert (snapshot.capacity, snapshot.active, snapshot.free) == (2, 0, 2)
asyncio.run(_run())
def test_queue_limit_scales_with_the_serving_slots():
# The wait line follows --parallel: 16 per slot, floored at 64 so a 1-slot
# backend keeps the depth it had before scaling existed.
config = LlamaAdmissionConfig()
assert config.queue_limit(4) == 64 # --parallel 4 (the default)
assert config.queue_limit(8) == 128 # --parallel 8
assert config.queue_limit(16) == 256
assert config.queue_limit(1) == 64 # floor, not 16
assert config.queue_limit(2) == 64 # floor, not 32
# An explicit cap wins, and a None multiplier means an unbounded line.
assert LlamaAdmissionConfig(max_queue = 5).queue_limit(8) == 5
assert LlamaAdmissionConfig(queue_per_slot = None).queue_limit(8) is None
# Non-positive settings mean unbounded, never "reject everything".
assert LlamaAdmissionConfig(max_queue = 0).queue_limit(4) is None
assert LlamaAdmissionConfig(max_queue = -1).queue_limit(4) is None
assert LlamaAdmissionConfig(queue_per_slot = 0).queue_limit(4) is None
assert LlamaAdmissionConfig(queue_per_slot = -3).queue_limit(4) is None
def test_queue_limit_rejects_only_once_the_line_is_full():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
# Explicit cap, so the test drives rejection without standing up the 64
# waiters the scaled floor would otherwise require.
config = LlamaAdmissionConfig(max_queue = 4)
held = [queue.reserve(capacity = 2, config = config).lease_nowait() for _ in range(2)]
parked = [queue.reserve(capacity = 2, config = config) for _ in range(4)]
assert queue.snapshot().queued == 4
with pytest.raises(LlamaAdmissionQueueFull):
queue.reserve(capacity = 2, config = config)
for reservation in parked:
reservation.cancel()
for lease in held:
lease.release()
asyncio.run(_run())
def test_waiting_is_never_timed_out_by_default():
# "Wait forever": the default config sets no queue timeout at all.
assert llama_admission_config_from_env().queue_timeout_s is None
assert LlamaAdmissionConfig().queue_timeout_s is None
def test_single_request_at_a_time_never_queues_or_allocates_waiters():
# The common serving case: one request in flight at a time must take a slot
# straight away and never touch the wait line.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
for _ in range(50):
reservation = queue.reserve(capacity = 4, config = config)
lease = reservation.lease_nowait()
assert lease is not None # admitted immediately
assert queue.snapshot().queued == 0 # nobody ever lined up
lease.release()
snapshot = queue.snapshot()
assert (snapshot.active, snapshot.free, snapshot.queued) == (0, 4, 0)
asyncio.run(_run())
def test_unbounded_queue_keeps_waiting_instead_of_rejecting():
# queue_per_slot None is the "pool + unbounded wait line" mode: nothing is
# ever rejected, callers just line up for the next free slot.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(max_queue = None, queue_per_slot = None)
held = queue.reserve(capacity = 1, config = config).lease_nowait()
waiters = [queue.reserve(capacity = 1, config = config) for _ in range(200)]
assert queue.snapshot().queued == 200 # no LlamaAdmissionQueueFull
held.release()
first = await waiters[0].wait(0.1)
assert first is not None
first.release()
for waiter in waiters[1:]:
waiter.cancel()
asyncio.run(_run())
def test_queue_full_rejects_excess_waiter():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
@ -288,6 +487,105 @@ def test_lease_release_is_idempotent_under_concurrent_calls():
asyncio.run(_run())
def test_releasing_a_stale_lease_does_not_free_someone_elses_slot():
# The concurrent test above passes without the _released guard: the racing
# calls all target a still-live slot, which the bitmask already absorbs. The
# case the guard exists for is a slot released twice with a reuse in between.
# It is live: _wait_for_openai_admission_non_streaming releases and re-raises,
# then the caller's finally cancels the reservation and releases the same
# lease again, by which point the slot can belong to another request.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
stale = queue.reserve(capacity = 1, config = config).lease_nowait()
stale.release()
other = queue.reserve(capacity = 1, config = config).lease_nowait()
assert other.slot == stale.slot # the slot got reused
stale.release()
assert queue.snapshot().active == 1, "stale release handed back a live slot"
other.release()
assert queue.snapshot().active == 0
asyncio.run(_run())
def test_grant_reclaims_the_slot_when_the_waiters_loop_is_gone():
# _grant_waiters_locked takes the slot before scheduling delivery, so if the
# schedule fails the bit is already set. Leaving it set strands the slot for
# good, because _free is rebuilt from the bitmask.
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
held = None
dead = asyncio.new_event_loop()
try:
async def _fill_and_queue():
nonlocal held
held = queue.reserve(capacity = 1, config = config).lease_nowait()
assert queue.reserve(capacity = 1, config = config).lease_nowait() is None
dead.run_until_complete(_fill_and_queue())
finally:
dead.close()
held.release() # grant path now hits the closed loop
assert queue.snapshot().active == 0
assert queue.is_idle()
def test_cancel_returns_the_granted_slot_when_the_waiters_loop_is_gone():
# Routes cancel() from finally blocks, so a raise here would mask their
# exception and skip the release that hands the granted slot back.
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
held = reservation = None
dead = asyncio.new_event_loop()
try:
async def _fill_and_queue():
nonlocal held, reservation
held = queue.reserve(capacity = 1, config = config).lease_nowait()
reservation = queue.reserve(capacity = 1, config = config)
dead.run_until_complete(_fill_and_queue())
held.release() # promotes the waiter, so cancel() has a lease to return
finally:
dead.close()
reservation.cancel()
assert queue.snapshot().active == 0
assert queue.is_idle()
def test_delivery_to_an_already_finished_waiter_releases_the_slot():
# A slot is taken before delivery is scheduled, so if the waiter finishes in
# that window someone has to hand it back. _deliver_lease does it twice over,
# in the dead-waiter branch and in the InvalidStateError backstop; this pins
# the outcome, not which one. Reaches into the waiter because no public call
# leaves that window open: queue.cancel() reclaims granted_lease itself.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
held = queue.reserve(capacity = 1, config = config).lease_nowait()
reservation = queue.reserve(capacity = 1, config = config)
waiter = reservation._waiter
held.release() # schedules _deliver_lease, sets granted_lease
waiter.future.cancel() # finishes the future before the callback runs
assert waiter.granted_lease is not None
await asyncio.sleep(0) # let the callback run
assert queue.snapshot().active == 0
assert queue.is_idle()
asyncio.run(_run())
def test_new_key_evicts_idle_prior_load_queues():
# Each model load carries a fresh ephemeral port, so a new base_url key must
# not leave the drained queues from earlier loads accumulating forever.
@ -318,3 +616,234 @@ def test_new_key_retains_in_flight_prior_load_queue():
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"}
asyncio.run(_run())
def test_capacity_shrink_never_admits_past_the_new_ceiling():
# A load that downshifts --parallel (or an unload resetting it to 1) shrinks the
# pool while slots are still held. Those holdovers keep occupying the backend, so
# they must count against the ceiling; sizing on free ids alone over-admits.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
held = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
assert all(lease is not None for lease in held)
waiter = queue.reserve(capacity = 4, config = config)
queue.reserve(capacity = 1, config = config) # capacity collapses to 1
# Release the one id that still falls inside the shrunk pool, so it goes
# back on the free list; ids at or above capacity retire instead.
low = min(held, key = lambda lease: lease.slot)
assert low.slot == 0
low.release()
# The other 3 holdovers are still generating, which already meets the new
# ceiling, so the freed id must not be handed on. Gating on "is an id free"
# alone grants it here and puts 4 generations on a 1-slot backend.
with pytest.raises(asyncio.TimeoutError):
await waiter.wait(0.2)
assert queue.snapshot().active == 3
waiter.cancel()
for lease in held:
if lease is not low:
lease.release()
asyncio.run(_run())
def test_queue_per_slot_env_is_parsed(monkeypatch):
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "4")
assert llama_admission_config_from_env().queue_limit(32) == 128
# Non-positive asks for an unbounded line rather than rejecting everything.
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "0")
assert llama_admission_config_from_env().queue_limit(32) is None
def test_max_queue_zero_from_env_is_unbounded_end_to_end(monkeypatch):
# Guards the whole env path, not just the parsed field: a regression that let
# queue_per_slot survive MAX_QUEUE=0 would silently re-bound the line.
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0")
config = llama_admission_config_from_env()
assert config.max_queue is None and config.queue_per_slot is None
assert config.queue_limit(1) is None and config.queue_limit(64) is None
def test_legacy_env_fallback_covers_every_setting(monkeypatch):
for canonical, legacy in llama_admission._LEGACY_ENV.items():
monkeypatch.delenv(canonical, raising = False)
monkeypatch.setenv(legacy, "0" if "CONTROL" in canonical else "7")
config = llama_admission_config_from_env()
assert config.enabled is False
assert config.queue_timeout_s == 7.0
assert config.keepalive_interval_s == 7.0
assert config.max_queue == 7
def test_empty_canonical_env_falls_through_to_legacy(monkeypatch):
# The branch _raw_env exists for: set but blank must not mask the legacy name.
monkeypatch.setenv(ADMISSION_CONTROL_ENV, " ")
monkeypatch.setenv(llama_admission._LEGACY_ENV[ADMISSION_CONTROL_ENV], "0")
assert llama_admission_config_from_env().enabled is False
def test_explicit_queue_per_slot_is_not_floored(monkeypatch):
# The floor exists so a 1-slot backend keeps its old depth by default, not to
# override an operator who asked for a shallow line.
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, "2")
config = llama_admission_config_from_env()
assert config.queue_limit(1) == 2
assert config.queue_limit(8) == 16
# Unset, the default multiplier is floored instead.
monkeypatch.delenv(ADMISSION_QUEUE_PER_SLOT_ENV, raising = False)
assert llama_admission_config_from_env().queue_limit(1) == 64
# A value that does not parse falls back to the default multiplier, so it has
# to keep the default's floor. Otherwise a typo quietly shrinks the line 4x.
for garbage in ("abc", "1e3", "16.0"):
monkeypatch.setenv(ADMISSION_QUEUE_PER_SLOT_ENV, garbage)
assert llama_admission_config_from_env().queue_limit(1) == 64, garbage
def test_module_imports_on_python_39(monkeypatch):
"""No 3.10+ API on an import path. The package declares >=3.9 but CI only
runs 3.12, so a regression here would ship broken."""
import ast
import pathlib
src = pathlib.Path(llama_admission.__file__).read_text(encoding = "utf-8")
tree = ast.parse(src)
# int.bit_count() (3.10+)
assert not [
n
for n in ast.walk(tree)
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "bit_count"
]
# dataclass(slots = ...) is 3.10+, so every dataclass must take it through
# the version gate instead of naming it. A new one that forgets the gate
# loses slots silently, so require the **_SLOTS unpack rather than allow it.
seen = 0
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
name = getattr(node.func, "id", None) or getattr(node.func, "attr", None)
if name != "dataclass":
continue
seen += 1
assert "slots" not in {kw.arg for kw in node.keywords}
assert [
kw
for kw in node.keywords
if kw.arg is None and getattr(kw.value, "id", None) == "_SLOTS"
], ast.dump(node)
assert seen
def test_slots_gate_matches_the_running_interpreter():
"""The gate is only worth having if it actually applies where it can."""
import sys
gated = (LlamaAdmissionConfig, llama_admission.LlamaAdmissionSnapshot, llama_admission._Waiter)
if sys.version_info >= (3, 10):
assert llama_admission._SLOTS == {"slots": True}
for cls in gated:
assert getattr(cls, "__slots__", None), cls
else:
assert llama_admission._SLOTS == {}
# Construct through the gate either way: slots=True rebuilds the class, so a
# field it cannot carry over would only show up on instantiation.
config = LlamaAdmissionConfig(max_queue = 7)
assert config.max_queue == 7 and config.queue_limit(4) == 7
assert llama_admission.LlamaAdmissionSnapshot("k", 1, 1, 0).capacity == 1
def test_held_count_tracks_the_bitmask():
# _held replaces int.bit_count(); the two must never drift apart.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
popcount = lambda: bin(queue._in_use).count("1")
leases = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
assert queue._held == popcount() == 4
leases[1].release()
assert queue._held == popcount() == 3
shrunk = queue.reserve(capacity = 2, config = config) # shrink with slots held
assert queue._held == popcount() == 3
shrunk.cancel() # else it is granted a slot as the others drain
for lease in leases:
lease.release()
assert queue._held == popcount() == 0
asyncio.run(_run())
def test_snapshot_free_never_exceeds_what_can_be_admitted():
# After a shrink, low ids can sit in _free while holdovers fill the ceiling.
# Reporting them as free made the admission log contradict itself.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
held = [queue.reserve(capacity = 4, config = config).lease_nowait() for _ in range(4)]
queue.reserve(capacity = 1, config = config) # capacity collapses to 1
min(held, key = lambda lease: lease.slot).release()
snapshot = queue.snapshot()
assert snapshot.free == 0, snapshot # nothing is actually takeable
assert snapshot.active == 3
for lease in held:
lease.release()
asyncio.run(_run())
def test_a_newcomer_does_not_barge_past_a_parked_waiter():
# Anti-starvation, pinned as behaviour rather than as the `if not self._waiters`
# check: _take_slot_locked consults _can_admit_locked anyway, so either alone
# refuses the newcomer. This fails if both ever go.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
held = queue.reserve(capacity = 1, config = config).lease_nowait()
parked = queue.reserve(capacity = 1, config = config)
assert parked.lease_nowait() is None
held.release()
newcomer = queue.reserve(capacity = 1, config = config)
assert newcomer.lease_nowait() is None, "newcomer barged past the parked waiter"
assert (await parked.wait(0.1)) is not None
asyncio.run(_run())
def test_dead_waiters_stop_counting_against_the_queue_limit():
# A future cancelled out of band leaves the entry in the deque: cancel() is not
# called, so only the prune drops it. Without that, depth, is_idle() and the
# queue-full limit all drift for the life of the queue.
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(max_queue = 2)
held = queue.reserve(capacity = 1, config = config).lease_nowait()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
assert queue.snapshot().queued == 2
with pytest.raises(LlamaAdmissionQueueFull):
queue.reserve(capacity = 1, config = config)
first._waiter.future.cancel()
second._waiter.future.cancel()
assert queue.snapshot().queued == 0, "dead waiters still occupy the line"
# The freed depth is usable again, and an idle queue is evictable.
queue.reserve(capacity = 1, config = config).cancel()
held.release()
assert queue.is_idle()
asyncio.run(_run())

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"

1798
studio/backend/tests/test_openai_auto_download.py Normal file
View file

File diff suppressed because it is too large Load diff

553
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,19 @@ 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 without this a
test that exercises the hook can publish its own fixture's scan and, inside the
TTL, hand it to the next test that expects a fresh one.
"""
resolver.invalidate_index()
yield
resolver.invalidate_index()
class _FakeBackend:
effective_parallel_slots = 1
_slot_save_binary = None
@ -94,7 +108,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 +130,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 +405,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 +594,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 +1936,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 +3009,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 +3356,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 +3381,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 +4057,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():
@ -4485,3 +4759,236 @@ def test_removal_of_a_path_still_only_touches_the_exact_key(monkeypatch):
)
# A different file must survive its neighbour being forgotten.
assert settings.get_model_override("/models/foo.gguf")["max_seq_length"] == 8192
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 path could not load it: with auto-download on it
# probed the Hub and 404d on a quant that was never a quant, and with it off it
# refused without switching. 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 request was answered
# by the resident model instead. 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 it
# with no evidence about any local model 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
# starting an F16 on a box sized for the Q4 sitting right next to it, 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):
# Retaining the old index covers what was already known, but nothing covers the
# model that just landed until the next scan finishes. A bare request for it in
# that window was answered by the unrelated resident model.
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 while another model is
# resident, instead of letting a foreign id fall through, and would kick off a
# multi-directory model 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 counted as satisfied by a
# resident /srv/models/Foo.gguf and returned before the case-preserving compare
# further down ever ran. A repo alias must 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

163
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,157 @@ 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 separately 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 and carries a GGUF quant. A
# Transformers model live from a directory that also holds GGUF exports is not one
# of them, and marking the alias loaded had the usage examples pin alias:quant that
# nothing can serve while switching is 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)

103
studio/backend/tests/test_public_check_optout.py Normal file
View file

@ -0,0 +1,103 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Coverage for UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK (#7307 Problem 8).
A wildcard bind asks ifconfig.me for the public IP and check-host.net whether the
port is reachable. Both stay on by default; setting the var skips both, which is
what lab and privacy-sensitive deployments asked for.
"""
import socket
import urllib.request
import pytest
import run
from run import (
DISABLE_PUBLIC_CHECK_ENV,
_resolve_external_ip,
_verify_global_reachability,
public_check_disabled,
)
IFCONFIG = "https://ifconfig.me"
CHECK_HOST = "check-host.net"
class _FakeSocket:
"""Stand-in for the step 3 UDP route lookup."""
def connect(self, addr):
pass
def getsockname(self):
return ("192.168.1.50", 0)
def close(self):
pass
@pytest.fixture
def calls(monkeypatch):
"""Record every outbound URL and fail it, so resolution reaches the LAN step."""
seen = []
def _urlopen(req, *args, **kwargs):
seen.append(req if isinstance(req, str) else req.full_url)
raise OSError("no network in this test")
monkeypatch.setattr(urllib.request, "urlopen", _urlopen)
monkeypatch.setattr(socket, "socket", lambda *a, **k: _FakeSocket())
monkeypatch.delenv(DISABLE_PUBLIC_CHECK_ENV, raising = False)
return seen
# ── public_check_disabled ───────────────────────────────────────────
def test_enabled_by_default(monkeypatch):
monkeypatch.delenv(DISABLE_PUBLIC_CHECK_ENV, raising = False)
assert public_check_disabled() is False
@pytest.mark.parametrize("raw", ["1", "true", "TRUE", "Yes", " 1 "])
def test_disabling_values(monkeypatch, raw):
monkeypatch.setenv(DISABLE_PUBLIC_CHECK_ENV, raw)
assert public_check_disabled() is True
@pytest.mark.parametrize("raw", ["0", "false", "no", "off", "", " ", "ture"])
def test_anything_else_leaves_it_on(monkeypatch, raw):
monkeypatch.setenv(DISABLE_PUBLIC_CHECK_ENV, raw)
assert public_check_disabled() is False
# ── the two lookups ─────────────────────────────────────────────────
def test_public_ip_lookup_runs_by_default(calls):
assert _resolve_external_ip() == "192.168.1.50"
assert IFCONFIG in calls
def test_public_ip_lookup_skipped_when_disabled(monkeypatch, calls):
monkeypatch.setenv(DISABLE_PUBLIC_CHECK_ENV, "1")
assert _resolve_external_ip() == "192.168.1.50", "the LAN address still resolves"
assert IFCONFIG not in calls
def test_reachability_probe_runs_by_default(calls):
_verify_global_reachability("95.216.11.2", 8888)
assert any(CHECK_HOST in url for url in calls)
def test_reachability_probe_skipped_when_disabled(monkeypatch, calls, capsys):
monkeypatch.setenv(DISABLE_PUBLIC_CHECK_ENV, "1")
_verify_global_reachability("95.216.11.2", 8888)
capsys.readouterr()
assert not any(CHECK_HOST in url for url in calls)
assert run._public_reachable is None, "skipping must not claim a reachability result"

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

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

@ -209,7 +209,27 @@ 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 for modeling
files importing `from fla.`, and `models/qwen3_5/` only exists from
transformers 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
passes through a no-op instead of exercising the install path. Pinning keeps
these tests hermetic across the whole supported transformers range, the same
way test_hook_does_not_install_tilelang_for_model_outside_allowlist pins it
against newly added FLA model_types.
"""
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 +335,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 +353,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 +380,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 +425,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 +466,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 +497,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 +559,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 +613,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 +677,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 +745,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 +954,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 +984,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 +1097,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",
}

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:

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

@ -3,10 +3,14 @@
"""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 (and auto-switch is too), a
``/v1`` request naming a GGUF repo that is *not* downloaded 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 +33,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 +101,25 @@ 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.
The settings UI reads this so toggling auto-switch off displays and
round-trips the saved value rather than 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: auto-switch is what loads the model once it lands, so
downloading without it would fetch gigabytes nothing can then 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 +195,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 +216,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 +228,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 {

5
studio/frontend/src/features/api-monitor/api-monitor-overlay.tsx
View file

@ -24,6 +24,7 @@ import {
useRef,
useState,
} from "react";
import { isLifecycleEntry, lifecycleLabel } from "./api-monitor-page";
import { useApiMonitorOverlayStore } from "./overlay-store";
import { computeStats } from "./use-api-monitor";
@ -344,7 +345,9 @@ export function ApiMonitorOverlay(): ReactElement | null {
aria-hidden={true}
/>
<span className="shrink-0 truncate text-ui-12p5 font-medium tracking-nav text-nav-fg">
{compactEndpoint(entry.endpoint)}
{isLifecycleEntry(entry)
? lifecycleLabel(entry)
: compactEndpoint(entry.endpoint)}
</span>
<span
className={cn(

121
studio/frontend/src/features/api-monitor/api-monitor-page.tsx
View file

@ -19,6 +19,9 @@ import {
} from "@/components/ui/select";
import { Skeleton } from "@/components/ui/skeleton";
import { usePlatformStore } from "@/config/env";
import { getInferenceStatus, unloadModel } from "@/features/chat/api/chat-api";
import { resolveInferenceCheckpointId } from "@/features/chat/lib/apply-inference-status-to-store";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import type { ApiMonitorEntry } from "@/features/chat/types/api";
import { useSettingsDialogStore } from "@/features/settings";
import { getApiBase, isTauri } from "@/lib/api-base";
@ -31,6 +34,7 @@ import {
Globe02Icon,
PauseIcon,
PlayIcon,
PowerSocket01Icon,
RefreshIcon,
Settings02Icon,
} from "@hugeicons/core-free-icons";
@ -82,6 +86,40 @@ function compactEndpoint(endpoint: string): string {
.replace(V1_PREFIX_RE, "/");
}
// A lifecycle row is a model load/unload/download, not an HTTP call: it carries
// an event and reason instead of a prompt, so there is no payload to expand.
export function isLifecycleEntry(entry: ApiMonitorEntry): boolean {
return entry.kind === "lifecycle";
}
export 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";
}
function statusDotClass(status: ApiMonitorEntry["status"]): string {
switch (status) {
case "running":
@ -234,6 +272,37 @@ function RequestRow({
entry.reply_preview ||
entry.prompt_preview ||
(entry.status === "running" ? "Waiting for output…" : "No preview");
// A load, unload or download has no prompt or reply, so it reads as a status
// line rather than a request with a payload behind it.
if (isLifecycleEntry(entry)) {
return (
<div className="flex w-full min-w-0 flex-col gap-1 border-b border-border/50 bg-muted/25 px-4 py-3 last:border-b-0">
<div className="flex min-w-0 items-center gap-2">
<span
className={cn(
"size-2 shrink-0 rounded-full",
statusDotClass(entry.status),
)}
aria-hidden={true}
/>
<span className="truncate text-ui-13 font-medium text-foreground">
{lifecycleLabel(entry)}
</span>
<span className="ml-auto shrink-0 text-ui-11 tabular-nums text-muted-foreground">
{formatTime(entry.started_at)}
</span>
</div>
<div className="min-w-0 truncate pl-4 text-ui-11 text-muted-foreground">
{entry.model}
</div>
{entry.error ? (
<div className="min-w-0 break-words pl-4 text-ui-11 text-red-600 dark:text-red-400">
{entry.error}
</div>
) : null}
</div>
);
}
return (
<button
type="button"
@ -465,6 +534,34 @@ export function ApiMonitorPage(): ReactElement {
requestDetail,
} = useApiMonitor();
const serverUrl = usePlatformStore((s) => s.serverUrl);
const [unloading, setUnloading] = useState(false);
const [unloadError, setUnloadError] = useState<string | null>(null);
// Manual release of the loaded model, so VRAM can be freed without waiting for
// the idle timer. /unload matches on the internal id, which the monitor does
// not carry (it advertises a host path), so read it from status.
const unloadActiveModel = async (): Promise<void> => {
setUnloading(true);
try {
const status = await getInferenceStatus();
const checkpoint = resolveInferenceCheckpointId(status);
if (!checkpoint) {
setUnloadError(null);
return;
}
await unloadModel({ model_path: checkpoint });
// Same as the chat eject flow: the store still holds the freed checkpoint.
useChatRuntimeStore.getState().clearCheckpoint();
setUnloadError(null);
refresh();
} catch (err: unknown) {
setUnloadError(
err instanceof Error ? err.message : "Failed to unload the model",
);
} finally {
setUnloading(false);
}
};
const [statusFilter, setStatusFilter] = useState<MonitorStatusFilter>("all");
const [query, setQuery] = useState("");
const [selectedId, setSelectedId] = useState<string | null>(null);
@ -555,6 +652,26 @@ export function ApiMonitorPage(): ReactElement {
/>
{paused ? "Resume" : "Pause"}
</Button>
<Button
type="button"
variant="outline"
size="sm"
onClick={() => void unloadActiveModel()}
disabled={unloading || !data?.active_model}
title={
data?.active_model
? `Unload ${data.active_model} and free its VRAM`
: "No model is loaded"
}
className="h-9 gap-1.5 rounded-full"
>
<HugeiconsIcon
icon={PowerSocket01Icon}
strokeWidth={1.75}
className="size-4"
/>
{unloading ? "Unloading" : "Unload"}
</Button>
<Button
type="button"
variant="outline"
@ -665,9 +782,9 @@ export function ApiMonitorPage(): ReactElement {
) : null}
</section>
{error ? (
{error || unloadError ? (
<div className="rounded-xl border border-red-500/40 bg-red-500/5 px-4 py-3 text-sm text-red-600 dark:text-red-400">
{error}
{error || unloadError}
</div>
) : null}

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

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

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(

4
studio/frontend/src/features/settings/components/monitor-link.tsx
View file

@ -5,7 +5,9 @@
// reached from the floating panel; this card is the way in from Settings.
import { Switch } from "@/components/ui/switch";
import { useApiMonitorOverlayStore } from "@/features/api-monitor";
// Direct path, not the barrel: the barrel re-exports the page, which would
// pull it into this chunk and defeat the route's dynamic import.
import { useApiMonitorOverlayStore } from "@/features/api-monitor/overlay-store";
import { getApiMonitor } from "@/features/chat/api/chat-api";
import type { ApiMonitorResponse } from "@/features/chat/types/api";
import { cn } from "@/lib/utils";

330
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: 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,112 @@ 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 running on its own (UNSLOTH_MODEL_IDLE_TTL, switching off) still
// reloads exactly what it freed on the next request. That restores the stored
// checkpoint only, never an arbitrary catalog entry, so it is tracked apart.
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 not evidence that the
// server holds nothing, and feeding those negatives in 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 of
// the same repo 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 +531,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 +608,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 +632,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 +644,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 +664,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 +774,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

@ -170,10 +170,10 @@ export function ApiKeysTab() {
<MonitorLink />
<UsageExamples apiKey={revealed} />
<ModelAutoSwitchSection />
<UsageExamples apiKey={revealed} />
<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",

18
studio/install_llama_prebuilt.py
View file

@ -2504,11 +2504,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).
@ -3050,7 +3050,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 +4264,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 +4680,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"):
@ -5526,7 +5526,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,13 +5539,13 @@ 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")

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

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

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

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

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

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

21
tests/studio/test_model_picker_contracts.py
View file

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

30
tests/studio/test_settings_compact_overflow_contract.py
View file

@ -5,7 +5,10 @@ from pathlib import Path
REPO = Path(__file__).resolve().parents[2]
SETTINGS_DIALOG = REPO / "studio/frontend/src/features/settings/settings-dialog.tsx"
API_MONITOR = REPO / "studio/frontend/src/features/settings/components/api-monitor-console.tsx"
# The monitor moved out of the settings dialog and onto its own page; Settings
# now links to it. The shrink contract still applies to both surfaces.
API_MONITOR_PAGE = REPO / "studio/frontend/src/features/api-monitor/api-monitor-page.tsx"
MONITOR_LINK = REPO / "studio/frontend/src/features/settings/components/monitor-link.tsx"
GENERAL_TAB = REPO / "studio/frontend/src/features/settings/tabs/general-tab.tsx"
@ -16,15 +19,26 @@ def test_dialog_content_can_shrink_inside_the_dialog_grid():
def test_api_monitor_entries_and_expanded_text_can_shrink():
source = API_MONITOR.read_text(encoding = "utf-8")
source = API_MONITOR_PAGE.read_text(encoding = "utf-8")
# Rows and the detail pane sit in flex parents, so they need min-w-0 or a long
# model id or endpoint pushes the layout wider than the viewport.
assert '"flex w-full min-w-0 flex-col gap-1 border-b border-border/50' in source
assert '<section className="flex min-w-0 flex-col gap-1.5">' in source
# Prompt and reply are unbounded user text: they must be height-capped,
# scrollable, and wrap rather than stretch the pane.
assert (
'<article className="min-w-0 rounded-lg border border-border/70 bg-background">' in source
"max-h-72 overflow-auto whitespace-pre-wrap break-words rounded-lg bg-muted/50" in source
)
assert (
'<section className="flex min-w-0 flex-col rounded-lg border border-border/70 bg-background">'
in source
)
assert source.count('className="max-h-44 overflow-auto whitespace-pre-wrap break-words') == 2
# A model id or path has no spaces to wrap on, so it needs break-all.
assert 'className="min-w-0 break-all font-mono' in source
def test_settings_monitor_link_can_shrink():
source = MONITOR_LINK.read_text(encoding = "utf-8")
assert "flex w-full min-w-0 items-center gap-3" in source
# The summary line carries a model id, so it truncates instead of widening
# the settings dialog.
assert '<span className="truncate text-xs text-muted-foreground">' in source
def test_embedding_model_controls_stack_on_the_narrowest_viewports():

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

@ -0,0 +1,214 @@
# 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 that misses these never re-ran, so a finished load left the first
# fetch's name. It must not be gated on having no checkpoint either: 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 the model being
# deleted. Preferring it on the switch setting alone kept naming one /v1/models
# had already proved absent, so 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 after an idle unload 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" until the next poll.
# The catalog is deliberately tri-state, and a failure must stay the unknown one.
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 of the same repo remains, and pinning
# it emitted repo:deleted-quant, a missing-quant 404 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
# The monitor moved out of the settings dialog onto its own page. Settings keeps
# configuration and links across; these contracts follow the behaviour, not the
# old file.
API_MONITOR_TSX = REPO / "studio/frontend/src/features/api-monitor/api-monitor-page.tsx"
MONITOR_LINK_TSX = SETTINGS / "components/monitor-link.tsx"
def test_api_monitor_history_does_not_reorder_under_the_reader():
# The backend retains 50 terminal entries and moves an entry to the front when
# it finishes. The console froze ids while paging; the full page pauses the
# poll instead, which holds the whole list still while a payload is read.
src = API_MONITOR_TSX.read_text(encoding = "utf-8")
assert "paused" in src
assert "setPaused" in src
# Filters and search are what keep 50 rows usable without paging.
assert "filterEntries(" in src
assert "STATUS_FILTERS" in src
def test_api_monitor_renders_lifecycle_rows():
src = API_MONITOR_TSX.read_text(encoding = "utf-8")
assert "export function isLifecycleEntry(" in src
assert 'entry.kind === "lifecycle"' in src
for label in ("Loading model", "Model loaded", "Model unloaded"):
assert label in src
# A lifecycle row has no prompt or reply, so it is not selectable for detail.
assert "if (isLifecycleEntry(entry)) {" in src
def test_auto_switch_section_sits_above_the_usage_examples():
tab = API_KEYS_TAB_TSX.read_text(encoding = "utf-8")
# The console became a link out to the monitor page; ordering still puts
# configuration ahead of the examples that depend on it.
assert tab.index("<MonitorLink />") < tab.index("<ModelAutoSwitchSection />")
assert tab.index("<ModelAutoSwitchSection />") < 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_settings_still_reaches_the_monitor():
# The console is gone, so Settings must still have a way through to it.
link = MONITOR_LINK_TSX.read_text(encoding = "utf-8")
assert 'to: "/api-monitor"' in link
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

8
unsloth/import_fixes.py
View file

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

2
unsloth/models/_utils.py
View file

@ -2471,7 +2471,7 @@ def _get_statistics(statistics = None, force_download = True):
for vendor_file in vendor_files:
path = Path(vendor_file)
if path.is_file():
file_content = path.read_text().lower()
file_content = path.read_text(encoding = "utf-8").lower()
if "amazon" in file_content:
return "aws"
elif "microsoft corporation" in file_content:

2
unsloth/models/loader.py
View file

@ -1585,7 +1585,7 @@ class FastModel(FastBaseModel):
if do_logging:
redirector = contextlib.nullcontext()
else:
redirector = contextlib.redirect_stdout(open(os.devnull, "w"))
redirector = contextlib.redirect_stdout(open(os.devnull, "w", encoding = "utf-8"))
model_types = ["siglip"] + model_types
# Set forced float32 env flag

56
unsloth/models/loader_utils.py
View file

@ -191,19 +191,43 @@ def _get_new_mapper():
.replace("MAP_TO_UNSLOTH_16bit", "NEW_MAP_TO_UNSLOTH_16bit")
)
exec(new_mapper, globals())
# Exec into a throwaway namespace, never globals(). The slice also carries
# FLOAT_TO_FP8_BLOCK_MAPPER / FLOAT_TO_FP8_ROW_MAPPER, the _add_* helpers
# and the builder's loop variables, so exec'ing into globals() would swap
# the FP8 tables this module imported from the installed mapper for the
# ones on GitHub main. This is only a probe for "would a newer Unsloth
# support this name?", so it must not change what the installed version
# resolves; the fetched FP8 tables are returned for the probe to use
# instead of being written over the installed ones.
namespace = {}
exec(new_mapper, namespace)
return (
NEW_INT_TO_FLOAT_MAPPER,
NEW_FLOAT_TO_INT_MAPPER,
NEW_MAP_TO_UNSLOTH_16bit,
namespace["NEW_INT_TO_FLOAT_MAPPER"],
namespace["NEW_FLOAT_TO_INT_MAPPER"],
namespace["NEW_MAP_TO_UNSLOTH_16bit"],
# .get, not []: these two come from the fetched file under its own names (unlike
# the NEW_ names above, renamed here), so an older or renamed mapper.py would
# KeyError into the bare except and take the 4bit half of the probe down too.
# {} is safe: the probe runs only after the installed tables already missed.
namespace.get("FLOAT_TO_FP8_BLOCK_MAPPER", {}),
namespace.get("FLOAT_TO_FP8_ROW_MAPPER", {}),
)
except:
return {}, {}, {}
return {}, {}, {}, {}, {}
def _resolve_with_mappers(
model_name, load_in_4bit, load_in_fp8, int_to_float, float_to_int, map_to_unsloth_16bit
model_name,
load_in_4bit,
load_in_fp8,
int_to_float,
float_to_int,
map_to_unsloth_16bit,
fp8_block = None,
fp8_row = None,
):
# fp8_block/fp8_row default to the installed tables; the newer-mapper probe passes the
# fetched ones so it can answer for new FP8 repos without rebinding the installed ones.
return __get_model_name(
model_name = model_name,
load_in_4bit = load_in_4bit,
@ -211,8 +235,8 @@ def _resolve_with_mappers(
FLOAT_TO_INT_MAPPER = float_to_int,
MAP_TO_UNSLOTH_16bit = map_to_unsloth_16bit,
load_in_fp8 = load_in_fp8,
FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER,
FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER,
FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER if fp8_block is None else fp8_block,
FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER if fp8_row is None else fp8_row,
)
@ -252,9 +276,13 @@ def get_model_name(
and not _env_says_offline() # offline: skip the remote (raw GitHub) mapper refresh
):
# Try checking if a new Unsloth version allows it!
NEW_INT_TO_FLOAT_MAPPER, NEW_FLOAT_TO_INT_MAPPER, NEW_MAP_TO_UNSLOTH_16bit = (
_get_new_mapper()
)
(
NEW_INT_TO_FLOAT_MAPPER,
NEW_FLOAT_TO_INT_MAPPER,
NEW_MAP_TO_UNSLOTH_16bit,
NEW_FP8_BLOCK_MAPPER,
NEW_FP8_ROW_MAPPER,
) = _get_new_mapper()
upgraded_model_name = _resolve_with_mappers(
model_name = model_name,
load_in_4bit = load_in_4bit,
@ -262,6 +290,10 @@ def get_model_name(
int_to_float = NEW_INT_TO_FLOAT_MAPPER,
float_to_int = NEW_FLOAT_TO_INT_MAPPER,
map_to_unsloth_16bit = NEW_MAP_TO_UNSLOTH_16bit,
# the fp8 probe has to look at the FETCHED tables too, or a new fp8 repo would
# miss both here and in the installed tables and skip the upgrade message
fp8_block = NEW_FP8_BLOCK_MAPPER,
fp8_row = NEW_FP8_ROW_MAPPER,
)
if upgraded_model_name is not None:
raise NotImplementedError(
@ -450,7 +482,7 @@ def _load_fp8_weight_map(
index_path = None
if index_path is not None:
import json
with open(index_path, "r") as f:
with open(index_path, "r", encoding = "utf-8") as f:
return json.load(f).get("weight_map", None)
# Unsharded single file: map every tensor to it.

47
unsloth/models/sentence_transformer.py
View file

@ -186,32 +186,23 @@ def _save_pretrained_gguf(
if tokenizer is None:
tokenizer = self.tokenizer
# 4. Patch environment so Unsloth treats this embedding model correctly
@contextlib.contextmanager
def patch_unsloth_gguf_save():
# Prevent deletion of the directory self.save_pretrained just created
original_rmtree = shutil.rmtree
try:
yield
finally:
shutil.rmtree = original_rmtree
# 4. Call Unsloth's GGUF saver on the inner model targeting the transformer subdirectory
# No rmtree guard here: the merge cleanup that deletes save_directory is gated on
# push_to_hub, which is forced False below.
result = unsloth_save_pretrained_gguf(
inner_model,
save_directory = transformer_dir,
tokenizer = tokenizer,
quantization_method = quantization_method,
first_conversion = first_conversion,
push_to_hub = False, # Force local first to move files
token = token,
max_shard_size = max_shard_size,
temporary_location = temporary_location,
maximum_memory_usage = maximum_memory_usage,
)
# 5. Call Unsloth's GGUF saver on the inner model targeting the transformer subdirectory
with patch_unsloth_gguf_save():
result = unsloth_save_pretrained_gguf(
inner_model,
save_directory = transformer_dir,
tokenizer = tokenizer,
quantization_method = quantization_method,
first_conversion = first_conversion,
push_to_hub = False, # Force local first to move files
token = token,
max_shard_size = max_shard_size,
temporary_location = temporary_location,
maximum_memory_usage = maximum_memory_usage,
)
# 6. Move GGUF files from the subdirectory (0_Transformer) to the root save_directory
# 5. Move GGUF files from the subdirectory (0_Transformer) to the root save_directory
gguf_files = result.get("gguf_files", [])
new_gguf_locations = []
@ -241,7 +232,7 @@ def _save_pretrained_gguf(
result["gguf_files"] = new_gguf_locations
# 7. Add branding
# 6. Add branding
try:
FastSentenceTransformer._add_unsloth_branding(save_directory)
@ -256,7 +247,7 @@ def _save_pretrained_gguf(
except:
pass
# 8. Handle Push to Hub if requested
# 7. Handle Push to Hub if requested
if push_to_hub:
if token is None:
token = get_token()
@ -2329,7 +2320,7 @@ def _patch_st_trainer_load_from_checkpoint():
if not os.path.isfile(modules_json):
raise RuntimeError("Unsloth: PEFT checkpoint is missing modules.json.")
try:
with open(modules_json, "r") as f:
with open(modules_json, "r", encoding = "utf-8") as f:
module_configs = json.load(f)
except Exception as e:
raise RuntimeError("Unsloth: Cannot parse checkpoint modules.json.") from e

6
unsloth_cli/_inference.py
View file

@ -98,9 +98,9 @@ def _json_rank_count_from_env(name: str) -> Optional[int]:
if value.lstrip().startswith(("[", "{")):
data = json.loads(value)
else:
with open(value, "r") as f:
with open(value, "r", encoding = "utf-8") as f:
data = json.load(f)
except (OSError, json.JSONDecodeError):
except (json.JSONDecodeError, OSError, UnicodeDecodeError):
return None
if isinstance(data, list):
return len(data)
@ -158,7 +158,7 @@ def quiet_if_nonzero_mlx_rank():
sys.stderr.flush()
saved_stdout_fd = os.dup(1)
saved_stderr_fd = os.dup(2)
with open(os.devnull, "w") as devnull:
with open(os.devnull, "w", encoding = "utf-8") as devnull:
try:
os.dup2(devnull.fileno(), 1)
os.dup2(devnull.fileno(), 2)

36
unsloth_cli/claude_subagent_mcp.py
View file

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

73
unsloth_cli/commands/start.py
View file

@ -4,6 +4,7 @@
"""`unsloth start` — launch a coding agent against a running Unsloth server."""
import atexit
import base64
import contextlib
import json
import os
@ -2052,6 +2053,32 @@ def _opencode_subagent_inline_config(path: Path, permission: dict) -> dict:
return inline
def _b64_path(path: Path) -> str:
"""Path as base64, so it can cross a shell without being expanded."""
return base64.b64encode(str(path).encode("utf-8")).decode("ascii")
_CLAUDE_PLAN_GATE_SCRIPT = '''\
"""Deny the editing agent while the parent session is in plan mode."""
import json, sys
try:
mode = (json.load(sys.stdin) or {}).get("permission_mode")
except Exception:
sys.exit(0) # fail open: a hook error must never block the parent session
if mode == "plan":
print(json.dumps({"hookSpecificOutput": {
"hookEventName": "PreToolUse",
"permissionDecision": "deny",
"permissionDecisionReason": (
"Plan mode is active. Call the read-only Unsloth plan agent "
"(unsloth_plan_agent) instead of unsloth_agent."
),
}}))
sys.exit(0)
'''
def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path:
"""Write a session plugin that exposes the local Claude child through MCP."""
plugin = path / "unsloth-local-agent"
@ -2094,6 +2121,52 @@ def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path:
}
},
)
# Claude already refuses the editing tool in plan mode, since it advertises
# readOnlyHint false. This PreToolUse hook replaces that dead end with a reason
# naming the read-only tool to call instead. Skipped under the WSL bridge, where
# the gate is a Linux path but the hook would run beside the Windows claude.
gate = plugin / "hooks" / "plan_gate.py"
if command == "wsl.exe":
# A persisted plugin dir may still hold a gate from an earlier non-WSL run.
for stale in (gate, plugin / "hooks" / "hooks.json"):
stale.unlink(missing_ok = True)
else:
_write_private_text(gate, _CLAUDE_PLAN_GATE_SCRIPT)
_write_private_json(
plugin / "hooks" / "hooks.json",
{
"hooks": {
"PreToolUse": [
{
"matcher": _CLAUDE_SUBAGENT_TOOL,
"hooks": [
{
"type": "command",
# Run through runpy rather than handing the path to
# the interpreter: a missing gate is then an
# ordinary traceback (exit 1, fails open) instead
# of exit 2, which Claude treats as a blocking
# error and would deny the tool in every mode.
# The path is base64'd because this string goes
# through a shell: a temp root holding $(..) or a
# backtick expands under sh, %VAR% under cmd, and
# the gate then silently fails open. base64's
# alphabet has no metacharacter in either.
"command": (
f'"{sys.executable}" -c '
f'"import base64,runpy; runpy.run_path('
f"base64.b64decode('{_b64_path(gate)}').decode())\""
),
# A hook with no timeout stalls the parent for as
# long as it hangs; measured unbounded past 400s.
"timeout": 10,
}
],
}
]
}
},
)
skill = plugin / "skills" / "local-agent" / "SKILL.md"
skill.parent.mkdir(parents = True, exist_ok = True, mode = 0o700)
skill.write_text(

8
unsloth_cli/commands/studio.py
View file

@ -335,7 +335,7 @@ def _iter_editable_studio_source_roots(venv_dir: 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- or multi-line dict literals; [^}]* still
# rejects nested dicts, which the setuptools template never
@ -719,7 +719,7 @@ def _cli_update_password(conn: sqlite3.Connection, username: str, new_password:
# credential after a later reset-password deletes auth.db. Mirrors
# backend clear_bootstrap_password().
try:
stale_path.write_text("")
stale_path.write_text("", encoding = "utf-8")
cleared = True
except OSError:
cleared = False
@ -2406,7 +2406,7 @@ def stop():
typer.echo("No running Unsloth server found (no PID file).")
raise typer.Exit(0)
pid_text = _PID_FILE.read_text().strip()
pid_text = _PID_FILE.read_text(encoding = "utf-8").strip()
if not pid_text.isdigit():
typer.echo(f"Invalid PID file contents: {pid_text}")
_PID_FILE.unlink(missing_ok = True)
@ -2863,7 +2863,7 @@ def reset_password():
path.unlink(missing_ok = True)
except OSError:
try:
path.write_text("")
path.write_text("", encoding = "utf-8")
except OSError as exc:
typer.echo(
f"Error: could not remove or clear {path.name} ({exc}); delete "

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

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

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