diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index b394b308e4..6e53a290cf 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -222,6 +222,9 @@ jobs: for s in \ tests/sh/test_get_torch_index_url.sh \ tests/sh/test_mac_intel_compat.sh \ + tests/sh/test_node_decision.sh \ + tests/sh/test_studio_home_node_dir.sh \ + tests/sh/test_system_node_readonly.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index a6f1401067..ceae8e049d 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -79,6 +79,8 @@ jobs: } pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1 + pwsh -NoProfile -File tests/studio/test_node_decision.ps1 + pwsh -NoProfile -File tests/studio/test_node_probe_guard.ps1 - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 4a4806cfb1..1a2a7df493 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -79,15 +79,15 @@ jobs: # Two surgical fixes against measured Windows-only install # waste (vs Mac/Linux on the same SHA): # - # (1) npm. setup.ps1 line 1109-1145 requires Node 22.12+ (or - # 20.19+ / 23+) AND npm >=11 because Vite 8 needs both. + # (1) npm. setup.ps1's Get-NodeDecision requires Node 22.12+ + # (or 20.19+ / 23+) AND npm >=11 because Vite 8 needs both. # actions/setup-node@v4 with `node-version: '22'` lands - # Node 22.22.2 + the npm 10.9.7 it bundles, so the npm - # check fails and setup.ps1 falls through to the - # "winget install Node.js LTS" branch -- a ~35 s reinstall - # of Node we don't need. `npm install -g npm@^11` updates - # the bundled npm in-place in ~5 s, which makes setup.ps1 - # short-circuit on the existing Node. + # Node 22.22.2 + the npm 10.9.7 it bundles, so the decision + # is "bundled" and setup.ps1 downloads an isolated Node (~30 + # MB) we don't need on a runner that already has a fine Node. + # `npm install -g npm@^11` updates the runner's npm in-place + # in ~5 s, flipping the decision to "system" so setup.ps1 + # reuses the existing Node with no download. # # (2) Defender. windows-latest's real-time scan opens / hashes # every file Studio writes during install (Vite output = diff --git a/README.md b/README.md index b6a4b836a4..6656033523 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,8 @@ unsloth studio -p 8888 ``` For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally. +For a secure HTTPS link instead of a raw network port, use `unsloth studio --secure`. Studio stays bound to localhost and is served only through a free Cloudflare HTTPS tunnel (it fails closed if the tunnel can't start, so the raw port is never exposed). + #### Docker Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run: ```bash @@ -162,13 +164,19 @@ Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Ad ## 📥 Advanced Installation The below advanced instructions are for Unsloth Studio. For Unsloth Core advanced installation, [view our docs](https://unsloth.ai/docs/get-started/install/pip-install#advanced-pip-installation). -#### Developer installs: macOS, Linux, WSL: +#### Developer / Nightly / Experimental installs: macOS, Linux, WSL: +The developer install builds from the `main` branch, which is the latest (nightly) source. ```bash git clone https://github.com/unslothai/unsloth cd unsloth ./install.sh --local unsloth studio -p 8888 ``` +To install into an isolated location (its own virtual env, `auth/`, `studio.db`, cache and llama.cpp build), set `UNSLOTH_STUDIO_HOME` and pass it again at launch: +```bash +UNSLOTH_STUDIO_HOME="$PWD/.studio" ./install.sh --local +UNSLOTH_STUDIO_HOME="$PWD/.studio" unsloth studio -p 8888 +``` Then to update : ```bash cd unsloth && git pull @@ -176,7 +184,8 @@ cd unsloth && git pull unsloth studio -p 8888 ``` -#### Developer installs: Windows PowerShell: +#### Developer / Nightly / Experimental installs: Windows PowerShell: +The developer install builds from the `main` branch, which is the latest (nightly) source. ```powershell git clone https://github.com/unslothai/unsloth.git cd unsloth @@ -184,40 +193,31 @@ Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass .\install.ps1 --local unsloth studio -p 8888 ``` +To install into an isolated location (its own virtual env, `auth/`, `studio.db`, cache and llama.cpp build), set `UNSLOTH_STUDIO_HOME` and pass it again at launch: +```powershell +$env:UNSLOTH_STUDIO_HOME="$PWD\.studio"; .\install.ps1 --local +$env:UNSLOTH_STUDIO_HOME="$PWD\.studio"; unsloth studio -p 8888 +``` Then to update : -```bash -cd unsloth && git pull -./install.sh --local -unsloth studio -p 8888 -``` - -#### Nightly: MacOS, Linux, WSL: -```bash -git clone https://github.com/unslothai/unsloth -cd unsloth -git checkout nightly -./install.sh --local -unsloth studio -p 8888 -``` -Then to launch every time: -```bash -unsloth studio -p 8888 -``` - -#### Nightly: Windows: -Run in Windows Powershell: ```powershell -git clone https://github.com/unslothai/unsloth.git -cd unsloth -git checkout nightly -Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass +cd unsloth; git pull .\install.ps1 --local unsloth studio -p 8888 ``` -Then to launch every time: + +#### Remote access: `--secure` (HTTPS tunnel) vs raw port +By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach it from another device, pick one of: + +- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Studio stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed. ```bash -unsloth studio -p 8888 +unsloth studio --secure -p 8888 ``` +- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network. +```bash +unsloth studio -H 0.0.0.0 -p 8888 +``` + +Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Studio. #### Advanced launch options Installer options can be passed as environment variables. On macOS, Linux and WSL place the variable after the pipe so the shell passes it to `sh`; on Windows set it with `$env:` before piping to `iex`. diff --git a/install.ps1 b/install.ps1 index c241a04166..d48c9fa4f3 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2298,7 +2298,9 @@ exit 0 # ── Run studio setup ── # setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools, - # CUDA Toolkit, Node.js, and other dependencies automatically via winget. + # CUDA Toolkit, and other dependencies automatically via winget. Node.js is + # NOT installed via winget -- setup.ps1 uses an isolated Node it manages and + # never touches the system Node/npm. Write-TauriLog "STEP" "Running studio setup" step "setup" "running unsloth studio setup..." $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index ee6678d9b9..796d03ff68 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -5,6 +5,7 @@ import hashlib import hmac +import ipaddress import os import secrets import sqlite3 @@ -100,6 +101,14 @@ def get_connection() -> sqlite3.Connection: """Get a connection to the auth database, creating tables if needed.""" ensure_dir(DB_PATH.parent) conn = sqlite3.connect(DB_PATH) + # Keep the auth dir + DB private (they hold the JWT/identity secrets and + # password hashes); sqlite3.connect would otherwise create the DB 0644 under + # a 022 umask, letting another OS user read the identity secret and forge proofs. + for _path, _mode in ((DB_PATH.parent, 0o700), (DB_PATH, 0o600)): + try: + os.chmod(_path, _mode) + except OSError: + pass conn.row_factory = sqlite3.Row conn.execute( """ @@ -210,6 +219,57 @@ def _get_or_create_api_key_pbkdf2_salt() -> bytes: return salt +# Secret answering the /api/auth/identity challenge (HMAC(secret, nonce)). Lives +# in this same-user DB so a port squatter or remote/fake server can't forge a +# proof. Separate from the per-user JWT secret. +_IDENTITY_SECRET_DB_KEY = "studio_identity_secret" +_identity_secret_cache: Optional[bytes] = None + + +def get_or_create_identity_secret() -> bytes: + """Return the identity secret (hex 32-byte row in app_secrets), creating it once.""" + global _identity_secret_cache + if _identity_secret_cache is not None: + return _identity_secret_cache + + conn = get_connection() + try: + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + if row is None: + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + (_IDENTITY_SECRET_DB_KEY, secrets.token_hex(32)), + ) + conn.commit() + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_IDENTITY_SECRET_DB_KEY,), + ).fetchone() + secret = bytes.fromhex(row["value"]) + finally: + conn.close() + + _identity_secret_cache = secret + return secret + + +def compute_identity_proof(nonce: bytes, host: str, port: int) -> str: + """HMAC-SHA256 proof that the caller holds this install's identity secret, + bound to the loopback address and port the connection landed on. A proof + relayed from a Studio on a different address/port (a squatter proxying to the + real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was + computed for that other endpoint and won't match the one the client dialed.""" + try: + host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms + except ValueError: + host = (host or "").lower() + msg = b"|".join([nonce, host.encode(), str(int(port)).encode()]) + return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest() + + _API_KEY_PBKDF2_ITERATIONS = 100_000 DESKTOP_SECRET_PREFIX = "desktop-" _DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash" diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index fbe847f9ce..ebb1d39dfb 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -14,6 +14,7 @@ from pathlib import Path from typing import Any from loggers import get_logger +from utils.node_runtime import resolve_node_executable from utils.paths import ensure_dir, oxc_validator_tmp_root logger = get_logger(__name__) @@ -231,6 +232,14 @@ def _run_oxc_batch( "code_shape": code_shape, "codes": code_values, } + # Resolve a usable Node (system or the isolated install, which is not on the + # user's PATH); a bare "node" would fail for isolated-Node users. + node_executable = resolve_node_executable() + if not node_executable: + return _fallback_results( + len(code_values), + "Node.js not found (install Node >= 20.19, or re-run Studio setup to provision it).", + ) try: tmp_dir = ensure_dir(oxc_validator_tmp_root()) env = child_env_without_native_path_secret() @@ -238,8 +247,13 @@ def _run_oxc_batch( env["TMPDIR"] = tmp_dir_str env["TMP"] = tmp_dir_str env["TEMP"] = tmp_dir_str + # Resolved node's dir first on the child PATH so it finds its own npm/npx. + node_bin_dir = os.path.dirname(node_executable) + if node_bin_dir: + env["PATH"] = node_bin_dir + os.pathsep + env.get("PATH", "") + env.pop("NODE_PATH", None) proc = subprocess.run( - ["node", str(_OXC_RUNNER_PATH)], + [node_executable, str(_OXC_RUNNER_PATH)], cwd = str(_OXC_TOOL_DIR), input = json.dumps(payload), text = True, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 16b312d8f6..b6f83b760c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -13,7 +13,6 @@ import json import os import re import struct -import structlog from loggers import get_logger import shutil import signal @@ -40,15 +39,7 @@ from core.inference.llama_server_args import ( strip_split_mode_only, ) from core.tool_healing import ( - _TC_END_TAG_RE, - _TC_FUNC_CLOSE_RE, - _TC_FUNC_START_RE, - _TC_JSON_START_RE, - _TC_PARAM_CLOSE_RE, - _TC_PARAM_START_RE, _TOOL_ALL_PATS, - _TOOL_CLOSED_PATS, - parse_tool_calls_from_text, strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret @@ -379,6 +370,13 @@ def _period_from_layer_types(layer_types: list) -> Optional[int]: return None +def _swa_entry_from_layer_types(lt) -> Optional[object]: + """Period int, or per-layer bool mask, from a transformers ``layer_types`` list.""" + if isinstance(lt, list) and lt: + return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] + return None + + def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: try: from huggingface_hub import hf_hub_download @@ -392,10 +390,7 @@ def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: period = src.get("sliding_window_pattern") if isinstance(period, int) and period > 0: return period - lt = src.get("layer_types") - if isinstance(lt, list) and lt: - return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] - return None + return _swa_entry_from_layer_types(src.get("layer_types")) def _arch_aliases(arch: str) -> tuple: @@ -412,10 +407,7 @@ def _swa_entry_from_config_obj(cfg) -> Optional[object]: period = getattr(src, "sliding_window_pattern", None) if isinstance(period, int) and period > 0: return period - lt = getattr(src, "layer_types", None) - if isinstance(lt, list) and lt: - return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt] - return None + return _swa_entry_from_layer_types(getattr(src, "layer_types", None)) _SWA_PATTERN_SOURCE_RE = re.compile(r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)") @@ -581,12 +573,14 @@ def detect_reasoning_flags( ) -> dict: """Classify a chat template's reasoning and tool-calling capabilities. - Returns the same five keys as the GGUF sniffer: ``supports_reasoning``, - ``reasoning_style`` (``"enable_thinking"`` | ``"reasoning_effort"``), - ``reasoning_always_on``, ``supports_preserve_thinking``, - ``supports_tools``. Used by both the llama-server backend at load time - and the safetensors/transformers paths in ``routes/inference`` so they - agree on what the frontend sees. + Returns the same six keys as the GGUF sniffer: ``supports_reasoning``, + ``reasoning_style`` (``"enable_thinking"`` | ``"reasoning_effort"`` | + ``"enable_thinking_effort"``), ``reasoning_always_on``, + ``reasoning_effort_levels``, ``supports_preserve_thinking``, + ``supports_tools``. A falsy ``chat_template`` yields the all-default dict. + Used by both the llama-server backend at load time and the + safetensors/transformers paths in ``routes/inference`` so they agree on + what the frontend sees. """ flags = { "supports_reasoning": False, @@ -606,7 +600,7 @@ def detect_reasoning_flags( if ("reasoning_effort" in tpl and "enable_thinking" in tpl) else [] ) - if "enable_thinking" in tpl and "reasoning_effort" in tpl and effort_levels: + if effort_levels: # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort # level among a discrete set (e.g. 'high' | 'max'). Distinct from # gpt-oss (reasoning_effort only, no on/off gate) and Qwen @@ -657,11 +651,32 @@ def detect_reasoning_flags( return flags +# Gemma 4 ships MTP as a separate drafter (no "-mtp" in the name). Gemma 3n +# ships no drafter, so it is excluded -- it takes the normal non-MTP path. +_GEMMA_MTP_FAMILY_RE = re.compile(r"gemma[-_]?4[-_]", re.IGNORECASE) + + +def _is_gemma_mtp_family(name: Optional[str]) -> bool: + """Match Gemma 4 by name.""" + return bool(name) and bool(_GEMMA_MTP_FAMILY_RE.search(name)) + + +def _is_gemma_mtp_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: + """Match Gemma 4 by id or GGUF filename.""" + return _is_gemma_mtp_family(model_identifier) or _is_gemma_mtp_family( + Path(gguf_path).name if gguf_path else None + ) + + def _is_mtp_model_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: """Name-based MTP detector. Fallback for the metadata signal.""" for cand in (model_identifier, Path(gguf_path).name if gguf_path else None): if cand and "-mtp" in cand.lower(): return True + # Recognise Gemma 4 too, so a failed drafter download surfaces a + # fallback reason instead of silently defaulting. + if cand and _is_gemma_mtp_family(cand): + return True return False @@ -1353,6 +1368,12 @@ class LlamaCppBackend: def base_url(self) -> str: return f"http://127.0.0.1:{self._port}" + @property + def _auth_headers(self) -> "Optional[dict[str, str]]": + """Bearer header matching the --api-key direct-stream mode uses, else + None (so unauthenticated llama-server calls don't get a spurious 401).""" + return {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + @property def model_identifier(self) -> Optional[str]: return self._model_identifier @@ -1374,6 +1395,11 @@ class LlamaCppBackend: def gguf_path(self) -> Optional[str]: return self._gguf_path + @property + def hf_repo(self) -> Optional[str]: + """HF repo of the loaded model, or None for local/native file loads.""" + return self._hf_repo + @property def mtp_draft_path(self) -> Optional[str]: return self._mtp_draft_path @@ -1626,6 +1652,30 @@ class LlamaCppBackend: # ── Binary discovery ────────────────────────────────────────── + @staticmethod + def _resolved_studio_root_and_is_legacy() -> "tuple[Optional[Path], bool]": + """Resolve the Studio install root and classify it as the legacy + ~/.unsloth/studio root vs. a custom (env/venv-inferred) root. + + Returns (resolved_root, is_legacy). On any import/resolution failure the + root is treated as legacy and resolved_root is None -- callers must read + resolved_root only when is_legacy is False. Shared by + _find_llama_server_binary (discovery) and _kill_orphaned_servers + (cleanup) so the two never disagree on which root is legacy. + """ + try: + from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 + + resolved = _sr() + legacy_studio = Path.home() / ".unsloth" / "studio" + try: + is_legacy = resolved.resolve() == legacy_studio.resolve() + except (OSError, ValueError): + is_legacy = resolved == legacy_studio + return (None if is_legacy else resolved), is_legacy + except (ImportError, OSError, ValueError): + return None, True + @staticmethod def _find_llama_server_binary(*, include_denied: bool = False) -> Optional[str]: """ @@ -1712,33 +1762,16 @@ class LlamaCppBackend: # 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp; # default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio). legacy_llama = Path.home() / ".unsloth" / "llama.cpp" - try: - from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 - - _resolved_sr = _sr() - _legacy_studio = Path.home() / ".unsloth" / "studio" - try: - _is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve() - except (OSError, ValueError): - _is_legacy = _resolved_sr == _legacy_studio - if _is_legacy: - search_roots = [legacy_llama] - else: - # _kill_orphaned_servers excludes the legacy root in custom - # mode; discovery must match so we never spawn a server we - # then refuse to clean up. UNSLOTH_LLAMA_CPP_PATH (handled - # earlier) is the explicit way to share a build across roots. - search_roots = [_resolved_sr / "llama.cpp"] - except (ImportError, OSError, ValueError): + _resolved_sr, _is_legacy = LlamaCppBackend._resolved_studio_root_and_is_legacy() + if _is_legacy: search_roots = [legacy_llama] - _seen_roots: set[str] = set() - _unique_roots: list[Path] = [] - for r in search_roots: - k = str(r) - if k not in _seen_roots: - _seen_roots.add(k) - _unique_roots.append(r) - for unsloth_home in _unique_roots: + else: + # _kill_orphaned_servers excludes the legacy root in custom mode; + # discovery must match so we never spawn a server we then refuse to + # clean up. UNSLOTH_LLAMA_CPP_PATH (handled earlier) is the explicit + # way to share a build across roots. + search_roots = [_resolved_sr / "llama.cpp"] + for unsloth_home in search_roots: hit, locked = _scan_pinned(_layout_candidates(unsloth_home)) if locked is not None: return _unavailable(locked) @@ -1960,6 +1993,36 @@ class LlamaCppBackend: return total + @staticmethod + def _resolve_visible_physical_ids() -> Optional[list[int]]: + """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on + ROCm, CUDA otherwise). None when no mask is set; empty list for an empty + mask. Shared by the APU / datacenter / free-memory probes so they agree + on the ordinal->physical mapping.""" + try: + import torch + is_rocm = getattr(torch.version, "hip", None) is not None + except Exception: + is_rocm = False + if is_rocm: + hip_v = os.environ.get("HIP_VISIBLE_DEVICES") + rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") + cvd = ( + hip_v + if hip_v is not None + else rocr_v + if rocr_v is not None + else os.environ.get("CUDA_VISIBLE_DEVICES") + ) + else: + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is None: + return None + try: + return [int(x.strip()) for x in cvd.split(",") if x.strip()] + except ValueError: + return None + @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: """True only for AMD unified-memory APUs (gfx1150/gfx1151), where @@ -1976,21 +2039,7 @@ class LlamaCppBackend: return False # Map visible ordinal -> physical id via the active ROCm mask (HIP, # then ROCR, then CUDA), mirroring _get_gpu_memory's ROCm branch. - physical_ids: Optional[list[int]] = None - hip_v = os.environ.get("HIP_VISIBLE_DEVICES") - rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") - cvd = ( - hip_v - if hip_v is not None - else rocr_v - if rocr_v is not None - else os.environ.get("CUDA_VISIBLE_DEVICES") - ) - if cvd is not None: - try: - physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] - except ValueError: - physical_ids = None + physical_ids = LlamaCppBackend._resolve_visible_physical_ids() arch_by_id: dict[int, str] = {} for ordinal in range(torch.cuda.device_count()): try: @@ -2042,13 +2091,7 @@ class LlamaCppBackend: # Mirror _get_gpu_free_memory: map visible ordinal -> physical id via # CUDA_VISIBLE_DEVICES; unset/unparsable leaves physical id == ordinal. - physical_ids: Optional[list[int]] = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] - except ValueError: - physical_ids = None + physical_ids = LlamaCppBackend._resolve_visible_physical_ids() pattern = LlamaCppBackend._DATACENTER_GPU_RE names_by_id: dict[int, str] = {} @@ -2202,29 +2245,12 @@ class LlamaCppBackend: # feed these IDs back into the subprocess as CVD, so visible ordinals # must be translated to physical indices first; otherwise CVD=2,3 # gets rewritten to 0,1 and targets the wrong GPUs. - physical_ids: Optional[list[int]] = None # Match utils/hardware/hardware.py::_get_parent_visible_gpu_spec: # treat an empty mask (HIP_VISIBLE_DEVICES="") as "no GPUs" rather # than falling through. ``or`` would coerce "" to the wrong source. - if getattr(torch.version, "hip", None) is not None: - hip_v = os.environ.get("HIP_VISIBLE_DEVICES") - rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") - cvd = ( - hip_v - if hip_v is not None - else rocr_v - if rocr_v is not None - else os.environ.get("CUDA_VISIBLE_DEVICES") - ) - else: - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - # Empty mask (CVD="") yields an empty list -> no GPUs, - # consistent with the nvidia-smi path. - physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] - except ValueError: - physical_ids = None + # Empty mask (CVD="") yields an empty list -> no GPUs, consistent + # with the nvidia-smi path. + physical_ids = LlamaCppBackend._resolve_visible_physical_ids() gpus = [] for ordinal in range(torch.cuda.device_count()): free_bytes, total_bytes = torch.cuda.mem_get_info(ordinal) @@ -2463,8 +2489,7 @@ class LlamaCppBackend: lib_dirs = [] # WSL: system HIP before the bundle's (which segfaults on /dev/dxg). - for _wsl_rocm in _wsl_system_rocm_lib_dirs(): - lib_dirs.append(_wsl_rocm) + lib_dirs.extend(_wsl_system_rocm_lib_dirs()) if lib_dirs: env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") lib_dirs.append(binary_dir) @@ -2476,33 +2501,8 @@ class LlamaCppBackend: import glob as _glob for _nv_pattern in [ - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cu*", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cudnn", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "nvjitlink", - "lib", - ), + os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", _sub, "lib") + for _sub in ("cu*", "cudnn", "nvjitlink") ]: for _nv_dir in _glob.glob(_nv_pattern): if os.path.isdir(_nv_dir): @@ -2617,6 +2617,12 @@ class LlamaCppBackend: return self._n_kv_heads_by_layer[layer_idx] return fallback + def _legacy_head_dim(self) -> int: + """Head-dim fallback for GGUFs without explicit key/value dims. Reached + only via the legacy branch of _can_estimate_kv(), so _embedding_length + is non-None here.""" + return self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + def _estimate_kv_cache_bytes( self, n_ctx: int, @@ -2680,7 +2686,7 @@ class LlamaCppBackend: n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division if key_len is not None and val_len is not None: return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe) - head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + head_dim = self._legacy_head_dim() return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern @@ -2747,7 +2753,7 @@ class LlamaCppBackend: return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe) # Path 5: Legacy fallback (old GGUFs without explicit dimensions) - head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + head_dim = self._legacy_head_dim() return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]: @@ -3311,8 +3317,8 @@ class LlamaCppBackend: attr = arch_keys.get(key) if attr == "n_kv_heads" and val_a is not None: self._n_kv_heads_by_layer = [int(x) for x in val_a] - if self._n_kv_heads is None and val_a: - self._n_kv_heads = max(int(x) for x in val_a) + if self._n_kv_heads is None and self._n_kv_heads_by_layer: + self._n_kv_heads = max(self._n_kv_heads_by_layer) elif attr == "sliding_window_pattern" and val_a is not None: self._sliding_window_pattern = [bool(x) for x in val_a] sliding_window_pattern_period = None @@ -3411,8 +3417,6 @@ class LlamaCppBackend: alongside llama-server. Returns None if neither can be found. """ import importlib.util - import os - import sys # Visual-server binary: env override, else next to llama-server or in the # install's build/bin (where the prebuilt/installer puts it). .exe on Windows. @@ -3472,8 +3476,6 @@ class LlamaCppBackend: visual decoder) and wait for health. Presents the same /v1 + /health interface as llama-server, so the rest of Studio is unchanged. """ - import os - assets = self._find_diffusion_assets() if assets is None: raise RuntimeError( @@ -3582,9 +3584,8 @@ class LlamaCppBackend: # (auto-sized to VRAM). Read it back so the UI context bar shows the real budget. chosen = maxtok try: - import re as _re for _ln in reversed(self._stdout_lines): - _m = _re.search(r"MAXTOK=(\d+)", _ln) + _m = re.search(r"MAXTOK=(\d+)", _ln) if _m: chosen = int(_m.group(1)) break @@ -3776,13 +3777,9 @@ class LlamaCppBackend: hf_token, cancel_event = self._cancel_event, ) - except RuntimeError as e: - if "Cancelled" in str(e): - raise - raise RuntimeError( - f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" - ) except Exception as e: + if isinstance(e, RuntimeError) and "Cancelled" in str(e): + raise raise RuntimeError( f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" ) @@ -3813,11 +3810,31 @@ class LlamaCppBackend: return None target: Optional[str] = None - try: - from huggingface_hub import list_repo_files - target = pick(list_repo_files(hf_repo, token = hf_token)) - except Exception as e: - logger.debug(f"Could not list repo files for {label}: {e}") + from huggingface_hub import list_repo_files + + # Retry a transient listing blip; permanent repo/auth errors and offline + # mode are not retried (offline raises at once -> fall through to cache). + for attempt in range(3): + if self._cancel_event.is_set(): + return None + try: + target = pick(list_repo_files(hf_repo, token = hf_token)) + break + except Exception as e: + if type(e).__name__ in ( + "RepositoryNotFoundError", + "GatedRepoError", + "RevisionNotFoundError", + "EntryNotFoundError", + "OfflineModeIsEnabled", + ): + logger.debug(f"Could not list repo files for {label}: {e}") + break + logger.debug( + f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}" + ) + if attempt < 2: + self._cancel_event.wait(2**attempt) if target is None: try: @@ -4341,6 +4358,16 @@ class LlamaCppBackend: out.append(tok) return out + @staticmethod + def _redacted_cmd_for_log(cmd: "list[str]") -> "list[str]": + """Copy of cmd with the value after --api-key replaced by .""" + out = list(cmd) + if "--api-key" in out: + ki = out.index("--api-key") + 1 + if ki < len(out): + out[ki] = "" + return out + def _start_llama_process(self, cmd: list[str], env: dict) -> None: """Spawn llama-server from cmd and start draining its output. @@ -4377,12 +4404,7 @@ class LlamaCppBackend: # Log the argv per attempt (the text-only mmproj retry re-enters here # with --mmproj stripped), redacting the API key. - _log_cmd = list(cmd) - if "--api-key" in _log_cmd: - _ki = _log_cmd.index("--api-key") + 1 - if _ki < len(_log_cmd): - _log_cmd[_ki] = "" - logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") + logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") self._process = subprocess.Popen( cmd, @@ -4495,36 +4517,8 @@ class LlamaCppBackend: except Exception as exc: logger.debug("Fast-path audio probe failed: %s", exc) detected = None - if detected in ("snac", "bicodec", "dac"): - with self._lock: - if not self._healthy: - return False - try: - self.init_audio_codec(detected) - self._is_audio = True - self._audio_type = detected - except Exception as exc: - logger.warning( - "Failed to init audio codec '%s': %s", - detected, - exc, - ) - self._audio_probed = False - return False - elif detected: - # csm / whisper / audio_vlm: track type but keep - # _is_audio False -- GGUF TTS routing only fires for - # snac/bicodec/dac. - with self._lock: - if not self._healthy: - return False - self._audio_type = detected - # Re-derive after a retried probe (_mmproj_has_audio persists). - from utils.models.model_config import is_audio_input_type - - self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( - self._mmproj_has_audio - ) + if not self._apply_detected_audio(detected): + return False if not self._healthy: return False return True @@ -4831,6 +4825,12 @@ class LlamaCppBackend: self._nextn_predict_layers or _is_mtp_model_name(model_identifier, model_path) or bool(mtp_draft_path) + ) and not ( + # Drafterless Gemma falls back to ngram-mod; reserve no + # drafter VRAM for it (mirrors the launch resolver). + _is_gemma_mtp_name(model_identifier, model_path) + and not mtp_draft_path + and not self._nextn_predict_layers ) _mtp_binary_ok = True _mtp_probe_raised = False @@ -5024,13 +5024,12 @@ class LlamaCppBackend: ) _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve - if tensor_parallel and effective_is_vision: - logger.info( - "Tensor parallelism skipped for vision model: " - "--split-mode tensor is incompatible with --mmproj " - "in the current llama.cpp build; using layer split." - ) - tensor_parallel = False + def _restore_after_tensor_downgrade(): + # Tensor mode dropped a quantized KV and stripped the cache + # extras (it rejects quantized); layer split supports them, so + # restore the original type + extras (minus --split-mode) and + # clear the env flag so the layer launch re-emits them. + nonlocal cache_type_kv, _cache_type_from_env, extra_args if _tensor_dropped_cache_type_kv is not None: cache_type_kv = _tensor_dropped_cache_type_kv _cache_type_from_env = False @@ -5040,6 +5039,15 @@ class LlamaCppBackend: else extra_args ) + if tensor_parallel and effective_is_vision: + logger.info( + "Tensor parallelism skipped for vision model: " + "--split-mode tensor is incompatible with --mmproj " + "in the current llama.cpp build; using layer split." + ) + tensor_parallel = False + _restore_after_tensor_downgrade() + # Tensor mode replicates a compute buffer on every GPU, so drop # GPUs below that reserve from the set up front (gpu_indices # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). @@ -5078,21 +5086,9 @@ class LlamaCppBackend: ) tensor_parallel = False # Layer split supports a quantized KV the tensor attempt - # dropped; restore it and re-emit it (clear the env flag the - # tensor re-adoption may have set, so the restored type wins - # over a stale inherited env on the layer launch). - if _tensor_dropped_cache_type_kv is not None: - cache_type_kv = _tensor_dropped_cache_type_kv - _cache_type_from_env = False - # Restore the original extras (with the real, possibly - # asymmetric, --cache-type-k/-v the tensor attempt stripped), - # then drop the user --split-mode tensor so the downgrade - # actually applies (extras are appended last). - extra_args = strip_split_mode_only( - _tensor_dropped_extra_args - if _tensor_dropped_extra_args is not None - else extra_args - ) + # dropped; restore the original cache type + extras (minus + # --split-mode) so the layer launch re-emits them. + _restore_after_tensor_downgrade() if tensor_parallel and tp_gpus: # Pooled usable budget (after each device's compute buffer) @@ -5125,18 +5121,9 @@ class LlamaCppBackend: "per-device compute buffers; falling back to layer split." ) tensor_parallel = False - # Restore the dropped quantized KV (layer split supports - # it); clear the env flag so the restored type is emitted. - if _tensor_dropped_cache_type_kv is not None: - cache_type_kv = _tensor_dropped_cache_type_kv - _cache_type_from_env = False - # Restore the original (possibly asymmetric) cache extras - # too, dropping only the user --split-mode tensor. - extra_args = strip_split_mode_only( - _tensor_dropped_extra_args - if _tensor_dropped_extra_args is not None - else extra_args - ) + # Restore the dropped quantized KV + original cache extras + # (minus --split-mode); layer split supports them. + _restore_after_tensor_downgrade() if tensor_parallel and tp_gpus: # Tensor-parallel allocation; see _plan_tensor_parallel. @@ -5585,10 +5572,9 @@ class LlamaCppBackend: logger.info(f"Using mmproj for vision: {launch_mmproj_path}") # Option C: --api-key for direct client access when enabled - import os as _os import secrets as _secrets - if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": + if os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": self._api_key = _secrets.token_urlsafe(32) cmd.extend(["--api-key", self._api_key]) logger.info("llama-server started with --api-key for direct streaming") @@ -5624,12 +5610,7 @@ class LlamaCppBackend: cmd.extend(str(a) for a in extra_args) logger.info(f"Appending user extra args to llama-server: {list(extra_args)}") - _log_cmd = list(cmd) - if "--api-key" in _log_cmd: - _ki = _log_cmd.index("--api-key") + 1 - if _ki < len(_log_cmd): - _log_cmd[_ki] = "" - logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") + logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") # Library paths so llama-server finds its shared libs and CUDA DLLs. env = self._llama_server_env_for_binary(binary) @@ -6087,37 +6068,8 @@ class LlamaCppBackend: except Exception as exc: logger.debug("Audio probe failed: %s", exc) detected = None - if detected in ("snac", "bicodec", "dac"): - with self._lock: - if not self._healthy: - return False - try: - self.init_audio_codec(detected) - self._is_audio = True - self._audio_type = detected - except Exception as exc: - # Surface as HTTP 500 (matches pre-PR contract). - logger.warning( - "Failed to init audio codec '%s': %s", - detected, - exc, - ) - self._audio_probed = False - return False - elif detected: - # csm / whisper / audio_vlm: track type but keep _is_audio - # False -- GGUF TTS routing only fires for snac/bicodec/dac. - with self._lock: - if not self._healthy: - return False - self._audio_type = detected - - # Audio input = token probe (audio_vlm/whisper) OR mmproj encoder. - from utils.models.model_config import is_audio_input_type - - self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( - self._mmproj_has_audio - ) + if not self._apply_detected_audio(detected): + return False if not self._healthy: return False @@ -6202,6 +6154,13 @@ class LlamaCppBackend: _mtp_too_small = ( _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) ) + # Drafterless Gemma (name-only MTP, no embedded head): emitting MTP + # would abort llama-server, so every mode below falls back instead. + _mtp_drafter_missing = ( + _is_gemma_mtp_name(model_identifier, model_path) + and not mtp_draft_path + and not self._nextn_predict_layers + ) # Embedded MTP head on an MLA model (GLM-5.2/DeepSeek/Kimi, detected by # kv_lora_rank): llama.cpp's MLA/DSA MTP path is ~2x slower than no spec, # so Auto drops it (override via the Settings dropdown / forced mtp, or @@ -6259,6 +6218,8 @@ class LlamaCppBackend: if mtp_draft_path: flags.extend(["--model-draft", mtp_draft_path]) logger.info(f"Using separate MTP drafter: {mtp_draft_path}") + spec_value = mtp_token + ngram_knobs: list[str] = [] if chain_ngram: ngram_knobs = _build_ngram_mod_flags(caps) if ngram_knobs: @@ -6268,25 +6229,8 @@ class LlamaCppBackend: "llama-server lacks ngram-mod tuning " "flags; loading MTP only (no ngram chain)" ) - spec_value = mtp_token - flags.extend( - [ - "--spec-type", - spec_value, - n_max_flag, - str(draft_n_max), - ] - ) - flags.extend(ngram_knobs) - else: - flags.extend( - [ - "--spec-type", - mtp_token, - n_max_flag, - str(draft_n_max), - ] - ) + flags.extend(["--spec-type", spec_value, n_max_flag, str(draft_n_max)]) + flags.extend(ngram_knobs) self._speculative_type = "draft-mtp" chain_label = "chained ngram-mod" if chain_ngram else "MTP-only" logger.info(f"Spec decoding: {mtp_token} ({chain_label})") @@ -6307,6 +6251,20 @@ class LlamaCppBackend: logger.info("Spec decoding: ngram-mod") return True + def _fallback_drafter_not_found() -> None: + """Drafterless Gemma: use ngram-mod (or spec-default) and record why.""" + logger.warning( + "Model %s is MTP-capable but no drafter or head was found; " + "falling back. Check network or run `unsloth studio update`.", + model_identifier, + ) + if self.probe_server_capabilities(binary).get("supports_ngram_mod"): + _emit_ngram_mod() + else: + flags.append("--spec-default") + self._speculative_type = "default" + self._spec_fallback_reason = "drafter_not_found" + if effective_mode == "off": return flags # nothing to emit if effective_mode == "ngram-simple": @@ -6327,6 +6285,10 @@ class LlamaCppBackend: flags.append("--spec-default") self._speculative_type = "default" return flags + if _mtp_drafter_missing: + # Drafterless: draft-mtp would abort llama-server, so fall back. + _fallback_drafter_not_found() + return flags if _mtp_too_small: logger.warning( f"Forcing MTP on a {_mtp_size_b:.1f}B model; " @@ -6345,6 +6307,10 @@ class LlamaCppBackend: ) _emit_ngram_mod() return flags + if _mtp_drafter_missing: + # No head/drafter: keep ngram-mod, drop the draft-mtp chain. + _fallback_drafter_not_found() + return flags if _mtp_too_small: logger.warning( f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; " @@ -6380,13 +6346,18 @@ class LlamaCppBackend: ) # spec-off: emit nothing, mirroring the sub-3B no-ngram path. elif is_mtp_model and not _mtp_too_small: - # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP. - _emit_mtp(chain_ngram = not gpus) + if _mtp_drafter_missing: + # Name-only MTP, drafter did not resolve (download failed/absent). + _fallback_drafter_not_found() + else: + # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP. + _emit_mtp(chain_ngram = not gpus) elif is_mtp_model and _mtp_too_small: # Sub-3B fallback: drop the MTP draft head, keep ngram-mod when # the binary supports it. - _small_caps = self.probe_server_capabilities(binary) - if _small_caps.get("supports_ngram_mod"): + if _mtp_drafter_missing: + _fallback_drafter_not_found() + elif self.probe_server_capabilities(binary).get("supports_ngram_mod"): logger.info( f"MTP GGUF detected but model size {_mtp_size_b:.1f}B " "is below the 3B speedup threshold; using ngram-mod " @@ -6476,6 +6447,16 @@ class LlamaCppBackend: if req_mode != backend_mode: return False + # Prior HF load fell back with drafter_not_found; a same-settings reload + # must retry the download in load_model, not dedupe to the stale fallback + # (HF loads resolve the drafter there, so gguf_path is None here). + if ( + self._spec_fallback_reason == "drafter_not_found" + and gguf_path is None + and req_mode in ("auto", "mtp", "mtp+ngram") + ): + return False + # spec_draft_n_max only matters when an MTP variant is engaged. Compare # on the resolved spec so an Auto request promoted to draft-mtp still # bounces a reload when n_max changes. @@ -6590,7 +6571,6 @@ class LlamaCppBackend: # Clean up temp chat template file. if hasattr(self, "_chat_template_file") and self._chat_template_file: try: - import os os.unlink(self._chat_template_file.name) except Exception: pass @@ -6868,20 +6848,10 @@ class LlamaCppBackend: install_roots: list[Path] = [] # Env-mode custom root (mirrors _find_llama_server_binary). - _is_custom_root = False - try: - from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 - - _resolved_sr = _sr() - _legacy_studio = Path.home() / ".unsloth" / "studio" - try: - _is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve() - except (OSError, ValueError): - _is_custom_root = _resolved_sr != _legacy_studio - if _is_custom_root: - install_roots.append(_resolved_sr / "llama.cpp") - except (ImportError, OSError, ValueError): - pass + _resolved_sr, _is_legacy = LlamaCppBackend._resolved_studio_root_and_is_legacy() + _is_custom_root = not _is_legacy + if _is_custom_root: + install_roots.append(_resolved_sr / "llama.cpp") # Primary install dir (default mode only). Env-mode skips this so a # custom-root Studio can't kill a default-install Studio's server. @@ -7043,12 +7013,10 @@ class LlamaCppBackend: tokens generate (e.g. under --split-mode tensor). False on any error so the caller can drop MTP and retry. """ - url = f"http://127.0.0.1:{self._port}/completion" + url = f"{self.base_url}/completion" payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False} - # Match the --api-key auth direct-stream mode uses, else a spurious 401. - headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None try: - resp = httpx.post(url, json = payload, timeout = timeout, headers = headers) + resp = httpx.post(url, json = payload, timeout = timeout, headers = self._auth_headers) except Exception as e: logger.debug(f"MTP decode probe failed: {e}") return False @@ -7176,7 +7144,7 @@ class LlamaCppBackend: ) -> bool: """Poll llama-server's /health until 200; also detect early exit/crash.""" deadline = time.monotonic() + timeout - url = f"http://127.0.0.1:{self._port}/health" + url = f"{self.base_url}/health" while time.monotonic() < deadline: # Process crashed? @@ -7245,7 +7213,7 @@ class LlamaCppBackend: The memory-fit step or ``--parallel`` slot split can leave this below the requested ``-c``; requests are validated against this value. """ - url = f"http://127.0.0.1:{self._port}/props" + url = f"{self.base_url}/props" try: resp = httpx.get(url, timeout = 5.0) if resp.status_code != 200: @@ -7318,6 +7286,33 @@ class LlamaCppBackend: # ── Generation (proxy to llama-server) ──────────────────────── + @contextlib.contextmanager + def _open_stream(self, url: str, payload: dict, cancel_event): + """Open a streaming POST to llama-server, retrying through prefill, and + yield ``(response, first_token_deadline)`` once a 200 lands. Owns the + httpx.Client + auth headers for the stream's lifetime; raises + RuntimeError on a non-200. Shared scaffold for the streaming consumers, + which differ only in how they parse the SSE body.""" + stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) + with httpx.Client( + timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) + ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + with self._stream_with_retry( + client, + url, + payload, + cancel_event, + headers = self._auth_headers, + first_token_deadline = first_token_deadline, + ) as response: + if response.status_code != 200: + error_body = response.read().decode() + raise RuntimeError( + f"llama-server returned {response.status_code}: {error_body}" + ) + yield response, first_token_deadline + @staticmethod def _iter_text_cancellable( response: "httpx.Response", @@ -7567,119 +7562,102 @@ class LlamaCppBackend: _metadata_finish_reason = None try: - # Prefill can use the long first-token timeout; body reads are lowered after headers. - stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client( - timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) - ) as client: - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - with self._stream_with_retry( - client, - url, - payload, + with self._open_stream(url, payload, cancel_event) as ( + response, + first_token_deadline, + ): + buffer = "" + has_content_tokens = False + reasoning_text = "" + for raw_chunk in self._iter_text_cancellable( + response, cancel_event, - headers = _auth_headers, first_token_deadline = first_token_deadline, - ) as response: - if response.status_code != 200: - error_body = response.read().decode() - raise RuntimeError( - f"llama-server returned {response.status_code}: {error_body}" - ) + ): + buffer += raw_chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() - buffer = "" - has_content_tokens = False - reasoning_text = "" - for raw_chunk in self._iter_text_cancellable( - response, - cancel_event, - first_token_deadline = first_token_deadline, - ): - buffer += raw_chunk - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) - line = line.strip() + if not line: + continue + if line == "data: [DONE]": + if in_thinking: + if has_content_tokens: + # Real thinking + content: close the tag + cumulative += "" + yield cumulative + else: + # Only reasoning_content, no content: + # model put its whole reply in reasoning + # (e.g. Qwen3 always-think). Show it as + # the main response, not a thinking block. + cumulative = reasoning_text + yield cumulative + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue - if not line: + try: + data = json.loads(line[6:]) + # Diffusion frame (per-step canvas) from the shim: forward untouched so + # the frontend renders it in place. No assistant text, so it never enters + # the cumulative content. + if data.get("type") == "diffusion_frame": + yield data continue - if line == "data: [DONE]": - if in_thinking: - if has_content_tokens: - # Real thinking + content: close the tag + # Capture server timings/usage from final chunks. + _chunk_timings = data.get("timings") + if _chunk_timings: + _metadata_timings = _chunk_timings + _chunk_usage = data.get("usage") + if _chunk_usage: + _metadata_usage = _chunk_usage + choices = data.get("choices", []) + if choices: + delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _metadata_finish_reason = _fr + + # Reasoning/thinking tokens: llama-server + # sends these as "reasoning_content"; wrap + # in tags for the frontend parser. + reasoning = delta.get("reasoning_content", "") + if reasoning: + reasoning_text += reasoning + if not in_thinking: + cumulative += "" + in_thinking = True + cumulative += reasoning + yield cumulative + + token = delta.get("content", "") + if token: + has_content_tokens = True + if in_thinking: cumulative += "" - yield cumulative - else: - # Only reasoning_content, no content: - # model put its whole reply in reasoning - # (e.g. Qwen3 always-think). Show it as - # the main response, not a thinking block. - cumulative = reasoning_text - yield cumulative - _stream_done = True - break # exit inner while - if not line.startswith("data: "): - continue - - try: - data = json.loads(line[6:]) - # Diffusion frame (per-step canvas) from the shim: forward untouched so - # the frontend renders it in place. No assistant text, so it never enters - # the cumulative content. - if data.get("type") == "diffusion_frame": - yield data - continue - # Capture server timings/usage from final chunks. - _chunk_timings = data.get("timings") - if _chunk_timings: - _metadata_timings = _chunk_timings - _chunk_usage = data.get("usage") - if _chunk_usage: - _metadata_usage = _chunk_usage - choices = data.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - _fr = choices[0].get("finish_reason") - if _fr: - _metadata_finish_reason = _fr - - # Reasoning/thinking tokens: llama-server - # sends these as "reasoning_content"; wrap - # in tags for the frontend parser. - reasoning = delta.get("reasoning_content", "") - if reasoning: - reasoning_text += reasoning - if not in_thinking: - cumulative += "" - in_thinking = True - cumulative += reasoning - yield cumulative - - token = delta.get("content", "") - if token: - has_content_tokens = True - if in_thinking: - cumulative += "" - in_thinking = False - cumulative += token - yield cumulative - except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") - if _stream_done: - break # exit outer for - if _metadata_usage or _metadata_timings or _metadata_finish_reason: - _metadata_usage = _backfill_usage_from_timings( - _metadata_usage, _metadata_timings - ) - yield { - "type": "metadata", - # Never None: a finish-only metadata event (no usage, - # no timings) would otherwise crash consumers that do - # usage.get(...) on the non-streaming paths. - "usage": _metadata_usage or {}, - "timings": _metadata_timings, - "finish_reason": _metadata_finish_reason, - } + in_thinking = False + cumulative += token + yield cumulative + except json.JSONDecodeError: + logger.debug(f"Skipping malformed SSE line: {line[:100]}") + if _stream_done: + break # exit outer for + if _metadata_usage or _metadata_timings or _metadata_finish_reason: + _metadata_usage = _backfill_usage_from_timings( + _metadata_usage, _metadata_timings + ) + yield { + "type": "metadata", + # Never None: a finish-only metadata event (no usage, + # no timings) would otherwise crash consumers that do + # usage.get(...) on the non-streaming paths. + "usage": _metadata_usage or {}, + "timings": _metadata_timings, + "finish_reason": _metadata_finish_reason, + } except httpx.ConnectError as e: # Server already down. If this was an MTP+tensor crash, recover by @@ -7793,6 +7771,42 @@ class LlamaCppBackend: text = pat.sub("", text) return text + def _build_metadata_event(usage, timings, finish_reason): + """Final usage+timings metadata event for the given pass, merging its + usage/timings with the running cross-iteration accumulators. None when + there is nothing to report.""" + _fu = _backfill_usage_from_timings(usage, timings) or {} + _fp = _fu.get("prompt_tokens", 0) + _tc = _fu.get("completion_tokens", 0) + _accumulated_completion_tokens + if not (usage or timings or _accumulated_completion_tokens or finish_reason): + return None + _mt = dict(timings) if timings else {} + if _accumulated_predicted_ms or _accumulated_predicted_n: + _mt["predicted_ms"] = _mt.get("predicted_ms", 0) + _accumulated_predicted_ms + _mt["predicted_n"] = _mt.get("predicted_n", 0) + _accumulated_predicted_n + if _mt["predicted_ms"] > 0: + _mt["predicted_per_second"] = _mt["predicted_n"] / ( + _mt["predicted_ms"] / 1000.0 + ) + return { + "type": "metadata", + "usage": { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + }, + "timings": _mt, + "finish_reason": finish_reason, + } + + def _flush_reasoning_and_buffer(): + """Append buffered reasoning (as a block) then the held + content_buffer to the cumulative display text.""" + nonlocal cumulative_display + if reasoning_accum: + cumulative_display += "" + reasoning_accum + "" + cumulative_display += content_buffer + tool_controller = ToolLoopController( tools = tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -7831,7 +7845,7 @@ class LlamaCppBackend: if not active_tools: _append_budget_exhausted_nudge = False break - _tool_xml_signals = TOOL_XML_SIGNALS if active_tools else () + _tool_xml_signals = TOOL_XML_SIGNALS # Build payload -- stream: True so we detect tool signals # in the first 1-2 chunks without a non-streaming penalty. @@ -7864,10 +7878,6 @@ class LlamaCppBackend: payload["seed"] = seed try: - _auth_headers = ( - {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - ) - # ── Speculative buffer state machine ────────────────── # BUFFERING: accumulate content, check for tool signals # STREAMING: no tool detected, yield tokens to caller @@ -7893,194 +7903,220 @@ class LlamaCppBackend: provisional_render_html_tool_call_ids = set() _suppress_visible_output = _forced_tool_call_pending - stream_timeout = httpx.Timeout( - connect = 10, - read = 0.5, - write = 10, - pool = 10, - ) - with httpx.Client( - timeout = stream_timeout, - limits = httpx.Limits(max_keepalive_connections = 0), - ) as client: - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - with self._stream_with_retry( - client, - url, - payload, + with self._open_stream(url, payload, cancel_event) as ( + response, + first_token_deadline, + ): + raw_buf = "" + for raw_chunk in self._iter_text_cancellable( + response, cancel_event, - headers = _auth_headers, first_token_deadline = first_token_deadline, - ) as response: - if response.status_code != 200: - error_body = response.read().decode() - raise RuntimeError( - f"llama-server returned {response.status_code}: {error_body}" - ) + ): + raw_buf += raw_chunk + while "\n" in raw_buf: + line, raw_buf = raw_buf.split("\n", 1) + line = line.strip() - raw_buf = "" - for raw_chunk in self._iter_text_cancellable( - response, - cancel_event, - first_token_deadline = first_token_deadline, - ): - raw_buf += raw_chunk - while "\n" in raw_buf: - line, raw_buf = raw_buf.split("\n", 1) - line = line.strip() + if not line: + continue + if line == "data: [DONE]": + # Flush thinking state for STREAMING + if detect_state == _S_STREAMING and in_thinking: + if has_content_tokens: + cumulative_display += "" + if not _suppress_visible_output: + yield { + "type": "content", + "text": _strip_tool_markup( + cumulative_display, + final = True, + ), + } + else: + cumulative_display = reasoning_accum + if not _suppress_visible_output: + yield { + "type": "content", + "text": cumulative_display, + } + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue - if not line: + try: + chunk_data = json.loads(line[6:]) + _ct = chunk_data.get("timings") + if _ct: + _iter_timings = _ct + _cu = chunk_data.get("usage") + if _cu: + _iter_usage = _cu + + choices = chunk_data.get("choices", []) + if not choices: continue - if line == "data: [DONE]": - # Flush thinking state for STREAMING - if detect_state == _S_STREAMING and in_thinking: - if has_content_tokens: + + delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _iter_finish_reason = _fr + + # ── Structured tool_calls ── + tc_deltas = delta.get("tool_calls") + if tc_deltas: + # llama-server can emit visible assistant + # preface content before native structured + # tool_calls. Preserve content_accum as + # the assistant pre-tool text and still + # drain/execute the structured call. + has_structured_tc = True + detect_state = _S_DRAINING + for tc_d in tc_deltas: + idx = tc_d.get("index", 0) + if idx not in tool_calls_acc: + tool_calls_acc[idx] = { + "id": tc_d.get("id", f"call_{idx}"), + "type": "function", + "function": { + "name": "", + "arguments": "", + }, + } + elif tc_d.get("id"): + # Update ID if a real one + # arrives on a later delta. + tool_calls_acc[idx]["id"] = tc_d["id"] + func = tc_d.get("function", {}) + if func.get("name"): + tool_calls_acc[idx]["function"]["name"] += func["name"] + if func.get("arguments"): + tool_calls_acc[idx]["function"]["arguments"] += func[ + "arguments" + ] + current_name = tool_calls_acc[idx]["function"].get( + "name", "" + ) + fallback_id = f"call_{idx}" + current_id = tool_calls_acc[idx].get("id", fallback_id) + already_started = ( + current_id in provisional_render_html_tool_call_ids + ) + has_real_id = current_id != fallback_id + if ( + current_name == "render_html" + and not _tool_succeeded("render_html") + and any( + ( + (tool.get("function") or {}).get("name") + == "render_html" + ) + for tool in active_tools + ) + and not already_started + and not provisional_render_html_tool_call_ids + and has_real_id + ): + provisional_render_html_tool_call_ids.add(current_id) + yield { + "type": "tool_start", + "tool_name": "render_html", + "tool_call_id": current_id, + "arguments": {}, + "provenance": tool_event_provenance( + provisional = True, + ), + } + continue + + # ── Reasoning tokens ── + # Yield only in STREAMING. In BUFFERING and + # DRAINING, accumulate silently so we don't + # corrupt the consumer's prev_text tracker + # (routes/inference.py never resets it + # between tool iterations). + reasoning = delta.get("reasoning_content", "") + if reasoning: + reasoning_accum += reasoning + if detect_state == _S_STREAMING: + if not in_thinking: + cumulative_display += "" + in_thinking = True + cumulative_display += reasoning + if not _suppress_visible_output: + yield { + "type": "content", + "text": cumulative_display, + } + + # ── Content tokens ── + token = delta.get("content", "") + if token: + has_content_tokens = True + content_accum += token + + if detect_state == _S_DRAINING: + pass # accumulate silently + + elif detect_state == _S_STREAMING: + if in_thinking: cumulative_display += "" + in_thinking = False + cumulative_display += token + cleaned = _strip_tool_markup_streaming(cumulative_display) + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned if not _suppress_visible_output: yield { "type": "content", - "text": _strip_tool_markup( - cumulative_display, - final = True, - ), - } - else: - cumulative_display = reasoning_accum - if not _suppress_visible_output: - yield { - "type": "content", - "text": cumulative_display, - } - _stream_done = True - break # exit inner while - if not line.startswith("data: "): - continue - - try: - chunk_data = json.loads(line[6:]) - _ct = chunk_data.get("timings") - if _ct: - _iter_timings = _ct - _cu = chunk_data.get("usage") - if _cu: - _iter_usage = _cu - - choices = chunk_data.get("choices", []) - if not choices: - continue - - delta = choices[0].get("delta", {}) - _fr = choices[0].get("finish_reason") - if _fr: - _iter_finish_reason = _fr - - # ── Structured tool_calls ── - tc_deltas = delta.get("tool_calls") - if tc_deltas: - # llama-server can emit visible assistant - # preface content before native structured - # tool_calls. Preserve content_accum as - # the assistant pre-tool text and still - # drain/execute the structured call. - has_structured_tc = True - detect_state = _S_DRAINING - for tc_d in tc_deltas: - idx = tc_d.get("index", 0) - if idx not in tool_calls_acc: - tool_calls_acc[idx] = { - "id": tc_d.get("id", f"call_{idx}"), - "type": "function", - "function": { - "name": "", - "arguments": "", - }, - } - elif tc_d.get("id"): - # Update ID if a real one - # arrives on a later delta. - tool_calls_acc[idx]["id"] = tc_d["id"] - func = tc_d.get("function", {}) - if func.get("name"): - tool_calls_acc[idx]["function"]["name"] += func[ - "name" - ] - if func.get("arguments"): - tool_calls_acc[idx]["function"]["arguments"] += ( - func["arguments"] - ) - current_name = tool_calls_acc[idx]["function"].get( - "name", "" - ) - fallback_id = f"call_{idx}" - current_id = tool_calls_acc[idx].get("id", fallback_id) - already_started = ( - current_id in provisional_render_html_tool_call_ids - ) - has_real_id = current_id != fallback_id - if ( - current_name == "render_html" - and not _tool_succeeded("render_html") - and any( - ( - (tool.get("function") or {}).get("name") - == "render_html" - ) - for tool in active_tools - ) - and not already_started - and not provisional_render_html_tool_call_ids - and has_real_id - ): - provisional_render_html_tool_call_ids.add( - current_id - ) - yield { - "type": "tool_start", - "tool_name": "render_html", - "tool_call_id": current_id, - "arguments": {}, - "provenance": tool_event_provenance( - provisional = True, - ), - } - continue - - # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). - reasoning = delta.get("reasoning_content", "") - if reasoning: - reasoning_accum += reasoning - if detect_state == _S_STREAMING: - if not in_thinking: - cumulative_display += "" - in_thinking = True - cumulative_display += reasoning - if not _suppress_visible_output: - yield { - "type": "content", - "text": cumulative_display, + "text": cleaned, } - # ── Content tokens ── - token = delta.get("content", "") - if token: - has_content_tokens = True - content_accum += token + elif detect_state == _S_BUFFERING: + content_buffer += token + stripped_buf = content_buffer.lstrip() + if not stripped_buf: + continue - if detect_state == _S_DRAINING: - pass # accumulate silently + # Check tool signal prefixes. + is_prefix = False + is_match = False + for sig in _tool_xml_signals: + if stripped_buf.startswith(sig): + is_match = True + break + if sig.startswith(stripped_buf): + is_prefix = True + break - elif detect_state == _S_STREAMING: - if in_thinking: - cumulative_display += "" - in_thinking = False - cumulative_display += token + if is_match: + # Tool signal -- flush any visible + # prefix before DRAINING so the + # route sends it before tool_start. + _flush_reasoning_and_buffer() cleaned = _strip_tool_markup_streaming( - cumulative_display + cumulative_display, + force = True, + ) + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned + if not _suppress_visible_output: + yield { + "type": "content", + "text": cleaned, + } + detect_state = _S_DRAINING + elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS: + pass # keep buffering + else: + # Not a tool -- flush buffer + detect_state = _S_STREAMING + # Flush reasoning accumulated + # during BUFFERING. + _flush_reasoning_and_buffer() + cleaned = _strip_tool_markup( + cumulative_display, ) if len(cleaned) > len(_last_emitted): _last_emitted = cleaned @@ -8090,73 +8126,10 @@ class LlamaCppBackend: "text": cleaned, } - elif detect_state == _S_BUFFERING: - content_buffer += token - stripped_buf = content_buffer.lstrip() - if not stripped_buf: - continue - - # Check tool signal prefixes. - is_prefix = False - is_match = False - for sig in _tool_xml_signals: - if stripped_buf.startswith(sig): - is_match = True - break - if sig.startswith(stripped_buf): - is_prefix = True - break - - if is_match: - # Tool signal -- flush any visible - # prefix before DRAINING so the - # route sends it before tool_start. - if reasoning_accum: - cumulative_display += "" - cumulative_display += reasoning_accum - cumulative_display += "" - cumulative_display += content_buffer - cleaned = _strip_tool_markup_streaming( - cumulative_display, - force = True, - ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned - if not _suppress_visible_output: - yield { - "type": "content", - "text": cleaned, - } - detect_state = _S_DRAINING - elif ( - is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS - ): - pass # keep buffering - else: - # Not a tool -- flush buffer - detect_state = _S_STREAMING - # Flush reasoning accumulated - # during BUFFERING. - if reasoning_accum: - cumulative_display += "" - cumulative_display += reasoning_accum - cumulative_display += "" - cumulative_display += content_buffer - cleaned = _strip_tool_markup( - cumulative_display, - ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned - if not _suppress_visible_output: - yield { - "type": "content", - "text": cleaned, - } - - except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") - if _stream_done: - break # exit outer for + except json.JSONDecodeError: + logger.debug(f"Skipping malformed SSE line: {line[:100]}") + if _stream_done: + break # exit outer for # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: @@ -8167,11 +8140,7 @@ class LlamaCppBackend: detect_state = _S_STREAMING if content_buffer: # Flush reasoning first. - if reasoning_accum: - cumulative_display += "" - cumulative_display += reasoning_accum - cumulative_display += "" - cumulative_display += content_buffer + _flush_reasoning_and_buffer() if not _suppress_visible_output: yield { "type": "content", @@ -8292,31 +8261,11 @@ class LlamaCppBackend: # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} - _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} - _fc = _fu.get("completion_tokens", 0) - _fp = _fu.get("prompt_tokens", 0) - _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: - _mt = dict(_iter_timings) if _iter_timings else {} - if _accumulated_predicted_ms or _accumulated_predicted_n: - _mt["predicted_ms"] = ( - _mt.get("predicted_ms", 0) + _accumulated_predicted_ms - ) - _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n - _mt["predicted_n"] = _tn - _tms = _mt["predicted_ms"] - if _tms > 0: - _mt["predicted_per_second"] = _tn / (_tms / 1000.0) - yield { - "type": "metadata", - "usage": { - "prompt_tokens": _fp, - "completion_tokens": _tc, - "total_tokens": _fp + _tc, - }, - "timings": _mt, - "finish_reason": _iter_finish_reason, - } + _meta = _build_metadata_event( + _iter_usage, _iter_timings, _iter_finish_reason + ) + if _meta is not None: + yield _meta return # Safety net caught tool XML -- treat as tool call. @@ -8367,31 +8316,11 @@ class LlamaCppBackend: content_accum = _strip_tool_markup(content_accum, final = True) if content_accum: yield {"type": "content", "text": content_accum} - _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {} - _fc = _fu.get("completion_tokens", 0) - _fp = _fu.get("prompt_tokens", 0) - _tc = _fc + _accumulated_completion_tokens - if _iter_usage or _iter_timings or _accumulated_completion_tokens: - _mt = dict(_iter_timings) if _iter_timings else {} - if _accumulated_predicted_ms or _accumulated_predicted_n: - _mt["predicted_ms"] = ( - _mt.get("predicted_ms", 0) + _accumulated_predicted_ms - ) - _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n - _mt["predicted_n"] = _tn - _tms = _mt["predicted_ms"] - if _tms > 0: - _mt["predicted_per_second"] = _tn / (_tms / 1000.0) - yield { - "type": "metadata", - "usage": { - "prompt_tokens": _fp, - "completion_tokens": _tc, - "total_tokens": _fp + _tc, - }, - "timings": _mt, - "finish_reason": _iter_finish_reason, - } + _meta = _build_metadata_event( + _iter_usage, _iter_timings, _iter_finish_reason + ) + if _meta is not None: + yield _meta return # ── Execute tool calls ── @@ -8592,125 +8521,85 @@ class LlamaCppBackend: _stream_done = False try: - stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client( - timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) - ) as client: - first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - with self._stream_with_retry( - client, - url, - stream_payload, + with self._open_stream(url, stream_payload, cancel_event) as ( + response, + first_token_deadline, + ): + buffer = "" + for raw_chunk in self._iter_text_cancellable( + response, cancel_event, - headers = _auth_headers, first_token_deadline = first_token_deadline, - ) as response: - if response.status_code != 200: - error_body = response.read().decode() - raise RuntimeError( - f"llama-server returned {response.status_code}: {error_body}" - ) + ): + buffer += raw_chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() - buffer = "" - for raw_chunk in self._iter_text_cancellable( - response, - cancel_event, - first_token_deadline = first_token_deadline, - ): - buffer += raw_chunk - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) - line = line.strip() + if not line: + continue + if line == "data: [DONE]": + if in_thinking: + if has_content_tokens: + cumulative += "" + yield { + "type": "content", + "text": _strip_tool_markup(cumulative, final = True), + } + else: + cumulative = reasoning_text + yield {"type": "content", "text": cumulative} + _stream_done = True + break # exit inner while + if not line.startswith("data: "): + continue - if not line: - continue - if line == "data: [DONE]": - if in_thinking: - if has_content_tokens: + try: + chunk_data = json.loads(line[6:]) + # Capture server timings/usage from final chunks. + _chunk_timings = chunk_data.get("timings") + if _chunk_timings: + _metadata_timings = _chunk_timings + _chunk_usage = chunk_data.get("usage") + if _chunk_usage: + _metadata_usage = _chunk_usage + choices = chunk_data.get("choices", []) + if choices: + delta = choices[0].get("delta", {}) + _fr = choices[0].get("finish_reason") + if _fr: + _metadata_finish_reason = _fr + + reasoning = delta.get("reasoning_content", "") + if reasoning: + reasoning_text += reasoning + if not in_thinking: + cumulative += "" + in_thinking = True + cumulative += reasoning + yield {"type": "content", "text": cumulative} + + token = delta.get("content", "") + if token: + has_content_tokens = True + if in_thinking: cumulative += "" - yield { - "type": "content", - "text": _strip_tool_markup(cumulative, final = True), - } - else: - cumulative = reasoning_text - yield {"type": "content", "text": cumulative} - _stream_done = True - break # exit inner while - if not line.startswith("data: "): - continue - - try: - chunk_data = json.loads(line[6:]) - # Capture server timings/usage from final chunks. - _chunk_timings = chunk_data.get("timings") - if _chunk_timings: - _metadata_timings = _chunk_timings - _chunk_usage = chunk_data.get("usage") - if _chunk_usage: - _metadata_usage = _chunk_usage - choices = chunk_data.get("choices", []) - if choices: - delta = choices[0].get("delta", {}) - _fr = choices[0].get("finish_reason") - if _fr: - _metadata_finish_reason = _fr - - reasoning = delta.get("reasoning_content", "") - if reasoning: - reasoning_text += reasoning - if not in_thinking: - cumulative += "" - in_thinking = True - cumulative += reasoning - yield {"type": "content", "text": cumulative} - - token = delta.get("content", "") - if token: - has_content_tokens = True - if in_thinking: - cumulative += "" - in_thinking = False - cumulative += token - cleaned = _strip_tool_markup(cumulative) - # Emit only when cleaned text grows (monotonic). - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned - yield {"type": "content", "text": cleaned} - except json.JSONDecodeError: - logger.debug(f"Skipping malformed SSE line: {line[:100]}") - if _stream_done: - break # exit outer for - _final_usage = _metadata_usage or {} - _final_completion = _final_usage.get("completion_tokens", 0) - _final_prompt = _final_usage.get("prompt_tokens", 0) - _total_completion = _final_completion + _accumulated_completion_tokens - if _metadata_usage or _metadata_timings or _metadata_finish_reason: - _merged_timings = dict(_metadata_timings) if _metadata_timings else {} - if _accumulated_predicted_ms or _accumulated_predicted_n: - _merged_timings["predicted_ms"] = ( - _merged_timings.get("predicted_ms", 0) + _accumulated_predicted_ms - ) - _total_predicted_n = ( - _merged_timings.get("predicted_n", 0) + _accumulated_predicted_n - ) - _merged_timings["predicted_n"] = _total_predicted_n - _total_predicted_ms = _merged_timings["predicted_ms"] - if _total_predicted_ms > 0: - _merged_timings["predicted_per_second"] = _total_predicted_n / ( - _total_predicted_ms / 1000.0 - ) - yield { - "type": "metadata", - "usage": { - "prompt_tokens": _final_prompt, - "completion_tokens": _total_completion, - "total_tokens": _final_prompt + _total_completion, - }, - "timings": _merged_timings, - "finish_reason": _metadata_finish_reason, - } + in_thinking = False + cumulative += token + cleaned = _strip_tool_markup(cumulative) + # Emit only when cleaned text grows (monotonic). + if len(cleaned) > len(_last_emitted): + _last_emitted = cleaned + yield {"type": "content", "text": cleaned} + except json.JSONDecodeError: + logger.debug(f"Skipping malformed SSE line: {line[:100]}") + if _stream_done: + break # exit outer for + _meta = _build_metadata_event( + _metadata_usage, _metadata_timings, _metadata_finish_reason + ) + if _meta is not None: + yield _meta except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") @@ -8748,8 +8637,6 @@ class LlamaCppBackend: continue if not isinstance(block, dict): return True - if block.get("type") == "text" and isinstance(block.get("text"), str): - continue if isinstance(block.get("text"), str): continue return True @@ -8770,9 +8657,7 @@ class LlamaCppBackend: parts = [] for block in content: if isinstance(block, dict): - if block.get("type") == "text" and isinstance(block.get("text"), str): - parts.append(block["text"]) - elif isinstance(block.get("text"), str): + if isinstance(block.get("text"), str): parts.append(block["text"]) elif isinstance(block, str): parts.append(block) @@ -8787,8 +8672,7 @@ class LlamaCppBackend: system_text = _block_text(system) try: - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client(timeout = 10, headers = _auth_headers) as client: + with httpx.Client(timeout = 10, headers = self._auth_headers) as client: def _tokenize(text: str) -> int: r = client.post( @@ -8867,12 +8751,44 @@ class LlamaCppBackend: logger.debug(f"Audio type detection failed: {e}") return None + def _apply_detected_audio(self, detected: Optional[str]) -> bool: + """Apply a probed audio codec under self._lock. Returns True to continue + the load (codec inited OK, or nothing to init), False to abort (server + unhealthy or codec init failed). Shared by the fast-path retry and the + main load path.""" + if detected in ("snac", "bicodec", "dac"): + with self._lock: + if not self._healthy: + return False + try: + self.init_audio_codec(detected) + self._is_audio = True + self._audio_type = detected + except Exception as exc: + # Surface as HTTP 500 (matches pre-PR contract). + logger.warning("Failed to init audio codec '%s': %s", detected, exc) + self._audio_probed = False + return False + elif detected: + # csm / whisper / audio_vlm: track type but keep _is_audio False -- + # GGUF TTS routing only fires for snac/bicodec/dac. + with self._lock: + if not self._healthy: + return False + self._audio_type = detected + # Audio input = token probe (audio_vlm/whisper) OR mmproj encoder. + from utils.models.model_config import is_audio_input_type + + self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( + self._mmproj_has_audio + ) + return True + def _detect_audio_type_strict(self) -> Optional[str]: """Codec name on match, None on non-audio, raises on transport/JSON errors.""" if not self.is_loaded: return None - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client(timeout = 10, headers = _auth_headers) as client: + with httpx.Client(timeout = 10, headers = self._auth_headers) as client: def _detok(tid: int) -> str: # Non-200 means "marker not in vocab" -- keep probing. @@ -8986,8 +8902,9 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 - _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None - with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client: + with httpx.Client( + timeout = httpx.Timeout(300, connect = 10), headers = self._auth_headers + ) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}") diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d520896798..26825a472e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -412,12 +412,13 @@ class InferenceStatusResponse(BaseModel): "(auto on an MTP model, or forced mtp / mtp+ngram). " "'binary_no_mtp' / 'binary_outdated' -> a newer prebuilt would " "re-enable it (show the update affordance); 'runtime_error' -> the " - "current build could not run it. 'mla_mtp_disabled' -> an Auto-mode " - "policy downgrade: the model is MLA (GLM-5.2 et al.) whose llama.cpp " - "MTP path runs slower than no speculation, so Auto used ngram-mod or " - "spec-off instead -- updating won't help; choose MTP in Settings (or " - "set UNSLOTH_MLA_MTP_ENABLED=1) to force it. None when MTP engaged or " - "was not requested." + "current build could not run it; 'drafter_not_found' -> the model's " + "separate MTP drafter could not be resolved; 'mla_mtp_disabled' -> " + "an Auto-mode policy downgrade: the model is MLA (GLM-5.2 et al.) " + "whose llama.cpp MTP path runs slower than no speculation, so Auto " + "used ngram-mod or spec-off instead -- updating won't help; choose " + "MTP in Settings (or set UNSLOTH_MLA_MTP_ENABLED=1) to force it. " + "None when MTP engaged or was not requested." ), ) llama_cpp_prebuilt_stale: bool = Field( diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index 3ac85380fe..d2b3bf94e9 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -5,6 +5,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +import base64 import ipaddress import os import shlex @@ -214,6 +215,34 @@ def _clear_login_bucket(key: tuple[str, str]) -> None: _LOGIN_IP_BUCKETS.pop(ip, None) +# Sync def (not async): compute_identity_proof touches SQLite on the first call, +# so FastAPI runs it in the threadpool rather than blocking the event loop. +@router.get("/identity") +def identity(nonce: str, request: Request) -> dict: + """Challenge-response proof this is the real local Studio: caller sends a nonce, + gets HMAC(install identity secret, nonce, connection address + port). + Unauthenticated and side-effect free; a process that can't read the same-user + secret can't forge a proof, and binding to the address/port the connection + landed on stops a squatter relaying a proof from the real Studio elsewhere.""" + try: + raw = base64.urlsafe_b64decode(nonce) + except Exception: + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, detail = "nonce must be base64url" + ) + if not 16 <= len(raw) <= 128: + raise HTTPException( + status_code = status.HTTP_400_BAD_REQUEST, detail = "nonce must decode to 16-128 bytes" + ) + # The address + port the connection actually landed on, from the socket + # (request.scope is getsockname, so it is the real local address even when + # bound to 0.0.0.0), never the client-controlled Host header. + server = request.scope.get("server") or ("", 0) + host = server[0] or "" + port = server[1] if server[1] is not None else 0 + return {"proof": storage.compute_identity_proof(raw, host, port)} + + @router.get("/status", response_model = AuthStatusResponse) async def auth_status() -> AuthStatusResponse: """Auth initialization state; ``default_username`` is exposed for first-boot UI prefill only.""" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 17bcc084ee..046968cb71 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -16,7 +16,6 @@ from starlette.requests import ClientDisconnect from typing import Any, List, Optional, Union import json import httpx -import structlog from loggers import get_logger import asyncio import threading @@ -120,6 +119,11 @@ def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Op return candidate if candidate and candidate in chat_template else None +_LOST_CONNECTION_MSG = ( + "Lost connection to the model server. It may have crashed -- try reloading the model." +) + + def _friendly_error(exc: Exception) -> str: """Extract a user-friendly message from known llama-server errors.""" if isinstance(exc, httpx.ReadTimeout): @@ -140,9 +144,7 @@ def _friendly_error(exc: Exception) -> str: # WriteError, PoolTimeout, ...) means the llama-server subprocess is # unreachable -- crashed or still coming up. if isinstance(exc, httpx.RequestError): - return ( - "Lost connection to the model server. It may have crashed -- try reloading the model." - ) + return _LOST_CONNECTION_MSG msg = str(exc) m = _re.search( r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)", @@ -155,9 +157,7 @@ def _friendly_error(exc: Exception) -> str: f"or shorten the conversation." ) if "Lost connection to llama-server" in msg: - return ( - "Lost connection to the model server. It may have crashed -- try reloading the model." - ) + return _LOST_CONNECTION_MSG template_msg = _template_raise_message(msg, _loaded_chat_template()) if template_msg: return f"An internal error occurred: {template_msg}" @@ -230,6 +230,31 @@ def _raise_unsupported_n(path_label: str) -> None: _raise_unsupported_openai_parameter("n", f"n > 1 is not supported for {path_label}.") +def _sse_streaming_response(content) -> StreamingResponse: + """A ``text/event-stream`` response with the standard SSE headers used by + every streaming path here: no client/proxy caching, no proxy buffering, and + a one-shot connection. Two callers build their response inline instead: the + external-provider proxy omits ``Connection: close``, and the OpenAI + passthrough returns an empty ``keep-alive`` stream when the request is + cancelled before the upstream response starts. + + Built on ``_SameTaskStreamingResponse`` (not Starlette's stock + ``StreamingResponse``) so the SSE generator runs in the request task. The + legacy AnyIO task-group wrapper trips "Attempted to exit a cancel scope in a + different task" on Python 3.13 + httpx, which surfaced as a mid-stream + ``response.failed``. The streaming paths that take their response inline use + ``_SameTaskStreamingResponse`` directly for the same reason.""" + return _SameTaskStreamingResponse( + content, + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) + + def _openai_stream_error_chunk(exc) -> dict: """Build an in-band OpenAI error chunk for a mid-stream failure. Once the stream's 200 headers are flushed the status can't change, so the error must @@ -404,12 +429,11 @@ def _apply_overflow_truncation(body: dict, err_text: str) -> bool: if counts: n_prompt, n_ctx = counts keep_ratio = min(0.95, (_OVERFLOW_PROMPT_TARGET_FRACTION * n_ctx) / max(1, n_prompt)) - # Scale the server-token target into char-estimate units. - target_est = int(total_est * keep_ratio) else: n_ctx = None keep_ratio = 0.6 # no counts in the error; cut conservatively - target_est = int(total_est * keep_ratio) + # Scale the server-token target into char-estimate units. + target_est = int(total_est * keep_ratio) new_messages, dropped = _truncate_middle_messages(messages, keep_ratio) if dropped: @@ -523,6 +547,52 @@ def _openai_stream_usage_chunk( return f"data: {usage_chunk.model_dump_json(exclude_none = True)}\n\n" +def _chat_chunk_sse(completion_id, created, model_name, *, delta, finish_reason) -> str: + """One ``ChatCompletionChunk`` as an SSE ``data:`` line. The role / content / + final chunks every in-process streamer emits differ only in their ``delta`` + and ``finish_reason``.""" + chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ChunkChoice(delta = delta, finish_reason = finish_reason)], + ) + return f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + + +def _chat_role_chunk(completion_id, created, model_name) -> str: + """Opening assistant-role chunk for a chat stream.""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(role = "assistant"), + finish_reason = None, + ) + + +def _chat_content_chunk(completion_id, created, model_name, text) -> str: + """A content-delta chunk carrying ``text``.""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(content = text), + finish_reason = None, + ) + + +def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: + """Terminal stop chunk (empty delta) carrying the finish reason.""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(), + finish_reason = finish_reason, + ) + + def _rewrite_cmpl_id(raw: bytes) -> bytes: """Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-`` prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key @@ -670,21 +740,11 @@ except ImportError: def _llama_non_streaming_generation_timeout() -> httpx.Timeout: - return httpx.Timeout( - connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - ) + return httpx.Timeout(_DEFAULT_FIRST_TOKEN_TIMEOUT_S) def _llama_streaming_generation_timeout() -> httpx.Timeout: - return httpx.Timeout( - connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, - ) + return httpx.Timeout(_DEFAULT_FIRST_TOKEN_TIMEOUT_S) def _set_stream_response_read_timeout( @@ -758,6 +818,39 @@ class _SameTaskStreamingResponse(StreamingResponse): raise ClientDisconnect() if self.background is not None: await self.background() +async def _aclose_stream_resources( + *, + watchers = (), + iterator = None, + resp = None, + client = None, +) -> None: + """Tear down an httpx streaming generator's resources in the required order: + cancel + await each watcher task, then aclose() the byte/line iterator, the + response, and the client. Each step swallows its own exceptions so teardown + always completes. See _anthropic_passthrough_stream for the ordering rationale.""" + for watcher in watchers: + if watcher is not None: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + if iterator is not None: + try: + await iterator.aclose() + except Exception: + pass + if resp is not None: + try: + await resp.aclose() + except Exception: + pass + if client is not None: + try: + await client.aclose() + except Exception: + pass async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: @@ -921,11 +1014,9 @@ from models.inference import ( ImageContentPart, ImageUrl, ResponsesRequest, - ResponsesInputMessage, ResponsesInputTextPart, ResponsesInputImagePart, ResponsesOutputTextPart, - ResponsesUnknownContentPart, ResponsesUnknownInputItem, ResponsesFunctionCallInputItem, ResponsesFunctionCallOutputInputItem, @@ -963,14 +1054,13 @@ from state.tool_approvals import resolve_tool_decision from core.inference.key_exchange import decrypt_api_key from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client -from core.inference.providers import get_provider_info, get_base_url +from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override from storage import providers_db from utils.utils import safe_error_detail, log_and_http_error import io -import wave import base64 import numpy as np from datetime import date as _date @@ -1068,6 +1158,26 @@ _ARTIFACT_PREVIEW_FRAME_HTML = """ """ +async def _authenticate_header_or_query(request: Request, token: Optional[str]) -> str: + """Resolve the bearer token from the Authorization header or the ``?token=`` + query param (needed for /