diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index d1bea819eb..489ee4ca08 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -7,7 +7,7 @@ # # Why a separate workflow: # - studio-backend-ci.yml's "Repo tests (CPU)" job already auto-discovers -# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 16 +# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 17 # Bucket-A tests below live inside those --ignore dirs (CPU-runnable but # historically excluded with their GPU siblings); pulling them out into # a sibling job keeps the existing 760-passed baseline stable while we @@ -274,6 +274,7 @@ jobs: tests/saving/test_export_dispatch.py \ tests/saving/test_imatrix_export.py \ tests/saving/test_gguf_single_pass_export.py \ + tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py @@ -365,6 +366,7 @@ jobs: tests/saving/test_export_dispatch.py \ tests/saving/test_imatrix_export.py \ tests/saving/test_gguf_single_pass_export.py \ + tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ @@ -2129,7 +2131,7 @@ jobs: pip show unsloth_zoo echo "::endgroup::" echo "Consolidated job done. Coverage:" - echo " - 16 unsloth Bucket-A tests under tests/saving/ + tests/utils/" + echo " - 17 unsloth Bucket-A tests under tests/saving/ + tests/utils/" echo " - unsloth_zoo @ ${UNSLOTH_ZOO_REF} pytest tests/ (5 GPU cases deselected)" echo " - unsloth_zoo.compiler.test_apply_fused_lm_head" diff --git a/.gitignore b/.gitignore index 39ca2226ca..fafd17aa95 100644 --- a/.gitignore +++ b/.gitignore @@ -238,4 +238,5 @@ package-lock.json !studio/package-lock.json llama.cpp/ # Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo. -/~/ +~/ +/temp/ diff --git a/README.md b/README.md index 514454f985..e0fc8ee44c 100644 --- a/README.md +++ b/README.md @@ -103,7 +103,7 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**. * **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more * **macOS:** Training, MLX and GGUF inference are ALL supported. * **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd). -* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819). +* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819). Vulkan accelerates GGUF inference only; training still requires a supported PyTorch or MLX backend. * **Multi-GPU:** Available now, with a major upgrade on the way #### macOS, Linux, WSL: @@ -112,12 +112,28 @@ curl -fsSL https://unsloth.ai/install.sh | sh ``` Use the same command to update. +To force the Vulkan llama.cpp backend, set `UNSLOTH_FORCE_VULKAN=1` **before installing or updating**. The setting selects the llama.cpp binary bundle, so setting it only when launching Studio cannot replace an existing CPU bundle: + +```bash +export UNSLOTH_FORCE_VULKAN=1 +curl -fsSL https://unsloth.ai/install.sh | sh +``` + #### Windows: ```powershell irm https://unsloth.ai/install.ps1 | iex ``` Use the same command to update. +To force the Vulkan llama.cpp backend, set the environment variable before running the installer or updater: + +```powershell +$env:UNSLOTH_FORCE_VULKAN=1 +irm https://unsloth.ai/install.ps1 | iex +``` + +Re-running the current installer replaces a previously selected CPU bundle when the backend differs. A separate Vulkan SDK is not required; the GPU driver must provide a working Vulkan runtime. + #### Launch ```bash unsloth studio -p 8888 @@ -263,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`): diff --git a/install.sh b/install.sh index 4cf14a5743..3a1454d742 100755 --- a/install.sh +++ b/install.sh @@ -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 @@ -3307,10 +3352,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 @@ -3338,6 +3393,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) @@ -3564,6 +3670,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 @@ -3605,13 +3712,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) @@ -3802,8 +3914,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)..." @@ -3882,6 +3999,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 diff --git a/scripts/lint_workflow_triggers.py b/scripts/lint_workflow_triggers.py index 0688f6c65c..8f22fcaf45 100644 --- a/scripts/lint_workflow_triggers.py +++ b/scripts/lint_workflow_triggers.py @@ -52,14 +52,14 @@ def _normalise_on(on_field): def _load_workflow(path: Path): try: - return yaml.safe_load(path.read_text()) + return yaml.safe_load(path.read_text(encoding = "utf-8")) except Exception as exc: print(f"ERROR: failed to parse {path}: {exc}", file = sys.stderr) sys.exit(2) def _extract_cache_keys(path: Path) -> list[str]: - text = path.read_text() + text = path.read_text(encoding = "utf-8") keys: list[str] = [] for m in re.finditer(r"(?:^|\n)\s*key:\s*([^\n]+)", text): keys.append(m.group(1).strip()) @@ -104,7 +104,7 @@ def main() -> int: for t in RESTRICTED_TRIGGERS: if t in triggers: - text = path.read_text() + text = path.read_text(encoding = "utf-8") if "lint:workflow_triggers-allow-workflow_run" not in text: findings.append( f"{path.name}: RESTRICTED trigger '{t}' requires an " diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 1c21f8da86..58b7f95ab1 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -98,6 +98,14 @@ "evidence": "L587: while True: sha256:06c2c7f15d73bf192e5e3272c5ff5fcaeff7f6774fef5f4eca6ef473ae50e2b3", "evidence_hash": "57acd497f404c203e4450d0580ad85aa8a33406e8d64ad06fbac6cf47d97b24d" }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L592: while True: sha256:84283c09277ded3296998b2a6a838744457b606829cf5ab5d0da6f222ff020a0", + "evidence_hash": "a7295004315e26a8f3c64fb837521e9fdd7268219bb43e000fb0236ab0259223" + }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index dfb8fc513e..94df994928 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -164,6 +164,21 @@ async def get_current_subject_allow_password_change( ) +# The literal the examples ship with; pasted unedited more often than a revoked key. +API_KEY_PLACEHOLDER = f"{API_KEY_PREFIX}YOUR_KEY" + + +def _invalid_api_key_detail(token: str) -> str: + """Why the key failed. Only the example placeholder is called out; every real + key gets one indistinguishable message, so this leaks no key existence.""" + if token == API_KEY_PLACEHOLDER: + return ( + "This is the placeholder key from the example. Create an API key in " + f"Unsloth Studio under Settings > API and use it in place of {API_KEY_PLACEHOLDER}." + ) + return "Invalid or expired API key" + + async def _get_current_subject( credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool ) -> str: @@ -176,7 +191,7 @@ async def _get_current_subject( if username is None: raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, - detail = "Invalid or expired API key", + detail = _invalid_api_key_detail(token), ) return username diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 39fa691304..5f80ad89a3 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -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 diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 051d80abfe..bf4a6a44b5 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -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 diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index e364ea4f3a..4979ebd48d 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -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: diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py index 706346daad..4bfefc21ce 100644 --- a/studio/backend/core/inference/_vulkan_probe.py +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -6,12 +6,14 @@ Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the Vulkan instance never lives in the long-running backend process. Loads the bundled ggml Vulkan backend from ```` and prints one -``\\t\\t\\t`` 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. +``\\t\\t\\t\\t`` 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 diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py index f76a38576f..ce32a6d3ef 100644 --- a/studio/backend/core/inference/api_monitor.py +++ b/studio/backend/core/inference/api_monitor.py @@ -52,6 +52,13 @@ class ApiMonitorEntry: total_tokens: Optional[int] = None total_tokens_authoritative: bool = False error: Optional[str] = None + # "request" (HTTP call) or "lifecycle" (model load/unload: event/reason, not a prompt; shared). + kind: str = "request" + event: Optional[str] = None + reason: Optional[str] = None + shared: bool = False + # 0-100 for a running download row; None when not applicable. + progress: Optional[float] = None def snapshot(self, *, include_details: bool = True) -> dict[str, Any]: duration_ms = None @@ -85,6 +92,10 @@ class ApiMonitorEntry: "completion_tokens": self.completion_tokens, "total_tokens": self.total_tokens, "error": self.error, + "kind": self.kind, + "event": self.event, + "reason": self.reason, + "progress": self.progress, } if include_details: payload["prompt"] = self.prompt @@ -127,6 +138,73 @@ class ApiMonitor: self._trim_terminal_locked() return entry.id + def record_lifecycle( + self, + *, + event: str, + model: str, + reason: Optional[str] = None, + running: bool = False, + ) -> str: + """Record a model load/unload alongside the request traffic that caused it. + + ``running=True`` opens the row for the caller to close with :meth:`finish` / + :meth:`fail`; an unload is terminal on arrival. Rows are shared (visible to + every subject) and share the request retention budget. + """ + now = time.time() + entry = ApiMonitorEntry( + id = f"apievt_{uuid.uuid4().hex[:12]}", + endpoint = f"model.{event}", + method = "", + model = model or "default", + prompt = "", + status = "running" if running else "completed", + started_at = now, + updated_at = now, + started_monotonic = time.monotonic(), + finished_at = None if running else now, + finished_monotonic = None if running else time.monotonic(), + kind = "lifecycle", + event = event, + reason = reason, + shared = True, + ) + with self._lock: + self._entries.appendleft(entry) + self._trim_terminal_locked() + return entry.id + + def relabel(self, entry_id: Optional[str], model: str) -> None: + """Rename an open lifecycle row once the load resolves its real id: up front + the caller only has the load path, which may be an HF snapshot dir.""" + if not entry_id or not model: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is not None: + entry.model = model + entry.updated_at = time.time() + + def set_progress(self, entry_id: Optional[str], progress: Optional[float]) -> None: + """Update an open download row's percentage (clamped to 0-100).""" + if not entry_id or progress is None: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is not None and entry.status == "running": + entry.progress = min(100.0, max(0.0, float(progress))) + entry.updated_at = time.time() + + def discard(self, entry_id: Optional[str]) -> None: + """Drop a row that turned out not to be an event (an already-satisfied load).""" + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is not None: + self._entries.remove(entry) + def append_reply(self, entry_id: Optional[str], text: str) -> None: if not entry_id or not text: return @@ -212,6 +290,18 @@ class ApiMonitor: self._entries.appendleft(entry) self._trim_terminal_locked() + def fail_open(self, entry_id: Optional[str], error: str) -> None: + """Fail only a still-open row: unlike :meth:`fail`, a catch-all in a + ``finally`` cannot stamp an error onto a request that already succeeded.""" + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None or entry.finished_at is not None: + return + # Same lock as the check, so a finish() cannot land in between. + self._fail_locked(entry, error) + def fail(self, entry_id: Optional[str], error: str) -> None: if not entry_id: return @@ -224,15 +314,18 @@ class ApiMonitor: if error: entry.error = _trim(error, 1000) return - now = time.time() - entry.status = "error" - entry.error = _trim(error, 1000) - entry.updated_at = now - entry.finished_at = now - entry.finished_monotonic = time.monotonic() - self._entries.remove(entry) - self._entries.appendleft(entry) - self._trim_terminal_locked() + self._fail_locked(entry, error) + + def _fail_locked(self, entry: ApiMonitorEntry, error: str) -> None: + now = time.time() + entry.status = "error" + entry.error = _trim(error, 1000) + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() def snapshot( self, @@ -244,7 +337,7 @@ class ApiMonitor: return [ entry.snapshot(include_details = include_details) for entry in self._entries - if subject is None or entry.subject == subject + if self._visible(entry, subject) ] def get( @@ -257,22 +350,29 @@ class ApiMonitor: entry = self._find_locked(entry_id) if entry is None: return None - if subject is not None and entry.subject != subject: + if not self._visible(entry, subject): return None return entry.snapshot(include_details = True) def active_count(self, *, subject: Optional[str] = None) -> int: + # Lifecycle rows show as "running" while loading but are not in-flight API requests. with self._lock: return sum( 1 for entry in self._entries - if entry.status == "running" and (subject is None or entry.subject == subject) + if entry.status == "running" + and entry.kind != "lifecycle" + and (subject is None or entry.subject == subject) ) def clear(self) -> None: with self._lock: self._entries.clear() + @staticmethod + def _visible(entry: ApiMonitorEntry, subject: Optional[str]) -> bool: + return subject is None or entry.subject == subject or entry.shared + def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]: for entry in self._entries: if entry.id == entry_id: diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 2f46470091..563a6732a1 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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: diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py index b6a939c87b..1a9ae04b0e 100644 --- a/studio/backend/core/inference/llama_admission.py +++ b/studio/backend/core/inference/llama_admission.py @@ -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)), ) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 397cba8842..ff966e8446 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -127,6 +127,15 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = ( "then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)" ) +# Shared by the route, pre-teardown and post-metadata rejections (#7205). +_VULKAN_DIFFUSION_GPU_IDS_ERROR = ( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." +) + # llama-server can serve HTTP 200 while running a model entirely on CPU when a # GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so @@ -237,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"): @@ -560,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 @@ -574,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 @@ -611,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 @@ -3492,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`` 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`` 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: @@ -3528,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(), @@ -3545,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``. + """ + 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( @@ -3568,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: " @@ -3587,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 @@ -4710,6 +4755,13 @@ class LlamaCppBackend: probe._read_gguf_metadata(gguf_path) return probe._is_diffusion + def _reject_vulkan_diffusion_gpu_ids_before_teardown( + self, gguf_path: str, model_identifier: str + ) -> None: + """Reject Vulkan + gpu_ids for diffusion GGUFs before Phase 1 teardown.""" + if self._gguf_path_is_diffusion(gguf_path, model_identifier): + raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR) + def _read_gguf_metadata(self, gguf_path: str) -> None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -5122,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 @@ -6339,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 @@ -6515,12 +6567,7 @@ class LlamaCppBackend: f"present. Available Vulkan devices: {sorted(_pf_probed)}." ) - # A remote uncached GGUF may only reveal that it needs the - # single-device diffusion runner after download. On Vulkan, an - # explicit gpu_ids request cannot be mapped from ggml ordinals to - # that runner's CUDA physical index. Download and classify the main - # file before killing the healthy server so this late rejection is - # non-destructive. The Phase 2 call below reuses this cached path. + # Classify before killing the healthy server (#7205); Phase 2 reuses this path. _preflight_model_path = None if is_vulkan_backend and gpu_ids and hf_repo: _resolved_repo = _resolve_repo_id_casing(hf_repo) @@ -6537,14 +6584,17 @@ class LlamaCppBackend: hf_variant = hf_variant, hf_token = hf_token, ) - if self._gguf_path_is_diffusion(_preflight_model_path, model_identifier): - raise ValueError( - "GPU selection (gpu_ids) is not supported for a DiffusionGemma " - "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " - "its device by CUDA physical index, which has no defined mapping " - "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " - "device." - ) + self._reject_vulkan_diffusion_gpu_ids_before_teardown( + _preflight_model_path, + model_identifier, + ) + elif is_vulkan_backend and gpu_ids and gguf_path and not hf_repo: + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + self._reject_vulkan_diffusion_gpu_ids_before_teardown( + gguf_path, + model_identifier, + ) # ── Phase 1: kill old process (under lock, fast) ────────── with self._lock: @@ -6621,21 +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: - # The diffusion runner pins its child by CUDA visibility mask, so a - # ggml Vulkan ordinal cannot be honored (wrong GPU / CPU fallback). - # Route and remote-download preflights reject before teardown; keep - # this as a final defense if classification ever disagrees. - if is_vulkan_backend and gpu_ids: - raise ValueError( - "GPU selection (gpu_ids) is not supported for a DiffusionGemma " - "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " - "its device by CUDA physical index, which has no defined mapping " - "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " - "device." - ) # 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") @@ -8297,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 @@ -9472,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}") @@ -9606,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 @@ -11049,16 +11101,20 @@ class LlamaCppBackend: build_rag_autoinject, execute_tool, is_always_safe_tool, - is_potentially_unsafe_tool_call, + is_high_risk_tool_call, ) - # Normalize the mode: "full" and bypass_permissions are the same - # switch, whichever arrives first wins toward the permissive side. - # "off" keeps the sandbox but never prompts. + # "full" and bypass_permissions are the same switch, whichever arrives + # first wins. "off" keeps the sandbox but never prompts. Unset defaults to + # "auto"; unknown falls back to the stricter "ask". An explicit + # confirm_tool_calls=True with no mode is already resolved to "ask" at the + # request layer, so it never arrives here as an ambiguous unset. if permission_mode == "full": bypass_permissions = True elif bypass_permissions: permission_mode = "full" + elif permission_mode is None: + permission_mode = "auto" elif permission_mode not in ("ask", "auto", "off"): permission_mode = "ask" @@ -12152,18 +12208,16 @@ class LlamaCppBackend: decision.as_assistant_tool_call() ) - # Bypass wins over the confirm gate at the loop level too, - # so a direct internal caller with both flags never prompts. - # In "auto" mode only calls detected as potentially unsafe - # pause; read-only calls run straight through. "off" never - # prompts (sandbox stays on). + # Bypass wins here too, so a direct internal caller with both + # flags never prompts. "auto" pauses only high-risk calls; + # "off" never prompts (sandbox stays on). needs_confirm = ( bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off" ) if needs_confirm and permission_mode == "auto": - needs_confirm = is_potentially_unsafe_tool_call( + needs_confirm = is_high_risk_tool_call( decision.tool_name, decision.arguments ) approval_id = new_approval_id() if needs_confirm else "" diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py index 3380ebf5f5..05b1271b27 100644 --- a/studio/backend/core/inference/llama_keepwarm.py +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -345,6 +345,22 @@ def _loaded_identity(backend): return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised) +def _note_idle_unload_event(freed) -> None: + """Monitor row for an idle auto-unload. Best-effort; uses the stash's + advertised repo id so the row never shows the on-disk load path.""" + try: + from core.inference.api_monitor import api_monitor + from core.inference.model_ids import public_model_id + + identifier, variant, advertised = (list(freed) + [None, None, None])[:3] + label = public_model_id(advertised or identifier) or "model" + if variant and ":" not in label: + label = f"{label}:{variant}" + api_monitor.record_lifecycle(event = "unload", model = label, reason = "idle") + except Exception as exc: + logger.debug("idle unload monitor event failed: %s", exc) + + async def idle_unload_loop(poll_seconds: float = 15.0) -> None: """Unload the loaded GGUF once idle past the configured TTL. Inert when off.""" from utils.openai_auto_switch_settings import ( @@ -407,6 +423,8 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None: elif manifest: _delete_resume_files(manifest) logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) + # An idle unload stashes for reload and skips note_model_unloaded. + _note_idle_unload_event(freed) seen_model = None except Exception as exc: logger.debug("idle_unload_loop iteration failed: %s", exc) diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 7b42d2f40d..6f1b931a7f 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -118,6 +118,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: parse_ctx_override(out) parse_cache_override(out) parse_split_mode_override(out) + parse_gpu_layers_override(out) return out @@ -193,9 +194,8 @@ _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS # inherited -ngl is respected (the offload_overridden path), so this group is # opt-in, not default. Layer flags are shared with llama_cpp's override # detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them). -_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset( - {"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"} -) +_GPU_LAYER_FLAGS: frozenset[str] = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers"}) +_LAYER_OFFLOAD_FLAGS: frozenset[str] = _GPU_LAYER_FLAGS | frozenset({"-fit", "--fit"}) _MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"}) _OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS @@ -306,6 +306,26 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return _last_flag_value(args, _CACHE_FLAGS) +def parse_gpu_layers_override(args: Optional[Iterable[str]]) -> Optional[int]: + """Return the last user-supplied GPU layer count from extras. + + Manual GPU memory mode strips llama.cpp offload flags because the + first-class load fields own them. Callers use this parser first to preserve + an explicit ``-ngl`` / ``--gpu-layers`` / ``--n-gpu-layers`` value when + translating the extras into those fields. + """ + raw_value = _last_flag_value(args, _GPU_LAYER_FLAGS) + if raw_value is None: + return None + try: + value = int(raw_value) + except ValueError as exc: + raise ValueError("llama-server GPU layers flag requires an integer value") from exc + if value < -1: + raise ValueError("llama-server GPU layers flag requires an integer value of at least -1") + return value + + def parse_cache_override_per_axis( args: Optional[Iterable[str]], ) -> tuple[Optional[str], Optional[str]]: diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py index 9e3eaeda3f..5d2a9e9c87 100644 --- a/studio/backend/core/inference/local_model_resolver.py +++ b/studio/backend/core/inference/local_model_resolver.py @@ -34,6 +34,15 @@ class _LocalGgufEntry: _CACHE_TTL_S = 5.0 _lock = threading.Lock() _scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {}) +# Not _lock: that is held for the whole scan, so the request path would wait on it. +_warm_lock = threading.Lock() +# Repos that finished downloading but are not in the published index yet: nothing +# else covers them until the next scan, and the request path must not call them absent. +_just_downloaded: set[str] = set() +_warming = False +_last_scan_s = 0.0 +# Rescan at most a tenth of the time: on the TTL alone a slow scan would run continuously. +_WARM_DUTY = 10.0 def _is_abs_path_id(value: str) -> bool: @@ -103,17 +112,26 @@ def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]: load_dir = _resolve_load_dir(p) variants, _ = list_local_gguf_variants(str(load_dir)) quants = tuple(v.quant for v in variants if getattr(v, "quant", None)) - return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None + if not quants: + return None + # That call orders by descending size, so the head is the biggest quant (often + # F16). Downstream reads [0], and a bare id must mean whichever quant a plain + # load would take: answering with the largest can evict a model and then OOM. + from core.inference.openai_auto_download import preferred_quant + + best = preferred_quant(quants) + if best and quants[0] != best: + quants = (best, *(q for q in quants if q != best)) + return _LocalGgufEntry(loader_id, str(load_dir), quants) except Exception: return None -def info_has_local_gguf(info) -> bool: - """True when *info* (a LocalModelInfo) points to on-disk GGUF weights the - auto-switch path can load. Read from the files, not ``info.model_format``: the - HF-cache scanner leaves model_format unset for GGUF snapshots, so a - model_format filter would drop every cached GGUF. Lets /v1/models advertise - exactly what /v1 can serve.""" +def local_gguf_quants(info) -> Optional[tuple[str, ...]]: + """On-disk quant labels for *info*, or None when it is not a servable local + GGUF. Read from the files, not ``info.model_format``: the HF-cache scanner + leaves that unset for GGUF snapshots, so filtering on it drops every cached + GGUF. One scan tells /v1/models what it can serve and which quant to name.""" from pathlib import Path path = getattr(info, "path", None) @@ -123,8 +141,14 @@ def info_has_local_gguf(info) -> bool: if isinstance(path, str) and any( seg in (".studio_links", "ollama_links") for seg in Path(path).parts ): - return False - return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None + return None + entry = _local_gguf_entry(getattr(info, "id", "") or "", info) + return entry.variants if entry is not None else None + + +def info_has_local_gguf(info) -> bool: + """True when *info* points to on-disk GGUF weights the auto-switch path can load.""" + return local_gguf_quants(info) is not None def _build_index() -> dict[str, _LocalGgufEntry]: @@ -147,10 +171,16 @@ def _build_index() -> dict[str, _LocalGgufEntry]: ) from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs from utils.hf_cache_settings import known_hf_hub_caches + from core.inference.model_ids import public_model_id index: dict[str, _LocalGgufEntry] = {} seen_hf: set[str] = set() + try: + active_root = str(Path(_resolve_hf_cache_dir()).resolve()) + except Exception: + active_root = None + def _scan_hf_once(directory) -> list: if directory is None: return [] @@ -162,7 +192,13 @@ def _build_index() -> dict[str, _LocalGgufEntry]: if rp in seen_hf: return [] seen_hf.add(rp) - return _scan_hf_cache(directory) + # Only the active cache loads by repo id. Say so, or an inactive repo is + # indexed under an id it cannot load by, and its snapshot basename (what + # /v1/models advertises once loaded by path) is never a key at all. + # No format classification here: nothing on this path reads model_format, + # and its recursive walk would duplicate the one _local_gguf_entry already + # does per snapshot, on the request path. + return _scan_hf_cache(directory, active_cache = rp == active_root, classify_format = False) except Exception as exc: # a missing/malformed root must skip, never crash the index logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc) return [] @@ -220,12 +256,91 @@ def _build_index() -> dict[str, _LocalGgufEntry]: continue # Index every alias (including the path) so a client can resolve by any of # them, even though only the non-path loader_id is advertised. - for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)): + for key in ( + raw_id, + getattr(info, "model_id", None), + getattr(info, "display_name", None), + public_model_id(raw_id), + ): if key: index.setdefault(key.strip().lower(), entry) + # Other revisions of the same repo resolve to their own weights, so a pin on + # one keeps working after Hugging Face writes a newer snapshot. + for name, sibling_entry in _sibling_revision_entries(raw_id, loader_id): + index.setdefault(name.strip().lower(), sibling_entry) return index +def _sibling_revision_entries(raw_id: str, loader_id: str): + """Yield ``(revision_name, entry)`` for the repo's OTHER cached revisions. + + An inactive-cache repo carries its snapshot path as the id, and /v1/models + advertises only that directory's basename once loaded, so anything durable + pinned to it (a subagent config) holds one revision hash. Hugging Face writes a + new snapshot dir on every update, and the scan emits a single entry per repo + pointed at the newest one, so that pin would otherwise stop resolving and drop + through to whatever model is loaded. + + Each revision gets an entry for its OWN directory rather than an alias onto the + scanned one: aliasing would redirect a pin that names an older complete revision + onto a newer half-downloaded snapshot and break a request that works today. + Incomplete revisions are skipped for the same reason. + + Sibling names are only revisions inside a real cache repo + (``/models--org--name/snapshots/``). A scan folder that merely happens + to be called ``snapshots`` holds unrelated models, and treating those as + revisions would silently serve one model in place of another. + """ + from pathlib import Path + from types import SimpleNamespace + + snapshots = Path(raw_id).parent + if snapshots.name != "snapshots" or not snapshots.parent.name.startswith("models--"): + return + from routes.models import snapshot_variants_all_complete + + try: + siblings = [p for p in snapshots.iterdir() if p.is_dir() and p.name != Path(raw_id).name] + except OSError: + return + for sibling in siblings: + if not snapshot_variants_all_complete(str(sibling)): + continue + entry = _local_gguf_entry(loader_id, SimpleNamespace(path = str(sibling))) + if entry is not None: + yield sibling.name, entry + + +def note_downloaded(repo_id: Optional[str]) -> None: + """Record a repo as present ahead of the scan that will index it.""" + if not repo_id: + return + with _lock: + _just_downloaded.add(repo_id.strip().lower()) + + +def recently_downloaded(repo_id: str) -> bool: + """Whether *repo_id* finished downloading since the last completed scan.""" + if not isinstance(repo_id, str) or not repo_id.strip(): + return False + return repo_id.strip().lower() in _just_downloaded + + +def invalidate_index() -> None: + """Mark the cached scan stale so the next resolve sees a just-finished download + instead of waiting out the TTL. + + Keeps the entries: the request path reads this cache without scanning, so + emptying it would leave it with no evidence about any local model until the + rebuild lands, and a bare request for one would be answered by whatever is + resident. Only a completed download invalidates, and that only adds, so the + retained entries stay true. + """ + global _scan + with _lock: + _scan = (0.0, _scan[1]) + + def _index() -> dict[str, _LocalGgufEntry]: global _scan # Build under the lock so concurrent callers with an expired cache don't all @@ -240,23 +355,74 @@ def _index() -> dict[str, _LocalGgufEntry]: # an install with many local models can itself exceed the TTL, which would # store the cache already expired and make every request rebuild the index. _scan = (time.monotonic(), fresh) + # The scan supersedes the notes: whatever landed is in the index now. + _just_downloaded.clear() return fresh -def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]: +def index_is_built() -> bool: + """Whether a scan has ever completed, freshness aside. + + Lock-free on purpose: ``_lock`` is held for the whole scan, so taking it would + park the request path on the scan it is trying to stay off. Safe because + ``_scan`` is only ever rebound, never mutated. + """ + return bool(_scan[0]) + + +def warm_index_soon() -> None: + """(Re)build the index off the request path when it is missing or past its TTL. + + The only refresh for callers using ``allow_scan=False``. Covers a stale index, + not just an absent one: a model downloaded through the Hub UI or dropped into a + scan folder has no invalidation hook and would otherwise stay invisible to them + for the life of the process. Never blocks, and never touches ``_lock``. + """ + global _warming + if time.monotonic() - _scan[0] < max(_CACHE_TTL_S, _last_scan_s * _WARM_DUTY): + return + with _warm_lock: + if _warming: + return + _warming = True + + def _run() -> None: + global _warming, _last_scan_s + started = time.monotonic() + try: + _index() + except Exception: + pass + finally: + _last_scan_s = time.monotonic() - started + with _warm_lock: + _warming = False + + threading.Thread(target = _run, name = "local-model-index-warm", daemon = True).start() + + +def resolve_local_gguf( + requested: str, *, allow_scan: bool = True +) -> Optional[tuple[str, Optional[str], str]]: """Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None. ``load_path`` is the concrete on-disk path to hand /load (so it never fetches a remote), ``loader_id`` is the advertised id used as the launch-override key. ``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first (so ids containing a colon still resolve); else the last ``:VARIANT`` is split - off and resolves only when that quant is on disk. + off and resolves only when that quant is on disk, unless it names no quant at + all (an Ollama-style ":latest"), which means the repo. + + ``allow_scan=False`` answers from the last built index and never rebuilds, for + the request path: the scan walks several model dirs and HF caches, takes seconds + on a large install, and holds a lock everyone queues behind. Stale is fine there, + since disk barely moves and a finished download calls :func:`invalidate_index`. """ if not isinstance(requested, str) or not requested.strip(): return None requested = requested.strip() try: - index = _index() + index = _index() if allow_scan else _scan[1] entry = index.get(requested.lower()) if entry is not None: variant = entry.variants[0] if entry.variants else None @@ -272,8 +438,44 @@ def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str for v in entry.variants: if v.lower() == wanted: return entry.load_path, v, entry.loader_id - return None + from core.inference.openai_auto_download import looks_like_quant + + if looks_like_quant(variant): + return None + # ":latest" or ":8b" names no file, so it means the repo; a real quant that + # is not on disk still misses, or a swap would serve the wrong weights. + return entry.load_path, (entry.variants[0] if entry.variants else None), entry.loader_id except Exception: # Best-effort: any resolver failure falls through to the loaded model, # so a malformed name can never turn a servable request into a 500. return None + + +MISS_MODEL_NOT_FOUND = "model_not_found" +MISS_VARIANT_NOT_FOUND = "variant_not_found" + + +def describe_local_miss(requested: str) -> tuple[str, tuple[str, ...]]: + """Why :func:`resolve_local_gguf` missed, so an error can say "wrong quant" + instead of "no such model". + + ``(MISS_VARIANT_NOT_FOUND, )`` when the repo is downloaded but the + requested ``:VARIANT`` is not, else ``(MISS_MODEL_NOT_FOUND, ())``. Fail-safe: a + scan failure reports the generic miss rather than raising into the handler. + """ + if not isinstance(requested, str) or not requested.strip(): + return MISS_MODEL_NOT_FOUND, () + base, sep, variant = requested.strip().rpartition(":") + from core.inference.openai_auto_download import looks_like_quant + + # Split like the resolver or the two disagree: a tag naming no quant means the + # repo there, so reporting a missing quant for it would name one nobody asked for. + if not sep or not looks_like_quant(variant): + return MISS_MODEL_NOT_FOUND, () + try: + entry = _index().get(base.strip().lower()) + except Exception: + return MISS_MODEL_NOT_FOUND, () + if entry is None or not entry.variants: + return MISS_MODEL_NOT_FOUND, () + return MISS_VARIANT_NOT_FOUND, entry.variants diff --git a/studio/backend/core/inference/model_ids.py b/studio/backend/core/inference/model_ids.py index 548cc60f94..3886307ae2 100644 --- a/studio/backend/core/inference/model_ids.py +++ b/studio/backend/core/inference/model_ids.py @@ -39,10 +39,29 @@ def _looks_like_path(identifier: str) -> bool: return False +def hf_cache_repo_id(path: Optional[str]) -> Optional[str]: + """``.../models--org--name/snapshots/`` -> ``org/name``, else None. + + A model loaded from the HF cache is identified by its snapshot dir, whose + basename is a commit hash; recover the repo id so callers don't show that. + """ + if not path: + return None + parts = str(path).replace("\\", "/").split("/") + for index, part in enumerate(parts): + # Only inside the real cache layout: a "models--" name alone is not a repo id. + if part.startswith("models--") and parts[index + 1 : index + 2] == ["snapshots"]: + return part[len("models--") :].replace("--", "/") + return None + + def public_model_id(identifier: Optional[str]) -> Optional[str]: """Return a clean, path-free public id for *identifier*. - - Local GGUF path -> the file stem with ``.gguf`` stripped, e.g. + - HF cache path -> the repo id it came from, e.g. + ``~/.cache/huggingface/hub/models--unsloth--X-GGUF/snapshots/`` -> + ``unsloth/X-GGUF``. + - Other local GGUF path -> the file stem with ``.gguf`` stripped, e.g. ``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``. - HF repo id (``org/model``) and already-clean names -> returned unchanged. - ``None`` / empty -> returned unchanged. @@ -51,6 +70,9 @@ def public_model_id(identifier: Optional[str]) -> Optional[str]: return identifier if not _looks_like_path(identifier): return identifier + repo_id = hf_cache_repo_id(identifier) + if repo_id: + return repo_id name = os.path.basename(identifier.replace("\\", "/").rstrip("/")) if name.lower().endswith(_GGUF_SUFFIX): name = name[: -len(_GGUF_SUFFIX)] diff --git a/studio/backend/core/inference/openai_auto_download.py b/studio/backend/core/inference/openai_auto_download.py new file mode 100644 index 0000000000..cad5e40d14 --- /dev/null +++ b/studio/backend/core/inference/openai_auto_download.py @@ -0,0 +1,812 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in: fetch a GGUF a /v1 request names but this server doesn't have. + +Auto-switch only loads models already on disk. With +``openai_api_auto_download_model`` on, a miss that looks like a real Hub repo is +fetched in the background and the request is told to retry rather than held +open: a quant is routinely tens of GB, far longer than any client (or the +Cloudflare edge on ``--secure``) will wait, and the inference lifecycle gate must +not be held meanwhile. The resident model keeps serving, and the retry that lands +after the download goes through the ordinary auto-switch path. + +Admission is deliberately narrow, since a request only needs an API key: +- ``namespace/name`` only, and only when the Hub confirms GGUF weights. A + namespace is not evidence of intent (LiteLLM and OpenRouter address every + provider that way), so ``gpt-4`` and ``anthropic/claude-3.5-sonnet`` alike + fall through to the resident model as before. +- GGUF only, decided from the remote file list, not the repo name: GGUF runs + under llama.cpp, which never imports repo Python. +- ``auto_map`` is refused, so ``trust_remote_code`` is only ever granted + deliberately in the UI, never by an API call. +- One download at a time, so a key holder cannot fan out fetches. +""" + +from __future__ import annotations + +import asyncio +import shutil +import threading +import time +from dataclasses import dataclass +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +# Keep the Hub probe short so a slow Hub can't stall the request path. +_MODEL_INFO_TIMEOUT_S = 8.0 +# auth_check and hf_hub_download take no timeout of their own and run while the +# provisional slot is held, so an unresponsive Hub would pin the single flight. The +# code probe fetches up to three configs, so it gets more room than the auth call. +_CODE_PROBE_TIMEOUT_S = 20.0 +# Headroom left free after the download, so filling the disk can't wedge the box. +_DISK_RESERVE_BYTES = 5 * 1024**3 +_WATCH_POLL_S = 2.0 +# A stalled watcher must not pin the single-flight slot forever. +_MAX_WATCH_S = 24 * 60 * 60 +# Past the watch window the row is resolved, so poll only to see whether the +# worker is still alive and still owns the slot. +_TIMED_OUT_POLL_S = 60.0 +_RETRY_AFTER_S = 30 +# Long enough for a client honouring Retry-After to come back and be told, short +# enough that one that never returns cannot hold the slot. +_FAILED_HOLD_S = 3 * _RETRY_AFTER_S +_MAX_LISTED_VARIANTS = 8 + + +@dataclass(frozen = True) +class AutoDownloadRefusal: + """Why this request cannot be served yet; the route raises it in the + surface's own error envelope.""" + + status: int + code: str + message: str + retry_after: Optional[int] = None + + +@dataclass +class _Active: + repo_id: str + # None while the Hub probe is still deciding which quant to fetch. + variant: Optional[str] = None + expected_bytes: int = 0 + monitor_id: Optional[str] = None + started_at: float = 0.0 + # Set when the worker failed. Held until a retry surfaces it: Retry-After is far + # longer than the watcher poll, so the client would restart the same failing download. + error: Optional[str] = None + failed_at: float = 0.0 + + +_lock = threading.Lock() +_active: Optional[_Active] = None + +# Repos the Hub says are not servable, so a "vendor/model" miss doesn't re-probe every request. +_NOT_SERVABLE_TTL_S = 10 * 60 +_NOT_SERVABLE_MAX = 256 +_cache_lock = threading.Lock() +_not_servable: dict[str, float] = {} + + +def _public_label(repo_id: str, variant: Optional[str]) -> str: + return f"{repo_id}:{variant}" if variant else repo_id + + +def split_model_ref(requested: str) -> tuple[str, Optional[str]]: + """``org/repo:QUANT`` -> ``("org/repo", "QUANT")``; no suffix -> variant None. + + Splits on the last colon. A slash-bearing suffix is only a variant when a real + Hub repo precedes it: "build/llama-13b" is a subdirectory GGUF key the catalog + advertises, while "C:/models/x.gguf" leaves a drive letter that is no repo id. + """ + text = (requested or "").strip() + base, sep, suffix = text.rpartition(":") + if not sep or not base or not suffix: + return text, None + stripped = base.strip() + if "/" in suffix: + from hub.utils.paths import is_valid_repo_id + if "/" not in stripped or not is_valid_repo_id(stripped): + return text, None + return stripped, suffix.strip() + + +def is_downloadable_ref(requested: str) -> bool: + """Whether *requested* is shaped like a Hub repo we may fetch. + + Requires an explicit namespace: keeps ``gpt-4`` and other foreign ids falling + through, and stops ModelConfig.from_identifier's bare-name ``unsloth/`` + prefixing from turning an unrelated label into a real repo. + """ + from hub.utils.paths import is_valid_repo_id + + repo_id, variant = split_model_ref(requested) + if "/" not in repo_id or not is_valid_repo_id(repo_id): + return False + if variant is not None: + from hub.utils.paths import is_valid_gguf_variant + return is_valid_gguf_variant(variant) + return True + + +def looks_like_quant(variant: Optional[str]) -> bool: + """Whether a ``:suffix`` names a GGUF quant rather than a foreign tag. + + Neither a namespace nor a colon proves a request was meant for this server + (``vendor/model`` is LiteLLM/OpenRouter, ``name:latest`` is Ollama). A real + quant label does. + """ + import re + + from utils.models.model_config import _GGUF_KNOWN_QUANT_RE + + if not variant: + return False + # _extract_quant_label can append a bpw modifier (IQ4_XS-3.53bpw); still a quant. + label = re.sub(r"-[0-9]+(?:\.[0-9]+)?bpw$", "", variant.strip(), flags = re.IGNORECASE) + return _GGUF_KNOWN_QUANT_RE.fullmatch(label) is not None + + +def _hub_token(hf_token: Optional[str]): + """The caller's token, or an explicit False. None makes huggingface_hub fall + back to a cached login (here the server owner's); only False is anonymous.""" + return hf_token or False + + +def _servable_key(repo_id: str, hf_token: Optional[str]) -> str: + """Cache key, per credential. + + The Hub 404s a private repo the caller cannot see, so a tokenless verdict says + nothing about a caller who has one. Digested, so no token is held here. + """ + import hashlib + + seen_as = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else "anon" + return f"{repo_id.lower()}\n{seen_as}" + + +def _mark_not_servable(repo_id: str, hf_token: Optional[str]) -> None: + with _cache_lock: + if len(_not_servable) >= _NOT_SERVABLE_MAX: + _not_servable.clear() + _not_servable[_servable_key(repo_id, hf_token)] = time.monotonic() + _NOT_SERVABLE_TTL_S + + +def _is_not_servable(repo_id: str, hf_token: Optional[str]) -> bool: + key = _servable_key(repo_id, hf_token) + with _cache_lock: + expires = _not_servable.get(key) + if expires is None: + return False + if expires <= time.monotonic(): + del _not_servable[key] + return False + return True + + +def _gated_refusal(repo_id: str) -> AutoDownloadRefusal: + return AutoDownloadRefusal( + status = 403, + code = "model_access_denied", + message = ( + f"'{repo_id}' is gated on Hugging Face. Accept its licence, then retry with " + "your own token in the X-Unsloth-HF-Token header: automatic download never " + "uses this server's Hugging Face identity." + ), + ) + + +async def _bounded_probe(fn, *args, timeout: float, default): + """Run a blocking Hub probe off the loop, bounding only the wait. + + The thread is left to finish (a blocking socket read cannot be cancelled); the + caller takes *default*, chosen per call site so a timeout errs the safe way. + """ + try: + return await asyncio.wait_for(asyncio.to_thread(fn, *args), timeout) + except (TimeoutError, asyncio.TimeoutError): + logger.debug("hub probe %s timed out after %ss", getattr(fn, "__name__", fn), timeout) + return default + + +def _auth_denied(repo_id: str, hf_token: Optional[str]) -> bool: + """Whether this token lacks file access to a gated repo. False when the + check is inconclusive: the download's own auth is the real gate.""" + from hub.utils.hf_errors import hf_error_status + + try: + from huggingface_hub import auth_check + auth_check(repo_id, token = _hub_token(hf_token)) + except Exception as exc: + return hf_error_status(exc) in (401, 403) + return False + + +def _gguf_variants(siblings) -> dict[str, int]: + """Quant label -> bytes the download will actually fetch. + + Mirrors list_gguf_variants for the selectable labels: companions (mmproj/MTP) + and big-endian builds are not quants, and sharded quants sum across shards. + Bytes come from the download plan, which folds companions back into every + quant, so the disk reserve is measured against what the worker fetches. + """ + from hub.utils.gguf import extract_quant_label as canonical_quant_label + from hub.utils.gguf_plan import build_gguf_variant_plans + from utils.models.model_config import ( + _extract_quant_label, + _is_big_endian_gguf_path, + _is_mmproj, + _is_mtp_drafter, + ) + + siblings = list(siblings or []) + plans = build_gguf_variant_plans(siblings) + sizes: dict[str, int] = {} + for sibling in siblings: + name = getattr(sibling, "rfilename", "") or "" + if not name.lower().endswith(".gguf"): + continue + quant = _extract_quant_label(name) + if not looks_like_quant(quant): + # With no recognized quant token the extractors part ways: this one takes + # the last hyphenated segment ("7b" of llama-7b) while the plan and worker + # key the whole stem, so advertising ours dispatches an unresolvable variant. + quant = canonical_quant_label(name) or quant + if _is_mmproj(name) or _is_mtp_drafter(name) or _is_big_endian_gguf_path(name, quant): + continue + plan = plans.get(quant.lower()) + if plan is not None: + sizes[quant] = plan.download_size_bytes + else: + sizes[quant] = sizes.get(quant, 0) + int(getattr(sibling, "size", 0) or 0) + return sizes + + +def _remaining_bytes(repo_id: str, plan, expected_bytes: int) -> int: + """Bytes still to fetch: a resumed quant or a companion shared with another + quant is already on disk, and charging for it can 507 a download that fits.""" + try: + from hub.utils.download_registry import existing_blob_bytes + + hashes = frozenset( + file.sha256 for file in getattr(plan, "expected_files", ()) or () if file.sha256 + ) + if not hashes: + return expected_bytes + return max(0, expected_bytes - existing_blob_bytes("model", repo_id, hashes)) + except Exception: + return expected_bytes + + +def _enough_disk(need_bytes: int) -> tuple[bool, int]: + """(fits, free_bytes). Fail-open on an unreadable cache root: the download + worker runs its own preflight, this only adds the reserve margin.""" + try: + from hub.utils.hf_cache_state import hf_cache_root + + root = hf_cache_root(create = True) + if root is None: + return True, 0 + free = shutil.disk_usage(root).free + except Exception: + return True, 0 + return free >= need_bytes + _DISK_RESERVE_BYTES, free + + +def _gb(num_bytes: int) -> str: + return f"{num_bytes / 1024**3:.1f} GB" + + +async def _job_state(repo_id: str, variant: Optional[str]) -> tuple[str, Optional[str]]: + from hub.services.models import downloads + try: + status = await downloads.get_download_status_response(repo_id, variant or "") + return status.state, status.error + except Exception as exc: + # "unknown", not "idle": idle ends the watch, and a failed probe proves nothing. + logger.debug("auto-download: status probe failed for %r: %s", repo_id, exc) + return "unknown", None + + +async def _progress_percent( + repo_id: str, variant: Optional[str], expected_bytes: int, hf_token: Optional[str] +) -> Optional[float]: + """0-100, or None. The hub service reports a 0-1 fraction, so scale it.""" + from hub.services.models import downloads + try: + payload = await downloads.get_gguf_download_progress_response( + repo_id, variant or "", expected_bytes, hf_token + ) + fraction = payload.get("progress") + if not isinstance(fraction, (int, float)): + return None + return min(100.0, max(0.0, float(fraction) * 100.0)) + except Exception: + return None + + +def _release(active: Optional[_Active]) -> None: + """Free the single-flight slot, but only while *active* still owns it. + + Keying on ``repo_id`` alone let a stale operation clear a newer one: variant A + errors, an adopting request frees the slot, a retry starts B, then A's watcher + matches the repo and clears B, admitting a second download alongside it. + """ + global _active + if active is None: + return + with _lock: + if _active is active: + _active = None + + +async def _watch(active: _Active, hf_token: Optional[str]) -> None: + """Poll a dispatched job so the monitor row resolves and the resolver cache + is dropped the moment the weights land.""" + from core.inference import api_monitor as monitor_module + + api_monitor = monitor_module.api_monitor + deadline = time.monotonic() + _MAX_WATCH_S + timed_out = False + try: + while True: + await asyncio.sleep(_TIMED_OUT_POLL_S if timed_out else _WATCH_POLL_S) + state, error = await _job_state(active.repo_id, active.variant) + if state in ("running", "cancelling", "unknown"): + if timed_out: + # A running worker still owns the slot: releasing on the clock alone + # would admit a second multi-GB download beside it. "unknown" cannot + # confirm it is alive, so release then, or a broken probe wedges us. + if state == "unknown": + return + continue + if time.monotonic() >= deadline: + api_monitor.fail_open(active.monitor_id, "Download timed out") + timed_out = True + continue + # Only "running" has progress; the others are still in flight, so keep the slot. + if state == "running": + api_monitor.set_progress( + active.monitor_id, + await _progress_percent( + active.repo_id, active.variant, active.expected_bytes, hf_token + ), + ) + continue + if state == "cancelled": + api_monitor.finish(active.monitor_id, status = "cancelled") + return + if state == "complete": + # No invalidate here: finalize_worker_exit already dropped the cache and + # warmed it; a second would mark that fresh scan stale and push a + # synchronous rescan onto the client's retry. + api_monitor.finish(active.monitor_id, status = "completed") + elif state == "idle": + # The job vanished without a terminal state (worker killed). + api_monitor.fail_open(active.monitor_id, "Download did not complete") + else: + api_monitor.fail_open(active.monitor_id, error or f"Download {state}") + # Keep the slot so the next retry is told it failed instead of + # silently restarting the same download. + active.error = error or f"Download {state}" + active.failed_at = time.monotonic() + return + return + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("auto-download: watcher failed for %r: %s", active.repo_id, exc) + api_monitor.fail_open(active.monitor_id, "Download tracking failed") + finally: + if not active.failed_at: + _release(active) + + +def _downloading_refusal(label: str, percent: Optional[float]) -> AutoDownloadRefusal: + progress = f" ({percent:.0f}% done)" if percent is not None else "" + return AutoDownloadRefusal( + status = 503, + code = "model_downloading", + message = (f"Downloading '{label}'{progress}. Retry shortly. Track it in Unsloth Studio."), + retry_after = _RETRY_AFTER_S, + ) + + +async def _is_downloadable_model(repo_id: str, hf_token: Optional[str]) -> bool: + """Whether the Hub has this repo with GGUF weights we could fetch. + + Only asked while another download holds the slot, to tell a second download + apart from an ordinary foreign label. Any failure answers False: refusing + would strand normal traffic for the length of the download. + """ + if _is_not_servable(repo_id, hf_token): + return False + + def _probe(): + from huggingface_hub import HfApi + return HfApi(token = _hub_token(hf_token)).model_info(repo_id, timeout = _MODEL_INFO_TIMEOUT_S) + + try: + info = await asyncio.to_thread(_probe) + except Exception: + return False + # The same filter admission uses, not a bare .gguf test: mmproj, MTP drafters and + # big-endian builds are companions, not quants. Answering otherwise would hold an + # ordinary foreign label at model_download_busy for an unrelated download. + servable = bool(_gguf_variants(getattr(info, "siblings", None))) + if not servable: + _mark_not_servable(repo_id, hf_token) + return servable + + +async def maybe_auto_download( + requested_model: str, + *, + hf_token: Optional[str] = None, + require_vision: bool = False, +) -> Optional[AutoDownloadRefusal]: + """Start (or report on) a background fetch of *requested_model*. + + Returns None when the request should carry on unchanged, or a refusal the + caller must raise. Only called after the local resolver has already missed. + + ``require_vision`` refuses a target with no mmproj companion rather than spend + gigabytes on weights that cannot answer the request; the local capability guard + only ever sees an already-downloaded model. + """ + global _active + + repo_id, wanted_variant = split_model_ref(requested_model) + if not is_downloadable_ref(requested_model): + return None + if _is_not_servable(repo_id, hf_token) and not looks_like_quant(wanted_variant): + return None + + # Settle the single-flight slot before the network, so retries during a download stay cheap. + busy: Optional[_Active] = None + with _lock: + current = _active + if current is not None and current.failed_at: + # A held failure only owns the slot until someone is told about it. + if current.repo_id != repo_id and time.monotonic() - current.failed_at > _FAILED_HOLD_S: + _active = current = None + if current is not None and current.repo_id == repo_id: + adopted = current + elif current is not None: + adopted = None + busy = current + else: + adopted = None + provisional = _Active(repo_id = repo_id, started_at = time.time()) + _active = provisional + + if busy is not None: + # Refusing before the probe blocks ordinary drop-in traffic: a namespaced label + # that is no downloadable GGUF repo (LiteLLM/OpenRouter style) would be told to + # wait out a multi-hour download. Only a downloadable label is a 2nd download. + if not await _is_downloadable_model(repo_id, hf_token): + return None + return AutoDownloadRefusal( + status = 503, + code = "model_download_busy", + message = ( + f"Already downloading '{_public_label(busy.repo_id, busy.variant)}'. " + f"Retry '{requested_model}' once it finishes." + ), + retry_after = _RETRY_AFTER_S, + ) + + if adopted is not None: + if adopted.variant is None: + # Still probing: no job yet, and a stale whole-repo error would free the probe's slot. + return _downloading_refusal(adopted.repo_id, None) + state, error = await _job_state(adopted.repo_id, adopted.variant) + if state in ("running", "cancelling", "unknown"): + return _downloading_refusal( + _public_label(adopted.repo_id, adopted.variant), + await _progress_percent( + adopted.repo_id, adopted.variant, adopted.expected_bytes, hf_token + ), + ) + if state == "error" or adopted.error: + error = error or adopted.error + # Surface once, then free the slot so a retry can start over. + _release(adopted) + return AutoDownloadRefusal( + status = 502, + code = "model_download_failed", + message = f"Downloading '{requested_model}' failed: {error or 'unknown error'}", + ) + # complete/idle/cancelled: the watcher is about to free the slot, so retry once more. + return _downloading_refusal( + _public_label(adopted.repo_id, adopted.variant), + 100.0 if state == "complete" else None, + ) + + try: + return await _admit_and_start( + repo_id, wanted_variant, requested_model, hf_token, provisional, require_vision + ) + except BaseException: + # Not `except Exception`: a cancel mid-probe would otherwise wedge the provisional slot. + _release(provisional) + raise + + +async def _admit_and_start( + repo_id: str, + wanted_variant: Optional[str], + requested_model: str, + hf_token: Optional[str], + active: _Active, + require_vision: bool = False, +) -> Optional[AutoDownloadRefusal]: + from hub.utils.hf_errors import hf_error_status + + def _probe(): + from huggingface_hub import HfApi + return HfApi(token = _hub_token(hf_token)).model_info( + repo_id, files_metadata = True, timeout = _MODEL_INFO_TIMEOUT_S + ) + + try: + info = await asyncio.to_thread(_probe) + except Exception as exc: + _release(active) + status = hf_error_status(exc) + if status == 401: + return AutoDownloadRefusal( + status = 401, + code = "model_access_denied", + message = ( + f"Hugging Face rejected the token sent for '{repo_id}'. Replace the " + "X-Unsloth-HF-Token header with a valid token; retrying will not help." + ), + ) + if status == 403: + return _gated_refusal(repo_id) + if status == 404: + _mark_not_servable(repo_id, hf_token) + # Unknown to the Hub reads as a foreign label; only an explicit quant makes it ours. + if not looks_like_quant(wanted_variant): + return None + # A private repo reads as absent without a token; don't confirm either way. + return AutoDownloadRefusal( + status = 404, + code = "model_not_found", + message = ( + f"'{repo_id}' was not found on Hugging Face, or is not accessible. " + "If it is private, send a token in the X-Unsloth-HF-Token header." + ), + ) + logger.warning("auto-download: Hub lookup failed for %r: %s", repo_id, exc) + return AutoDownloadRefusal( + status = 503, + code = "model_lookup_failed", + message = f"Could not reach Hugging Face to look up '{repo_id}'. Retry shortly.", + retry_after = _RETRY_AFTER_S, + ) + + # Inconclusive on timeout: the download's own auth is the real gate. + if getattr(info, "gated", False) and await _bounded_probe( + _auth_denied, repo_id, hf_token, timeout = _MODEL_INFO_TIMEOUT_S, default = False + ): + # Metadata for a gated repo is not file access; unchecked, the config read below lies. + _release(active) + return _gated_refusal(repo_id) + + variants = _gguf_variants(getattr(info, "siblings", None)) + if not variants: + _release(active) + _mark_not_servable(repo_id, hf_token) + if not looks_like_quant(wanted_variant): + return None + return AutoDownloadRefusal( + status = 400, + code = "model_not_supported", + message = ( + f"'{repo_id}' has no GGUF weights. Automatic download serves GGUF only; " + "load other formats from Unsloth Studio." + ), + ) + + # trust_remote_code gate: _config_has_auto_map is tri-state, so refuse on True and on None. + from utils.security.consent import _config_has_auto_map + + # _hub_token, not the raw token: None lets huggingface_hub fall back to a cached + # server login, so a caller-named repo would be probed with this server's identity. + # Defaults to None on timeout, which refuses: unchecked is not cleared. + has_auto_map = await _bounded_probe( + _config_has_auto_map, + repo_id, + _hub_token(hf_token), + timeout = _CODE_PROBE_TIMEOUT_S, + default = None, + ) + if has_auto_map is not False: + _release(active) + unknown = has_auto_map is None + return AutoDownloadRefusal( + status = 403, + code = "remote_code_consent_required", + message = ( + f"'{repo_id}' " + + ( + "could not be checked for custom code" + if unknown + else "ships custom code that runs on load" + ) + + ". Load it once in Unsloth Studio to review and approve it, then retry." + ), + ) + + variant = _match_variant(wanted_variant, variants) + if variant is None: + _release(active) + listed = sorted(variants) + shown = ", ".join(listed[:_MAX_LISTED_VARIANTS]) + extra = len(listed) - _MAX_LISTED_VARIANTS + return AutoDownloadRefusal( + status = 404, + code = "model_not_found", + message = ( + f"'{repo_id}' has no quant '{wanted_variant}'. Available quants: " + f"{shown}{f' and {extra} more' if extra > 0 else ''}." + ), + ) + + expected_bytes = variants[variant] + from hub.utils.gguf_plan import build_gguf_variant_plans + + plan = build_gguf_variant_plans(list(getattr(info, "siblings", None) or [])).get( + variant.lower() + ) + if require_vision and not (plan and plan.mmproj_filenames): + _release(active) + return AutoDownloadRefusal( + status = 400, + code = "invalid_value", + message = ( + f"'{_public_label(repo_id, variant)}' ships no mmproj companion, so it " + "cannot answer the image or audio input in this request. It was not " + "downloaded." + ), + ) + + need_bytes = _remaining_bytes(repo_id, plan, expected_bytes) + fits, free = _enough_disk(need_bytes) + if not fits: + _release(active) + return AutoDownloadRefusal( + status = 507, + code = "insufficient_disk_space", + message = ( + f"'{_public_label(repo_id, variant)}' needs {_gb(need_bytes)} plus " + f"{_gb(_DISK_RESERVE_BYTES)} headroom, but only {_gb(free)} is free." + ), + ) + + return await _dispatch(repo_id, variant, expected_bytes, requested_model, hf_token, active) + + +def preferred_quant(labels) -> Optional[str]: + """The quant a plain load would pick from *labels*, or None. + + The one ranking for "which quant did they mean": local resolution, remote + admission and /v1/models must agree, or a bare id means a different quant + depending on which of them answered it. + """ + from utils.models.model_config import _pick_best_gguf + + # _pick_best_gguf ranks filenames and matches upper-case tokens, so feed "