diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py new file mode 100644 index 0000000000..f76a38576f --- /dev/null +++ b/studio/backend/core/inference/api_monitor.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Small in-memory monitor for OpenAI-compatible API traffic.""" + +from __future__ import annotations + +import threading +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Optional + + +_MAX_ENTRIES = 50 +_MAX_PROMPT_CHARS = 12000 +_MAX_REPLY_CHARS = 12000 +_PREVIEW_CHARS = 360 + + +def _trim(text: Optional[str], limit: int) -> str: + if not text: + return "" + if len(text) <= limit: + return text + # Guard against limit < 3 (slice would underflow). + if limit <= 3: + return "..."[:limit] + return text[: limit - 3] + "..." + + +@dataclass +class ApiMonitorEntry: + id: str + endpoint: str + method: str + model: str + prompt: str + status: str + started_at: float + updated_at: float + subject: Optional[str] = None + # Monotonic anchors so duration math survives wall-clock steps (NTP). + started_monotonic: float = 0.0 + finished_monotonic: Optional[float] = None + reply: str = "" + finished_at: Optional[float] = None + context_length: Optional[int] = None + prompt_tokens: Optional[int] = None + completion_tokens: Optional[int] = None + total_tokens: Optional[int] = None + total_tokens_authoritative: bool = False + error: Optional[str] = None + + def snapshot(self, *, include_details: bool = True) -> dict[str, Any]: + duration_ms = None + if self.finished_monotonic is not None: + duration_ms = max( + 0, + int((self.finished_monotonic - self.started_monotonic) * 1000), + ) + elif self.finished_at is not None: + duration_ms = max(0, int((self.finished_at - self.started_at) * 1000)) + context_usage = None + if self.total_tokens is not None and self.context_length: + context_usage = min(1.0, max(0.0, self.total_tokens / self.context_length)) + payload = { + "id": self.id, + "endpoint": self.endpoint, + "method": self.method, + "model": self.model, + "prompt_preview": _trim(self.prompt, _PREVIEW_CHARS), + "reply_preview": _trim(self.reply, _PREVIEW_CHARS), + "prompt_truncated": len(self.prompt) > _PREVIEW_CHARS, + "reply_truncated": len(self.reply) > _PREVIEW_CHARS, + "status": self.status, + "started_at": self.started_at, + "updated_at": self.updated_at, + "finished_at": self.finished_at, + "duration_ms": duration_ms, + "context_length": self.context_length, + "context_usage": context_usage, + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, + "total_tokens": self.total_tokens, + "error": self.error, + } + if include_details: + payload["prompt"] = self.prompt + payload["reply"] = self.reply + return payload + + +class ApiMonitor: + def __init__(self, max_entries: int = _MAX_ENTRIES): + self._entries: deque[ApiMonitorEntry] = deque() + self._max_entries = max(0, max_entries) + self._lock = threading.Lock() + + def start( + self, + *, + endpoint: str, + method: str, + model: str, + prompt: str, + context_length: Optional[int] = None, + subject: Optional[str] = None, + ) -> str: + now = time.time() + entry = ApiMonitorEntry( + id = f"apireq_{uuid.uuid4().hex[:12]}", + endpoint = endpoint, + method = method, + model = model or "default", + prompt = _trim(prompt, _MAX_PROMPT_CHARS), + status = "running", + started_at = now, + updated_at = now, + subject = subject, + started_monotonic = time.monotonic(), + context_length = context_length, + ) + with self._lock: + self._entries.appendleft(entry) + self._trim_terminal_locked() + return entry.id + + def append_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id or not text: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + # Preview is capped: once the "..." marker is present the head is + # frozen, so skip the per-chunk re-concat (avoids O(n^2) on long + # generations). A reply that landed exactly on the cap has no marker + # yet, so let one more append record the truncation before freezing. + if len(entry.reply) >= _MAX_REPLY_CHARS: + if not entry.reply.endswith("..."): + entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + return + entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + entry.reply = _trim(text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_usage( + self, + entry_id: Optional[str], + *, + prompt_tokens: Optional[int] = None, + completion_tokens: Optional[int] = None, + total_tokens: Optional[int] = None, + context_length: Optional[int] = None, + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if prompt_tokens is not None: + entry.prompt_tokens = prompt_tokens + if completion_tokens is not None: + entry.completion_tokens = completion_tokens + if total_tokens is not None: + entry.total_tokens = total_tokens + entry.total_tokens_authoritative = True + elif not entry.total_tokens_authoritative and ( + prompt_tokens is not None or completion_tokens is not None + ): + # Derive only when no authoritative total has been set; + # a later partial chunk must not clobber a provider total. + entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0) + if context_length is not None: + entry.context_length = context_length + entry.updated_at = time.time() + + def finish( + self, + entry_id: Optional[str], + status: str = "completed", + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + # Idempotent: second call (e.g. [DONE] after the finally block + # already ran) must not move finished_*. + if entry.finished_at is not None: + return + now = time.time() + entry.status = status + 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 fail(self, entry_id: Optional[str], error: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if entry.finished_at is not None: + # Already terminal; refresh error text only. + 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() + + def snapshot( + self, + *, + include_details: bool = True, + subject: Optional[str] = None, + ) -> list[dict[str, Any]]: + with self._lock: + return [ + entry.snapshot(include_details = include_details) + for entry in self._entries + if subject is None or entry.subject == subject + ] + + def get( + self, + entry_id: str, + *, + subject: Optional[str] = None, + ) -> Optional[dict[str, Any]]: + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return None + if subject is not None and entry.subject != subject: + return None + return entry.snapshot(include_details = True) + + def active_count(self, *, subject: Optional[str] = None) -> int: + with self._lock: + return sum( + 1 + for entry in self._entries + if entry.status == "running" and (subject is None or entry.subject == subject) + ) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + + def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]: + for entry in self._entries: + if entry.id == entry_id: + return entry + return None + + def _trim_terminal_locked(self) -> None: + terminal_seen = 0 + kept: deque[ApiMonitorEntry] = deque() + for entry in self._entries: + if entry.status == "running": + kept.append(entry) + continue + if terminal_seen < self._max_entries: + kept.append(entry) + terminal_seen += 1 + self._entries = kept + + +api_monitor = ApiMonitor() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 0f4abe04f5..062944264a 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -23,18 +23,19 @@ import sys import threading import time from pathlib import Path -from typing import Callable, Collection, Generator, Iterable, List, Optional, Union +from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union import httpx from core.inference.llama_server_args import ( + _effective_tensor_parallel, + _tensor_parallel_matches_loaded, extra_args_disable_mmproj, parse_cache_override, + parse_cache_override_per_axis, parse_ctx_override, parse_split_mode_override, - resolve_cache_type_kv, resolve_requested_ctx, - resolve_tensor_parallel, strip_shadowing_flags, strip_split_mode_only, ) @@ -51,6 +52,7 @@ from core.tool_healing import ( strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret +from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) @@ -723,18 +725,63 @@ def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]: # _build_speculative_flags); auto mode drops MTP under it. _MTP_MIN_SIZE_B = 3.0 -# Context-fit VRAM budget: tighter than _GPU_PIN_VRAM_FRACTION (0.95) on -# purpose -- over-promising context OOMs at runtime (#5106). -_CTX_FIT_VRAM_FRACTION = 0.90 +# Cap total GPU occupancy at this fraction of the card. The fit reserves an +# absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction +# of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve. +_CTX_FIT_VRAM_FRACTION = 0.95 -# Extra VRAM fraction reserved when MTP will engage: the draft model's -# weights, KV cache, and compute buffers live outside the main model's -# estimate. Applied to BOTH the ctx-fit budget and the GPU pin thresholds -- -# tightening only the fit lets a load whose weights land between the two -# fractions pin without any room for the drafter. +# Flat MTP reserve, used only when GGUF dims are too sparse for the byte-accurate +# reserve (_estimate_mtp_overhead_bytes). Applied to both the fit budget and pin. _MTP_VRAM_RESERVE_FRAC = 0.05 +def _kv_bytes_per_elem(cache_type: Optional[str]) -> float: + """Bytes per KV-cache element for a llama.cpp cache type (f16 default).""" + return { + "f32": 4.0, + "f16": 2.0, + "bf16": 2.0, + "q8_0": 34 / 32, + "q5_1": 0.75, + "q5_0": 0.6875, + "q4_1": 0.625, + "q4_0": 0.5625, + "iq4_nl": 0.5625, + }.get((cache_type or "f16").strip().lower(), 2.0) + + +def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]: + """Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it + exceeds the f16 default, else None. Studio emits --cache-type only for the + param/extras path, so a heavier env (f32) would otherwise reach the child + unbudgeted; quantized env types stay over-reserved by f16 (-> None).""" + e = os.environ if env is None else env + f16_bpe = _kv_bytes_per_elem("f16") + heaviest: Optional[str] = None + heaviest_bpe = f16_bpe + for var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + raw = (e.get(var) or "").strip().lower() + if not raw: + continue + bpe = _kv_bytes_per_elem(raw) + if bpe > heaviest_bpe: + heaviest, heaviest_bpe = raw, bpe + return heaviest + + +def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) -> Optional[str]: + """Heavier (max bytes/elem) of the explicit --cache-type-k/-v extras, or None. + + Extras are appended last and win per axis, so an asymmetric K=f32,V=f16 must be + budgeted by its heavier axis. resolve_cache_type_kv returns only the last-wins + single type, which under-reserves the heavier axis when the lighter one is last.""" + k, v = parse_cache_override_per_axis(extra_args) + candidates = [c for c in (k, v) if c] + if not candidates: + return None + return max(candidates, key = _kv_bytes_per_elem) + + def _auto_mode_drops_mtp( req_mode: Optional[str], size_b: Optional[float], @@ -778,6 +825,196 @@ def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collect return False +def _effective_spec_type( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """The --spec-type llama-server will use: the last CLI --spec-type (or + --spec-default, which resolves non-MTP), else LLAMA_ARG_SPEC_TYPE. A CLI flag + overrides the env (matching llama.cpp), so a stale MTP env can't make the + budget reserve a drafter the launch won't load. None if neither sets it.""" + args = [str(a) for a in extra_args] if extra_args else [] + cli_present = False + cli_value: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag == "--spec-default": + cli_present = True + cli_value = "default" + continue + if flag != "--spec-type": + continue + cli_present = True + cli_value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if cli_present: + return cli_value + return (os.environ if env is None else env).get("LLAMA_ARG_SPEC_TYPE") + + +def _extra_args_requests_mtp( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects MTP (mtp/draft-mtp), so the + budget must reserve for it.""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("mtp", "draft-mtp") for p in value.split(",")) + + +def _extra_args_requests_separate_draft( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects a non-MTP model draft mode + (draft-simple/draft-eagle3), which loads a separate draft model the budget + must reserve (draft-mtp -> _extra_args_requests_mtp; ngram-* load no model).""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(",")) + + +def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optional[int]: + """Draft depth from extras (``--spec-draft-n-max`` or legacy ``--draft-max``), else None.""" + if not extra_args: + return None + args = [str(a) for a in extra_args] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--spec-draft-n-max", "--draft-max"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + return found + + +def _extra_args_mtp_draft_path( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """Separate drafter path from extras (local --model-draft/-md or HF + --spec-draft-hf/-hfd/...), else the LLAMA_ARG_SPEC_DRAFT_MODEL/_HF_REPO env, + else None. An HF repo isn't a local file, so the budget can't size it (falls + back to the flat reserve), but recognizing it avoids sizing the wrong one.""" + flags = { + "--model-draft", + "--spec-draft-model", + "-md", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", + } + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if value and not value.startswith("-"): + found = value + if found is not None: + return found + e = os.environ if env is None else env + return e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") or None + + +def _extra_args_draft_cache_types( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> tuple[Optional[str], Optional[str]]: + """Draft KV cache types (k_type, v_type), each from extras else the + LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V env, else None (f16). K and V are + independent: a one-sided override must not apply to both.""" + args = [str(a) for a in extra_args] if extra_args else [] + k_flags = {"--cache-type-k-draft", "--spec-draft-type-k", "-ctkd"} + v_flags = {"--cache-type-v-draft", "--spec-draft-type-v", "-ctvd"} + k_type: Optional[str] = None + v_type: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in k_flags and flag not in v_flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if not value or value.startswith("-"): + continue + if flag in k_flags: + k_type = value + else: + v_type = value + e = os.environ if env is None else env + if k_type is None: + k_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K") or None + if v_type is None: + v_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V") or None + return k_type, v_type + + +def _extra_args_draft_offloaded_to_cpu( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the SEPARATE draft model is on CPU (so the budget must not charge + its weights+KV): --spec-draft-ngl 0, or --spec-draft-device naming only + cpu/none, else the LLAMA_ARG_N_GPU_LAYERS_DRAFT env the child honors (the + device flag has no env). An embedded MTP head follows the main -ngl, so these + draft-only flags don't move it. Last-wins, so only each flag's final value counts.""" + ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"} + dev_flags = {"--spec-draft-device", "-devd", "--device-draft"} + args = [str(a) for a in extra_args] if extra_args else [] + last_ngl: Optional[str] = None + last_dev: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if flag in ngl_flags: + last_ngl = value + elif flag in dev_flags: + last_dev = value + if last_ngl is None: + last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT") + if last_ngl is not None: + try: + if int(last_ngl) == 0: + return True + except (TypeError, ValueError): + pass + if last_dev is not None: + devs = [d.strip().lower() for d in last_dev.split(",") if d.strip()] + if devs and all(d in ("cpu", "none") for d in devs): + return True + return False + + +def _extra_args_n_ubatch( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[int]: + """Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH + env, else None. It sizes the compute-graph buffer, so an override must reach + the VRAM reserve.""" + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--ubatch-size", "-ub"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + if found is not None: + return found + raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH") + if raw: + try: + return int(raw) + except (TypeError, ValueError): + pass + return None + + def _build_ngram_mod_flags( caps: Optional[dict], n_match: int = 24, @@ -923,6 +1160,7 @@ class LlamaCppBackend: self._is_diffusion: bool = False self._diffusion_visual_bin: Optional[str] = None self._healthy = False + self._stats_logger = None # vLLM-style engine-stats poller, set on load # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None self._context_length: Optional[int] = None @@ -953,6 +1191,9 @@ class LlamaCppBackend: self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None self._embedding_length: Optional[int] = None + # For the compute-graph buffer estimate; vocab from the tokens array len. + self._feed_forward_length: Optional[int] = None + self._vocab_size: Optional[int] = None # Architecture-aware KV fields for 5-path estimation self._kv_key_length: Optional[int] = None self._kv_value_length: Optional[int] = None @@ -1449,6 +1690,7 @@ class LlamaCppBackend: "supports_cache_ram": False, "supports_ctx_checkpoints": False, "supports_no_cache_prompt": False, + "supports_metrics": False, } try: mtime = int(Path(bin_path).stat().st_mtime) @@ -1467,6 +1709,7 @@ class LlamaCppBackend: supports_cache_ram = False supports_ctx_checkpoints = False supports_no_cache_prompt = False + supports_metrics = False try: probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( @@ -1562,6 +1805,7 @@ class LlamaCppBackend: supports_cache_ram = _is_real("--cache-ram") supports_ctx_checkpoints = _is_real("--ctx-checkpoints") supports_no_cache_prompt = _is_real("--no-cache-prompt") + supports_metrics = _is_real("--metrics") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") @@ -1577,6 +1821,7 @@ class LlamaCppBackend: "supports_cache_ram": supports_cache_ram, "supports_ctx_checkpoints": supports_ctx_checkpoints, "supports_no_cache_prompt": supports_no_cache_prompt, + "supports_metrics": supports_metrics, } cls._capability_cache[cache_key] = info return info @@ -1728,7 +1973,14 @@ class LlamaCppBackend: @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: - """Query free memory per GPU. + """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by + index; empty if no supported GPU is reachable. Thin wrapper over + ``_get_gpu_memory`` for callers that only need free VRAM.""" + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + + @staticmethod + def _get_gpu_memory() -> list[tuple[int, int, int]]: + """Query free AND total memory per GPU. Order: 1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects @@ -1739,15 +1991,15 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. - Returns list of (gpu_index, free_mib) sorted by index; empty if no - supported GPU is reachable. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no + supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. """ # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( [ "nvidia-smi", - "--query-gpu=index,memory.free", + "--query-gpu=index,memory.free,memory.total", "--format=csv,noheader,nounits", ], capture_output = True, @@ -1767,15 +2019,30 @@ class LlamaCppBackend: allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) except ValueError: pass - gpus: list[tuple[int, int]] = [] + gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): - parts = line.split(",") - if len(parts) == 2: - idx = int(parts[0].strip()) - free_mib = int(parts[1].strip()) - if allowed is not None and idx not in allowed: - continue - gpus.append((idx, free_mib)) + parts = [p.strip() for p in line.split(",")] + if len(parts) < 2: + continue + # Index and free required; skip a bad line rather than abandon + # the probe to the torch fallback. + try: + idx = int(parts[0]) + free_mib = int(parts[1]) + except ValueError: + continue + # Total parsed separately: a two-column line or a non-integer + # total ("N/A" on MIG/vGPU) keeps the GPU at total 0 (fit uses + # the free*frac fallback) instead of dropping it. + total_mib = 0 + if len(parts) >= 3 and parts[2]: + try: + total_mib = int(parts[2]) + except ValueError: + total_mib = 0 + if allowed is not None and idx not in allowed: + continue + gpus.append((idx, free_mib, total_mib)) # Match the docstring's sort-by-id guarantee (driver order isn't). gpus.sort(key = lambda g: g[0]) if gpus: @@ -1820,13 +2087,13 @@ class LlamaCppBackend: physical_ids = None gpus = [] for ordinal in range(torch.cuda.device_count()): - free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal) + free_bytes, total_bytes = torch.cuda.mem_get_info(ordinal) idx = ( physical_ids[ordinal] if physical_ids is not None and ordinal < len(physical_ids) else ordinal ) - gpus.append((idx, free_bytes // (1024 * 1024))) + gpus.append((idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024))) # Match the nvidia-smi path's docstring guarantee of sorted-by-id. return sorted(gpus, key = lambda g: g[0]) except Exception as e: @@ -1904,19 +2171,17 @@ class LlamaCppBackend: # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). _GPU_PIN_VRAM_FRACTION = 0.95 - # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the - # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes - # via graph_reserve -- it is roughly EQUAL on every device (not proportional - # to the tensor split) and independent of context. Measured ~2.3 GB - # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a - # conservative headroom above that. It is (a) subtracted from each GPU's free - # VRAM before computing --tensor-split, so the roomier GPU absorbs more - # weight and the smallest GPU keeps room for KV, and (b) reserved per device - # when capping context. The auto-fallback to layer split covers any - # underestimate. NOTE: scales with the model's vocab / batch size; tune if a - # large-vocab model OOMs at load. + # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF + # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived + # path) returns 0. _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120 + # Fixed per-device overhead on every GPU of a LAYER split (CUDA context + + # scratch), beyond the conserved slot-scaling buffer. ~0.9 GB/device measured + # (Qwen3.6-27B, b9625), independent of --parallel; reserved per extra GPU so a + # tight layer split can't advertise a context that OOMs at load. + _PIPELINE_PER_DEVICE_OVERHEAD_MIB = 1024 + # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) @@ -2084,6 +2349,8 @@ class LlamaCppBackend: model_size_bytes: int, gpus: list[tuple[int, int]], usable_fraction: Optional[float] = None, + total_by_idx: Optional[dict[int, int]] = None, + per_device_overhead_bytes: int = 0, ) -> tuple[Optional[list[int]], bool]: """Pick GPU(s) for a model from estimated VRAM and free memory. @@ -2091,6 +2358,11 @@ class LlamaCppBackend: ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides headroom for compute buffers, CUDA context, and other runtime overhead; callers lower it when MTP reserves VRAM for a draft model. + ``total_by_idx`` (index -> total MiB) makes the headroom an ABSOLUTE + ``(1 - fraction) * total`` per GPU instead of a fraction of free. + ``per_device_overhead_bytes`` is the fixed layer-split cost per GPU beyond + the first; a k-GPU pin must hold ``model + (k-1) * overhead`` or it can OOM + a device after -ngl -1 (no --fit fallback). Single-GPU adds none. Returns (gpu_indices, use_fit): - ([1], False) fits on 1 GPU at the headroom threshold @@ -2104,20 +2376,31 @@ class LlamaCppBackend: if usable_fraction is None: usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION - # Sort GPUs by free memory descending - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # Per-GPU usable budget: free - (1-frac)*total when total is known, else + # the legacy free*frac (also covers a total-0 two-column probe). + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - usable_fraction) * t) + return free_mib * usable_fraction + + # Rank by usable budget (free - reserve), not raw free: a more-used large + # card can have less usable room than a less-used small one. + ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True) # Try 1 GPU at the usable-VRAM threshold. - if ranked[0][1] * usable_fraction >= model_size_mib: + if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: return [ranked[0][0]], False - # Try N GPUs (accumulate free memory from most-free) - cumulative = 0 + # Try N GPUs (accumulate usable memory from most-free). Each GPU past the + # first adds a fixed per-device overhead the pool must hold. + overhead_mib = per_device_overhead_bytes / (1024 * 1024) + cumulative = 0.0 selected = [] for idx, free_mib in ranked: selected.append(idx) - cumulative += free_mib * usable_fraction - if cumulative >= model_size_mib: + cumulative += _usable(idx, free_mib) + if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib: return sorted(selected), False # Too large even for all GPUs; let --fit handle it @@ -2172,14 +2455,10 @@ class LlamaCppBackend: 5. Legacy -- fallback using embed // n_heads Server-flag knobs (mirror llama-server's CLI): - swa_full -- ``--swa-full``: force SWA layers to cache full - ``n_ctx`` (collapses path 3 to path 4 for them). - n_parallel -- ``--parallel`` slots: non-SWA layers stay constant - (cells split across slots), SWA layers scale linearly. - kv_unified -- ``--kv-unified`` (default on): no-op for memory math; - kept for API forward-compat. - ctx_checkpoints -- ``--ctx-checkpoints`` (PR #15293): N SWA snapshots - per slot, one sliding-window of state per SWA layer. + swa_full -- --swa-full: SWA layers cache full n_ctx (path 3->4). + n_parallel -- --parallel slots: non-SWA constant, SWA scale linearly. + kv_unified -- --kv-unified: memory no-op (API forward-compat). + ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot. Returns 0 if metadata is insufficient. """ @@ -2194,17 +2473,7 @@ class LlamaCppBackend: n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization - bpe = { - "f32": 4.0, - "f16": 2.0, - "bf16": 2.0, - "q8_0": 34 / 32, - "q5_1": 0.75, - "q5_0": 0.6875, - "q4_1": 0.625, - "q4_0": 0.5625, - "iq4_nl": 0.5625, - }.get(cache_type_kv or "f16", 2.0) + bpe = _kv_bytes_per_elem(cache_type_kv) slots = max(1, n_parallel) @@ -2232,15 +2501,12 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] 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 filled by the resolver at parse time; if absent, falls through - # to the legacy 1/4-global heuristic below. Per-layer-type --parallel N - # accounting (verified against llama-server): - # * non-SWA layers: total cells = n_ctx split across slots -> CONSTANT. - # * SWA layers: per-slot cells = 2*sliding_window (capped at n_ctx - # and per_slot_ctx) -> grows LINEARLY in slots. - # --swa-full forces full n_ctx for SWA layers; --ctx-checkpoints N adds - # N snapshots per SWA layer per slot. + # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern + # from the resolver; if absent, falls through to the legacy 1/4-global + # heuristic. --parallel N accounting (verified against llama-server): + # non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells + # = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots. + # --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -2249,8 +2515,7 @@ class LlamaCppBackend: ): swa = self._sliding_window per_slot_ctx = max(1, n_ctx // slots) - # --swa-full caches full context like non-SWA (per-slot cells = - # per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA + # --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA # caches 2*sliding_window per slot, clamped at per-slot ctx. swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) key_len_swa = self._kv_key_length_swa or key_len @@ -2303,6 +2568,149 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) + def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]: + """Lightweight backend with a drafter GGUF's metadata, to size its own KV + via _estimate_kv_cache_bytes. Cached per path; None if unreadable.""" + cache = getattr(self, "_draft_backend_cache", None) + if cache is not None and cache[0] == drafter_path: + return cache[1] + db: Optional[LlamaCppBackend] = None + try: + db = LlamaCppBackend.__new__(LlamaCppBackend) + for attr in ( + "_context_length", + "_n_layers", + "_n_kv_heads", + "_n_heads", + "_embedding_length", + "_kv_key_length", + "_kv_value_length", + "_kv_lora_rank", + "_sliding_window", + "_sliding_window_pattern", + "_ssm_inner_size", + "_full_attention_interval", + "_key_length_mla", + "_n_kv_heads_by_layer", + "_kv_key_length_swa", + "_kv_value_length_swa", + "_shared_kv_layers", + "_nextn_predict_layers", + ): + setattr(db, attr, None) + db._model_identifier = "mtp-draft" + db._read_gguf_metadata(drafter_path) + except Exception as e: # unreadable drafter -> caller falls back + logger.debug(f"Could not read drafter GGUF for MTP budget: {e}") + db = None + self._draft_backend_cache = (drafter_path, db) + return db + + def _mtp_draft_kv_bytes( + self, + n_ctx: int, + *, + drafter_path: Optional[str] = None, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + n_parallel: int = 1, + ) -> Optional[int]: + """Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are + independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes + at the heavier type. Embedded head (Qwen): nextn_predict_layers attention + layers from the main dims. None when dims are missing (flat fallback).""" + if n_ctx <= 0: + return None + bpe_k = _kv_bytes_per_elem(draft_cache_type_k) + bpe_v = _kv_bytes_per_elem(draft_cache_type_v) + if drafter_path: + db = self._draft_backend_for(drafter_path) + if db is None or not db._can_estimate_kv(): + return None + heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v + # The drafter is served under the same --parallel slot count as the + # main model, so price its KV per slot too: a sliding-window drafter + # (Gemma) grows KV with slots and would otherwise be under-reserved. + kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel) + return kv or None + nextn = self._nextn_predict_layers or 0 + n_kv = self._n_kv_heads or self._n_heads + k_len = self._kv_key_length + v_len = self._kv_value_length + if not (nextn and n_kv and k_len and v_len): + return None + # The embedded MTP head is one draft layer, so a quantized draft KV can't + # amortize its overhead and fits *less* context than f16 (llama.cpp#24102). + # Floor it at f16: a quantized override is priced as f16, f32 keeps its 4 + # bytes. The separate-drafter branch is multi-layer, so it keeps its type. + f16_bpe = _kv_bytes_per_elem("f16") + bpe_k = max(bpe_k, f16_bpe) + bpe_v = max(bpe_v, f16_bpe) + return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx) + + def _estimate_mtp_overhead_bytes( + self, + n_ctx: int, + *, + spec_draft_n_max: int = 0, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + drafter_path: Optional[str] = None, + draft_weights_bytes: int = 0, + n_parallel: int = 1, + ) -> Optional[int]: + """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate- + drafter weights. The verify buffer rides in the ctx-fit headroom (no tuned + constant). None when the draft KV can't be sized (caller keeps the flat + fallback). ``draft_weights_bytes`` is the drafter file size (0 for embedded).""" + draft_kv = self._mtp_draft_kv_bytes( + n_ctx, + drafter_path = drafter_path, + draft_cache_type_k = draft_cache_type_k, + draft_cache_type_v = draft_cache_type_v, + n_parallel = n_parallel, + ) + weights = max(0, draft_weights_bytes) + if draft_kv is None: + # KV unsized (exotic/remote drafter): still reserve known weights so a + # large drafter can't launch over budget (the small unsized KV rides in + # the cushion). Nothing known -> None, so the caller keeps the flat + # fallback. + return weights if weights > 0 else None + return draft_kv + weights + + _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it + _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + + def _estimate_compute_buffer_bytes( + self, + *, + n_ubatch: Optional[int] = None, + n_parallel: int = 1, + per_device_tensor: bool = False, + ) -> int: + """Per-device compute-graph buffer (bytes) from GGUF dims: a vocab-width + output buffer + activation scratch. Context-independent; scales with + ``--parallel`` (serving slots). Tensor mode materializes it on every device. + A slight upper bound over measured allocations; 0 when dims are missing.""" + n_vocab = self._vocab_size or 0 + n_embd = self._embedding_length or 0 + if n_vocab <= 0 or n_embd <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + par = max(1, int(n_parallel)) + out_buffer = n_vocab * ub * 4 # f32 output/logits buffer + act_scratch = 4 * n_embd * ub * 4 # a few resident hidden-width buffers + if per_device_tensor: + # Output + comm/staging materialized on every device, every slot. + compute = 2 * act_scratch + out_buffer * par + else: + # Each extra concurrent slot adds one output buffer (chat decode sizes + # ~one logit row per slot; would under-count embeddings/--logits-all, + # not run here). Matches measured {1:36,2:492,4:1388,8:3220} MiB. + compute = act_scratch + out_buffer * max(0, par - 1) + return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _fit_context_to_vram( self, requested_ctx: int, @@ -2317,13 +2725,15 @@ class LlamaCppBackend: ctx_checkpoints: int = 0, kv_on_gpu: bool = True, mtp_engaged: bool = False, + mtp_overhead_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, + total_mib: Optional[int] = None, ) -> int: """Return the largest context length that fits in GPU VRAM. - Uses 90% of available VRAM as the ctx-fit budget -- tighter than - ``_GPU_PIN_VRAM_FRACTION`` on purpose (over-promising context OOMs at - runtime). If the weights alone don't fit, returns ``requested_ctx``. + Budget caps occupancy at ``_CTX_FIT_VRAM_FRACTION`` of the card: an + absolute ``free - (1 - frac) * total`` when ``total_mib`` is given, else + ``free * frac``. Weights alone over budget returns ``requested_ctx``. ``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV cache lives in CPU RAM and the requested context is honored verbatim. @@ -2351,17 +2761,25 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - # MTP engaged: carve the drafter's reserve out of the fit budget. Callers - # can override outright (tensor-parallel mode passes a fatter margin), so - # only compute a default when none was supplied. + # byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback + # when dims can't size the draft KV); callers may override budget_frac. if budget_frac is None: - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) - budget_bytes = available_mib * 1024 * 1024 * budget_frac + flat_mtp = mtp_engaged and mtp_overhead_fn is None + budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0) + # Absolute reserve off total when known, else fraction-of-free; clamp >=0. + if total_mib is not None and total_mib > 0: + budget_mib = max(0.0, available_mib - (1.0 - budget_frac) * total_mib) + else: + budget_mib = available_mib * budget_frac + budget_bytes = budget_mib * 1024 * 1024 model_footprint = model_size_bytes + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: return requested_ctx # Weights alone exceed budget -- reducing ctx can't help; --fit handles it. @@ -2374,7 +2792,7 @@ class LlamaCppBackend: ) return requested_ctx - # Binary search for max context that fits + # Binary search for max context that fits (KV + MTP draft reserve at that ctx) remaining = budget_bytes - model_footprint effective_min = min(min_ctx, requested_ctx) lo, hi = effective_min, requested_ctx @@ -2382,7 +2800,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv <= remaining: + if kv + _mtp_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -2561,6 +2979,8 @@ class LlamaCppBackend: self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None + self._feed_forward_length = None + self._vocab_size = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None @@ -2582,6 +3002,8 @@ class LlamaCppBackend: WANTED = { "general.architecture", "tokenizer.chat_template", + # Vocab size = tokens array length (no vocab_size key in many GGUFs). + "tokenizer.ggml.tokens", # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner. "diffusion.canvas_length", # Source-repo hints for the SWA resolver's HF fallback. @@ -2645,6 +3067,7 @@ class LlamaCppBackend: f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", + f"{arch}.feed_forward_length": "feed_forward_length", f"{arch}.attention.key_length": "kv_key_length", f"{arch}.attention.value_length": "kv_value_length", f"{arch}.attention.sliding_window": "sliding_window", @@ -2678,6 +3101,9 @@ class LlamaCppBackend: elif vtype == 9: # ARRAY atype = struct.unpack(" tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -3464,22 +3899,41 @@ class LlamaCppBackend: Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``). + one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. - ``tensor_split`` is None (llama.cpp's even default, safe for every arch incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even - share fits the smallest GPU; otherwise it is weighted by - ``(free - buffer)`` so the roomier GPU absorbs more weight and the - smallest GPU keeps room for KV. + share fits the smallest GPU; otherwise it is weighted by usable budget + so the roomier GPU absorbs more weight and the smallest keeps room for KV. + ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes + the compute buffer. """ - # Drop GPUs that can't hold the per-device compute-graph buffer; they'd - # OOM in tensor mode. load_model already filters before calling, so this - # is defense-in-depth that also keeps the pure function self-contained - # (and unit-testable without a GPU). - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - usable_gpus = [g for g in gpus if g[1] >= reserve_mib] + + # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a + # two-column probe) the legacy free*frac. Mirrors _select_gpus and + # _gpu_usable so the 5% cushion is kept on every path, not dropped here. + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - _CTX_FIT_VRAM_FRACTION) * t) + return max(0.0, free_mib * _CTX_FIT_VRAM_FRACTION) + + # Drop GPUs whose usable budget can't hold the per-device compute-graph + # buffer; they'd OOM in tensor mode. Admitting on raw free would let a + # partly-used big card in with no budget left. Defense-in-depth (load_model + # gates too). Derived per-device reserve; flat fallback. + _reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = n_parallel, per_device_tensor = True + ) + reserve_mib = ( + _reserve_bytes // (1024 * 1024) + if _reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + usable_gpus = [g for g in gpus if _usable(g[0], g[1]) >= reserve_mib] gpu_indices = sorted(idx for idx, _ in usable_gpus) if len(gpu_indices) < 2: # Tensor parallelism is meaningless on <2 GPUs (the caller drops the @@ -3491,21 +3945,50 @@ class LlamaCppBackend: None, ) free_by_idx = {idx: free for idx, free in usable_gpus} - pool_mib = sum(free_by_idx.values()) - kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - if mtp_engaged: - # MTP keeps a draft model + its own KV cache on GPU. - kv_budget_b -= 2 * 1024**3 + usable_by_idx = {idx: _usable(idx, free_by_idx[idx]) for idx in gpu_indices} + pool_mib = sum(usable_by_idx.values()) + # MTP reserve: byte-accurate per-ctx inside _fit_ctx (mtp_overhead_fn) plus + # a flat cushion that the byte fn can't size -- 2 GiB when dims are wholly + # unavailable (no fn), or mtp_flat_reserve_bytes when the fn is weights-only + # because the draft KV couldn't be sized (_mtp_kv_unsized). Without this the + # binary search spends the unsized-KV cushion on main context and OOMs. + flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) + if mtp_engaged and mtp_overhead_fn is None: + flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + kv_budget_b = ( + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + ) + + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 def _fit_ctx(ctx: int) -> int: - # Largest context whose KV fits the pooled budget. Floors small, but - # never raises an explicit ctx above what was asked. + # Largest context whose KV (+ MTP draft reserve) fits the pooled + # budget. Floors small, but never raises an explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: # Weights + buffers exceed the pool -> floor; the load then # falls back to layer split. return ctx_floor + if mtp_overhead_fn is not None: + # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + def _consumer(c: int) -> int: + return self._estimate_kv_cache_bytes( + c, cache_type_kv, n_parallel = n_parallel + ) + _mtp_at(c) + + if _consumer(ctx) <= kv_budget_b: + return ctx + lo, hi, best = ctx_floor, ctx, ctx_floor + while lo <= hi: + mid = (lo + hi) // 2 + if _consumer(mid) <= kv_budget_b: + best = mid + lo = mid + 1 + else: + hi = mid - 1 + return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) if kv_at <= kv_budget_b: return ctx @@ -3520,16 +4003,19 @@ class LlamaCppBackend: max_available_ctx = _fit_ctx(max_ctx_target) effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) - min_free_mib = min(free_by_idx.values()) + min_usable_mib = min(usable_by_idx.values()) kv_bytes = ( self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) if (self._can_estimate_kv() and effective_ctx > 0) else 0 ) - even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024) + # The MTP reserve also has to fit the even split (mirror the pooled budget): + # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. + mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes + even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) tensor_split: Optional[list[int]] = None - if even_share_mib > (min_free_mib - reserve_mib): - adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices] + if even_share_mib > (min_usable_mib - reserve_mib): + adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -3856,27 +4342,58 @@ class LlamaCppBackend: ctx_override = parse_ctx_override(extra_args) requested_ctx = resolve_requested_ctx(extra_args, n_ctx) cache_override = parse_cache_override(extra_args) - cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) - # A user --split-mode in extras last-wins-overrides the - # toggle, so reconcile it back into tensor_parallel state. + # Budget the heavier of asymmetric --cache-type-k/-v extras (they + # win per axis at launch, appended last); resolve_cache_type_kv only + # returns the last-wins type, which under-reserves the heavier axis. + # The user's extras still set the real (possibly asymmetric) child + # cache, so this only affects the reserve, not the emitted command. + _extras_cache = _extra_args_main_cache_type_for_budget(extra_args) + cache_type_kv = _extras_cache if _extras_cache is not None else cache_type_kv + _cache_type_from_env = False + if cache_type_kv is None: + # Param/extras set nothing, so the child inherits + # LLAMA_ARG_CACHE_TYPE_K/_V. Adopt a heavier env type (f32) for + # the reserve only; the launch does NOT re-emit it (that would + # rewrite an asymmetric K=f32,V=f16 env into symmetric flags), + # so _cache_type_from_env keeps it out of the emitted flags. + cache_type_kv = _env_main_cache_type_for_budget() + _cache_type_from_env = cache_type_kv is not None + # A user --split-mode in extras last-wins-overrides the toggle, and + # an inherited tensor LLAMA_ARG_SPLIT_MODE flips it on (the child + # would run tensor unbudgeted otherwise). The duplicate-load matchers + # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) - tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel) + tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type - # that would re-impose it when appended last). The layer-split - # fallback re-runs with tensor_parallel False and keeps the type. - if ( - tensor_parallel - and cache_type_kv - and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES - ): + # that would re-impose it when appended last). Layer split does + # support it, so remember the dropped type and the original extras + # to restore (verbatim, incl. an asymmetric K/V) if we later fall + # back to layer split below. + _tensor_dropped_cache_type_kv: Optional[str] = None + _tensor_dropped_extra_args: Optional[list] = None + # Tensor mode rejects any quantized axis. cache_type_kv is the + # heavier-by-bytes budget type, which can mask a quantized axis (an + # f16 budget hides a paired q4_0), so also test each explicit + # --cache-type-k/-v extra, not just the budget type. + _ck_extra, _cv_extra = parse_cache_override_per_axis(extra_args) + _cache_non_tensor_safe = any( + c and c.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + for c in (cache_type_kv, _ck_extra, _cv_extra) + ) + if tensor_parallel and _cache_non_tensor_safe: logger.info( "Tensor parallelism requires a non-quantized KV cache; " "ignoring cache type %s for the tensor attempt.", cache_type_kv, ) + _tensor_dropped_cache_type_kv = cache_type_kv cache_type_kv = None if extra_args: + # Keep the originals so a layer downgrade restores the real + # (possibly asymmetric) --cache-type-k/-v the layer path + # supports, not just the scalar heavier type. + _tensor_dropped_extra_args = list(extra_args) extra_args = strip_shadowing_flags( extra_args, strip_context = False, @@ -3885,10 +4402,24 @@ class LlamaCppBackend: strip_template = False, strip_split_mode = False, ) + # The launch keeps an inherited tensor-safe env cache type (the + # env cleanup only pops quantized ones), so re-adopt a heavier + # env type (f32) for the budget here too -- mirrors the initial + # adoption, which was skipped because the param/extras set the + # (now-dropped) quantized type. Else the child allocates f32 KV + # against an f16 budget. + _env_tensor_cache = _env_main_cache_type_for_budget() + if _env_tensor_cache is not None: + cache_type_kv = _env_tensor_cache + _cache_type_from_env = True if ctx_override is not None and ctx_override > 0: logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") if cache_override is not None: - logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + _ck, _cv = parse_cache_override_per_axis(extra_args) + logger.info( + f"User --cache-type-k/-v (k={_ck}, v={_cv}) honored; " + "KV estimate budgets the heavier axis" + ) if split_mode_override is not None: logger.info( f"User --split-mode {split_mode_override} honored; " @@ -3920,7 +4451,27 @@ class LlamaCppBackend: self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0 ) model_size = gguf_size + mmproj_size - gpus = self._get_gpu_free_memory() + # 2-tuple gpus for existing logic + a total map for the absolute + # per-GPU headroom (correct when the GPU is already partly used). + _gpu_mem = self._get_gpu_memory() + gpus = [(idx, free) for idx, free, _t in _gpu_mem] + total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + + def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): + # Per-GPU usable budget for ranking: free - (1-frac)*total. + # Callers pass the ACTIVE fraction so the ranking matches the + # budget the fit then tests (else mixed totals mis-order). + idx, free = g + t = total_by_idx.get(idx, 0) + if t > 0: + return free - (1.0 - frac) * t + return free * frac + + def _pool_budget_mib(subset, frac): + # Sum each GPU's own usable budget. Pooling free and total + # separately would let an unknown-total GPU (MIG/vGPU/N/A) + # add full free with no cushion among known-total GPUs. + return sum(max(0.0, _gpu_usable(g, frac)) for g in subset) # Resolve effective context: 0 means let llama-server use # the model's native length. Only expand to a known native @@ -3936,12 +4487,10 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx - # Will MTP engage on this load? If so, auto-fit reserves - # extra VRAM for the draft model. Mirrors - # _build_speculative_flags' resolver: forced mtp / mtp+ngram - # always engage; auto only on an MTP model >= 3B; ngram / - # ngram-simple / off never engage MTP. A separate drafter - # (Gemma) counts as an MTP model just like a baked-in head. + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. + # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always + # engage; auto only on an MTP model >= 3B; ngram/off never. A + # separate drafter (Gemma) counts as an MTP model. _mtp_canonical = _canonicalize_spec_mode(speculative_type) _mtp_effective = _mtp_canonical or "auto" _mtp_size_for_fit = _extract_model_size_b(model_identifier) @@ -3952,49 +4501,229 @@ class LlamaCppBackend: and _mtp_size_for_fit < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) ) + # LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras + # nor Studio emit a spec flag (mode "off", no user --spec-type), + # since _build_speculative_flags emits one for every other mode. + # Consult the env for the reserve only then, else a stale MTP env + # would over-reserve. + _spec_env: Mapping[str, str] = ( + os.environ + if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off") + else {} + ) + # Extras can run MTP even when Studio suppresses its own emission. + _user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env) + # A non-MTP model-based draft mode (draft-simple/draft-eagle3) in + # extras also loads a separate draft model that needs reserving; + # engage only when extras actually name a drafter for it. + _user_draft_via_extras = _extra_args_requests_separate_draft( + extra_args, env = _spec_env + ) and bool(_extra_args_mtp_draft_path(extra_args)) + # Mirror _build_speculative_flags: reserve only for MTP the launch + # resolver will actually emit (needs a head/drafter and a binary + # that supports --spec-type mtp). + _mtp_model_for_fit = bool( + self._nextn_predict_layers + or _is_mtp_model_name(model_identifier, model_path) + or bool(mtp_draft_path) + ) + _mtp_binary_ok = True + if not _user_mtp_via_extras: + try: + _mtp_binary_ok = bool( + (self.probe_server_capabilities(binary) or {}).get("mtp_token") + ) + except Exception: + _mtp_binary_ok = False _mtp_will_engage = bool( - not _extra_args_set_spec_type(extra_args) - and ( - _mtp_effective in ("mtp", "mtp+ngram") - or ( - _mtp_effective == "auto" - and ( - bool(self._nextn_predict_layers) - or _is_mtp_model_name(model_identifier, model_path) - or bool(mtp_draft_path) - ) - and not _mtp_sub_3b_for_fit + _user_mtp_via_extras + or _user_draft_via_extras + or ( + not _extra_args_set_spec_type(extra_args) + and _mtp_binary_ok + and _mtp_model_for_fit + and ( + _mtp_effective in ("mtp", "mtp+ngram") + or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit) ) ) ) - # Auto-cap context to fit GPU VRAM and select GPUs. Two - # policies by whether the user set n_ctx: - # Explicit n_ctx: honor it. Try the full context with - # _select_gpus (as many GPUs as needed); cap only if it - # doesn't fit on any combination. - # Auto n_ctx=0 (native): prefer fewer GPUs with reduced - # context, since multi-GPU is slower. + # Effective draft depth: extras win (last-wins at launch), else + # the field, else the platform default (2 GPU / 3 CPU). + _extra_n_max = _extra_args_spec_draft_n_max(extra_args) + _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max + if _mtp_eff_n_max is None: + _mtp_eff_n_max = 2 if gpus else 3 + # Separate-drafter weights live on GPU (an embedded head is + # already in model_size). Size the drafter the launch loads, by + # precedence: extras --model-draft (last-wins), else Studio's + # emitted mtp_draft_path, else the env drafter. Sizing the wrong + # one would under-reserve and OOM. + _cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {}) + _studio_draft_for_budget = ( + mtp_draft_path + if ( + _mtp_will_engage + and mtp_draft_path + and not _extra_args_set_spec_type(extra_args) + ) + else None + ) + _env_draft_for_budget = _extra_args_mtp_draft_path([], env = os.environ) + _mtp_draft_for_budget = ( + _cli_draft_for_budget or _studio_draft_for_budget or _env_draft_for_budget + ) + # Drafter offloaded to CPU keeps its weights+KV off the GPU, so + # drop it from the budget (an embedded head stays in the model). + # Consult the env too: the child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT. + _draft_on_cpu = _extra_args_draft_offloaded_to_cpu(extra_args, env = os.environ) + if _draft_on_cpu: + _mtp_draft_for_budget = None + _mtp_draft_weights = 0 + if _mtp_draft_for_budget: + try: + _mtp_draft_weights = self._get_gguf_size_bytes(_mtp_draft_for_budget) + except Exception: + _mtp_draft_weights = 0 + # Draft K/V types (f16 by default; independent extras overrides). + _mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types(extra_args) + + # Byte-accurate reserve when dims allow, else None -> flat fallback. + mtp_overhead_fn: Optional[Callable[[int], int]] = None + # True when the byte reserve is the drafter weights ONLY because + # its KV couldn't be sized; the flat fraction must then stay on + # as the cushion for that unsized draft KV (it is not covered by + # the weights-only mtp_overhead_fn). + _mtp_kv_unsized = False + if _mtp_will_engage: + _probe_ctx = self._context_length or ( + effective_ctx if effective_ctx > 0 else 4096 + ) + _draft_kv_probe = self._mtp_draft_kv_bytes( + _probe_ctx, + drafter_path = _mtp_draft_for_budget, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + n_parallel = n_parallel, + ) + if ( + self._estimate_mtp_overhead_bytes( + _probe_ctx, + spec_draft_n_max = _mtp_eff_n_max, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + drafter_path = _mtp_draft_for_budget, + draft_weights_bytes = _mtp_draft_weights, + n_parallel = n_parallel, + ) + is not None + ): + # Reserve is weights-only when the draft KV is unsizable. + _mtp_kv_unsized = _draft_kv_probe is None + + # Closure binding this load's draft params; ctx varies. + def mtp_overhead_fn( + ctx: int, + _n: int = _mtp_eff_n_max, + _ck: Optional[str] = _mtp_draft_ck, + _cv: Optional[str] = _mtp_draft_cv, + _dp: Optional[str] = _mtp_draft_for_budget, + _w: int = _mtp_draft_weights, + _np: int = n_parallel, + ) -> int: + v = self._estimate_mtp_overhead_bytes( + ctx, + spec_draft_n_max = _n, + draft_cache_type_k = _ck, + draft_cache_type_v = _cv, + drafter_path = _dp, + draft_weights_bytes = _w, + n_parallel = _np, + ) + return v if v is not None else 0 + + def _mtp_bytes(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + + # Effective micro-batch (a user --ubatch override scales the + # compute buffer); None -> the 512 default in the estimate. + _effective_ubatch = _extra_args_n_ubatch(extra_args) + + # Layer-split compute buffer (one lump; tensor mode reserves it + # per device in _plan_tensor_parallel). Context-independent, so + # fold it into the model footprint for the branches below. Falls + # back to the flat reserve when dims are missing (returns 0), a + # safe upper bound since the tensor buffer >= the layer one. + _compute_buffer_pipeline = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = False, + ) + if _compute_buffer_pipeline <= 0: + _compute_buffer_pipeline = ( + self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + ) + model_size_fit = model_size + _compute_buffer_pipeline + + # Layer split adds a fixed per-device overhead on every GPU. The + # folded buffer covers one device; reserve the extra devices' + # share so a k-GPU split can't pin a context that OOMs a device + # (k=1 adds nothing). + _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + + # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: + # honor it, cap only if it fits no combination. Auto (native): + # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True # Per-GPU weight proportions for tensor mode (None = even). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 - # MTP draft model lives outside the main estimates; carve - # its reserve out of every fit budget and pin threshold so - # a load can't pin into the drafter's headroom. - _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0 - _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve + # Flat MTP reserve fraction: used only as the fallback when the + # byte-accurate mtp_overhead_fn can't size the draft KV (dims + # unavailable, or _mtp_kv_unsized = weights-only). A separate + # drafter on CPU uses no GPU (no reserve); an embedded head is on + # GPU regardless of draft-offload flags (keep its reserve). + _flat_mtp_engages = _mtp_will_engage and ( + mtp_overhead_fn is None or _mtp_kv_unsized + ) + _draft_cpu_no_embedded = _draft_on_cpu and not self._nextn_predict_layers + # MTP reserves GPU VRAM unless its only drafter is a separate + # CPU-offloaded one (an embedded head stays on GPU). The tensor + # path reserves like the layer path; gate both on this. + _mtp_reserves_gpu = _mtp_will_engage and not _draft_cpu_no_embedded + _flat_mtp_reserve = ( + _MTP_VRAM_RESERVE_FRAC + if (_flat_mtp_engages and not _draft_cpu_no_embedded) + else 0.0 + ) + _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve - # Tensor mode allocates a compute-graph buffer on every - # participating GPU, so a GPU with less free VRAM than that - # reserve can't host it and would OOM at load. Drop those - # from the tensor-parallel set up front (gpu_indices below - # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded - # from llama-server entirely, not just given zero weight). + # 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). tp_gpus = gpus if tensor_parallel: - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - tp_gpus = [g for g in gpus if g[1] >= reserve_mib] + # Deterministic per-device compute buffer (replicated on + # every device in tensor mode); flat fallback when dims + # are unavailable. _plan_tensor_parallel uses the same. + _tp_reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = True, + ) + reserve_mib = ( + _tp_reserve_bytes // (1024 * 1024) + if _tp_reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + # Admit by usable budget (free - (1-frac)*total), not raw + # free: a partly-used big card can clear the reserve on raw + # free yet have no budget left. + tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] if tensor_parallel and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single @@ -4010,21 +4739,81 @@ class LlamaCppBackend: len(gpus), ) tensor_parallel = False - # A user --split-mode tensor in extras is appended after - # Studio's flags, so it would still reach llama-server and - # fail here; strip it so the downgrade actually applies. - extra_args = strip_split_mode_only(extra_args) + # 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 + ) if tensor_parallel and tp_gpus: - # Tensor-parallel allocation: use all usable GPUs, weight - # the split by (free - buffer), and cap context to the - # pooled VRAM after weights + per-device compute-graph - # buffers. See _plan_tensor_parallel for the policy. + # Pooled usable budget (after each device's compute buffer) + # must hold the non-shrinkable footprint: weights + the MTP + # reserve. The planner can shrink ctx/KV, not these. + _tp_weight_budget_mib = ( + sum(_gpu_usable(g) for g in tp_gpus) - len(tp_gpus) * reserve_mib + ) + _tp_flat_mtp = 2 * 1024**3 # flat reserve when dims unavailable + if not _mtp_reserves_gpu: + # No MTP, or its only drafter is CPU-offloaded (no GPU). + _tp_mtp_floor = 0 + elif mtp_overhead_fn is not None and not _mtp_kv_unsized: + _tp_mtp_floor = _mtp_bytes( + min(2048, effective_ctx) if effective_ctx > 0 else 2048 + ) + else: + # Dims unavailable / weights-only: tensor mode has no + # --fit valve, so keep the flat reserve as the unsized-KV + # cushion, never below the known byte reserve. + _tp_mtp_floor = max( + _tp_flat_mtp, + _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), + ) + _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + if _tp_weight_budget_mib <= _tp_required_mib: + logger.info( + "Tensor parallelism requested but the pooled VRAM " + "budget cannot hold the weights, MTP reserve, and " + "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 + ) + + if tensor_parallel and tp_gpus: + # Tensor-parallel allocation; see _plan_tensor_parallel. target_ctx = ( effective_ctx if explicit_ctx else (self._context_length or effective_ctx) ) + # When the draft KV couldn't be sized (weights-only reserve), + # the planner's mtp_overhead_fn is non-None but covers only + # weights, so pass the flat cushion for the unsized KV (else + # the binary search spends it on context). + _tp_unsized_mtp_reserve = ( + 2 * 1024**3 if (_mtp_reserves_gpu and _mtp_kv_unsized) else 0 + ) ( effective_ctx, max_available_ctx, @@ -4036,10 +4825,14 @@ class LlamaCppBackend: target_ctx, cache_type_kv = cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + mtp_flat_reserve_bytes = _tp_unsized_mtp_reserve, # Report the UI ceiling from native ctx, not the # explicit small request. max_target_ctx = self._context_length or target_ctx, + total_by_idx = total_by_idx, + n_ubatch = _effective_ubatch, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -4048,24 +4841,38 @@ class LlamaCppBackend: # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: - ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + ranked_for_cap = sorted( + gpus, + key = lambda g: _gpu_usable( + g, _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve + ), + reverse = True, + ) best_cap = 0 + _cap_fraction = _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve for n_gpus in range(1, len(ranked_for_cap) + 1): subset = ranked_for_cap[:n_gpus] - pool_mib = sum(free for _, free in subset) + # Per-GPU-consistent pool budget (fixes mixed + # known/unknown totals); pass it as an absolute + # budget so the fit and the check below agree. + pool_budget = _pool_budget_mib(subset, _cap_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( native_ctx_for_cap, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * (_CTX_FIT_VRAM_FRACTION - _mtp_reserve): + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -4079,34 +4886,49 @@ class LlamaCppBackend: # Honor the requested context verbatim. If it fits, # pin GPUs and skip --fit; else ship -c --fit # on and let llama-server flex -ngl (CPU offload). - requested_total = model_size + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel + requested_total = ( + model_size_fit + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) + + _mtp_bytes(effective_ctx) ) gpu_indices, use_fit = self._select_gpus( - requested_total, gpus, usable_fraction = _pin_fraction + requested_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) # No silent shrink: effective_ctx stays == requested_ctx. else: # Auto context: prefer fewer GPUs, cap to fit. Same - # headroom threshold as _select_gpus (#5106). - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # headroom threshold as _select_gpus (#5106). Rank by the + # active pin fraction so the order matches the fit budget. pin_fraction = _pin_fraction + ranked = sorted( + gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True + ) for n_gpus in range(1, len(ranked) + 1): subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) + pool_budget = _pool_budget_mib(subset, pin_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( effective_ctx, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False @@ -4119,14 +4941,17 @@ class LlamaCppBackend: if effective_ctx > 0: for n_gpus in range(1, len(ranked) + 1): subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) kv = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel, ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = ( + _subset_model_size(n_gpus) + + kv + + _mtp_bytes(effective_ctx) + ) / (1024 * 1024) + if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) use_fit = False break @@ -4138,14 +4963,36 @@ class LlamaCppBackend: "Falling back to file-size-only GPU selection", model_size_gb = round(model_size / (1024**3), 2), ) + # Add the byte-accurate MTP reserve here too when it is + # available; otherwise _pin_fraction carries the flat + # fallback (the two are mutually exclusive by design). + _fs_total = model_size_fit + _mtp_bytes( + self._context_length or effective_ctx or 4096 + ) gpu_indices, use_fit = self._select_gpus( - model_size, gpus, usable_fraction = _pin_fraction + _fs_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 # so the slider isn't on an unusable native ctx. effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096 + # MTP reserve at the final context, for the logs below. + _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 + if _mtp_will_engage: + _mtp_note = ( + f"MTP reserve: {_mtp_reserve_bytes / (1024**3):.2f} GB " + f"(draft KV @ {effective_ctx} + verify n_max={_mtp_eff_n_max}" + + (", flat-frac fallback" if mtp_overhead_fn is None else "") + + "), " + ) + else: + _mtp_note = "" + if effective_ctx < original_ctx: kv_est = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel @@ -4153,7 +5000,9 @@ class LlamaCppBackend: logger.info( f"Context auto-reduced: {original_ctx} -> {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_est / (1024**3):.1f} GB)" + f"est. KV cache: {kv_est / (1024**3):.1f} GB, " + f"{_mtp_note}".rstrip(", ") + + ")" ) kv_cache_bytes = self._estimate_kv_cache_bytes( @@ -4166,6 +5015,7 @@ class LlamaCppBackend: f"GGUF size: {gguf_size / (1024**3):.1f} GB, " f"{mmproj_note}" f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " + f"{_mtp_note}" f"context: {effective_ctx}, " f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" ) @@ -4214,6 +5064,10 @@ class LlamaCppBackend: fully_gpu_offloaded = True server_caps = self.probe_server_capabilities(binary) + # Expose Prometheus /metrics for the engine-stats logger, only + # when the binary advertises it (older/custom binaries may not). + if server_caps.get("supports_metrics"): + cmd.append("--metrics") cmd.extend( self._ctx_integrity_flags( n_parallel, @@ -4260,7 +5114,11 @@ class LlamaCppBackend: "iq4_nl", "f32", } - if cache_type_kv and cache_type_kv in _valid_cache_types: + if ( + cache_type_kv + and cache_type_kv in _valid_cache_types + and not _cache_type_from_env + ): cmd.extend( [ "--cache-type-k", @@ -4272,6 +5130,8 @@ class LlamaCppBackend: self._cache_type_kv = cache_type_kv logger.info(f"KV cache type: {cache_type_kv}") else: + # An env-only type is left inherited (untouched) so an + # asymmetric K/V env reaches the child as set. self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor @@ -4429,6 +5289,31 @@ class LlamaCppBackend: if "--threads" not in cmd: env.pop("LLAMA_ARG_THREADS", None) + # Reconcile the inherited LLAMA_ARG_* env with Studio's final + # decision: stripping CLI extras on a tensor->layer downgrade + # can't remove env vars, so the child could run a mode/KV Studio + # didn't budget. + if not tensor_parallel: + # Layer split: clear a non-layer inherited split mode (and any + # paired tensor-split) so the child can't override the layer plan. + _inherited_sm = (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower() + if _inherited_sm and _inherited_sm != "layer": + env.pop("LLAMA_ARG_SPLIT_MODE", None) + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + else: + # Studio owns the tensor split: it emits --tensor-split when it + # picks an uneven one (CLI wins) and nothing when an even split + # is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even + # case can't be overridden by a stale env (the layer branch above + # clears it too). + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + # Tensor split aborts on a quantized KV; clear an inherited + # quantized cache type so the child uses the tensor-safe default. + for _ct_var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + _ct_raw = (env.get(_ct_var) or "").strip().lower() + if _ct_raw and _ct_raw not in self._TENSOR_PARALLEL_KV_TYPES: + env.pop(_ct_var, None) + # Windows + full offload: PASSIVE OMP + 2 threads stop # spin-wait burning CPU. CPU/partial offload keeps default # OMP parallelism. #5692. @@ -4721,6 +5606,18 @@ class LlamaCppBackend: logger.info( f"llama-server ready on port {self._port} for model '{model_identifier}'" ) + # Poll llama-server /metrics -> vLLM-style engine_stats logs + # (only when the binary exposes /metrics). + if server_caps.get("supports_metrics"): + try: + from core.inference.llama_stats import maybe_start_stats_logger + if self._stats_logger is not None: + self._stats_logger.stop() + self._stats_logger = maybe_start_stats_logger(self.base_url, logger) + except Exception as e: + logger.debug(f"engine-stats logger not started: {e}") + else: + self._stats_logger = None # Probe outside _lock (interruptible by /unload); init inside. self._is_audio = False @@ -4879,6 +5776,11 @@ class LlamaCppBackend: "run `unsloth studio update`. Loading without " "speculative decoding." ) + # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins + # over env) so the child matches the binary-capability gate and + # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. + flags.append("--spec-default") + self._speculative_type = "default" self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() @@ -5065,10 +5967,12 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False - # Reconcile a user --split-mode in extras (load_model does the same), so - # an extras-driven tensor load isn't seen as a mismatch that needlessly - # kills/reloads a healthy server. - if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + # Reconcile a user --split-mode in extras AND an inherited tensor + # LLAMA_ARG_SPLIT_MODE env, but only against a server that actually + # launched tensor: if load_model downgraded to layer split it scrubbed + # the child env, so the env must not force an endless reload of a healthy + # server. An identical request would downgrade the same way. + if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): return False # Compare on the canonical requested mode. With --spec-type in @@ -5221,6 +6125,11 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"Error killing llama-server process: {e}") finally: + # getattr: teardown must tolerate a partially-built backend (failed + # __init__ or a __new__-built instance), as with _llama_log_fh below. + if getattr(self, "_stats_logger", None) is not None: + self._stats_logger.stop() + self._stats_logger = None self._process = None # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). @@ -5228,8 +6137,9 @@ class LlamaCppBackend: # Drives _wait_for_vram_settle in the next load_model; set in finally # so both in-process and frontend Apply paths record the kill. self._last_kill_monotonic = time.monotonic() - if self._stdout_thread is not None: - self._stdout_thread.join(timeout = 2) + stdout_thread = getattr(self, "_stdout_thread", None) + if stdout_thread is not None: + stdout_thread.join(timeout = 2) self._stdout_thread = None fh = getattr(self, "_llama_log_fh", None) if fh is not None: diff --git a/studio/backend/core/inference/llama_http.py b/studio/backend/core/inference/llama_http.py new file mode 100644 index 0000000000..b554949c3e --- /dev/null +++ b/studio/backend/core/inference/llama_http.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared pooled httpx.AsyncClient for NON-streaming calls to the local llama-server. + +Streaming generation must NOT use this. It relies on ``Connection: close`` and +``max_keepalive_connections=0`` so a client disconnect tears down the upstream +socket and stops GPU decode (PR #5749). This pooled client is only for short +request/response proxy calls (non-streaming completions, embeddings) where +reusing a connection removes per-request setup cost. Per-request ``timeout`` is +still passed at each call site. +""" + +from __future__ import annotations + +import asyncio +import weakref + +import httpx + +_LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32) + + +def _new_client() -> httpx.AsyncClient: + try: + return httpx.AsyncClient(limits = _LIMITS) + except Exception: + # Mirror external_provider: an unsupported env proxy scheme can raise. + return httpx.AsyncClient(limits = _LIMITS, trust_env = False) + + +# One client per running event loop: an httpx client binds its transport to the +# loop it first runs on, so a single global instance breaks across a lifespan +# restart or a second test loop. Weak keys let a finished loop drop its client. +_clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]" = ( + weakref.WeakKeyDictionary() +) + + +def nonstreaming_client() -> httpx.AsyncClient: + loop = asyncio.get_running_loop() + client = _clients.get(loop) + if client is None or client.is_closed: + client = _new_client() + _clients[loop] = client + return client + + +async def aclose() -> None: + clients = list(_clients.values()) + _clients.clear() + for client in clients: + try: + await client.aclose() + except Exception: + pass diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 69a86fa3ba..b42be5ee0d 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -13,7 +13,8 @@ Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md from __future__ import annotations -from typing import Iterable, Optional +import os +from typing import Iterable, Mapping, Optional # Each group = every alias (short + long) of one hard-denied flag. # Extend the matching group when llama.cpp adds a new alias. @@ -124,7 +125,9 @@ def is_managed_flag(flag: str) -> bool: # from inherited extras so they can't last-wins-override an Apply that # re-sets the same field. _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) -_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}) +_CACHE_TYPE_K_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k"}) +_CACHE_TYPE_V_FLAGS: frozenset[str] = frozenset({"-ctv", "--cache-type-v"}) +_CACHE_FLAGS: frozenset[str] = _CACHE_TYPE_K_FLAGS | _CACHE_TYPE_V_FLAGS _SPEC_FLAGS: frozenset[str] = frozenset( { "--spec-default", @@ -133,13 +136,22 @@ _SPEC_FLAGS: frozenset[str] = frozenset( "--spec-ngram-size", "--draft-min", "--draft-max", - # MTP path (llama.cpp #22673). --model-draft and aliases are - # Studio-managed since the separate-drafter support (Gemma 4): an - # inherited copy must not last-wins-override the auto-detected - # drafter. Explicit extras for the current load are never stripped. + # MTP path (llama.cpp #22673). The drafter selectors (local --model-draft + # and HF --spec-draft-hf aliases) are Studio-managed since the separate- + # drafter support (Gemma 4): an inherited copy must not last-wins-override + # the auto-detected drafter. Explicit extras for the current load are never + # stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld, + # --spec-draft-device) are deliberately NOT stripped: the VRAM budget reads + # them via the same parsers the child honors, so they stay consistent on + # inherit, and stripping them would silently move a CPU-offloaded drafter + # back onto the GPU. "--model-draft", "-md", "--spec-draft-model", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", "--spec-draft-n-max", "--spec-draft-n-min", "--spec-draft-p-min", @@ -274,6 +286,20 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return _last_flag_value(args, _CACHE_FLAGS) +def parse_cache_override_per_axis( + args: Optional[Iterable[str]], +) -> tuple[Optional[str], Optional[str]]: + """Last-wins --cache-type-k / --cache-type-v values kept apart, as (k, v). + + parse_cache_override collapses both axes to one last-wins value; this keeps + them separate so an asymmetric K/V can be budgeted by its heavier axis. + """ + return ( + _last_flag_value(args, _CACHE_TYPE_K_FLAGS), + _last_flag_value(args, _CACHE_TYPE_V_FLAGS), + ) + + def resolve_cache_type_kv( args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str] ) -> Optional[str]: @@ -309,6 +335,60 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral return override.strip().lower() == "tensor" +def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool: + """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio + emits --split-mode only on its tensor branch, so a tensor env on the layer + path would run the child tensor-parallel unbudgeted; this flips the budget + to tensor. Only tensor is heavier, so other modes are ignored.""" + raw = (os.environ if env is None else env).get("LLAMA_ARG_SPLIT_MODE") + return bool(raw) and raw.strip().lower() == "tensor" + + +def _effective_tensor_parallel( + extra_args: Optional[Iterable[str]], + tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Tensor-parallel decision including the inherited LLAMA_ARG_SPLIT_MODE env. + + resolve_tensor_parallel (extras + toggle), flipped on when extras set no split + mode but the child inherits a tensor split env. Shared by load_model (which + budgets and launches it) and the tensor-fallback wrapper (so an env-only + tensor crash still retries layer split).""" + resolved = resolve_tensor_parallel(extra_args, tensor_parallel) + if ( + not resolved + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + return True + return resolved + + +def _tensor_parallel_matches_loaded( + extra_args: Optional[Iterable[str]], + requested_tensor_parallel: bool, + loaded_tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Whether a duplicate load request matches a loaded server's tensor state. + + Env-only tensor mode is a launch hint load_model may downgrade to layer split + (capacity/buffer), scrubbing the child env. So only let an inherited tensor env + raise a match against a server that *actually* launched tensor; on a downgraded + (layer) server the env is ignored, and an identical request would downgrade the + same way -- avoiding an endless reload of a healthy server.""" + requested = resolve_tensor_parallel(extra_args, requested_tensor_parallel) + if ( + loaded_tensor_parallel + and not requested + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + requested = True + return requested == loaded_tensor_parallel + + _MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) _MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) diff --git a/studio/backend/core/inference/llama_stats.py b/studio/backend/core/inference/llama_stats.py new file mode 100644 index 0000000000..6047aedbc0 --- /dev/null +++ b/studio/backend/core/inference/llama_stats.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Translate llama-server's Prometheus /metrics into a periodic, vLLM-style +engine-stats log line (generation/prompt throughput, requests in flight). + +llama-server already computes these (it needs `--metrics`); this lifts them +into Studio's structured log so the terminal shows serving health, not just +per-request access lines. Emitted only while there is activity. +""" + +import os +import re +import threading +import time +import urllib.request + +# Prometheus body lines: "llamacpp:[{labels}] " (skip "#" HELP/TYPE). +_METRIC_RE = re.compile(r"^llamacpp:(\w+)(?:\{[^}]*\})?\s+([0-9.eE+-]+)", re.MULTILINE) +_OFF = {"0", "false", "no", "off"} + + +class LlamaServerStatsLogger: + """Daemon poller that logs vLLM-style engine stats from llama-server. + + Keeps retrying through transient scrape failures; the backend stops it via + stop() on unload/reload, so a brief /metrics stall does not silence stats. + """ + + def __init__( + self, + base_url, + logger, + interval_s = 10.0, + ): + self._url = f"{base_url.rstrip('/')}/metrics" + self._log = logger + self._interval = max(1.0, float(interval_s)) + self._stop = threading.Event() + self._thread = None + + def start(self): + if self._thread is None: + self._thread = threading.Thread(target = self._run, name = "llama-stats", daemon = True) + self._thread.start() + + def stop(self): + self._stop.set() + + def _scrape(self): + try: + with urllib.request.urlopen(self._url, timeout = 3) as r: + if r.status != 200: + return None + body = r.read().decode("utf-8", "replace") + except Exception: + return None + out = {} + for k, v in _METRIC_RE.findall(body): + try: # a malformed value must not kill the daemon thread + out[k] = float(v) + except ValueError: + continue + return out + + def _run(self): + misses = 0 + prev = None # (monotonic_t, tokens_predicted_total, prompt_tokens_total) + while not self._stop.wait(self._interval): + m = self._scrape() + if not m: + misses += 1 + if misses == 3: # transient stall (load/GC); keep polling. + self._log.debug("engine_stats: /metrics scrape failing, still retrying") + continue # real shutdown is driven by stop() from _kill_process + misses = 0 + # Generation tokens come from tokens_predicted_total (counter) and + # predicted_tokens_seconds (gauge); n_decode_total counts + # llama_decode() calls, not tokens, so it must not feed tok/s. + now = time.monotonic() + predicted = m.get("tokens_predicted_total", 0.0) + prompt = m.get("prompt_tokens_total", 0.0) + gen_delta = prompt_delta = 0.0 + if prev is not None and now > prev[0]: + dt = now - prev[0] + gen_delta = max(0.0, (predicted - prev[1]) / dt) + prompt_delta = max(0.0, (prompt - prev[2]) / dt) + prev = (now, predicted, prompt) + # Prefer llama.cpp's own throughput gauges; fall back to the counter + # delta for binaries that expose only the counters. + gen_tps = m.get("predicted_tokens_seconds") or gen_delta + prompt_tps = m.get("prompt_tokens_seconds") or prompt_delta + running, waiting = ( + int(m.get("requests_processing", 0)), + int(m.get("requests_deferred", 0)), + ) + # Gate on real activity this tick so a stale gauge never logs at idle. + if running or waiting or gen_delta or prompt_delta: + self._log.info( + "engine_stats", + gen_tok_s = round(float(gen_tps), 1), + prompt_tok_s = round(float(prompt_tps), 1), + running = running, + waiting = waiting, + ) + + +def maybe_start_stats_logger(base_url, logger): + """Start a stats logger unless UNSLOTH_STUDIO_ENGINE_STATS disables it.""" + if (os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "").strip().lower() in _OFF: + return None + try: + interval = float(os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS_INTERVAL_S", "10")) + except ValueError: + interval = 10.0 + sl = LlamaServerStatsLogger(base_url, logger, interval) + sl.start() + return sl diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4fa48a2d88..6b6deb1265 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -30,15 +30,15 @@ from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union from utils.hardware import prepare_gpu_selection +# Re-exported from the shared helper so GGUF, training, and inference share one +# type; kept importable here for backwards compatibility. +from utils.hf_xet_fallback import DownloadStallError + logger = get_logger(__name__) _CTX = mp.get_context("spawn") -class DownloadStallError(RuntimeError): - """Raised when the worker reports no download progress for too long.""" - - # Dispatcher timeout constants (seconds) _DISPATCH_READ_TIMEOUT = 30.0 _DISPATCH_POLL_INTERVAL = 0.5 diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py index 73687165b8..3ceb1a268a 100644 --- a/studio/backend/core/inference/tensor_fallback.py +++ b/studio/backend/core/inference/tensor_fallback.py @@ -13,7 +13,7 @@ import logging from typing import Awaitable, Callable, Optional from core.inference.llama_server_args import ( - resolve_tensor_parallel, + _effective_tensor_parallel, strip_split_mode_only, ) @@ -34,18 +34,20 @@ async def load_with_tensor_fallback( True on success; it *raises* on a hard crash (llama-server aborts on some archs / older builds), which is treated the same as a False return. - Tensor mode can be requested by the toggle or by a ``--split-mode tensor`` - in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether - tensor mode is actually engaged, and it strips ``--split-mode`` from the - extras so the layer retry can't relaunch the same failing tensor load. A - non-tensor load keeps its original contract and propagates exceptions. + Tensor mode can be requested by the toggle, by a ``--split-mode tensor`` in + ``extra_args`` (an allowed shadow flag), or by an inherited + ``LLAMA_ARG_SPLIT_MODE=tensor`` env (load_model engages it the same way), so + the retry is keyed on whether tensor mode is actually engaged, and it forces + ``--split-mode layer`` on the retry so neither leftover extras nor the + inherited tensor env can relaunch the same failing tensor load. A non-tensor + load keeps its original contract and propagates exceptions. ``cancelled()`` distinguishes a real tensor-start failure from a user cancellation: ``attempt_load`` also returns False when the load was cancelled, so without this the helper would restart a load the user just cancelled. """ - tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor) + tensor_requested = _effective_tensor_parallel(extra_args, requested_tensor) try: success = await attempt_load(requested_tensor, extra_args) except Exception as exc: @@ -67,4 +69,8 @@ async def load_with_tensor_fallback( "(this model may not support tensor parallelism)", label, ) - return await attempt_load(False, strip_split_mode_only(extra_args)) + # Force --split-mode layer (CLI wins over env) so neither leftover extras nor + # an inherited LLAMA_ARG_SPLIT_MODE=tensor can re-engage tensor and re-crash + # the retry; load_model and the child both honor the explicit layer override. + layer_extras = strip_split_mode_only(extra_args) or [] + return await attempt_load(False, [*layer_extras, "--split-mode", "layer"]) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index d69246f584..0c3d42a2fe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -217,6 +217,12 @@ class TrainingBackend: self._db_config: Optional[dict] = None self._db_started_at: Optional[str] = None + # Xet -> HTTP model-load fallback state (config kept for the respawn). + self._last_full_config: Optional[dict] = None + self._in_model_load: bool = False + self._xet_fallback_used: bool = False + self._needs_xet_respawn: bool = False + logger.info("TrainingBackend initialized (subprocess mode)") # ------------------------------------------------------------------ @@ -311,6 +317,8 @@ class TrainingBackend: "trust_remote_code": kwargs.get("trust_remote_code", False), "gpu_ids": kwargs.get("gpu_ids"), "s3_config": kwargs.get("s3_config"), + # Flipped to True only by the HTTP-fallback respawn after a stall. + "disable_xet": kwargs.get("disable_xet", False), } # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. @@ -386,6 +394,11 @@ class TrainingBackend: self._db_total_steps_set = False self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() + # Start each job Xet-first; keep config so a stall can respawn over HTTP. + self._last_full_config = config + self._in_model_load = False + self._xet_fallback_used = False + self._needs_xet_respawn = False # Assign subprocess handles after state reset. self._event_queue = event_queue @@ -446,6 +459,97 @@ class TrainingBackend: output_dir, ) + def _handle_stall_event(self, event: dict) -> None: + """A worker reported a no-progress download stall. + + On the first model-load, terminate the worker so the pump loop respawns it + over HTTP. A later stall (already on HTTP, or outside model-load) surfaces + as an error instead. + """ + msg = event.get("message", "Download stalled") + with self._lock: + recover = self._in_model_load and not self._xet_fallback_used + proc = self._proc + if recover: + self._xet_fallback_used = True + self._needs_xet_respawn = True + self._progress.status_message = ( + "Model download stalled on Xet; retrying over HTTP..." + ) + else: + self._progress.error = self._progress.error or ( + "Model download stalled even over HTTP -- check your network connection" + ) + if recover: + logger.warning("Training model-load stalled on Xet; respawning over HTTP: %s", msg) + else: + logger.error("Training download stalled with no further fallback: %s", msg) + # Terminate either way so the pump loop proceeds (respawn or finalize). + if proc is not None and proc.is_alive(): + proc.terminate() + + def _respawn_worker_disable_xet(self) -> None: + """Respawn the worker once with HF_HUB_DISABLE_XET=1 after a model-load + stall. Runs on the exiting pump thread, reaps the terminated worker, and + starts a fresh worker + pump. DB/progress run-state is preserved so the + history row is not duplicated; the new worker re-formats and loads over HTTP. + """ + config = self._last_full_config + if config is None: + logger.error("Cannot respawn training worker: no stored config") + return + + with self._lock: + old_proc = self._proc + if old_proc is not None: + old_proc.join(timeout = 5.0) + if old_proc.is_alive(): + old_proc.kill() + old_proc.join(timeout = 2.0) + + config = {**config, "disable_xet": True} + self._last_full_config = config + logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall") + + from .worker import run_training_process + + try: + with native_path_secret_removed_for_child_start(): + event_queue = _CTX.Queue() + stop_queue = _CTX.Queue() + new_proc = _CTX.Process( + target = run_without_native_path_secret, + args = (run_training_process,), + kwargs = { + "event_queue": event_queue, + "stop_queue": stop_queue, + "config": config, + }, + daemon = True, + ) + new_proc.start() + except Exception: + logger.error("Failed to respawn training subprocess", exc_info = True) + with self._lock: + self._progress.is_training = False + self._progress.error = "Failed to recover stalled model download" + self._ensure_db_run_created() + self._finalize_run_in_db( + status = "error", + error_message = "Failed to recover stalled model download", + ) + return + + logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._in_model_load = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = new_proc + self._pump_thread = new_pump + new_pump.start() + def is_training_active(self) -> bool: """Check if training is currently active.""" with self._lock: @@ -566,6 +670,14 @@ class TrainingBackend: for e in self._drain_queue(self._event_queue): self._handle_event(e) + # Model-load stall: respawn over HTTP instead of finalizing as failure. + # Runs on THIS exiting pump thread and starts a fresh pump (never joins + # the current thread); DB run-state is preserved. + if self._needs_xet_respawn: + self._needs_xet_respawn = False + self._respawn_worker_disable_xet() + return + # Mark done if no explicit complete/error was received. with self._lock: if self._progress.is_training: @@ -597,6 +709,19 @@ class TrainingBackend: db_action: Optional[str] = None db_action_kwargs: dict = {} + # Model-load lifecycle + stall recovery (no DB metrics); handled first. + if etype == "model_load_started": + with self._lock: + self._in_model_load = True + return + if etype == "model_load_completed": + with self._lock: + self._in_model_load = False + return + if etype == "stall": + self._handle_stall_event(event) + return + with self._lock: if etype == "progress": self._progress.step = event.get("step", self._progress.step) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index b106aa1298..65533b6946 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1987,6 +1987,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # before imports + # HTTP-fallback respawn: disable Xet before any huggingface_hub import (the + # var is read at import time). Mirrors core/inference/worker.py. + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + print( + "Xet transport disabled for this training worker (HF_HUB_DISABLE_XET=1).", + file = sys.stderr, + flush = True, + ) + # Offline auto-detect: skip ~25s of HF retries per call when DNS is dead. if "HF_HUB_OFFLINE" not in os.environ: import socket as _socket @@ -2706,18 +2719,33 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> cpt_trains_embeddings = False # ── 4c. Load training model (uses VRAM — dataset already formatted) ── + # Watchdog lets the parent recover a stalled Xet download via respawn. _send_status(event_queue, "Loading model...") - success = trainer.load_model( - model_name = model_name, - max_seq_length = config["max_seq_length"], - load_in_4bit = config["load_in_4bit"], - full_finetuning = not use_lora, - hf_token = hf_token, - is_dataset_image = config.get("is_dataset_image", False), - is_dataset_audio = config.get("is_dataset_audio", False), - trust_remote_code = config.get("trust_remote_code", False), - gpu_ids = config.get("resolved_gpu_ids"), + from utils.hf_xet_fallback import start_watchdog + + event_queue.put({"type": "model_load_started", "ts": time.time()}) + _load_watchdog_stop = start_watchdog( + repo_ids = [model_name], + on_stall = lambda msg: event_queue.put( + {"type": "stall", "message": msg, "ts": time.time()} + ), + xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) + try: + success = trainer.load_model( + model_name = model_name, + max_seq_length = config["max_seq_length"], + load_in_4bit = config["load_in_4bit"], + full_finetuning = not use_lora, + hf_token = hf_token, + is_dataset_image = config.get("is_dataset_image", False), + is_dataset_audio = config.get("is_dataset_audio", False), + trust_remote_code = config.get("trust_remote_code", False), + gpu_ids = config.get("resolved_gpu_ids"), + ) + finally: + _load_watchdog_stop.set() + event_queue.put({"type": "model_load_completed", "ts": time.time()}) if not success or trainer.should_stop: if trainer.should_stop: event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py index 095e3de641..74c3ad6ce2 100644 --- a/studio/backend/hub/services/models/gguf_variants.py +++ b/studio/backend/hub/services/models/gguf_variants.py @@ -525,33 +525,51 @@ async def get_gguf_variants_response( return False def _any_mmproj_cached(filenames: frozenset[str]) -> bool: - return any( + if any( by_filename.get(name.lower()) is not None for by_filename in cached_filenames_by_snapshot for name in filenames + ): + return True + return any( + _is_mmproj_filename(name.rsplit("/", 1)[-1]) + for by_filename in cached_filenames_by_snapshot + for name in by_filename + ) + + def _quant_bytes_present(quant: str, size_bytes: int) -> bool: + # Small rounding tolerance for symlinks vs real sizes. + if size_bytes <= 0: + return False + return any( + by_quant.get(quant, 0) >= size_bytes * 0.99 + for by_quant in cached_quant_bytes_by_snapshot ) def _is_fully_downloaded(variant) -> bool: - requirement = requirements_by_quant.get(variant.quant.lower()) - if requirement is None: - if variant.size_bytes == 0: - return False - quant = variant.quant.lower() - # Allow small rounding tolerance (symlinks vs real sizes). - return any( - by_quant.get(quant, 0) >= variant.size_bytes * 0.99 - for by_quant in cached_quant_bytes_by_snapshot + quant = variant.quant.lower() + requirement = requirements_by_quant.get(quant) + # Vision repos ship an mmproj adapter; any precision on disk suffices. + if ( + requirement is not None + and _filenames_cached( + requirement.main_filenames, + requirement.main_size_bytes, + ) + and ( + not requirement.mmproj_filenames + or _any_mmproj_cached(requirement.mmproj_filenames) ) - if not _filenames_cached( - requirement.main_filenames, - requirement.main_size_bytes, ): + return True + # Byte fallback so a present quant isn't demoted by a filename mismatch; + # vision repos still need an mmproj cached (any precision). + if not _quant_bytes_present(quant, variant.size_bytes): return False - # Vision repos ship an mmproj adapter per variant. Any mmproj - # precision on disk suffices (the loader picks whichever is present); - # requiring the API-preferred one would falsely demote variants. - if requirement.mmproj_filenames and not _any_mmproj_cached( - requirement.mmproj_filenames, + if ( + requirement is not None + and requirement.mmproj_filenames + and not _any_mmproj_cached(requirement.mmproj_filenames) ): return False return True diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py index 4eb8180628..acd650bf42 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -353,6 +353,20 @@ def list_partial_gguf_variants_from_state( return variants, has_vision +def resolve_local_gguf_path(repo_id: str, gguf_variant: Optional[str]) -> Optional[str]: + """Absolute path to the (shard-1) GGUF file for ``repo_id`` + ``gguf_variant`` + if it is already downloaded in the HF cache, else ``None``. Read-only — never + triggers a download. Lets callers read header metadata before a load.""" + for snapshot in iter_hf_cache_snapshots(repo_id): + variants, _ = list_local_gguf_variants(str(snapshot)) + for variant in variants: + if gguf_variant is None or variant.quant == gguf_variant: + candidate = snapshot / variant.filename + if candidate.is_file(): + return str(candidate) + return None + + def list_gguf_variants( repo_id: str, hf_token: Optional[str] = None ) -> tuple[list[GgufVariantInfo], bool, Optional[list]]: diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py index b8691587ac..f471ceb300 100644 --- a/studio/backend/loggers/handlers.py +++ b/studio/backend/loggers/handlers.py @@ -8,6 +8,7 @@ filter_sensitive_data (structlog processor for sanitization), and get_logger (factory for structured loggers). """ +import os import re import time @@ -17,6 +18,31 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send from utils.native_path_leases import redact_native_paths logger = structlog.get_logger(__name__) + + +def _env_int(name: str, default: int) -> int: + try: + raw = (os.environ.get(name) or "").strip() + return int(raw) if raw else default + except ValueError: + return default + + +# Drop duplicate successful-GET access logs repeated within the window: the SPA +# fans one cache invalidation into many identical list fetches; only the first +# informs. Loading polls, mutations, and errors are unaffected. 0 = log all. +_ACCESS_LOG_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS", 300) +# Pure-liveness/UI polls whose access line carries no signal beyond "client still +# polling" (state changes are logged by their own modules). Collapsed to a longer +# heartbeat instead of one line per poll; first hit and any error still log. 0 = off. +_QUIET_POLL_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS", 10000) +_QUIET_POLL_PATHS = { + "/api/health", + "/api/auth/status", + "/api/inference/status", + "/api/inference/monitor", +} +_DEDUP_MAP_MAX = 4096 _NATIVE_PATH_LEASE_RE = re.compile( r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+" ) @@ -43,6 +69,30 @@ class LoggingMiddleware: def __init__(self, app: ASGIApp) -> None: self.app = app + # (method, path, query, status_code) -> monotonic ts of the last EMITTED log. + self._last_log: dict[tuple[str, str, bytes, int], float] = {} + + def _is_redundant_repeat( + self, method: str, path: str, query: bytes, status_code: int, now: float + ) -> bool: + """True if an identical GET/2xx log fired < window ago. The query string + is part of the identity, so distinct query-driven GETs are not collapsed. + Mutations and non-2xx are never deduped. Quiet-poll paths use a longer + heartbeat window. Stamps only on emit, so steady polls still log.""" + if method != "GET" or not (200 <= status_code < 300): + return False + window_ms = _QUIET_POLL_DEDUP_MS if path in _QUIET_POLL_PATHS else _ACCESS_LOG_DEDUP_MS + if window_ms <= 0: + return False + key = (method, path, query, status_code) + last = self._last_log.get(key) + if last is not None and (now - last) * 1000.0 < window_ms: + return True + self._last_log[key] = now + if len(self._last_log) > _DEDUP_MAP_MAX: + cutoff = now - (max(_ACCESS_LOG_DEDUP_MS, _QUIET_POLL_DEDUP_MS) / 1000.0) + self._last_log = {k: v for k, v in self._last_log.items() if v >= cutoff} + return False async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": @@ -78,13 +128,16 @@ class LoggingMiddleware: ) raise else: - if not excluded: + end_time = time.perf_counter() + if not excluded and not self._is_redundant_repeat( + scope["method"], path, scope.get("query_string", b""), status_code, end_time + ): logger.info( "request_completed", method = scope["method"], path = path, status_code = status_code, - process_time_ms = round((time.perf_counter() - start_time) * 1000, 2), + process_time_ms = round((end_time - start_time) * 1000, 2), ) diff --git a/studio/backend/main.py b/studio/backend/main.py index 08c7bc6826..54946fa992 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -15,6 +15,15 @@ from dataclasses import asdict # Suppress C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" +# Pin GPU index ordering to PCI bus id before any torch import creates a CUDA +# context. Without this, torch/CUDA default to FASTEST_FIRST while nvidia-smi +# (and Studio's VRAM probes) use PCI-bus order, so a GPU index chosen from +# nvidia-smi data can resolve to a different physical card via +# CUDA_VISIBLE_DEVICES. setdefault so an explicit user override wins. See +# utils/hardware/hardware.py for the full rationale; set here too so the entry +# process is covered before its heavy ML imports. +os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. @@ -465,6 +474,10 @@ async def lifespan(app: FastAPI): app.state.bootstrap_password = storage.get_bootstrap_password() yield + from core.inference.llama_http import aclose as _close_llama_http + + await _close_llama_http() + await run_lifespan_shutdown( terminate_hub_downloads, clear_unsloth_compiled_cache, @@ -492,8 +505,7 @@ app.add_middleware(LoggingMiddleware) # img/media-src allow any https origin so HF model-card assets render (mirrors # tauri.conf.json); scripts/frames/connect-src stay same-origin + HF. -from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402 -from starlette.requests import Request as _StarletteRequest # noqa: E402 +from starlette.datastructures import MutableHeaders # noqa: E402 _CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce" @@ -549,28 +561,51 @@ def _build_csp(script_nonce: "str | None" = None) -> str: ) -class SecurityHeadersMiddleware(BaseHTTPMiddleware): - """Set baseline security headers; splice per-response inline-script nonces into CSP.""" +class SecurityHeadersMiddleware: + """Set baseline security headers; splice per-response inline-script nonces into CSP. - async def dispatch(self, request: _StarletteRequest, call_next): - response = await call_next(request) - # Strip the internal nonce hand-off header so it never reaches the client - nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER) - if nonce is not None: - del response.headers[_CSP_SCRIPT_NONCE_HEADER] - response.headers.setdefault("Content-Security-Policy", _build_csp(nonce)) - # Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and - # DENY would block serve_kernel_port_as_iframe regardless of CSP. - if not _IS_COLAB and request.url.path != _ARTIFACT_PREVIEW_FRAME_PATH: - response.headers.setdefault("X-Frame-Options", "DENY") - response.headers.setdefault("X-Content-Type-Options", "nosniff") - response.headers.setdefault("Referrer-Policy", "no-referrer") - response.headers.setdefault( - "Permissions-Policy", - "camera=(), microphone=(self), geolocation=()", - ) - response.headers["server"] = "unsloth-studio" - return response + Pure ASGI (not BaseHTTPMiddleware) so streaming responses are not wrapped in + an anyio stream. Header logic mirrors the prior version exactly via + MutableHeaders on the response-start message. + """ + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + path = scope.get("path", "") + + async def send_wrapper(message): + if message["type"] == "http.response.start": + # ASGI headers are an iterable; coerce to a list so MutableHeaders + # can mutate in place even if a server sends a tuple or omits it. + raw = message.setdefault("headers", []) + if not isinstance(raw, list): + raw = list(raw) + message["headers"] = raw + headers = MutableHeaders(raw = raw) + # Strip the internal nonce hand-off header so it never reaches the client + nonce = headers.get(_CSP_SCRIPT_NONCE_HEADER) + if nonce is not None: + del headers[_CSP_SCRIPT_NONCE_HEADER] + headers.setdefault("Content-Security-Policy", _build_csp(nonce)) + # Omit X-Frame-Options in Colab: CSP frame-ancestors handles it, and + # DENY would block serve_kernel_port_as_iframe regardless of CSP. + if not _IS_COLAB and path != _ARTIFACT_PREVIEW_FRAME_PATH: + headers.setdefault("X-Frame-Options", "DENY") + headers.setdefault("X-Content-Type-Options", "nosniff") + headers.setdefault("Referrer-Policy", "no-referrer") + headers.setdefault( + "Permissions-Policy", + "camera=(), microphone=(self), geolocation=()", + ) + headers["server"] = "unsloth-studio" + await send(message) + + await self.app(scope, receive, send_wrapper) app.add_middleware(SecurityHeadersMiddleware) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 82949870da..fdc3bb25c6 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -125,6 +125,11 @@ class ValidateModelRequest(BaseModel): gguf_variant: Optional[str] = Field( None, description = "GGUF quantization variant (e.g. 'Q4_K_M')" ) + include_context_length: bool = Field( + False, + description = "Also read the native context length from the local GGUF header. " + "Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.", + ) class ValidateModelResponse(BaseModel): @@ -144,6 +149,11 @@ class ValidateModelResponse(BaseModel): False, description = "Whether the model defaults require trust_remote_code to be enabled for loading.", ) + context_length: Optional[int] = Field( + None, + description = "Native training context length, read from the GGUF header when the file " + "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", + ) class GenerateRequest(BaseModel): diff --git a/studio/backend/requirements/overrides.txt b/studio/backend/requirements/overrides.txt index 176c651b96..8df4089402 100644 --- a/studio/backend/requirements/overrides.txt +++ b/studio/backend/requirements/overrides.txt @@ -1,2 +1,5 @@ -# Torch AO overrides (installed with --force-reinstall --no-cache-dir) -torchao==0.14.0 +# torchao is installed by studio/install_python_stack.py, which selects the +# version matching the torch release actually installed in the venv (torchao's +# C++ extensions are built against one exact torch version, so a fixed pin here +# would skip them on a newer torch). See _select_torchao_spec / +# _probe_installed_torch_version in that file. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d7db0e6f8d..0b95ad3987 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -612,6 +612,7 @@ try: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _tensor_parallel_matches_loaded, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -646,6 +647,7 @@ except ImportError: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _tensor_parallel_matches_loaded, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -875,6 +877,8 @@ from auth.authentication import get_current_subject 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.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override @@ -1297,6 +1301,423 @@ def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str logger = get_logger(__name__) +def _monitor_content_text(content) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for part in content: + if isinstance(part, dict): + ptype = part.get("type") + if ptype in ("text", "input_text", "output_text"): + text = part.get("text") + if isinstance(text, str): + parts.append(text) + elif ptype in ("image_url", "input_image", "image"): + parts.append("[image]") + else: + parts.append(f"[{ptype or 'content'}]") + else: + ptype = getattr(part, "type", None) + text = getattr(part, "text", None) + if isinstance(text, str): + parts.append(text) + elif ptype in ("image_url", "input_image", "image"): + parts.append("[image]") + elif ptype: + parts.append(f"[{ptype}]") + return "\n".join(parts) + return str(content) + + +def _monitor_prompt_from_messages(messages) -> str: + lines: list[str] = [] + for msg in messages or []: + role = msg.get("role") if isinstance(msg, dict) else getattr(msg, "role", "") + content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "") + tool_calls = ( + msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) + ) + text = _monitor_content_text(content) + if tool_calls and not text: + text = "[tool calls]" + if text: + lines.append(f"{role or 'message'}: {text}") + return "\n\n".join(lines) + + +def _monitor_usage( + monitor_id: Optional[str], + usage: Optional[dict], + context_length = None, +): + if not usage: + return + api_monitor.set_usage( + monitor_id, + prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens"), + completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens"), + total_tokens = usage.get("total_tokens"), + context_length = context_length, + ) + + +def _monitor_call_text(name: Any, arguments: Any = None) -> str: + call_name = str(name or "tool") + if arguments is None or arguments == "": + return f"Tool call: {call_name}" + if not isinstance(arguments, str): + args_text = json.dumps(arguments, default = str) + else: + args_text = arguments + if len(args_text) > 500: + args_text = args_text[:497] + "..." + return f"Tool call: {call_name}({args_text})" + + +def _monitor_tool_calls_text(tool_calls: Any) -> str: + if not isinstance(tool_calls, list): + return "" + parts: list[str] = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + continue + fn = tool_call.get("function") or {} + if not isinstance(fn, dict): + fn = {} + name = fn.get("name") or tool_call.get("name") or "tool" + args = fn.get("arguments") + if args is None: + args = tool_call.get("arguments") + parts.append(_monitor_call_text(name, args)) + return "\n".join(parts) + + +def _monitor_openai_chunk( + monitor_id: Optional[str], + data: dict, + context_length = None, +): + if not monitor_id: + return + _monitor_usage(monitor_id, data.get("usage"), context_length) + # Defensive: ignore malformed shapes so the helper never raises into the + # streaming generator and aborts the user's response. + choices = data.get("choices") + if not isinstance(choices, list) or not choices: + return + reply_parts: list[tuple[int, str]] = [] + for idx, choice in enumerate(choices): + if not isinstance(choice, dict): + continue + delta = choice.get("delta") or {} + message = choice.get("message") or {} + content = delta.get("content") if isinstance(delta, dict) else None + if content: + api_monitor.append_reply(monitor_id, content) + continue + if isinstance(delta, dict): + tool_text = _monitor_tool_calls_text(delta.get("tool_calls")) + if tool_text: + api_monitor.append_reply(monitor_id, tool_text) + continue + if isinstance(choice.get("text"), str): + reply_parts.append((idx, choice["text"])) + elif isinstance(message, dict): + text = message.get("content") + if isinstance(text, str): + reply_parts.append((idx, text)) + else: + tool_text = _monitor_tool_calls_text(message.get("tool_calls")) + if tool_text: + reply_parts.append((idx, tool_text)) + if not reply_parts: + return + if len(choices) == 1: + api_monitor.append_reply(monitor_id, reply_parts[0][1]) + return + api_monitor.append_reply( + monitor_id, + "\n\n".join(f"Choice {idx + 1}:\n{text}" for idx, text in reply_parts), + ) + + +def _monitor_openai_error_message(data: dict) -> Optional[str]: + error = data.get("error") + if isinstance(error, dict): + message = error.get("message") + if isinstance(message, str) and message: + return message + return json.dumps(error) + if isinstance(error, str) and error: + return error + return None + + +def _monitor_openai_sse_line( + monitor_id: Optional[str], + raw_line: str, + context_length = None, +) -> Optional[str]: + if not monitor_id: + return None + # SSE spec allows `data:value` and `data: value`; accept both. + if not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + if data_str == "[DONE]": + api_monitor.finish(monitor_id) + return "done" + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + if isinstance(data, dict): + error_message = _monitor_openai_error_message(data) + if error_message: + api_monitor.fail(monitor_id, error_message) + return "error" + _monitor_openai_chunk(monitor_id, data, context_length) + return None + + +def _monitor_openai_sse_event( + monitor_id: Optional[str], + event: bytes, + context_length = None, +) -> None: + for line in event.decode("utf-8", errors = "ignore").splitlines(): + _monitor_openai_sse_line(monitor_id, line.strip(), context_length) + + +def _monitor_anthropic_usage( + monitor_id: Optional[str], + usage: Optional[dict], + context_length = None, +) -> None: + if not usage: + return + _monitor_usage( + monitor_id, + { + "prompt_tokens": usage.get("input_tokens") or usage.get("prompt_tokens"), + "completion_tokens": usage.get("output_tokens") or usage.get("completion_tokens"), + "total_tokens": usage.get("total_tokens"), + }, + context_length, + ) + + +_ANTHROPIC_MONITOR_TOOL_BLOCKS: dict[str, dict[int, bool]] = {} + + +def _monitor_anthropic_index(data: dict) -> int: + try: + return int(data.get("index") or 0) + except (TypeError, ValueError): + return 0 + + +def _monitor_anthropic_payload( + monitor_id: Optional[str], + data: dict, + context_length = None, +) -> Optional[str]: + if not monitor_id or not isinstance(data, dict): + return None + event_type = data.get("type") + if event_type == "message_start": + message = data.get("message") or {} + if isinstance(message, dict): + _monitor_anthropic_usage(monitor_id, message.get("usage"), context_length) + return None + if event_type == "content_block_start": + content_block = data.get("content_block") or {} + if isinstance(content_block, dict) and content_block.get("type") == "tool_use": + index = _monitor_anthropic_index(data) + _ANTHROPIC_MONITOR_TOOL_BLOCKS.setdefault(monitor_id, {})[index] = False + api_monitor.append_reply(monitor_id, _monitor_call_text(content_block.get("name"))) + return None + if event_type == "content_block_delta": + delta = data.get("delta") or {} + text = delta.get("text") if isinstance(delta, dict) else None + if isinstance(text, str) and text: + api_monitor.append_reply(monitor_id, text) + elif isinstance(delta, dict) and delta.get("type") == "input_json_delta": + index = _monitor_anthropic_index(data) + tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id) or {} + if index in tool_blocks: + if not tool_blocks[index]: + api_monitor.append_reply(monitor_id, "\nInput: ") + tool_blocks[index] = True + partial_json = delta.get("partial_json") + if isinstance(partial_json, str) and partial_json: + api_monitor.append_reply(monitor_id, partial_json) + return None + if event_type == "content_block_stop": + index = _monitor_anthropic_index(data) + tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id) + if tool_blocks is not None: + tool_blocks.pop(index, None) + if not tool_blocks: + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + return None + if event_type == "message_delta": + _monitor_anthropic_usage(monitor_id, data.get("usage"), context_length) + return None + if event_type == "error": + error = data.get("error") or {} + if isinstance(error, dict): + message = error.get("message") or json.dumps(error, default = str) + else: + message = str(error) + api_monitor.fail(monitor_id, message) + return "error" + return None + + +def _monitor_anthropic_sse_line( + monitor_id: Optional[str], + raw_line: str, + context_length = None, +) -> Optional[str]: + if not monitor_id or not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + return _monitor_anthropic_payload(monitor_id, data, context_length) + + +def _monitor_anthropic_content_blocks(content: Any) -> str: + if not isinstance(content, list): + return "" + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "text" and isinstance(block.get("text"), str): + parts.append(block["text"]) + elif block.get("type") == "tool_use": + parts.append(_monitor_call_text(block.get("name"), block.get("input"))) + return "".join(parts) + + +def _monitor_anthropic_json_response( + response, + monitor_id: Optional[str], + context_length = None, +) -> None: + if not monitor_id: + return + body = getattr(response, "body", b"") + try: + data = json.loads(body.decode("utf-8") if isinstance(body, bytes) else body) + except Exception: + api_monitor.finish(monitor_id) + return + if not isinstance(data, dict): + api_monitor.finish(monitor_id) + return + text = _monitor_anthropic_content_blocks(data.get("content")) + if text: + api_monitor.set_reply(monitor_id, text) + _monitor_anthropic_usage(monitor_id, data.get("usage"), context_length) + api_monitor.finish(monitor_id) + + +def _monitor_anthropic_response( + response, + monitor_id, + context_length = None, + cancel_event = None, +): + if not monitor_id: + return response + body_iterator = getattr(response, "body_iterator", None) + if body_iterator is None: + _monitor_anthropic_json_response(response, monitor_id, context_length) + return response + + async def _monitored_body(): + terminal = False + try: + async for chunk in body_iterator: + text = ( + chunk.decode("utf-8", errors = "ignore") + if isinstance(chunk, (bytes, bytearray)) + else str(chunk) + ) + for line in text.splitlines(): + if ( + _monitor_anthropic_sse_line( + monitor_id, + line.strip(), + context_length, + ) + == "error" + ): + terminal = True + yield chunk + if not terminal: + api_monitor.finish( + monitor_id, + "cancelled" + if cancel_event is not None and cancel_event.is_set() + else "completed", + ) + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + except asyncio.CancelledError: + if cancel_event is not None: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + raise + + response.body_iterator = _monitored_body() + return response + + +def _monitor_context_length() -> Optional[int]: + llama_backend = get_llama_cpp_backend() + if getattr(llama_backend, "is_loaded", False): + context_length = _positive_int_or_none(getattr(llama_backend, "context_length", None)) + if context_length is not None: + return context_length + backend = get_inference_backend() + if not backend.active_model_name: + return None + models = getattr(backend, "models", {}) or {} + model_info = models.get(backend.active_model_name, {}) if isinstance(models, dict) else {} + context_length = _positive_int_or_none(model_info.get("context_length")) + if context_length is not None: + return context_length + for candidate in ( + getattr(backend, "context_length", None), + getattr(backend, "max_seq_length", None), + ): + context_length = _positive_int_or_none(candidate) + if context_length is not None: + return context_length + return None + + +def _monitor_active_model() -> Optional[str]: + llama_backend = get_llama_cpp_backend() + if getattr(llama_backend, "is_loaded", False): + return getattr(llama_backend, "model_identifier", None) + backend = get_inference_backend() + return backend.active_model_name + + def _validate_native_gguf_companion( companion_path: str | None, gguf_path: str | None, label: str ) -> None: @@ -1395,9 +1816,8 @@ def _request_matches_loaded_settings( strip_split_mode = _should_strip_split_mode(request, backend_extra), ) ) - if ( - resolve_tensor_parallel(effective_extra, request.tensor_parallel) - != llama_backend.tensor_parallel + if not _tensor_parallel_matches_loaded( + effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False # Spec decoding works on vision models too (MTP is mmproj-compatible, @@ -1503,6 +1923,19 @@ def get_llama_cpp_backend() -> LlamaCppBackend: return _llama_cpp_backend +def _model_json_response(model, status_code: int = 200) -> Response: + """Serialize a pydantic response once via pydantic-core. + + Equivalent body to ``JSONResponse(content = model.model_dump())`` but + avoids the dict round-trip plus Starlette's second ``json.dumps``. + """ + return Response( + content = model.model_dump_json(), + media_type = "application/json", + status_code = status_code, + ) + + @router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, @@ -2117,6 +2550,33 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) + is_gguf = getattr(config, "is_gguf", False) + # Native context length, read from the local GGUF header when present. + # Lets the staged ("Load on selection" off) flow populate the context + # slider before the GPU load; None until the file is downloaded. + context_length: Optional[int] = None + if request.include_context_length and is_gguf: + from hub.utils.gguf import resolve_local_gguf_path + from utils.models.gguf_metadata import read_gguf_context_length + + # Best-effort: a header-read failure must never fail validation of an + # otherwise-valid model (the outer except turns it into a 400). + try: + if native_grant_backed: + # model_identifier is the resolved canonical .gguf path. + local_gguf = model_identifier + else: + # Local folder / exported GGUFs already have their file + # resolved on the config (gguf_file is None for HF repos, so + # those fall back to the HF-cache lookup). + local_gguf = config.gguf_file or resolve_local_gguf_path( + model_identifier, request.gguf_variant + ) + if local_gguf: + context_length = read_gguf_context_length(local_gguf) + except Exception as e: + logger.debug("Context-length probe failed for %s: %s", model_log_label, e) + return ValidateModelResponse( valid = True, message = "Model identifier is valid.", @@ -2124,12 +2584,13 @@ async def validate_model( display_name = model_log_label if native_grant_backed else getattr(config, "display_name", config.identifier), - is_gguf = getattr(config, "is_gguf", False), + is_gguf = is_gguf, is_lora = getattr(config, "is_lora", False), is_vision = getattr(config, "is_vision", False), requires_trust_remote_code = bool( load_inference_config(config.identifier).get("trust_remote_code", False) ), + context_length = context_length, ) except HTTPException: @@ -2159,6 +2620,20 @@ async def validate_model( f"Error validating model identifier '{request.model_path}': {e}", exc_info = True, ) + # RuntimeError / ValueError carry intentional, actionable messages here + # (e.g. "llama-server binary not found - cannot load GGUF models. Run + # setup.sh ..."), so surface them instead of a blank "Invalid model". + # Path-redact for safety and keep any other exception type generic so an + # unexpected internal error never leaks its details to the client. + if isinstance(e, (RuntimeError, ValueError)): + msg = redact_native_paths(str(e)).strip() + if msg: + if any(h.lower() in msg.lower() for h in not_supported_hints): + msg = f"This model is not supported yet. Try a different model. (Original error: {msg})" + raise HTTPException( + status_code = 400, + detail = msg, + ) raise HTTPException( status_code = 400, detail = "Invalid model", @@ -2247,6 +2722,35 @@ async def confirm_tool_call( return {"resolved": True} +@studio_router.get("/monitor") +async def get_api_monitor(current_subject: str = Depends(get_current_subject)): + """Return recent OpenAI-compatible API activity for Studio.""" + active_model = _monitor_active_model() + active_requests = api_monitor.active_count(subject = current_subject) + if active_requests: + operating_status = "generating" + elif active_model: + operating_status = "ready" + else: + operating_status = "idle" + return { + "status": operating_status, + "active_model": active_model, + "context_length": _monitor_context_length(), + "active_requests": active_requests, + "entries": api_monitor.snapshot(include_details = False, subject = current_subject), + } + + +@studio_router.get("/monitor/{entry_id}") +async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(get_current_subject)): + """Return full prompt/reply details for one OpenAI-compatible API request.""" + entry = api_monitor.get(entry_id, subject = current_subject) + if entry is None: + raise HTTPException(status_code = 404, detail = "Monitor entry not found") + return entry + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, current_subject: str = Depends(get_current_subject) @@ -3167,7 +3671,9 @@ def _build_external_messages( async def _proxy_to_external_provider( - payload: ChatCompletionRequest, request: Request + payload: ChatCompletionRequest, + request: Request, + current_subject: Optional[str] = None, ) -> StreamingResponse: """ Proxy a chat completion request to an external LLM provider. @@ -3238,6 +3744,16 @@ async def _proxy_to_external_provider( provider_type = provider_type, base_url = base_url, ) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = None, + subject = current_subject, + ) client = ExternalProviderClient( provider_type = provider_type, @@ -3279,14 +3795,29 @@ async def _proxy_to_external_provider( ) try: sent_done = False + stream_failed = False async for line in gen: + monitor_event = _monitor_openai_sse_line(monitor_id, line) + if monitor_event is None: + try: + _monitor_openai_chunk(monitor_id, json.loads(line)) + except Exception: + pass + if monitor_event == "error": + stream_failed = True yield f"{line}\n\n" - if "[DONE]" in line: + if monitor_event == "done": sent_done = True if not sent_done: + if not stream_failed: + api_monitor.finish(monitor_id) yield "data: [DONE]\n\n" + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as exc: logger.error("external_provider.stream_error", error = str(exc)) + api_monitor.fail(monitor_id, _friendly_error(exc)) finally: try: await gen.aclose() @@ -3546,7 +4077,7 @@ async def openai_chat_completions( ) if _wants_multiple_choices(payload): _raise_unsupported_n("external provider chat completions") - return await _proxy_to_external_provider(payload, request) + return await _proxy_to_external_provider(payload, request, current_subject) # Reject a malformed function tool here: it would otherwise reach # llama-server and surface as an opaque 500 "Failed to parse tools". @@ -3595,12 +4126,49 @@ async def openai_chat_completions( # ── Determine which backend is active ───────────────────── # Single-model server: any model name serves the loaded model (drop-in # OpenAI compat), so payload.model is only a fallback label here. + monitor_id = None + + async def _monitored_generate_audio(model_label: str, context_length: Optional[int] = None): + tts_monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + tts_monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model_label, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = context_length, + subject = current_subject, + ) + try: + response = await generate_audio(payload, request) + except asyncio.CancelledError: + api_monitor.finish(tts_monitor_id, "cancelled") + raise + except Exception as e: + api_monitor.fail(tts_monitor_id, _friendly_error(e)) + raise + if isinstance(response, JSONResponse): + try: + body = json.loads(response.body.decode()) + choices = body.get("choices") or [] + message = (choices[0].get("message") or {}) if choices else {} + content = message.get("content") + if isinstance(content, str): + api_monitor.set_reply(tts_monitor_id, content) + except Exception: + pass + api_monitor.finish(tts_monitor_id) + return response + if using_gguf: model_name = llama_backend.model_identifier or payload.model if getattr(llama_backend, "_is_audio", False): if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF audio chat completions") - return await generate_audio(payload, request) + return await _monitored_generate_audio( + model_name, + context_length = llama_backend.context_length, + ) else: backend = get_inference_backend() if not backend.active_model_name: @@ -3616,7 +4184,7 @@ async def openai_chat_completions( # (Whisper is ASR not TTS -- handled below in audio input path) model_info = backend.models.get(backend.active_model_name, {}) if model_info.get("is_audio") and model_info.get("audio_type") != "whisper": - return await generate_audio(payload, request) + return await _monitored_generate_audio(model_name) # ── Whisper without audio: return clear error ── if model_info.get("audio_type") == "whisper" and not payload.audio_base64: @@ -3625,10 +4193,25 @@ async def openai_chat_completions( detail = "Whisper models require audio input. Please upload an audio file.", ) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model_name, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + # ── Audio INPUT path: decode WAV and route to audio input generation ── if payload.audio_base64 and model_info.get("has_audio_input"): - audio_array = _decode_audio_base64(payload.audio_base64) - system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + try: + audio_array = _decode_audio_base64(payload.audio_base64) + system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + except Exception as e: + api_monitor.fail(monitor_id, _friendly_error(e)) + raise cancel_event = threading.Event() completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) @@ -3674,16 +4257,20 @@ async def openai_chat_completions( gen = audio_input_generate() _DONE = object() + cancelled = False while True: if cancel_event.is_set(): + cancelled = True break if await request.is_disconnected(): cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") return chunk_text = await asyncio.to_thread(next, gen, _DONE) if chunk_text is _DONE: break if chunk_text: + api_monitor.append_reply(monitor_id, chunk_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -3703,13 +4290,16 @@ async def openai_chat_completions( model = model_name, choices = [ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")], ) + api_monitor.finish(monitor_id, "cancelled" if cancelled else "completed") yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: logger.error(f"Error during audio input streaming: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" finally: _tracker.__exit__(None, None, None) @@ -3724,7 +4314,13 @@ async def openai_chat_completions( }, ) else: - full_text = "".join(audio_input_generate()) + try: + full_text = "".join(audio_input_generate()) + except Exception as e: + api_monitor.fail(monitor_id, _friendly_error(e)) + raise + api_monitor.set_reply(monitor_id, full_text) + api_monitor.finish(monitor_id) response = ChatCompletion( id = completion_id, created = created, @@ -3736,7 +4332,35 @@ async def openai_chat_completions( ) ], ) - return JSONResponse(content = response.model_dump()) + return _model_json_response(response) + + if monitor_id is None and not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model_name, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + + # Finalize the monitor entry on validation rejection before raising. + def _reject(status_code: int, detail: Any) -> "HTTPException": + if monitor_id is not None: + fail_detail = detail if isinstance(detail, str) else json.dumps(detail, default = str) + api_monitor.fail(monitor_id, fail_detail) + return HTTPException(status_code = status_code, detail = detail) + + def _reject_unsupported_n(path_label: str) -> "HTTPException": + return _reject( + 400, + openai_error_body( + f"n > 1 is not supported for {path_label}.", + status = 400, + code = "unsupported_parameter", + param = "n", + ), + ) # ── Standard OpenAI function-calling pass-through (GGUF only) ──── # When a client (opencode / Claude Code via OpenAI compat / Cursor / @@ -3769,14 +4393,14 @@ async def openai_chat_completions( and (_tools_passthrough or _has_response_format) ): if _wants_multiple_choices(payload): - _raise_unsupported_n("GGUF tool or response_format passthrough") + raise _reject_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: # This path forwards the request verbatim, so the transcoded audio # never gets injected. (The agentic tool loop below does support # audio.) - raise HTTPException( - status_code = 400, - detail = "Audio input is not supported together with guided decoding or client-supplied tools yet.", + raise _reject( + 400, + "Audio input is not supported together with guided decoding or client-supplied tools yet.", ) # Preserve the vision guard from the non-passthrough path below: @@ -3791,9 +4415,9 @@ async def openai_chat_completions( for m in payload.messages ) ): - raise HTTPException( - status_code = 400, - detail = "Image provided but current GGUF model does not support vision.", + raise _reject( + 400, + "Image provided but current GGUF model does not support vision.", ) cancel_event = threading.Event() @@ -3809,21 +4433,20 @@ async def openai_chat_completions( payload, model_name, completion_id, + monitor_id = monitor_id, ) return await _openai_passthrough_non_streaming( llama_backend, payload, model_name, + monitor_id = monitor_id, ) # ── Parse messages (handles multimodal content parts) ───── system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) if not chat_messages: - raise HTTPException( - status_code = 400, - detail = "At least one non-system message is required.", - ) + raise _reject(400, "At least one non-system message is required.") # ── GGUF path: proxy to llama-server /v1/chat/completions ── if using_gguf: @@ -3836,25 +4459,19 @@ async def openai_chat_completions( audio_format = "wav" if payload.audio_base64: if not getattr(llama_backend, "_has_audio_input", False): - raise HTTPException( - status_code = 400, - detail = "Audio provided but current GGUF model does not support audio input.", + raise _reject( + 400, + "Audio provided but current GGUF model does not support audio input.", ) if len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: - raise HTTPException( - status_code = 413, - detail = "Audio file is too large (max ~25 MB).", - ) + raise _reject(413, "Audio file is too large (max ~25 MB).") try: audio_b64, audio_format = await asyncio.to_thread( _prepare_audio_for_llama, payload.audio_base64 ) except Exception as e: logger.warning("Audio decode failed: %s", e, exc_info = True) - raise HTTPException( - status_code = 400, - detail = "Could not decode the provided audio file.", - ) + raise _reject(400, "Could not decode the provided audio file.") gguf_messages, _ = _openai_messages_for_gguf_chat( payload, @@ -3918,9 +4535,9 @@ async def openai_chat_completions( # Bypass Permissions suppresses confirm, so the stream requirement # (the gate needs streaming to prompt) no longer applies. if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream: - raise HTTPException( - status_code = 400, - detail = openai_error_body( + raise _reject( + 400, + openai_error_body( "confirm_tool_calls requires stream=true for local tool execution.", status = 400, code = "invalid_request_error", @@ -3928,7 +4545,7 @@ async def openai_chat_completions( ), ) if _wants_multiple_choices(payload): - _raise_unsupported_n("GGUF tool chat completions") + raise _reject_unsupported_n("GGUF tool chat completions") # ── Tool-use system prompt nudge ────────────────────── _nudge = _build_tool_action_nudge( tools = tools_to_use, @@ -4038,6 +4655,7 @@ async def openai_chat_completions( break if await request.is_disconnected(): cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") return event = await asyncio.to_thread(next, gen, _tool_sentinel) @@ -4086,6 +4704,7 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4121,16 +4740,22 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stream_usage, _monitor_context_length()) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: import traceback tb = traceback.format_exc() logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: @@ -4179,7 +4804,7 @@ async def openai_chat_completions( if payload.stream: if _wants_multiple_choices(payload): - _raise_unsupported_n("streaming GGUF chat completions") + raise _reject_unsupported_n("streaming GGUF chat completions") _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() @@ -4212,6 +4837,7 @@ async def openai_chat_completions( break if await request.is_disconnected(): cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") return cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) if cumulative is _gguf_sentinel: @@ -4236,6 +4862,7 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4272,13 +4899,19 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stream_usage, _monitor_context_length()) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: logger.error(f"Error during GGUF streaming: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: @@ -4300,6 +4933,7 @@ async def openai_chat_completions( _n = payload.n or 1 _choices = [] + _monitor_replies = [] _prompt_tokens = 0 _sum_completion = 0 _prompt_details = None @@ -4325,6 +4959,7 @@ async def openai_chat_completions( finish_reason = _clamp_finish_reason(completion_finish), ) ) + _monitor_replies.append(full_text) if completion_usage: # The prompt is shared across all n choices, so count its # tokens ONCE (OpenAI bills only generated tokens for each @@ -4346,10 +4981,27 @@ async def openai_chat_completions( prompt_tokens_details = _prompt_tokens_details(_prompt_details), ), ) - return JSONResponse(content = response.model_dump()) + monitor_reply = full_text + if _n > 1: + monitor_reply = "\n\n".join( + f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies) + ) + api_monitor.set_reply(monitor_id, monitor_reply) + _monitor_usage( + monitor_id, + { + "prompt_tokens": _prompt_tokens, + "completion_tokens": _sum_completion, + "total_tokens": _prompt_tokens + _sum_completion, + }, + _monitor_context_length(), + ) + api_monitor.finish(monitor_id) + return _model_json_response(response) except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) # An over-context prompt makes llama-server return 400; map any # upstream 4xx to a 400 client error rather than leaking a 500. _cls = _classify_llama_generation_error(e) @@ -4469,9 +5121,9 @@ async def openai_chat_completions( # Bypass Permissions suppresses confirm, so the stream requirement # (the gate needs streaming to prompt) no longer applies. if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream: - raise HTTPException( - status_code = 400, - detail = openai_error_body( + raise _reject( + 400, + openai_error_body( "confirm_tool_calls requires stream=true for local tool execution.", status = 400, code = "invalid_request_error", @@ -4593,6 +5245,7 @@ async def openai_chat_completions( if await request.is_disconnected(): cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") return event = await asyncio.to_thread(next, gen, _sf_tool_sentinel) @@ -4627,6 +5280,7 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4667,17 +5321,23 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") raise except Exception: backend.reset_generation_state() # Generic wire message; full trace stays in the log (CWE-209: # transformers/torch errors may leak paths). logger.exception("safetensors tool stream error") + api_monitor.fail(monitor_id, "An internal error occurred.") error_chunk = { "error": { "message": "An internal error occurred.", @@ -4721,6 +5381,11 @@ async def openai_chat_completions( return full_text content_text = await asyncio.to_thread(_drain_to_text) + api_monitor.set_reply(monitor_id, content_text) + _stats = _sf_stats_holder.get("stats") + if _stats: + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed") response = ChatCompletion( id = completion_id, created = created, @@ -4732,11 +5397,17 @@ async def openai_chat_completions( ) ], ) - return JSONResponse(content = response.model_dump()) + return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception: backend.reset_generation_state() # CWE-209: generic detail; full trace in log. logger.exception("safetensors tool completion error") + api_monitor.fail(monitor_id, "An internal error occurred.") raise HTTPException( status_code = 500, detail = "An internal error occurred.", @@ -4830,11 +5501,13 @@ async def openai_chat_completions( if await request.is_disconnected(): cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") return new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4876,15 +5549,21 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = { "error": { "message": _friendly_error(e), @@ -4923,11 +5602,17 @@ async def openai_chat_completions( ) ], ) - return JSONResponse(content = response.model_dump()) + api_monitor.set_reply(monitor_id, full_text) + _stats = stats_holder.get("stats") + if _stats: + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish(monitor_id) + return _model_json_response(response) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) @@ -5137,6 +5822,19 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) + raw_prompt = body.get("prompt", "") + if isinstance(raw_prompt, list): + prompt_text = "\n".join(str(part) for part in raw_prompt) + else: + prompt_text = str(raw_prompt) + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = str(body.get("model") or llama_backend.model_identifier or "default"), + prompt = prompt_text, + context_length = llama_backend.context_length, + subject = current_subject, + ) if is_stream: @@ -5164,10 +5862,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S resp = await _send_stream_with_preheader_cancel(client, req, request = request) if resp is None: + api_monitor.finish(monitor_id, "cancelled") return if resp.status_code != 200: err_bytes = await resp.aread() err_text = err_bytes.decode("utf-8", errors = "replace") + api_monitor.fail(monitor_id, err_text[:500]) raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}") disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, disconnect_event) @@ -5184,26 +5884,49 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge buffer += chunk while b"\n\n" in buffer: event, buffer = buffer.split(b"\n\n", 1) + _monitor_openai_sse_event( + monitor_id, + event, + llama_backend.context_length, + ) out = _cmpl_stream_event_out(event, _include_usage) if out is not None: yield out + b"\n\n" if not disconnect_event.is_set() and buffer: + _monitor_openai_sse_event( + monitor_id, + buffer, + llama_backend.context_length, + ) out = _cmpl_stream_event_out(buffer, _include_usage) if out is not None: # Re-add the SSE separator the split consumed, so a final # event arriving without a trailing blank line is still # terminated for the client's parser. yield out + b"\n\n" + if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") + return + api_monitor.finish(monitor_id) except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not disconnect_event.is_set(): logger.error("openai_completions stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") return + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + disconnect_event.set() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as e: if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return logger.error("openai_completions stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") return @@ -5231,15 +5954,27 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge return StreamingResponse(_stream(), media_type = "text/event-stream") else: - async with httpx.AsyncClient() as client: - resp = await client.post( + try: + resp = await nonstreaming_client().post( target_url, json = body, timeout = _llama_non_streaming_generation_timeout(), ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as e: + api_monitor.fail(monitor_id, _friendly_error(e)) + raise if resp.status_code != 200: + api_monitor.fail(monitor_id, resp.text[:500]) raise _openai_passthrough_error(resp.status_code, resp.text) + try: + _monitor_openai_chunk(monitor_id, resp.json(), llama_backend.context_length) + except Exception: + pass + api_monitor.finish(monitor_id) return Response( content = _rewrite_cmpl_id(resp.content), @@ -5272,9 +6007,42 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get body = await request.json() target_url = f"{llama_backend.base_url}/v1/embeddings" + raw_input = body.get("input", "") + if isinstance(raw_input, list): + prompt_text = "\n".join(str(part) for part in raw_input) + else: + prompt_text = str(raw_input) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = str(body.get("model") or llama_backend.model_identifier or "default"), + prompt = prompt_text, + context_length = llama_backend.context_length, + subject = current_subject, + ) - async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S) + try: + resp = await nonstreaming_client().post( + target_url, + json = body, + timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + if resp.status_code != 200: + api_monitor.fail(monitor_id, resp.text[:500]) + else: + try: + _monitor_usage(monitor_id, resp.json().get("usage"), _monitor_context_length()) + except Exception: + pass + api_monitor.finish(monitor_id) return Response( content = resp.content, status_code = resp.status_code, @@ -5788,84 +6556,128 @@ def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]: async def _responses_non_streaming( - payload: ResponsesRequest, messages: list[ChatMessage], request: Request + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, + current_subject: Optional[str] = None, ) -> JSONResponse: """Handle a non-streaming Responses API call.""" chat_req = _build_chat_request(payload, messages, stream = False) - result = await openai_chat_completions(chat_req, request) - - # openai_chat_completions returns a JSONResponse for non-streaming. - if isinstance(result, JSONResponse): - body = json.loads(result.body.decode()) - elif isinstance(result, Response): - body = json.loads(result.body.decode()) - else: - body = result - - choices = body.get("choices", []) - text = "" - reasoning_text = "" - tool_calls: list[dict] = [] - if choices: - msg = choices[0].get("message", {}) or {} - raw_content = msg.get("content", "") or "" - raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content) - llama_backend = get_llama_cpp_backend() - reasoning_text, text = _extract_responses_reasoning( - raw_text, - msg.get("reasoning_content"), - parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend), - ) - tool_calls = msg.get("tool_calls") or [] - - usage_data = body.get("usage", {}) - input_tokens = usage_data.get("prompt_tokens", 0) - output_tokens = usage_data.get("completion_tokens", 0) - - resp_id = f"resp_{uuid.uuid4().hex[:12]}" - - # Responses API emits each tool call as its own top-level output item, - # plus an optional assistant text message. Emit the text message only when - # the model produced content, so clients expecting a pure tool-call turn - # (finish_reason="tool_calls") don't see a spurious empty message item. - output_items: list[dict] = [] - if reasoning_text and not text and not tool_calls: - text = reasoning_text - if reasoning_text: - output_items.append(_responses_reasoning_output_item(reasoning_text)) - if text: - msg_id = f"msg_{uuid.uuid4().hex[:12]}" - output_items.append( - ResponsesOutputMessage( - id = msg_id, - status = "completed", - role = "assistant", - content = [ResponsesOutputTextContent(text = text)], - ).model_dump() - ) - output_items.extend(_chat_tool_calls_to_responses_output(tool_calls)) - - response = ResponsesResponse( - id = resp_id, - created_at = int(time.time()), - status = "completed", - model = body.get("model", payload.model), - output = output_items, - usage = ResponsesUsage( - input_tokens = input_tokens, - output_tokens = output_tokens, - total_tokens = input_tokens + output_tokens, - ), - temperature = payload.temperature, - top_p = payload.top_p, - max_output_tokens = payload.max_output_tokens, - instructions = payload.instructions, + request_state = getattr(request, "state", None) + if request_state is None: + request_state = type("_RequestState", (), {})() + try: + setattr(request, "state", request_state) + except Exception: + request_state = None + previous_skip_monitor = ( + bool(getattr(request_state, "skip_api_monitor", False)) + if request_state is not None + else False ) - return JSONResponse(content = response.model_dump()) + monitor_id = None + if not previous_skip_monitor: + monitor_id = api_monitor.start( + endpoint = getattr(getattr(request, "url", None), "path", "/v1/responses"), + method = getattr(request, "method", "POST"), + model = payload.model, + prompt = _monitor_prompt_from_messages(messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + if request_state is not None: + request_state.skip_api_monitor = True + + try: + result = await openai_chat_completions(chat_req, request) + + # openai_chat_completions returns a JSONResponse for non-streaming. + if isinstance(result, JSONResponse): + body = json.loads(result.body.decode()) + elif isinstance(result, Response): + body = json.loads(result.body.decode()) + else: + body = result + + choices = body.get("choices", []) + text = "" + reasoning_text = "" + tool_calls: list[dict] = [] + if choices: + msg = choices[0].get("message", {}) or {} + raw_content = msg.get("content", "") or "" + raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content) + llama_backend = get_llama_cpp_backend() + reasoning_text, text = _extract_responses_reasoning( + raw_text, + msg.get("reasoning_content"), + parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend), + ) + tool_calls = msg.get("tool_calls") or [] + + usage_data = body.get("usage", {}) + input_tokens = usage_data.get("prompt_tokens", 0) + output_tokens = usage_data.get("completion_tokens", 0) + + resp_id = f"resp_{uuid.uuid4().hex[:12]}" + + # Responses API emits each tool call as its own top-level output item, + # plus an optional assistant text message. Emit the text message only when + # the model produced content, so clients expecting a pure tool-call turn + # (finish_reason="tool_calls") don't see a spurious empty message item. + output_items: list[dict] = [] + if reasoning_text and not text and not tool_calls: + text = reasoning_text + if reasoning_text: + output_items.append(_responses_reasoning_output_item(reasoning_text)) + if text: + msg_id = f"msg_{uuid.uuid4().hex[:12]}" + output_items.append( + ResponsesOutputMessage( + id = msg_id, + status = "completed", + role = "assistant", + content = [ResponsesOutputTextContent(text = text)], + ).model_dump() + ) + output_items.extend(_chat_tool_calls_to_responses_output(tool_calls)) + + response = ResponsesResponse( + id = resp_id, + created_at = int(time.time()), + status = "completed", + model = body.get("model", payload.model), + output = output_items, + usage = ResponsesUsage( + input_tokens = input_tokens, + output_tokens = output_tokens, + total_tokens = input_tokens + output_tokens, + ), + temperature = payload.temperature, + top_p = payload.top_p, + max_output_tokens = payload.max_output_tokens, + instructions = payload.instructions, + ) + api_monitor.set_reply(monitor_id, text or _monitor_tool_calls_text(tool_calls)) + _monitor_usage(monitor_id, usage_data, _monitor_context_length()) + api_monitor.finish(monitor_id) + return _model_json_response(response) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + finally: + if request_state is not None: + request_state.skip_api_monitor = previous_skip_monitor async def _responses_stream( - payload: ResponsesRequest, messages: list[ChatMessage], request: Request + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, + monitor_id: Optional[str] = None, ): """Handle a streaming Responses API call, emitting named SSE events. @@ -6130,9 +6942,11 @@ async def _responses_stream( try: resp = await _send_stream_with_preheader_cancel(client, req, request = request) if resp is None: + api_monitor.finish(monitor_id, "cancelled") return except httpx.RequestError as e: logger.error("responses stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) yield _sse( "response.failed", { @@ -6158,6 +6972,7 @@ async def _responses_stream( resp.status_code, err_text[:500], ) + api_monitor.fail(monitor_id, err_text[:500]) yield _sse( "response.failed", { @@ -6209,6 +7024,11 @@ async def _responses_stream( if usage: input_tokens = usage.get("prompt_tokens", input_tokens) output_tokens = usage.get("completion_tokens", output_tokens) + _monitor_usage( + monitor_id, + usage, + llama_backend.context_length, + ) continue delta = choices[0].get("delta", {}) or {} @@ -6234,6 +7054,7 @@ async def _responses_stream( for event in _ensure_message_open(): yield event full_text += visible_delta + api_monitor.append_reply(monitor_id, visible_delta) yield _sse( "response.output_text.delta", { @@ -6305,9 +7126,18 @@ async def _responses_stream( if usage: input_tokens = usage.get("prompt_tokens", input_tokens) output_tokens = usage.get("completion_tokens", output_tokens) + _monitor_usage( + monitor_id, + usage, + llama_backend.context_length, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not disconnect_event.is_set(): logger.error("responses stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) status_code = 400 if _classify_llama_generation_error(e) is not None else 500 yield _sse( "response.failed", @@ -6316,8 +7146,10 @@ async def _responses_stream( return except Exception as e: if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return logger.error("responses stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) status_code = 400 if _classify_llama_generation_error(e) is not None else 500 yield _sse( "response.failed", @@ -6347,6 +7179,7 @@ async def _responses_stream( pass if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return final_reasoning, final_visible = extractor.finish() @@ -6368,6 +7201,7 @@ async def _responses_stream( for event in _ensure_message_open(): yield event full_text += final_visible + api_monitor.append_reply(monitor_id, final_visible) yield _sse( "response.output_text.delta", { @@ -6382,6 +7216,7 @@ async def _responses_stream( for event in _ensure_message_open(): yield event full_text = full_reasoning + api_monitor.set_reply(monitor_id, full_text) yield _sse( "response.output_text.delta", { @@ -6528,6 +7363,7 @@ async def _responses_stream( "arguments": st["arguments"], }, } + api_monitor.append_reply(monitor_id, _monitor_call_text(st["name"], st["arguments"])) yield _sse("response.output_item.done", item_done) # response.completed @@ -6548,6 +7384,7 @@ async def _responses_stream( }, }, } + api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) return StreamingResponse( @@ -6579,8 +7416,28 @@ async def openai_responses( raise HTTPException(status_code = 400, detail = "No input provided.") if payload.stream: - return await _responses_stream(payload, messages, request) - return await _responses_non_streaming(payload, messages, request) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = payload.model, + prompt = _monitor_prompt_from_messages(messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + try: + return await _responses_stream(payload, messages, request, monitor_id) + except HTTPException as exc: + detail = exc.detail + if not isinstance(detail, str): + detail = json.dumps(detail, default = str) + api_monitor.fail(monitor_id, detail) + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + return await _responses_non_streaming(payload, messages, request, current_subject) # ===================================================================== @@ -6900,14 +7757,67 @@ async def anthropic_messages( and payload.tool_choice.get("disable_parallel_tool_use") ) + monitor_id = None + monitor_context_length = _monitor_context_length() + request_state = getattr(request, "state", None) + if not getattr(request_state, "skip_api_monitor", False): + request_url = getattr(request, "url", None) + monitor_id = api_monitor.start( + endpoint = getattr(request_url, "path", "/v1/messages"), + method = getattr(request, "method", "POST"), + model = model_name, + prompt = _monitor_prompt_from_messages(openai_messages), + context_length = monitor_context_length, + subject = current_subject, + ) + + async def _monitored_anthropic(coro): + try: + response = await coro + except asyncio.CancelledError: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + return _monitor_anthropic_response( + response, + monitor_id, + monitor_context_length, + cancel_event, + ) + # ── Client-side pass-through path ───────────────────────── if client_tools: openai_tools = openai_client_tools if payload.stream: - return await _anthropic_passthrough_stream( - request, - cancel_event, + return await _monitored_anthropic( + _anthropic_passthrough_stream( + request, + cancel_event, + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + payload.max_tokens, + message_id, + model_name, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + tool_choice = openai_tool_choice, + session_id = payload.session_id, + cancel_id = payload.cancel_id, + disable_parallel_tool_use = _disable_parallel, + ) + ) + return await _monitored_anthropic( + _anthropic_passthrough_non_streaming( llama_backend, openai_messages, openai_tools, @@ -6922,26 +7832,8 @@ async def anthropic_messages( repetition_penalty = repetition_penalty, presence_penalty = presence_penalty, tool_choice = openai_tool_choice, - session_id = payload.session_id, - cancel_id = payload.cancel_id, disable_parallel_tool_use = _disable_parallel, ) - return await _anthropic_passthrough_non_streaming( - llama_backend, - openai_messages, - openai_tools, - temperature, - top_p, - top_k, - payload.max_tokens, - message_id, - model_name, - stop = stop, - min_p = min_p, - repetition_penalty = repetition_penalty, - presence_penalty = presence_penalty, - tool_choice = openai_tool_choice, - disable_parallel_tool_use = _disable_parallel, ) if server_tools: @@ -6949,6 +7841,10 @@ async def anthropic_messages( if bool(getattr(payload, "confirm_tool_calls", False)) and not bool( getattr(payload, "bypass_permissions", False) ): + api_monitor.fail( + monitor_id, + "confirm_tool_calls is not supported for Anthropic Messages server tools.", + ) raise HTTPException( status_code = 400, detail = anthropic_error_body( @@ -7009,22 +7905,26 @@ async def anthropic_messages( ) if payload.stream: - return await _anthropic_tool_stream( - request, - cancel_event, + return await _monitored_anthropic( + _anthropic_tool_stream( + request, + cancel_event, + _run_tool_gen, + message_id, + model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, + openai_tools = openai_tools, + disable_parallel_tool_use = _disable_parallel, + ) + ) + return await _monitored_anthropic( + _anthropic_tool_non_streaming( _run_tool_gen, message_id, model_name, - llama_backend = llama_backend, - openai_messages = openai_messages, - openai_tools = openai_tools, disable_parallel_tool_use = _disable_parallel, ) - return await _anthropic_tool_non_streaming( - _run_tool_gen, - message_id, - model_name, - disable_parallel_tool_use = _disable_parallel, ) # ── No-tool path ────────────────────────────────────────── @@ -7043,19 +7943,23 @@ async def anthropic_messages( ) if payload.stream: - return await _anthropic_plain_stream( - request, - cancel_event, + return await _monitored_anthropic( + _anthropic_plain_stream( + request, + cancel_event, + _run_plain_gen, + message_id, + model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, + ) + ) + return await _monitored_anthropic( + _anthropic_plain_non_streaming( _run_plain_gen, message_id, model_name, - llama_backend = llama_backend, - openai_messages = openai_messages, ) - return await _anthropic_plain_non_streaming( - _run_plain_gen, - message_id, - model_name, ) @@ -7361,7 +8265,7 @@ async def _anthropic_tool_non_streaming( output_tokens = usage.get("completion_tokens", 0), ), ) - return JSONResponse(content = resp.model_dump()) + return _model_json_response(resp) async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): @@ -7403,7 +8307,7 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): output_tokens = usage.get("completion_tokens", 0), ), ) - return JSONResponse(content = resp.model_dump()) + return _model_json_response(resp) # ===================================================================== @@ -7719,12 +8623,11 @@ async def _anthropic_passthrough_non_streaming( backend_ctx = llama_backend.context_length, ) - async with httpx.AsyncClient() as client: - resp = await client.post( - target_url, - json = body, - timeout = _llama_non_streaming_generation_timeout(), - ) + resp = await nonstreaming_client().post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise HTTPException( @@ -7775,7 +8678,7 @@ async def _anthropic_passthrough_non_streaming( output_tokens = usage.get("completion_tokens", 0), ), ) - return JSONResponse(content = resp_obj.model_dump()) + return _model_json_response(resp_obj) # ===================================================================== @@ -8106,7 +9009,13 @@ def _build_openai_passthrough_body( async def _openai_passthrough_stream( - request, cancel_event, llama_backend, payload, model_name, completion_id + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id: Optional[str] = None, ): """Streaming client-side pass-through for /v1/chat/completions. @@ -8150,6 +9059,7 @@ async def _openai_passthrough_stream( except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) if resp is not None: try: await resp.aclose() @@ -8164,6 +9074,7 @@ async def _openai_passthrough_stream( detail = _friendly_error(e), ) if resp is None: + api_monitor.finish(monitor_id, "cancelled") try: await client.aclose() except Exception: @@ -8205,6 +9116,7 @@ async def _openai_passthrough_stream( await client.aclose() except Exception: pass + api_monitor.fail(monitor_id, err_text[:500]) raise _openai_passthrough_error(upstream_status, err_text) async def _stream(): @@ -8218,6 +9130,7 @@ async def _openai_passthrough_stream( disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, cancel_event) ) + monitor_done = False try: lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( @@ -8237,21 +9150,39 @@ async def _openai_passthrough_stream( # relayed byte-for-byte. if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: raw_line = _cap_parallel_tool_calls_sse_line(raw_line) + monitor_event = _monitor_openai_sse_line( + monitor_id, + raw_line, + llama_backend.context_length, + ) # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" - if raw_line[6:].strip() == "[DONE]": + if monitor_event == "done" or raw_line[6:].strip() == "[DONE]": + monitor_done = True break + if not monitor_done: + api_monitor.finish( + monitor_id, + "cancelled" if cancel_event.is_set() else "completed", + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): # Watcher closed resp on cancel. Emit nothing extra; the client # initiated the cancel or already disconnected. if not cancel_event.is_set(): + api_monitor.fail(monitor_id, "Stream interrupted") raise + api_monitor.finish(monitor_id, "cancelled") except Exception as e: if cancel_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return # 200 headers already flushed; errors must go in the SSE body. logger.error("openai passthrough stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: @@ -8294,7 +9225,12 @@ async def _openai_passthrough_stream( raise -async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): +async def _openai_passthrough_non_streaming( + llama_backend, + payload, + model_name, + monitor_id: Optional[str] = None, +): """Non-streaming client-side pass-through for /v1/chat/completions. Returns llama-server's JSON response verbatim so the client sees the native @@ -8311,17 +9247,20 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): ) while True: try: - async with httpx.AsyncClient() as client: - resp = await client.post( - target_url, - json = body, - timeout = _llama_non_streaming_generation_timeout(), - ) + resp = await nonstreaming_client().post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. Surface the # same friendly message the sync chat path emits so operators don't see # a bare 500 with no diagnostic. logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) raise HTTPException( status_code = 502, detail = _friendly_error(e), @@ -8337,6 +9276,7 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): ): _truncate_budget -= 1 continue + api_monitor.fail(monitor_id, resp.text[:500]) raise _openai_passthrough_error(resp.status_code, resp.text) # The guided-decoding fence wraps each choice's JSON content in a @@ -8357,6 +9297,7 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): "openai passthrough non-streaming: response not JSON, relaying raw: %s", exc, ) + api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") changed = False @@ -8394,5 +9335,9 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): # Nothing mutated: relay the upstream bytes verbatim, skipping a redundant # parse + re-serialize round-trip. if not changed: + _monitor_openai_chunk(monitor_id, data, llama_backend.context_length) + api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") + _monitor_openai_chunk(monitor_id, data, llama_backend.context_length) + api_monitor.finish(monitor_id) return JSONResponse(content = data) diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 3aae6f4209..5559cc0404 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -55,6 +55,9 @@ class LlamaUpdateStatusResponse(BaseModel): source_build: bool = Field( False, description = "True when there is no marker (source build) but a prebuilt is offered." ) + update_size_bytes: Optional[int] = Field( + None, description = "Download size of the prebuilt Update would fetch, in bytes." + ) job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 0ed8de5532..92a87ce045 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -33,6 +33,7 @@ from core.inference.anthropic_compat import ( AnthropicStreamEmitter, AnthropicPassthroughEmitter, ) +from core.inference.api_monitor import ApiMonitor from routes.inference import ( _build_tool_action_nudge, _normalize_anthropic_openai_images, @@ -40,6 +41,7 @@ from routes.inference import ( _anthropic_requested_studio_tools, _anthropic_passthrough_stream, _anthropic_tool_non_streaming, + _monitor_anthropic_sse_line, anthropic_messages, ) from state.tool_policy import reset_tool_policy, set_tool_policy @@ -50,6 +52,46 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/messages", + method = "POST", + model = "m", + prompt = "hi", + ) + + for payload in ( + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_1", + "name": "lookup", + "input": {}, + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "input_json_delta", + "partial_json": '{"query":"weather"}', + }, + }, + {"type": "content_block_stop", "index": 0}, + ): + _monitor_anthropic_sse_line(monitor_id, f"data: {json.dumps(payload)}") + + entry = monitor.get(monitor_id) + assert entry is not None + assert entry["reply"] == 'Tool call: lookup\nInput: {"query":"weather"}' + + # ===================================================================== # Tool nudge tests # ===================================================================== @@ -1385,7 +1427,7 @@ def _mock_backend(monkeypatch, **overrides): def _gen_plain(**kwargs): calls.append(("plain", kwargs)) - yield {"type": "content", "text": "ok"} + yield "ok" def _gen_tools(**kwargs): calls.append(("tools", kwargs)) @@ -1396,6 +1438,8 @@ def _mock_backend(monkeypatch, **overrides): is_vision = False, supports_tools = True, model_identifier = "test-model", + context_length = 4096, + count_chat_tokens = lambda *args, **kwargs: 2, generate_chat_completion = _gen_plain, generate_chat_completion_with_tools = _gen_tools, calls = calls, @@ -1426,6 +1470,112 @@ def _reset_policy(): class TestAnthropicMessagesToolRouting: + class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/messages") + method = "POST" + + async def is_disconnected(self): + return False + + @staticmethod + def _consume_response(response): + async def _consume(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + return chunks + + return _drive(_consume()) + + def test_plain_non_streaming_records_api_monitor_entry(self, monkeypatch): + import routes.inference as inf_mod + + _mock_backend(monkeypatch, context_length = 2048) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + payload = _basic_payload() + + response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + assert response.status_code == 200 + [entry] = monitor.snapshot() + assert entry["endpoint"] == "/v1/messages" + assert entry["status"] == "completed" + assert entry["model"] == "test-model" + assert entry["prompt_preview"] == "user: hi" + assert entry["reply_preview"] == "ok" + assert entry["context_length"] == 2048 + assert monitor.active_count() == 0 + + def test_tool_use_non_streaming_records_api_monitor_reply(self, monkeypatch): + import routes.inference as inf_mod + + def _gen_tools(**_kwargs): + yield { + "type": "tool_start", + "tool_call_id": "call_1", + "tool_name": "lookup", + "arguments": {"query": "weather"}, + } + + _mock_backend( + monkeypatch, + context_length = 2048, + generate_chat_completion_with_tools = _gen_tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + payload = _basic_payload( + enable_tools = True, + tools = [{"type": "web_search_20250305", "name": "web_search"}], + ) + + response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + assert response.status_code == 200 + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply_preview"] == 'Tool call: lookup({"query": "weather"})' + + def test_plain_streaming_records_active_and_completed_monitor_entry(self, monkeypatch): + import routes.inference as inf_mod + + _mock_backend(monkeypatch, context_length = 2048) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + payload = _basic_payload(stream = True) + + response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + assert monitor.active_count() == 1 + self._consume_response(response) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply_preview"] == "ok" + assert entry["prompt_tokens"] == 2 + assert entry["context_length"] == 2048 + assert monitor.active_count() == 0 + + def test_plain_streaming_pre_response_cancel_finalizes_monitor(self, monkeypatch): + import routes.inference as inf_mod + + async def _cancelled_before_response(*_args, **_kwargs): + raise asyncio.CancelledError() + + _mock_backend(monkeypatch, context_length = 2048) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_anthropic_plain_stream", _cancelled_before_response) + payload = _basic_payload(stream = True) + + with pytest.raises(asyncio.CancelledError): + _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + def test_mixed_server_and_client_tools_rejected_with_400(self, monkeypatch): _mock_backend(monkeypatch) payload = _basic_payload( diff --git a/studio/backend/tests/test_api_monitor.py b/studio/backend/tests/test_api_monitor.py new file mode 100644 index 0000000000..56bc404350 --- /dev/null +++ b/studio/backend/tests/test_api_monitor.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from core.inference.api_monitor import ApiMonitor, _trim + + +def test_api_monitor_tracks_reply_usage_and_context(): + monitor = ApiMonitor(max_entries = 3) + + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "local-model", + prompt = "user: hello", + context_length = 100, + ) + monitor.append_reply(entry_id, "hi") + monitor.append_reply(entry_id, " there") + monitor.set_usage( + entry_id, + prompt_tokens = 4, + completion_tokens = 6, + ) + monitor.finish(entry_id) + + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hi there" + assert entry["total_tokens"] == 10 + assert entry["context_usage"] == 0.1 + assert entry["duration_ms"] is not None + + +def test_api_monitor_summary_omits_full_prompt_and_reply(): + monitor = ApiMonitor(max_entries = 3) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "local-model", + prompt = "p" * 500, + ) + monitor.set_reply(entry_id, "r" * 500) + + [summary] = monitor.snapshot(include_details = False) + assert "prompt" not in summary + assert "reply" not in summary + assert summary["prompt_preview"].endswith("...") + assert summary["reply_preview"].endswith("...") + assert summary["prompt_truncated"] is True + assert summary["reply_truncated"] is True + + detail = monitor.get(entry_id) + assert detail is not None + assert detail["prompt"] == "p" * 500 + assert detail["reply"] == "r" * 500 + + +def test_api_monitor_filters_entries_by_subject(): + monitor = ApiMonitor(max_entries = 3) + alice = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "alice prompt", + subject = "alice", + ) + bob = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "bob prompt", + subject = "bob", + ) + monitor.finish(bob) + + alice_entries = monitor.snapshot(subject = "alice") + assert [entry["id"] for entry in alice_entries] == [alice] + assert monitor.get(bob, subject = "alice") is None + assert monitor.get(bob, subject = "bob")["id"] == bob + assert monitor.active_count(subject = "alice") == 1 + assert monitor.active_count(subject = "bob") == 0 + + +def test_api_monitor_keeps_bounded_recent_history(): + monitor = ApiMonitor(max_entries = 2) + + first = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "first", + ) + second = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "second", + ) + third = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "third", + ) + monitor.finish(first) + monitor.finish(second) + monitor.finish(third) + + entries = monitor.snapshot() + ids = [entry["id"] for entry in entries] + assert ids[0] == third + assert [entry["prompt"] for entry in entries] == ["third", "second"] + assert first not in ids + assert monitor.active_count() == 0 + + +def test_api_monitor_keeps_running_entries_beyond_history_limit(): + monitor = ApiMonitor(max_entries = 1) + + running = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "running", + ) + for prompt in ("done-1", "done-2", "done-3"): + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = prompt, + ) + monitor.finish(entry_id) + + entries = monitor.snapshot() + ids = [entry["id"] for entry in entries] + assert running in ids + assert monitor.active_count() == 1 + + monitor.finish(running) + [entry] = monitor.snapshot() + assert entry["id"] == running + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + +def test_api_monitor_finish_is_idempotent(): + monitor = ApiMonitor(max_entries = 2) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "hi", + ) + monitor.finish(entry_id) + first = monitor.snapshot()[0] + monitor.finish(entry_id) + second = monitor.snapshot()[0] + assert first["finished_at"] == second["finished_at"] + assert first["duration_ms"] == second["duration_ms"] + + +def test_api_monitor_preserves_authoritative_total_tokens(): + monitor = ApiMonitor(max_entries = 2) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "hi", + ) + monitor.set_usage( + entry_id, + prompt_tokens = 10, + completion_tokens = 20, + total_tokens = 33, + ) + # A later partial chunk omitting `total_tokens` must not clobber 33. + monitor.set_usage(entry_id, prompt_tokens = 11) + assert monitor.snapshot()[0]["total_tokens"] == 33 + + +def test_api_monitor_recomputes_derived_total_tokens(): + monitor = ApiMonitor(max_entries = 2) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "hi", + ) + monitor.set_usage(entry_id, prompt_tokens = 10) + assert monitor.snapshot()[0]["total_tokens"] == 10 + + monitor.set_usage(entry_id, completion_tokens = 20) + entry = monitor.snapshot()[0] + assert entry["prompt_tokens"] == 10 + assert entry["completion_tokens"] == 20 + assert entry["total_tokens"] == 30 + + +def test_api_monitor_duration_non_negative_under_clock_step(monkeypatch): + import core.inference.api_monitor as m + + fake_now = [1000.0] + monkeypatch.setattr(m.time, "time", lambda: fake_now[0]) + monitor = ApiMonitor(max_entries = 1) + entry_id = monitor.start( + endpoint = "/x", + method = "POST", + model = "m", + prompt = "hi", + ) + fake_now[0] = 500.0 + monitor.finish(entry_id) + assert monitor.snapshot()[0]["duration_ms"] >= 0 + + +def test_api_monitor_trim_guards_tiny_limit(): + assert _trim("abcdefgh", 2) == ".." + assert _trim("abcdefgh", 0) == "" + assert _trim("abcdefgh", 3) == "..." + assert _trim("abcdefgh", 4) == "a..." + assert _trim("abcdefgh", 100) == "abcdefgh" + + +def test_api_monitor_append_reply_caps_without_regrowing(): + import core.inference.api_monitor as m + + monitor = ApiMonitor(max_entries = 1) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "go", + ) + monitor.append_reply(entry_id, "x" * (m._MAX_REPLY_CHARS + 500)) + capped = monitor.snapshot()[0]["reply"] + assert len(capped) == m._MAX_REPLY_CHARS and capped.endswith("...") + + # Chunks past the cap must not change or grow the stored preview. + monitor.append_reply(entry_id, "y" * 1000) + assert monitor.snapshot()[0]["reply"] == capped + + +def test_api_monitor_append_reply_exact_cap_then_more_marks_truncated(): + import core.inference.api_monitor as m + + monitor = ApiMonitor(max_entries = 1) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "go", + ) + # A reply landing exactly on the cap has no "..." marker yet. + monitor.append_reply(entry_id, "x" * m._MAX_REPLY_CHARS) + assert not monitor.snapshot()[0]["reply"].endswith("...") + # One more chunk must record the truncation, not silently freeze. + monitor.append_reply(entry_id, "y") + reply = monitor.snapshot()[0]["reply"] + assert len(reply) == m._MAX_REPLY_CHARS and reply.endswith("...") diff --git a/studio/backend/tests/test_api_perf_serialization.py b/studio/backend/tests/test_api_perf_serialization.py new file mode 100644 index 0000000000..f5ad53306d --- /dev/null +++ b/studio/backend/tests/test_api_perf_serialization.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""_model_json_response produces the same body as JSONResponse(model.model_dump()).""" + +import asyncio +import json +from typing import Optional + +from fastapi.responses import JSONResponse +from pydantic import BaseModel + +import routes.inference as inference_route +from core.inference import llama_http + + +class _Usage(BaseModel): + prompt_tokens: int = 3 + completion_tokens: int = 5 + details: Optional[dict] = None + + +class _Choice(BaseModel): + index: int = 0 + text: str = "hello" + logprobs: Optional[dict] = None + + +class _Resp(BaseModel): + id: str = "chatcmpl-abc" + object: str = "chat.completion" + created: int = 1700000000 + model: str = "unsloth/SmolLM2-135M-Instruct-GGUF" + choices: list[_Choice] = [_Choice()] + usage: _Usage = _Usage() + system_fingerprint: Optional[str] = None + + +def _old_body(model) -> bytes: + # What the previous code emitted: dict -> Starlette json.dumps. + return JSONResponse(content = model.model_dump()).body + + +def test_body_matches_old_jsonresponse(): + model = _Resp() + resp = inference_route._model_json_response(model) + # Same decoded JSON (key order is irrelevant once parsed), nulls preserved. + assert json.loads(resp.body) == json.loads(_old_body(model)) + assert json.loads(resp.body)["system_fingerprint"] is None # null kept, not dropped + + +def test_media_type_and_status(): + resp = inference_route._model_json_response(_Resp(), status_code = 200) + assert resp.media_type == "application/json" + assert resp.status_code == 200 + err = inference_route._model_json_response(_Resp(), status_code = 503) + assert err.status_code == 503 + + +def test_pooled_client_reused_within_loop_and_recreated_after_close(): + async def _scenario(): + a = llama_http.nonstreaming_client() + b = llama_http.nonstreaming_client() + assert a is b # reused within one loop + await llama_http.aclose() + assert a.is_closed + c = llama_http.nonstreaming_client() # must not return the closed client + assert c is not a and not c.is_closed + await llama_http.aclose() + + asyncio.run(_scenario()) + + +def test_pooled_client_is_per_event_loop(): + clients = [] + # Each asyncio.run uses a fresh loop; the pooled client must not leak across. + for _ in range(2): + + async def _grab(): + clients.append(llama_http.nonstreaming_client()) + await llama_http.aclose() + + asyncio.run(_grab()) + assert clients[0] is not clients[1] diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py new file mode 100644 index 0000000000..42c400383e --- /dev/null +++ b/studio/backend/tests/test_compute_buffer.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for ``_estimate_compute_buffer_bytes``: it scales with ``--parallel``, +tensor exceeds pipeline, and it is a safe upper bound on the allocations measured +on real hardware (Qwen3.6-27B-MTP: parallel 1/2/4/8 -> 36/492/1388/3220 MiB single +GPU, ~600 MiB/device tensor). No GPU, subprocess, or GGUF I/O.""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) +# httpx -- only stub when the real library is missing. Unconditional stubbing +# shadows HTTPError/Response that huggingface_hub.errors imports at load time, +# silently breaking the transformers introspection tier in tests collected after +# this one (the stub leaks via sys.modules for the whole session). +try: + import httpx as _httpx_real # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + sys.modules["httpx"] = _httpx_stub + +from core.inference.llama_cpp import LlamaCppBackend + +MIB = 1024 * 1024 + + +def _backend(vocab = 248320, embd = 5120): + """Backend with just the dims the compute-buffer estimate reads.""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._vocab_size = vocab + b._embedding_length = embd + return b + + +# Measured ground truth (MiB) the estimate must upper-bound. +_PIPELINE_MEASURED = {1: 36, 2: 492, 4: 1388, 8: 3220} +_TENSOR_MEASURED_PER_DEVICE = 600 + + +class TestSafeUpperBound: + """The estimate must be >= every measured allocation (never under-reserve).""" + + @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items())) + def test_pipeline_upper_bounds_measured(self, parallel, measured): + est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB + assert est >= measured, f"under-reserved at parallel={parallel}: {est:.0f} < {measured}" + + @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items())) + def test_pipeline_not_wildly_over(self, parallel, measured): + # Stay within ~2x of measured so we don't waste context (the point of + # replacing the flat reserve). parallel=1 is tiny in absolute terms. + est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB + assert est <= max(measured * 2.0, 128) + + def test_tensor_upper_bounds_measured(self): + est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB + assert est >= _TENSOR_MEASURED_PER_DEVICE + + def test_tensor_far_below_old_flat_reserve(self): + # The whole point: deterministic estimate << flat 5120 for this model. + est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB + assert est < LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + + +class TestScaling: + def test_grows_with_serving_slots(self): + b = _backend() + vals = [b._estimate_compute_buffer_bytes(n_parallel = p) for p in (1, 2, 4, 8)] + assert vals == sorted(vals) and vals[0] < vals[-1] + + def test_parallel_1_is_small(self): + # Single-token decode: a few tens of MiB, not gigabytes. + est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1) / MIB + assert est < 128 + + def test_tensor_exceeds_pipeline_at_same_parallel(self): + b = _backend() + pipe = b._estimate_compute_buffer_bytes(n_parallel = 1) + tens = b._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) + assert tens > pipe + + def test_scales_with_vocab(self): + small = _backend(vocab = 32000)._estimate_compute_buffer_bytes(n_parallel = 4) + big = _backend(vocab = 256000)._estimate_compute_buffer_bytes(n_parallel = 4) + assert big > small + + def test_scales_with_ubatch(self): + b = _backend() + lo = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 256) + hi = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 1024) + assert hi > lo + + +class TestFallback: + def test_zero_when_vocab_missing(self): + assert _backend(vocab = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0 + + def test_zero_when_embd_missing(self): + assert _backend(embd = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0 + + def test_zero_lets_tensor_plan_use_flat_fallback(self): + # When dims are missing, _plan_tensor_parallel must fall back to the flat + # reserve (defense-in-depth) rather than reserving 0 and OOMing. + b = _backend(vocab = None, embd = None) + b._n_layers = None # can't estimate KV -> floors ctx, still returns a plan + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, 48000)], 8 * 1024**3, 8192) + assert gi == [0, 1] # both GPUs usable under the flat fallback + + +class TestParallel1Default: + """At Studio's default --parallel 1 the buffer is negligible in pipeline.""" + + def test_default_n_parallel(self): + est = _backend()._estimate_compute_buffer_bytes() / MIB + assert est < 128 diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index d3d4387720..a5be07f8e3 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -13,6 +13,7 @@ from typing import Iterable, Mapping from utils.models.gguf_metadata import ( is_mmproj_by_metadata, pairing_score, + read_gguf_context_length, read_gguf_general_metadata, read_mmproj_audio_capability, ) @@ -21,6 +22,7 @@ from utils.models.gguf_metadata import ( _GGUF_MAGIC = 0x46554747 _VTYPE_STRING = 8 _VTYPE_UINT32 = 4 +_VTYPE_UINT64 = 10 _VTYPE_ARRAY = 9 _VTYPE_BOOL = 7 @@ -38,6 +40,10 @@ def _enc_kv_uint32(key: str, value: int) -> bytes: return _enc_string(key) + struct.pack(" bytes: + return _enc_string(key) + struct.pack(" bytes: return _enc_string(key) + struct.pack(" Path: """Minimal GGUF: header + KV body, no tensors.""" extra_uint32 = extra_uint32 or {} + extra_uint64 = extra_uint64 or {} extra_string_arrays = extra_string_arrays or {} extra_bools = extra_bools or {} kv_count = ( - len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + len(extra_bools) + len(general_strings) + + len(extra_uint32) + + len(extra_uint64) + + len(extra_string_arrays) + + len(extra_bools) ) body = b"" for k, v in general_strings.items(): body += _enc_kv_string(k, v) for k, v in extra_uint32.items(): body += _enc_kv_uint32(k, v) + for k, v in extra_uint64.items(): + body += _enc_kv_uint64(k, v) for k, v in extra_string_arrays.items(): body += _enc_kv_string_array(k, v) for k, v in extra_bools.items(): @@ -100,6 +114,66 @@ def test_returns_none_for_non_gguf(tmp_path: Path): assert read_gguf_general_metadata(str(p)) is None +def test_context_length_none_for_missing_file(tmp_path: Path): + assert read_gguf_context_length(str(tmp_path / "nope.gguf")) is None + + +def test_context_length_none_for_non_gguf(tmp_path: Path): + p = tmp_path / "garbage.gguf" + p.write_bytes(b"not a gguf file at all, just bytes") + assert read_gguf_context_length(str(p)) is None + + +def test_context_length_read_from_arch_namespaced_key(tmp_path: Path): + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"llama.context_length": 4096, "llama.block_count": 32}, + ) + assert read_gguf_context_length(str(p)) == 4096 + + +def test_context_length_none_when_absent(tmp_path: Path): + # Architecture present but no .context_length key. + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"llama.block_count": 32}, + ) + assert read_gguf_context_length(str(p)) is None + + +def test_context_length_ignores_foreign_arch_key(tmp_path: Path): + # A context_length under a different arch namespace must not match. + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"qwen2.context_length": 8192}, + ) + assert read_gguf_context_length(str(p)) is None + + +def test_context_length_read_from_uint64(tmp_path: Path): + # Some models store context_length as a uint64 (vtype 10). + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "qwen3"}, + extra_uint64 = {"qwen3.context_length": 262144}, + ) + assert read_gguf_context_length(str(p)) == 262144 + + +def test_context_length_zero_treated_as_absent(tmp_path: Path): + # A zero/garbage ceiling must read as None so the UI can't build a slider + # with max < min. + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"llama.context_length": 0}, + ) + assert read_gguf_context_length(str(p)) is None + + def test_extracts_general_string_fields(tmp_path: Path): p = _write_synthetic_gguf( tmp_path / "model.gguf", diff --git a/studio/backend/tests/test_gguf_xet_fallback_integration.py b/studio/backend/tests/test_gguf_xet_fallback_integration.py new file mode 100644 index 0000000000..cbcc73847e --- /dev/null +++ b/studio/backend/tests/test_gguf_xet_fallback_integration.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Integration: GGUF Chat-Mode downloads route through the Xet->HTTP helper, +preserving cancellation and the best-effort companion contract. No GPU, no +network, no real subprocess (the helper is patched). +""" + +from __future__ import annotations + +import sys +import threading +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Heavy-dep stubbing; prefer the real structlog so a bare stub never leaks to +# later modules that log at import time. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +try: + import structlog # noqa: F401 +except ImportError: + sys.modules["structlog"] = _types.ModuleType("structlog") +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + "HTTPStatusError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Request = type("Request", (), {}) + _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **k: None}) + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **k: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + sys.modules.setdefault("httpx", _httpx_stub) + +from huggingface_hub import constants as hf_constants + +from core.inference.llama_cpp import LlamaCppBackend +from utils.hf_xet_fallback import DownloadStallError + +REPO = "unsloth/vision-GGUF" + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + return tmp_path + + +def _build_cache( + root: Path, + repo_id: str, + files: dict[str, int], + sha: str = "a" * 40, +) -> Path: + repo_dir = root / f"models--{repo_id.replace('/', '--')}" + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + snap = repo_dir / "snapshots" / sha + snap.mkdir(parents = True, exist_ok = True) + for rel, size in files.items(): + (snap / rel).write_bytes(b"\0" * size) + return snap + + +def test_companion_routes_through_helper(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + captured = {} + + def fake_helper( + repo_id, + filename, + token = None, + **kwargs, + ): + captured["filename"] = filename + captured["cancel_event"] = kwargs.get("cancel_event") + return f"/fake/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out == "/fake/mmproj-vision-F16.gguf" + # _cancel_event must be threaded through so /unload can abort the download. + assert captured["cancel_event"] is backend._cancel_event + + +def test_companion_swallows_terminal_stall_to_none(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + + def stalling_helper( + repo_id, + filename, + token = None, + **kwargs, + ): + raise DownloadStallError("both transports stalled") + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", stalling_helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out is None, "a companion download is best-effort; a terminal stall must not raise" + + +def test_companion_cancelled_skips_download(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + backend._cancel_event.set() + called = {"n": 0} + + def helper( + repo_id, + filename, + token = None, + **kwargs, + ): + called["n"] += 1 + return "/should-not-happen" + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out is None + assert called["n"] == 0, "a cancelled load must not start a companion download" diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py new file mode 100644 index 0000000000..39ecebd328 --- /dev/null +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP +transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on. +CPU-only, no network, no real subprocess (the per-attempt download seam is +monkeypatched). +""" + +from __future__ import annotations + +import subprocess +import sys +import threading +import time +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub heavy/unavailable deps before importing the module under test. Use the +# real structlog when present; a bare stub left in sys.modules would break later +# modules that log at import time. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +try: + import structlog # noqa: F401 +except ImportError: + sys.modules["structlog"] = _types.ModuleType("structlog") + +import huggingface_hub +from huggingface_hub import constants as hf_constants + +import utils.hf_xet_fallback as xf + + +# --------------------------------------------------------------------------- # +# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total. +# --------------------------------------------------------------------------- # +REPO = "ztest/xet-watchdog" + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + return tmp_path + + +def _blobs_dir(root: Path, repo_id: str = REPO) -> Path: + d = root / f"models--{repo_id.replace('/', '--')}" / "blobs" + d.mkdir(parents = True, exist_ok = True) + return d + + +def _wait( + predicate, + timeout: float = 2.0, + step: float = 0.02, +) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(step) + return predicate() + + +def test_constant_incomplete_fires_stall(hf_cache): + blobs = _blobs_dir(hf_cache) + (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 + ) + try: + assert _wait( + lambda: len(calls) >= 1, timeout = 3.0 + ), "watchdog never fired on a constant-size .incomplete" + finally: + stop.set() + assert "stalled" in calls[0].lower() + + +def test_growing_incomplete_never_stalls(hf_cache): + blobs = _blobs_dir(hf_cache) + part = blobs / "growing.incomplete" + part.write_bytes(b"\0" * 1024) + + grow_stop = threading.Event() + + def _grow(): + size = 1024 + while not grow_stop.wait(0.05): + size += 4096 + part.write_bytes(b"\0" * size) + + grower = threading.Thread(target = _grow, daemon = True) + grower.start() + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 + ) + try: + time.sleep(1.0) # well past stall_timeout, but bytes keep growing + assert calls == [], "watchdog fired despite continuous progress" + finally: + stop.set() + grow_stop.set() + + +def test_no_incomplete_never_stalls(hf_cache): + blobs = _blobs_dir(hf_cache) + (blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 + ) + try: + time.sleep(0.8) + assert calls == [], "watchdog fired with no active .incomplete" + finally: + stop.set() + + +def test_stall_fires_at_most_once(hf_cache): + blobs = _blobs_dir(hf_cache) + (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048) + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2 + ) + try: + assert _wait(lambda: len(calls) >= 1, timeout = 3.0) + time.sleep(0.6) # keep ticking; must not fire again + assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1" + finally: + stop.set() + + +def test_get_state_empty_cache(hf_cache): + assert xf.get_hf_download_state([REPO]) == (0, False) + + +def test_get_state_absent_cache_root(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache")) + assert xf.get_hf_download_state([REPO]) == (0, False) + + +def test_get_state_skips_local_paths(hf_cache): + # Filesystem paths are not HF repo IDs and must be ignored without error. + assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False) + + +def test_get_state_sparse_aware(hf_cache): + blobs = _blobs_dir(hf_cache) + sparse = blobs / "sparse.incomplete" + with open(sparse, "wb") as f: + f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks + st = sparse.stat() + if getattr(st, "st_blocks", 0) == 0: + pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable") + total, has_incomplete = xf.get_hf_download_state([REPO]) + assert has_incomplete is True + assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks" + + +# --------------------------------------------------------------------------- # +# Transport policy: cached short-circuit, cancel, error propagation, and the +# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn. +# --------------------------------------------------------------------------- # +DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" + + +@pytest.fixture(autouse = True) +def _no_real_cache_hit(monkeypatch): + """Default: the cached probe misses; tests override it to force a hit.""" + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + + +class _FakeAttempt: + """Records calls to the download seam and returns scripted results.""" + + def __init__(self, results): + self._results = list(results) + self.calls = [] + + def __call__( + self, + repo_id, + filename, + token, + *, + repo_type, + disable_xet, + cancel_event, + stall_timeout, + interval, + grace_period, + on_status, + ): + self.calls.append( + _types.SimpleNamespace( + repo_id = repo_id, + filename = filename, + disable_xet = disable_xet, + repo_type = repo_type, + ) + ) + return self._results[len(self.calls) - 1] + + +def _install(monkeypatch, results): + fake = _FakeAttempt(results) + monkeypatch.setattr(xf, "_run_download_attempt", fake) + return fake + + +def test_cached_file_short_circuits(monkeypatch, tmp_path): + cached = tmp_path / "cached.gguf" + cached.write_bytes(b"\0" * 8) + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached)) + fake = _install(monkeypatch, []) # must not be called + + out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert out == str(cached) + assert fake.calls == [], "spawned a download for an already-cached file" + + +def test_cancel_before_start_raises_no_attempt(monkeypatch): + fake = _install(monkeypatch, []) + ev = threading.Event() + ev.set() + with pytest.raises(RuntimeError, match = "Cancelled"): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev) + assert fake.calls == [] + + +def test_nonstall_error_propagates_without_fallback(monkeypatch): + fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")]) + with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback" + assert fake.calls[0].disable_xet is False + + +def test_immediate_success_uses_xet_only(monkeypatch): + prepared = [] + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", + lambda *a, **k: prepared.append(a), + ) + fake = _install(monkeypatch, [("ok", "/cache/model.gguf")]) + out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert out == "/cache/model.gguf" + assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False + assert prepared == [], "no cache prep should run when Xet succeeds first try" + + +def test_stall_then_http_fallback_succeeds(monkeypatch): + prepared = [] + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", + lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), + ) + fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")]) + + out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert out == "/cache/model.gguf" + assert len(fake.calls) == 2 + assert fake.calls[0].disable_xet is False # Xet first + assert fake.calls[1].disable_xet is True # HTTP fallback + assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry" + + +def test_second_stall_raises_download_stall_error(monkeypatch): + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None + ) + fake = _install(monkeypatch, [("stall", None), ("stall", None)]) + with pytest.raises(xf.DownloadStallError): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert len(fake.calls) == 2 + + +def test_cancelled_midattempt_raises_no_fallback(monkeypatch): + fake = _install(monkeypatch, [("cancelled", None)]) + with pytest.raises(RuntimeError, match = "Cancelled"): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert len(fake.calls) == 1 + + +def test_per_file_independent_fallback(monkeypatch): + """A stalled shard falls back; a sibling shard that succeeds does not.""" + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None + ) + fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")]) + assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a" + assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b" + assert [c.disable_xet for c in fake.calls] == [False, False, True] + + +# --------------------------------------------------------------------------- # +# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect +# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it). +# --------------------------------------------------------------------------- # +def _safe_path() -> str: + import os + return os.environ.get("PATH", "") + + +def test_disable_xet_constant_set_in_fresh_interpreter(): + code = ( + "from huggingface_hub import constants as c; " + "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)" + ) + proc = subprocess.run( + [sys.executable, "-c", code], + env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()}, + capture_output = True, + text = True, + ) + assert proc.returncode == 0, ( + f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True " + f"(rc={proc.returncode}): {proc.stderr}" + ) + + +def test_default_leaves_xet_enabled(): + code = ( + "from huggingface_hub import constants as c; " + "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)" + ) + proc = subprocess.run( + [sys.executable, "-c", code], + env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET + capture_output = True, + text = True, + ) + assert proc.returncode == 0, ( + f"without the env var, constants.HF_HUB_DISABLE_XET was not False " + f"(rc={proc.returncode}): {proc.stderr}" + ) diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index cd834b345b..27e9d0f57a 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -71,7 +71,7 @@ except ImportError: ) sys.modules["httpx"] = _httpx_stub -from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend # Helpers @@ -1484,8 +1484,8 @@ class TestServerFlags: assert fitted < 32_768 def test_fit_mtp_engaged_returns_smaller_or_equal_context(self): - # MTP budget is 0.85 of available, non-MTP is 0.90; on a tight - # budget MTP must yield <= non-MTP. + # Flat MTP fallback budget is _CTX_FIT_VRAM_FRACTION - 0.05; non-MTP is + # the full fraction. On a tight budget MTP must yield <= non-MTP. b = self._gqa_backend() common = dict( requested_ctx = 32_768, @@ -1518,7 +1518,7 @@ class TestServerFlags: kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) assert kv_full > kv_default # Budget = model + kv_default (rounded up) -- swa_full must not fit. - budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / 0.90 + 1 + budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / _CTX_FIT_VRAM_FRACTION + 1 fitted_default = b._fit_context_to_vram( requested_ctx = ctx, available_mib = int(budget_mib), diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index abacc6e953..9866fc4ae1 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -67,7 +67,11 @@ _httpx_stub.Client = type( ) sys.modules.setdefault("httpx", _httpx_stub) -from core.inference.llama_cpp import LlamaCppBackend, classify_gpu_offload_lines +from core.inference.llama_cpp import ( + _CTX_FIT_VRAM_FRACTION, + LlamaCppBackend, + classify_gpu_offload_lines, +) from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx @@ -171,7 +175,7 @@ def _drive( ) kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: + if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -664,3 +668,39 @@ class TestClassifyGpuOffload: def test_module_level_no_signal_returns_none(self): assert classify_gpu_offload_lines(["INFO starting server"]) is None + + +def test_select_gpus_ranks_by_usable_not_raw_free(): + # 80 GB card (30 GB free -> 25.9 GB usable) vs 32 GB card (29 GB free -> 27.4 + # GB usable). A 27 GB model fits the 32 GB card alone; raw-free ranking would + # try the 80 GB card first and split across both. Usable ranking picks [1]. + gpus = [(0, 30000), (1, 29000)] + totals = {0: 81920, 1: 32607} + model = int(27000 * 1024 * 1024) + idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals) + assert idxs == [1] and use_fit is False + + +def test_select_gpus_reserves_per_device_overhead(): + # Two 16 GB cards, ~15181 MiB usable each at 0.95 -> 30362 MiB pooled. A 30000 + # MiB model fits the pool with no per-device overhead, but a layer split also + # pays ~1 GiB/extra-GPU; that pushes the 2-GPU need to 31024 MiB > pool, so a + # pin would OOM -> must fall back to --fit. Single-GPU fits add no overhead + # (Finding F1, the explicit/file-size multi-GPU pin gap). + gpus = [(0, 16000), (1, 16000)] + totals = {0: 16384, 1: 16384} + gib = 1024 * 1024 * 1024 + model = int(30000 * 1024 * 1024) + idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals) + assert idxs == [0, 1] and use_fit is False # fits 2 GPUs without overhead + idxs2, use_fit2 = LlamaCppBackend._select_gpus( + model, gpus, total_by_idx = totals, per_device_overhead_bytes = gib + ) + assert idxs2 is None and use_fit2 is True # overhead tips it past the pool + # A single-GPU fit is unchanged by the overhead (k=1 adds nothing). + small = int(15000 * 1024 * 1024) + a, _ = LlamaCppBackend._select_gpus(small, gpus, total_by_idx = totals) + b, _ = LlamaCppBackend._select_gpus( + small, gpus, total_by_idx = totals, per_device_overhead_bytes = gib + ) + assert a == [0] and b == [0] diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index f90c4ba0e7..2a2e113585 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -520,3 +520,114 @@ def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path) assert info["latest_tag"] == "b9596-mix-aaa" assert info["behind"] is True assert info["stale"] is True + + +# update_download_size_bytes (banner download-size lookup). + + +def _patch_assets(monkeypatch, mapping): + """Stub latest_release_assets with a per-repo {asset_name: size} lookup.""" + monkeypatch.setattr( + fr, + "latest_release_assets", + lambda repo, *, force_refresh = False: mapping.get(repo), + ) + + +def test_update_size_unsloth_prebuilt_exact_match(monkeypatch): + # The unsloth fork's own bundle (app--): the want= exact match + # on app-- wins. + marker = { + "asset": "app-b9190-linux-x64-cuda13-newer.tar.gz", + "published_repo": "unslothai/llama.cpp", + } + _patch_assets( + monkeypatch, + { + "unslothai/llama.cpp": { + "app-b9300-linux-x64-cuda13-newer.tar.gz": 123_456_789, + "app-b9300-windows-x64-cuda13-newer.zip": 999, + } + }, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 123_456_789 + + +def test_update_size_macos_fork_asset_suffix_fallback(monkeypatch): + # macOS bundles use the upstream-style llama--bin-macos-*, matched via the + # endswith fallback in the publish repo. + marker = { + "asset": "llama-b9190-bin-macos-arm64.tar.gz", + "published_repo": "unslothai/llama.cpp", + } + _patch_assets( + monkeypatch, + {"unslothai/llama.cpp": {"llama-b9300-bin-macos-arm64.tar.gz": 55_000_000}}, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 55_000_000 + + +def test_update_size_upstream_ubuntu_uses_binary_repo(monkeypatch): + # #6338 P2: ggml-org ubuntu-* prebuilt lives in binary_repo, not the fork + # publish repo. The size must still resolve. + marker = { + "asset": "llama-b9190-bin-ubuntu-x64.tar.gz", + "published_repo": "unslothai/llama.cpp", + "binary_repo": "ggml-org/llama.cpp", + } + _patch_assets( + monkeypatch, + { + "unslothai/llama.cpp": {"app-b9300-linux-x64-cuda13-newer.tar.gz": 1}, + "ggml-org/llama.cpp": { + "llama-b9673-bin-ubuntu-x64.tar.gz": 42_000_000, + "llama-b9673-bin-ubuntu-vulkan-x64.tar.gz": 7, + }, + }, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 42_000_000 + + +def test_update_size_upstream_windows_uses_binary_repo(monkeypatch): + # Regression (#6338 P2): the Windows upstream CPU prebuilt uses a win-* token. + marker = { + "asset": "llama-b9190-bin-win-cpu-x64.zip", + "published_repo": "unslothai/llama.cpp", + "binary_repo": "ggml-org/llama.cpp", + } + _patch_assets( + monkeypatch, + {"ggml-org/llama.cpp": {"llama-b9673-bin-win-cpu-x64.zip": 33_000_000}}, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 33_000_000 + + +def test_update_size_no_matching_asset_fails_open(monkeypatch): + # A ROCm version drift (installed 6.4 vs latest 7.2) leaves no suffix match; + # the helper fails open to None rather than guessing a wrong artifact. + marker = { + "asset": "llama-b9190-bin-ubuntu-rocm-6.4-x64.tar.gz", + "published_repo": "unslothai/llama.cpp", + "binary_repo": "ggml-org/llama.cpp", + } + _patch_assets( + monkeypatch, + {"ggml-org/llama.cpp": {"llama-b9673-bin-ubuntu-rocm-7.2-x64.tar.gz": 9}}, + ) + assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") is None + + +def test_update_size_missing_inputs_fail_open(monkeypatch): + _patch_assets( + monkeypatch, + {"unslothai/llama.cpp": {"app-b9300-linux-x64-cpu.tar.gz": 5}}, + ) + # No marker, no latest tag, or no asset string -> None (never raise). + assert fr.update_download_size_bytes(None, "b9300", "unslothai/llama.cpp") is None + assert ( + fr.update_download_size_bytes( + {"asset": "app-b9190-linux-x64-cpu.tar.gz"}, None, "unslothai/llama.cpp" + ) + is None + ) + assert fr.update_download_size_bytes({"asset": None}, "b9300", "unslothai/llama.cpp") is None diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py index 310aaf6c0f..545b9c4e7a 100644 --- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py +++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py @@ -74,7 +74,7 @@ _httpx_stub.Client = type( ) sys.modules.setdefault("httpx", _httpx_stub) -from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend # Helpers @@ -140,7 +140,7 @@ def _compute_max_available_ctx( ) kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: + if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 14c3848d33..063a9ce2cd 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -1268,9 +1268,10 @@ def test_build_speculative_flags_user_draft_n_max_override(monkeypatch): assert backend.spec_draft_n_max == 5 -def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch): - # Outdated llama-server with no MTP support: forced MTP must degrade - # to spec-off (warned) rather than emit a bad --spec-type. +def test_build_speculative_flags_mtp_token_missing_emits_spec_default(monkeypatch): + # Outdated llama-server with no MTP support: forced MTP must degrade (warned) + # and emit --spec-default so an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI + # wins over env) can't make the child attempt MTP the gate budgeted off. backend = _resolver_backend(monkeypatch, mtp_token = None) flags = backend._build_speculative_flags( speculative_type = "mtp", @@ -1282,10 +1283,11 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch): binary = "/fake/llama-server", ) assert "--spec-type" not in flags - # _speculative_type stays None (resolved emission was none); the user's - # choice is still reflected in _requested_spec_mode. + assert "--spec-default" in flags + # Degraded to non-speculative; the user's choice is still reflected. + assert backend.speculative_type == "default" assert backend.requested_spec_mode == "mtp" - assert backend.speculative_type is None + assert backend.spec_fallback_reason == "binary_no_mtp" def test_forced_mtp_on_non_mtp_model_defaults_back(monkeypatch): diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index c42c3e8a79..828d439e0a 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -901,3 +901,49 @@ def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path): res = upd.start_update() assert res["started"] is False assert res["reason"] == "up_to_date" + + +def test_status_update_available_includes_size(monkeypatch, tmp_path): + # Marker (prebuilt) update path attaches the download size of the asset the + # banner would fetch. + binary = _write_install(tmp_path, "b9493", asset = "app-b9493-linux-x64-cuda13-newer.tar.gz") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + monkeypatch.setattr( + freshness, + "latest_release_assets", + lambda repo, *, force_refresh = False: { + "app-b9518-linux-x64-cuda13-newer.tar.gz": 88_000_000 + }, + ) + st = upd.get_update_status(force_refresh = True) + assert st["update_available"] is True + assert st["update_size_bytes"] == 88_000_000 + + +def test_status_source_build_includes_update_size(monkeypatch, tmp_path): + # #6338 P3: a source build offered a prebuilt must carry the asset size too. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker -> source build + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt( + monkeypatch, + repo = "unslothai/llama.cpp", + release_tag = "b9585", + asset = "app-b9585-linux-x64-cpu.tar.gz", + ) + monkeypatch.setattr(upd, "_installed_build_number", lambda b: None) + monkeypatch.setattr( + upd, + "latest_release_assets", + lambda repo, *, force_refresh = False: ( + {"app-b9585-linux-x64-cpu.tar.gz": 77_000_000} + if repo == "unslothai/llama.cpp" + else None + ), + ) + st = upd.get_update_status() + assert st["source_build"] is True + assert st["update_available"] is True + assert st["update_size_bytes"] == 77_000_000 diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index 7f1d0c8a42..c103c6394a 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -281,6 +281,7 @@ def test_kill_process_records_timestamp_on_actual_kill(): backend = LlamaCppBackend.__new__(LlamaCppBackend) backend._process = None backend._healthy = False + backend._stats_logger = None # _kill_process stops it in finally backend._stdout_thread = None backend._llama_log_fh = None backend._last_kill_monotonic = 0.0 @@ -308,6 +309,26 @@ def test_kill_process_records_timestamp_on_actual_kill(): assert before <= backend._last_kill_monotonic <= after +def test_kill_process_tolerates_partially_constructed_backend(): + # Teardown must not AttributeError on a __new__-built backend that never ran + # __init__: _stats_logger / _stdout_thread / _llama_log_fh are left unset. + backend = LlamaCppBackend.__new__(LlamaCppBackend) + + class _FakeProcess: + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + backend._process = _FakeProcess() + backend._kill_process() + assert backend._process is None + + def test_helper_is_static_method_callable_off_class(): """Pin the @staticmethod binding so call sites can invoke off the class.""" ctx, _state = _patch_probe([[]]) diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py index bf0c4b731f..333f710b44 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -78,6 +78,27 @@ def test_status_response_exposes_source_build(): rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1}) +def test_status_response_exposes_update_size_bytes(): + payload = { + "supported": True, + "update_available": True, + "stale": False, + "installed_tag": "b9493", + "latest_tag": "b9518", + "published_repo": "unslothai/llama.cpp", + "installed_at_utc": None, + "age_days": None, + "source_build": False, + "update_size_bytes": 123_456_789, + "job": {"state": "idle"}, + } + model = rl.LlamaUpdateStatusResponse(**payload) + assert model.model_dump()["update_size_bytes"] == 123_456_789 + # Omitted -> defaults to None (the offline / no-matching-asset case). + without = {k: v for k, v in payload.items() if k != "update_size_bytes"} + assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None + + def test_status_handler_runs_off_event_loop(monkeypatch): seen = {} diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index f3a3ea1ec4..deeb228026 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -24,6 +24,7 @@ _lsa = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_lsa) is_managed_flag = _lsa.is_managed_flag parse_cache_override = _lsa.parse_cache_override +parse_cache_override_per_axis = _lsa.parse_cache_override_per_axis parse_ctx_override = _lsa.parse_ctx_override parse_split_mode_override = _lsa.parse_split_mode_override resolve_cache_type_kv = _lsa.resolve_cache_type_kv @@ -467,6 +468,25 @@ def test_parse_cache_override_rejects_malformed_values(args): parse_cache_override(args) +@pytest.mark.parametrize( + "args, expected", + [ + (["--cache-type-k", "f32", "--cache-type-v", "f16"], ("f32", "f16")), + (["-ctk", "q8_0", "-ctv", "q4_0"], ("q8_0", "q4_0")), + (["--cache-type-k=f32"], ("f32", None)), + (["--cache-type-v", "f16"], (None, "f16")), + (["-c", "4096"], (None, None)), + (None, (None, None)), + # Last-wins is kept per axis. + (["-ctk", "f16", "-ctk", "f32"], ("f32", None)), + ], +) +def test_parse_cache_override_per_axis(args, expected): + # Unlike parse_cache_override (collapses both axes to one last-wins value), + # this keeps K and V apart so an asymmetric cache can be budgeted per axis. + assert parse_cache_override_per_axis(args) == expected + + def test_resolve_cache_type_kv_uses_override_when_present(): assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0" @@ -649,6 +669,53 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec(): assert out == ["--top-k", "20"] +@pytest.mark.parametrize( + "selector", + [ + ["--spec-draft-hf", "org/repo"], + ["-hfd", "org/repo"], + ["-hfrd", "org/repo"], + ["--hf-repo-draft", "org/repo"], + ["--spec-draft-hf=org/repo"], + ], +) +def test_strip_shadowing_flags_drops_hf_drafter_selectors_with_spec(selector): + # HF drafter selectors must reset on inherit like local --model-draft, or a + # stale inherited HF drafter last-wins over Studio's re-derived spec choice. + out = strip_shadowing_flags( + selector + ["--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = True, + strip_template = False, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_draft_tuning_with_spec(): + # Per-drafter tuning knobs are deliberately preserved: the VRAM budget reads + # them via the same parsers the child honors (so they stay consistent on + # inherit), and stripping --spec-draft-ngl would move a CPU drafter to GPU. + keep = [ + "--spec-draft-type-k", + "q4_0", + "--spec-draft-type-v", + "q4_0", + "--spec-draft-ngl", + "0", + "--spec-draft-device", + "cpu", + ] + out = strip_shadowing_flags( + list(keep), + strip_context = False, + strip_cache = False, + strip_spec = True, + strip_template = False, + ) + assert out == keep + + def test_strip_shadowing_flags_keeps_split_mode_when_not_requested(): # No tensor_parallel field supplied on the Apply -> an inherited # --split-mode survives (mirrors the chat-template keep behavior). diff --git a/studio/backend/tests/test_llama_stats.py b/studio/backend/tests/test_llama_stats.py new file mode 100644 index 0000000000..5e43d80f25 --- /dev/null +++ b/studio/backend/tests/test_llama_stats.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the llama-server /metrics -> engine_stats translator: generation +throughput comes from generated-token metrics (not llama_decode() calls), and +the unexposed kv_cache_usage_ratio is never fabricated into the log line.""" + +from core.inference.llama_stats import LlamaServerStatsLogger + + +class _Capture: + def __init__(self): + self.events = [] + + def info(self, event, **kw): + self.events.append((event, dict(kw))) + + def debug(self, *a, **k): + pass + + +def _drive(snaps): + """Run _run() synchronously over `snaps`, then stop deterministically.""" + cap = _Capture() + lg = LlamaServerStatsLogger("http://127.0.0.1:0", cap) + lg._interval = 0.001 # bypass the 1s floor for a fast, synchronous run + state = {"i": 0} + + def fake_scrape(): + i = state["i"] + state["i"] += 1 + if i >= len(snaps): + lg.stop() + return None + return snaps[i] + + lg._scrape = fake_scrape + lg._run() + return [kw for ev, kw in cap.events if ev == "engine_stats"] + + +def test_gen_tok_s_uses_token_metrics_not_decode_calls(): + # tokens_predicted_total jumps 95 while n_decode_total only moves 9; the + # gauge reports 95 tok/s. Decode-call rate (9) must not be reported. + snaps = [ + { + "tokens_predicted_total": 0.0, + "prompt_tokens_total": 0.0, + "n_decode_total": 0.0, + "predicted_tokens_seconds": 95.0, + "prompt_tokens_seconds": 30.0, + "requests_processing": 1.0, + }, + { + "tokens_predicted_total": 95.0, + "prompt_tokens_total": 30.0, + "n_decode_total": 9.0, + "predicted_tokens_seconds": 95.0, + "prompt_tokens_seconds": 30.0, + "requests_processing": 1.0, + }, + ] + stats = _drive(snaps) + assert stats, "expected engine_stats while a request is processing" + assert all(s["gen_tok_s"] == 95.0 for s in stats) + assert all(s["prompt_tok_s"] == 30.0 for s in stats) + + +def test_kv_cache_pct_not_emitted_when_metric_absent(): + # llama.cpp does not expose kv_cache_usage_ratio, so it must not appear. + snaps = [ + { + "tokens_predicted_total": 0.0, + "prompt_tokens_total": 0.0, + "predicted_tokens_seconds": 10.0, + "requests_processing": 1.0, + }, + { + "tokens_predicted_total": 10.0, + "prompt_tokens_total": 5.0, + "predicted_tokens_seconds": 10.0, + "requests_processing": 1.0, + }, + ] + stats = _drive(snaps) + assert stats + assert all("kv_cache_pct" not in s for s in stats) + + +def test_scrape_parses_labelled_and_bare_metrics(monkeypatch): + # Prometheus samples may carry labels; both labelled and bare lines parse. + import core.inference.llama_stats as ls + + body = ( + 'llamacpp:tokens_predicted_total{model="m"} 20\n' + 'llamacpp:prompt_tokens_total{model="m"} 5\n' + "llamacpp:requests_processing 1\n" + "# HELP llamacpp:ignored ignored\n" + ) + + class _Resp: + status = 200 + + def read(self): + return body.encode() + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + monkeypatch.setattr(ls.urllib.request, "urlopen", lambda *a, **k: _Resp()) + m = ls.LlamaServerStatsLogger("http://127.0.0.1:0", _Capture())._scrape() + assert m["tokens_predicted_total"] == 20.0 + assert m["prompt_tokens_total"] == 5.0 + assert m["requests_processing"] == 1.0 + + +def test_counter_delta_fallback_without_gauges(): + # Older binaries expose only the counters; throughput falls back to deltas. + snaps = [ + {"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0}, + {"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0}, + ] + stats = _drive(snaps) + # running=1 keeps it emitting; gen_tok_s falls back to the (here zero) delta. + assert stats and all(s["gen_tok_s"] >= 0.0 for s in stats) diff --git a/studio/backend/tests/test_logging_middleware.py b/studio/backend/tests/test_logging_middleware.py index 34d43942ef..89061a5cbd 100644 --- a/studio/backend/tests/test_logging_middleware.py +++ b/studio/backend/tests/test_logging_middleware.py @@ -123,6 +123,100 @@ def test_non_http_scope_passes_through(logs): assert logs.events == [] +def test_duplicate_get_within_window_deduped(logs, monkeypatch): + monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000) + + async def app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + async def send(message): + pass + + mw = LoggingMiddleware(app) + for _ in range(3): + _run(mw(_http_scope("/api/chat/projects"), _noop_receive, send)) + + # Only the first of the identical GET/200 burst is logged. + assert len(logs.events) == 1 + assert logs.events[0][1] == "request_completed" + + +def test_mutations_and_errors_are_never_deduped(logs, monkeypatch): + monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000) + + async def post_ok(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + async def get_404(scope, receive, send): + await send({"type": "http.response.start", "status": 404, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + async def send(message): + pass + + mw = LoggingMiddleware(post_ok) + for _ in range(2): + _run(mw(_http_scope("/api/chat/threads", method = "POST"), _noop_receive, send)) + mw_404 = LoggingMiddleware(get_404) + for _ in range(2): + _run(mw_404(_http_scope("/api/models"), _noop_receive, send)) + + # 2 mutations + 2 errors all logged (dedup only touches GET/2xx). + assert len(logs.events) == 4 + + +def test_quiet_poll_paths_use_longer_heartbeat_window(logs, monkeypatch): + # Burst dedup off, quiet-poll heartbeat on: only liveness paths collapse. + monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 0) + monkeypatch.setattr(hmod, "_QUIET_POLL_DEDUP_MS", 1000) + + async def app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + async def send(message): + pass + + mw = LoggingMiddleware(app) + for _ in range(3): + _run(mw(_http_scope("/api/inference/monitor"), _noop_receive, send)) # quiet + for _ in range(3): + _run(mw(_http_scope("/api/chat/projects"), _noop_receive, send)) # normal + + paths = [e[2]["path"] for e in logs.events] + assert paths.count("/api/inference/monitor") == 1 # collapsed to one heartbeat + assert paths.count("/api/chat/projects") == 3 # base dedup off -> all logged + + +def test_distinct_query_strings_are_not_deduped(logs, monkeypatch): + monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000) + + async def app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok"}) + + async def send(message): + pass + + def scope(query): + return { + "type": "http", + "path": "/api/models/browse-folders", + "method": "GET", + "query_string": query, + } + + mw = LoggingMiddleware(app) + _run(mw(scope(b"path=/tmp/a"), _noop_receive, send)) + _run(mw(scope(b"path=/tmp/b"), _noop_receive, send)) # distinct query -> logs + _run(mw(scope(b"path=/tmp/a"), _noop_receive, send)) # repeat of first -> deduped + + # Two distinct query strings log; the immediate repeat of the first does not. + assert len(logs.events) == 2 + + def test_fastapi_static_asset_success_skips_log(tmp_path, logs): assets_dir = tmp_path / "assets" assets_dir.mkdir() diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 1005431926..11aeee6d77 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -293,6 +293,183 @@ class TestSecurityHeadersMiddleware: # not read directive-string `in` membership as URL sanitisation. assert any(src == "https:" for src in directives[name]) + def test_headers_applied_to_streaming_response(self, main_module): + # The ASGI middleware must set headers on streaming responses too. + from fastapi.responses import StreamingResponse + + app = FastAPI() + app.add_middleware(main_module.SecurityHeadersMiddleware) + + @app.get("/stream") + async def stream(): + async def gen(): + yield b"a" + yield b"b" + + return StreamingResponse(gen(), media_type = "text/plain") + + r = TestClient(app).get("/stream") + assert r.status_code == 200 + assert r.text == "ab" + assert r.headers["x-content-type-options"] == "nosniff" + assert r.headers["server"] == "unsloth-studio" + assert "content-security-policy" in r.headers + + def test_artifact_preview_frame_omits_x_frame_options(self, main_module): + app = FastAPI() + app.add_middleware(main_module.SecurityHeadersMiddleware) + + @app.get(main_module._ARTIFACT_PREVIEW_FRAME_PATH) + async def frame(): + return Response(content = b"", media_type = "text/html") + + r = TestClient(app).get(main_module._ARTIFACT_PREVIEW_FRAME_PATH) + assert r.status_code == 200 + assert "x-frame-options" not in {k.lower() for k in r.headers.keys()} + assert r.headers["referrer-policy"] == "no-referrer" + + def test_response_start_with_tuple_headers_is_hardened(self, main_module): + # An ASGI server may emit tuple-valued raw headers; the middleware must + # coerce to a list and still inject security headers without crashing. + import asyncio + + async def _inner_app(scope, receive, send): + await send( + { + "type": "http.response.start", + "status": 200, + "headers": ((b"content-type", b"text/plain"),), # tuple, not list + } + ) + await send({"type": "http.response.body", "body": b"ok"}) + + captured = {} + + async def _send(message): + if message["type"] == "http.response.start": + captured["headers"] = dict(message["headers"]) + + async def _receive(): + return {"type": "http.request"} + + mw = main_module.SecurityHeadersMiddleware(_inner_app) + asyncio.run(mw({"type": "http", "path": "/plain"}, _receive, _send)) + + hdrs = captured["headers"] + assert hdrs[b"server"] == b"unsloth-studio" + assert b"content-security-policy" in hdrs + assert hdrs[b"x-frame-options"] == b"DENY" + + def test_is_pure_asgi_not_basehttp_middleware(self, main_module): + # Regression: as a BaseHTTPMiddleware this wrapped the SSE stream in its + # own anyio task group, breaking disconnect detection (GPU stuck at 100%) + # and raising cancel scope errors. Must stay pure ASGI. + from starlette.middleware.base import BaseHTTPMiddleware + + cls = main_module.SecurityHeadersMiddleware + assert not issubclass(cls, BaseHTTPMiddleware) + assert not hasattr(cls, "dispatch") + + def test_forwards_receive_channel_unchanged(self, main_module): + # Must forward the ASGI receive channel untouched so client disconnects + # reach the streaming handler (BaseHTTPMiddleware swapped in its own). + seen = {} + + async def inner_app(scope, receive, send): + seen["receive"] = receive + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + mw = main_module.SecurityHeadersMiddleware(inner_app) + sentinel_receive = object() # forwarded verbatim, never wrapped/awaited + sent = [] + + async def send(message): + sent.append(message) + + async def run(): + await mw( + {"type": "http", "path": "/plain", "headers": []}, + sentinel_receive, + send, + ) + + asyncio.run(run()) + assert seen["receive"] is sentinel_receive + start = next(m for m in sent if m["type"] == "http.response.start") + names = {n.lower() for n, _ in start["headers"]} + assert b"content-security-policy" in names + assert b"server" in names + + def test_streaming_response_survives_client_disconnect(self, main_module): + # A StreamingResponse that polls is_disconnected() (like gguf_tool_stream) + # must unwind cleanly on client disconnect: no cancel scope error, the + # generator's finally runs, and security headers are still applied. + from fastapi import FastAPI, Request + from fastapi.responses import StreamingResponse + + state = {"cleaned_up": False} + app = FastAPI() + app.add_middleware(main_module.SecurityHeadersMiddleware) + + @app.get("/v1/chat/completions") + async def stream(request: Request): + async def gen(): + try: + for i in range(1000): + if await request.is_disconnected(): + break + yield f"data: {i}\n\n".encode() + await asyncio.sleep(0.01) + finally: + state["cleaned_up"] = True + + return StreamingResponse(gen(), media_type = "text/event-stream") + + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "GET", + "path": "/v1/chat/completions", + "raw_path": b"/v1/chat/completions", + "query_string": b"", + "root_path": "", + "scheme": "http", + "headers": [(b"host", b"testserver")], + "client": ("127.0.0.1", 50000), + "server": ("127.0.0.1", 80), + } + + async def run(): + body_started = asyncio.Event() + calls = {"n": 0} + + async def receive(): + calls["n"] += 1 + if calls["n"] == 1: + return {"type": "http.request", "body": b"", "more_body": False} + await body_started.wait() # client clicks Stop after tokens stream + return {"type": "http.disconnect"} + + sent = [] + + async def send(message): + sent.append(message) + if message["type"] == "http.response.body" and message.get("body"): + body_started.set() + + # Must return without raising the anyio cancel-scope RuntimeError. + await asyncio.wait_for(app(scope, receive, send), timeout = 5.0) + return sent + + sent = asyncio.run(run()) + assert state["cleaned_up"] is True + start = next(m for m in sent if m["type"] == "http.response.start") + names = {n.lower() for n, _ in start["headers"]} + assert b"content-security-policy" in names + assert b"server" in names + # /api/health auth gate diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py new file mode 100644 index 0000000000..f61ffa3c21 --- /dev/null +++ b/studio/backend/tests/test_mtp_vram_budget.py @@ -0,0 +1,1002 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the deterministic MTP VRAM reserve used by load-time auto-fit. + +reserve(ctx) = draft_KV(ctx, draft_cache_type) + separate_drafter_weights, sized +from GGUF dims (embedded head from the main model's dims; separate drafter from +its own KV). Anchors checked against real llama-server measurements. Pure: no +GPU, network, subprocess, or GGUF I/O.""" + +from __future__ import annotations + +import inspect +import os +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy/unavailable deps before importing the module under test. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +# httpx -- only stub when the real library is missing. Unconditional stubbing +# shadows HTTPError/Response that huggingface_hub.errors imports at load time, +# silently breaking the transformers introspection tier in tests collected after +# this one (the stub leaks via sys.modules for the whole session). +try: + import httpx as _httpx_real # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + ): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **kw: None}) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + sys.modules["httpx"] = _httpx_stub + +from core.inference.llama_cpp import ( # noqa: E402 + _CTX_FIT_VRAM_FRACTION, + LlamaCppBackend, + _extra_args_draft_cache_types, + _extra_args_draft_offloaded_to_cpu, + _extra_args_mtp_draft_path, + _extra_args_n_ubatch, + _extra_args_requests_mtp, + _extra_args_requests_separate_draft, + _extra_args_spec_draft_n_max, + _effective_tensor_parallel, + _env_main_cache_type_for_budget, + _extra_args_main_cache_type_for_budget, + _kv_bytes_per_elem, + _tensor_parallel_matches_loaded, +) +from core.inference.llama_server_args import _env_split_mode_is_tensor # noqa: E402 + +MIB = 1024 * 1024 +GIB = 1024**3 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_backend( + *, + nextn = 1, + n_kv_heads = 4, + n_heads = 24, + kv_key_length = 256, + kv_value_length = 256, + embedding_length = 5120, + n_layers = 65, + native_ctx = 262144, +): + """Qwen3.6-27B-MTP-class backend (embedded head) with the MTP-math dims.""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._nextn_predict_layers = nextn + b._n_kv_heads = n_kv_heads + b._n_heads = n_heads + b._kv_key_length = kv_key_length + b._kv_value_length = kv_value_length + b._embedding_length = embedding_length + b._n_layers = n_layers + b._context_length = native_ctx + # Hybrid attention/Mamba (qwen35 path) + remaining KV-estimator fields. + b._shared_kv_layers = 0 + b._kv_lora_rank = None + b._sliding_window = None + b._sliding_window_pattern = None + b._ssm_inner_size = 6144 + b._full_attention_interval = 4 + b._key_length_mla = None + b._n_kv_heads_by_layer = None + b._kv_key_length_swa = None + b._kv_value_length_swa = None + b._draft_backend_cache = None + return b + + +class _StubDrafter: + """Stand-in for a separate drafter backend (no GGUF I/O).""" + + def __init__(self, kv_per_token): + self._kv_per_token = kv_per_token + + def _can_estimate_kv(self): + return True + + def _estimate_kv_cache_bytes( + self, + n_ctx, + cache_type = None, + n_parallel = 1, + **_k, + ): + bpe = _kv_bytes_per_elem(cache_type) + # n_parallel scales like a sliding-window drafter's per-slot KV. + return 0 if n_ctx <= 0 else int(n_ctx * self._kv_per_token * bpe / 2.0 * n_parallel) + + +# --------------------------------------------------------------------------- +# Embedded draft KV: deterministic from nextn dims, scales with ctx + draft type +# --------------------------------------------------------------------------- + + +class TestEmbeddedDraftKv: + def test_scales_linearly_with_context(self): + b = _make_backend() + kv_8k = b._mtp_draft_kv_bytes(8192) + kv_16k = b._mtp_draft_kv_bytes(16384) + kv_64k = b._mtp_draft_kv_bytes(65536) + assert kv_8k and kv_16k and kv_64k + assert kv_16k == pytest.approx(2 * kv_8k) + assert kv_64k == pytest.approx(8 * kv_8k) + + def test_value_matches_dim_formula_f16(self): + # nextn(1) * n_kv(4) * (256+256) * 2(f16) * ctx -- no magic safety factor. + b = _make_backend() + ctx = 131072 + expected = int(1 * 4 * 512 * 2.0 * ctx) + assert b._mtp_draft_kv_bytes(ctx) == expected + # And that is 512 MiB, matching the measured 27B draft-KV slope (~4 MiB/1k). + assert b._mtp_draft_kv_bytes(ctx) / MIB == pytest.approx(512, abs = 1) + + def test_scales_with_nextn_predict_layers(self): + one = _make_backend(nextn = 1)._mtp_draft_kv_bytes(65536) + two = _make_backend(nextn = 2)._mtp_draft_kv_bytes(65536) + assert two == pytest.approx(2 * one) + + def test_embedded_draft_kv_floored_at_f16(self): + # The embedded MTP head is one layer, so llama.cpp's quantized-KV + # overhead is not amortized: a quantized draft KV fits LESS context than + # f16, not more (ggml-org/llama.cpp#24102). The embedded reserve floors a + # quantized draft type at f16 (never under-reserved); f32 still costs more. + b = _make_backend() + f16 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f16", draft_cache_type_v = "f16") + q8 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q8_0", draft_cache_type_v = "q8_0") + q4 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0") + f32 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f32", draft_cache_type_v = "f32") + assert q8 == f16 and q4 == f16 # quantized draft KV priced as f16, not less + assert f32 == pytest.approx(f16 * 2.0) # f32 genuinely larger, not floored + + def test_draft_kv_split_axes_no_under_reserve(self): + # A quantized draft type on either or both axes never reserves below the + # all-f16 value for the single-layer embedded head (the f16 floor; #24102). + b = _make_backend() + both_q4 = b._mtp_draft_kv_bytes( + 131072, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0" + ) + k_only = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "q4_0") # V defaults f16 + both_f16 = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "f16", draft_cache_type_v = "f16") + assert both_q4 == k_only == both_f16 # floored at f16, never under-reserved + + def test_none_when_dims_missing(self): + assert _make_backend(nextn = 0)._mtp_draft_kv_bytes(65536) is None + assert _make_backend(kv_key_length = None)._mtp_draft_kv_bytes(65536) is None + assert _make_backend()._mtp_draft_kv_bytes(0) is None + + +# --------------------------------------------------------------------------- +# Separate drafter (Gemma): sized from the drafter GGUF's own dims + weights +# --------------------------------------------------------------------------- + + +class TestSeparateDrafter: + def test_uses_drafter_kv_and_weights(self, monkeypatch): + b = _make_backend(nextn = None) # main has no embedded head + stub = _StubDrafter(kv_per_token = 2000) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: stub) + ctx = 65536 + kv = b._mtp_draft_kv_bytes(ctx, drafter_path = "/m/draft.gguf") + assert kv == stub._estimate_kv_cache_bytes(ctx) + total = b._estimate_mtp_overhead_bytes( + ctx, drafter_path = "/m/draft.gguf", draft_weights_bytes = GIB + ) + assert total == kv + GIB + + def test_drafter_kv_scales_with_context(self, monkeypatch): + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: _StubDrafter(2000)) + a = b._mtp_draft_kv_bytes(16384, drafter_path = "/m/d.gguf") + c = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") + assert c == pytest.approx(4 * a) + + def test_drafter_kv_scales_with_parallel_slots(self, monkeypatch): + # The drafter is served under the same --parallel slots as the main model, + # so a sliding-window drafter's KV grows per slot; the reserve must thread + # n_parallel or it under-reserves (Finding G1). + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: _StubDrafter(2000)) + one = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf", n_parallel = 1) + four = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf", n_parallel = 4) + assert four == pytest.approx(4 * one) + # And it threads through the overhead estimate too. + ov1 = b._estimate_mtp_overhead_bytes( + 65536, drafter_path = "/m/d.gguf", draft_weights_bytes = GIB, n_parallel = 1 + ) + ov4 = b._estimate_mtp_overhead_bytes( + 65536, drafter_path = "/m/d.gguf", draft_weights_bytes = GIB, n_parallel = 4 + ) + assert (ov4 - GIB) == pytest.approx(4 * (ov1 - GIB)) # KV scales, weights flat + + def test_none_when_drafter_unreadable(self, monkeypatch): + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: None) + assert b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") is None + assert b._estimate_mtp_overhead_bytes(65536, drafter_path = "/m/d.gguf") is None + + def test_keeps_weights_when_drafter_kv_unsizable(self, monkeypatch): + # KV can't be sized (exotic/remote drafter), but the local weights are + # known: reserve the weights so a drafter larger than the flat fallback + # cushion can't slip through and OOM (Finding C). Nothing known -> None. + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: None) + assert b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") is None + assert ( + b._estimate_mtp_overhead_bytes( + 65536, drafter_path = "/m/d.gguf", draft_weights_bytes = 3 * GIB + ) + == 3 * GIB + ) + assert b._estimate_mtp_overhead_bytes(65536, drafter_path = "/m/d.gguf") is None + + +# --------------------------------------------------------------------------- +# Total overhead = draft KV (+ separate drafter weights); no verify constant +# --------------------------------------------------------------------------- + + +class TestOverheadTotal: + def test_equals_draft_kv_for_embedded(self): + b = _make_backend() + for ctx in (16384, 65536, 131072): + assert b._estimate_mtp_overhead_bytes(ctx) == b._mtp_draft_kv_bytes(ctx) + + def test_does_not_depend_on_n_max(self): + # The verify buffer (the only n_max-dependent term) rides in headroom now. + b = _make_backend() + assert b._estimate_mtp_overhead_bytes( + 65536, spec_draft_n_max = 2 + ) == b._estimate_mtp_overhead_bytes(65536, spec_draft_n_max = 6) + + def test_none_when_draft_kv_unsizable(self): + assert _make_backend(nextn = 0)._estimate_mtp_overhead_bytes(65536) is None + + def test_includes_separate_drafter_weights(self): + b = _make_backend() + base = b._estimate_mtp_overhead_bytes(65536) + with_w = b._estimate_mtp_overhead_bytes(65536, draft_weights_bytes = GIB) + assert with_w - base == GIB + + @pytest.mark.parametrize( + "ctx,measured_draft_kv_mib", + # Measured 27B MTP delta minus the (headroom-covered) ~500 MiB verify + # buffer leaves the draft KV; the deterministic estimate must match it. + [(16384, 64), (65536, 256), (131072, 512)], + ) + def test_draft_kv_matches_measured(self, ctx, measured_draft_kv_mib): + b = _make_backend() + pred = b._estimate_mtp_overhead_bytes(ctx) / MIB + assert pred == pytest.approx(measured_draft_kv_mib, abs = 2) + + +# --------------------------------------------------------------------------- +# _fit_context_to_vram: MTP reserve lowers the chosen context +# --------------------------------------------------------------------------- + + +class TestFitContextWithMtp: + def _fit_backend(self, kv_per_token = 325_000): + b = _make_backend() + b._can_estimate_kv = lambda: True + b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else n * kv_per_token) + return b + + def test_overhead_fn_lowers_context(self): + b = self._fit_backend() + avail_mib = 24_000 + model = 8 * GIB + without = b._fit_context_to_vram(131072, avail_mib, model) + with_mtp = b._fit_context_to_vram( + 131072, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0, + ) + assert 0 < with_mtp < without + + def test_quantized_embedded_draft_kv_does_not_inflate_context(self): + # For the single-layer embedded head, quantizing the draft KV does NOT + # buy more context (it fits less in practice; ggml-org/llama.cpp#24102), + # so the f16-floored reserve advertises the same context as f16 -- never + # a larger one off a smaller (unsafe) reserve. + b = self._fit_backend() + avail_mib, model = 24_000, 8 * GIB + f16 = b._fit_context_to_vram( + 131072, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes( + c, draft_cache_type_k = "f16", draft_cache_type_v = "f16" + ) + or 0, + ) + q4 = b._fit_context_to_vram( + 131072, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes( + c, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0" + ) + or 0, + ) + assert 0 < q4 == f16 + + def test_no_mtp_unchanged(self): + b = self._fit_backend() + avail_mib, model = 24_000, 8 * GIB + a = b._fit_context_to_vram(131072, avail_mib, model) + bb = b._fit_context_to_vram( + 131072, avail_mib, model, mtp_engaged = False, mtp_overhead_fn = None + ) + assert a == bb + + def test_chosen_context_actually_fits_budget(self): + b = self._fit_backend() + avail_mib, model = 24_000, 8 * GIB + fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0 # noqa: E731 + ctx = b._fit_context_to_vram(131072, avail_mib, model, mtp_overhead_fn = fn) + budget = avail_mib * MIB * _CTX_FIT_VRAM_FRACTION + assert model + b._estimate_kv_cache_bytes(ctx) + fn(ctx) <= budget + + +# --------------------------------------------------------------------------- +# extra_args parsing: detect user-enabled MTP + draft depth + draft KV type +# --------------------------------------------------------------------------- + + +class TestExtraArgsMtpDetection: + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-type", "draft-mtp"], True), + (["--spec-type", "mtp"], True), + (["--spec-type", "ngram-mod,draft-mtp"], True), + (["--spec-type=draft-mtp"], True), + (["--spec-type", "ngram-mod"], False), + (["--spec-default"], False), + (["-c", "131072"], False), + (None, False), + ([], False), + ], + ) + def test_requests_mtp(self, args, expected): + assert _extra_args_requests_mtp(args, env = {}) is expected + + def test_requests_mtp_env(self): + # The child honors LLAMA_ARG_SPEC_TYPE; env-requested MTP must reserve too. + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) is True + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "ngram-mod,mtp"}) is True + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) is False + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "none"}) is False + + def test_requests_mtp_effective_spec_type(self): + # llama.cpp uses the LAST CLI --spec-type and ignores the env when any CLI + # --spec-type is present. The reserve must track that effective value, not + # any earlier/MTP-ish one, or it over-reserves a drafter the launch won't + # load (Finding B). + env_mtp = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"} + # Later CLI value overrides an earlier MTP one (last-wins). + assert ( + _extra_args_requests_mtp( + ["--spec-type", "draft-mtp", "--spec-type", "ngram-mod"], env = {} + ) + is False + ) + # A non-MTP CLI flag overrides a stale MTP env. + assert _extra_args_requests_mtp(["--spec-type", "ngram-mod"], env = env_mtp) is False + assert _extra_args_requests_mtp(["--spec-type", "none"], env = env_mtp) is False + # A later MTP CLI value still engages. + assert ( + _extra_args_requests_mtp( + ["--spec-type", "ngram-mod", "--spec-type", "draft-mtp"], env = {} + ) + is True + ) + # Same precedence for separate (draft-simple/eagle3) detection. + assert ( + _extra_args_requests_separate_draft( + ["--spec-type", "draft-simple", "--spec-type", "ngram-mod"], env = {} + ) + is False + ) + assert ( + _extra_args_requests_separate_draft( + ["--spec-type", "ngram-mod"], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"} + ) + is False + ) + + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-type", "draft-simple"], True), + (["--spec-type", "draft-eagle3"], True), + (["--spec-type=draft-eagle3"], True), + (["--spec-type", "draft-mtp"], False), # MTP path handles this one + (["--spec-type", "ngram-mod"], False), # loads no draft model + (["-c", "4096"], False), + (None, False), + ], + ) + def test_requests_separate_draft(self, args, expected): + assert _extra_args_requests_separate_draft(args, env = {}) is expected + + def test_requests_separate_draft_env(self): + assert ( + _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) + is True + ) + assert ( + _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) + is False + ) + + def test_load_model_reserves_for_non_mtp_draft_modes(self): + # load_model engages the draft reserve for a non-MTP model-based draft mode + # only when extras also name a drafter (else nothing is loaded to reserve). + # Strip all whitespace so the check survives any line-wrapping the + # formatter applies to the call (pre-commit black wraps long lines). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_user_draft_via_extras" in compact + # called with extra_args (an env kwarg may follow); prefix match stays + # robust to that and to any formatter line-wrapping. + assert "_extra_args_requests_separate_draft(extra_args" in compact + assert "or_user_draft_via_extras" in compact # OR'd into the reserve gate + # The drafter check must NOT force extras-only (env={}); the default + # env=None lets it see an env LLAMA_ARG_SPEC_DRAFT_MODEL, so an env-only + # drafter still engages the reserve (codex review 4507014299). + assert "bool(_extra_args_mtp_draft_path(extra_args))" in compact + + def test_env_only_drafter_engages_separate_draft_reserve(self, monkeypatch): + # An env-provided drafter (no --model-draft in extras) must still engage + # the draft reserve, or auto-fit spends the drafter's VRAM and OOMs. Mirror + # load_model's _user_draft_via_extras gate (codex review 4507014299). + monkeypatch.setenv("LLAMA_ARG_SPEC_DRAFT_MODEL", "/large.gguf") + monkeypatch.delenv("LLAMA_ARG_SPEC_DRAFT_HF_REPO", raising = False) + ea = ["--spec-type", "draft-simple"] # _spec_env is {} (extras set spec-type) + assert _extra_args_requests_separate_draft(ea, env = {}) is True + assert _extra_args_mtp_draft_path(ea) == "/large.gguf" # env=None -> os.environ + # -> _user_draft_via_extras True; _env_draft_for_budget sizes the drafter. + assert _extra_args_mtp_draft_path([], env = dict(os.environ)) == "/large.gguf" + + def test_load_model_gates_env_spec_type_on_off_mode(self): + # LLAMA_ARG_SPEC_TYPE only reaches the child when Studio emits no spec + # flag (UI mode "off", no user --spec-type); otherwise the emitted + # --spec-type/--spec-default overrides the env, so the reserve must not + # consult it or a stale MTP env over-reserves (Finding F3). Whitespace- + # stripped so the check survives formatter line-wrapping. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert '_mtp_canonical=="off"' in compact # the env-reaches-child gate + assert "_extra_args_requests_mtp(extra_args,env=_spec_env)" in compact + + def test_spec_default_overrides_env_mtp(self): + # --spec-default is a CLI spec flag (resolves to the model default, + # non-MTP) that overrides a stale LLAMA_ARG_SPEC_TYPE env, so the reserve + # must not treat it as MTP (reviewer.py R4). + env_mtp = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"} + assert _extra_args_requests_mtp(["--spec-default"], env = env_mtp) is False + assert ( + _extra_args_requests_separate_draft( + ["--spec-default"], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"} + ) + is False + ) + # A later --spec-type still wins over an earlier --spec-default. + assert ( + _extra_args_requests_mtp(["--spec-default", "--spec-type", "draft-mtp"], env = {}) is True + ) + + def test_load_model_drafter_budget_precedence(self): + # The budget sizes the drafter the launch actually loads: CLI extras win, + # then Studio's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL), + # then the env drafter -- not the env before Studio's (reviewer.py R3). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact + assert "_env_draft_for_budget=_extra_args_mtp_draft_path([],env=os.environ)" in compact + assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact + + def test_load_model_drops_cpu_offloaded_drafter_from_budget(self): + # A SEPARATE drafter offloaded to CPU (--spec-draft-ngl 0 / + # --spec-draft-device none) consumes no GPU, so it must be dropped from the + # budget and get no flat reserve (Finding F2). But an embedded head is on + # GPU regardless of those draft-only flags, so the flat reserve is only + # suppressed when there is no embedded head (Finding G5). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + # env-aware: also honors the inherited LLAMA_ARG_N_GPU_LAYERS_DRAFT. + assert ( + "_draft_on_cpu=_extra_args_draft_offloaded_to_cpu(extra_args,env=os.environ)" in compact + ) + assert "if_draft_on_cpu:_mtp_draft_for_budget=None" in compact + # flat reserve suppressed only for a CPU drafter with no embedded head + assert "_draft_cpu_no_embedded=_draft_on_cpuandnotself._nextn_predict_layers" in compact + assert "not_draft_cpu_no_embedded" in compact + + def test_load_model_keeps_flat_reserve_for_unsized_draft_kv(self): + # When only the drafter weights could be sized (KV unsizable), the flat + # fraction stays on as the cushion for the still-unsized draft KV, on top + # of the byte-accurate weights reserve (Finding G3). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_mtp_kv_unsized" in compact + assert "mtp_overhead_fnisNoneor_mtp_kv_unsized" in compact + + def test_load_model_ranks_subsets_by_active_pin_fraction(self): + # Auto/cap subset ranking uses the active budget fraction (lowered by the + # flat MTP reserve), not a hard-coded 0.95, so the ranking order matches + # the fit budget that is then tested (Finding G4). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_gpu_usable(g,pin_fraction)" in compact + assert "_gpu_usable(g,_CTX_FIT_VRAM_FRACTION-_flat_mtp_reserve)" in compact + + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-draft-ngl", "0"], True), + (["-ngld", "0"], True), + (["--spec-draft-ngl=0"], True), + (["--n-gpu-layers-draft", "0"], True), + (["--spec-draft-ngl", "20"], False), + (["--spec-draft-device", "none"], True), + (["--spec-draft-device", "CPU"], True), + (["-devd", "cpu,none"], True), + (["--spec-draft-device", "CUDA0"], False), + (["--spec-draft-device", "CUDA0,CPU"], False), # any GPU -> on GPU + (["-c", "4096"], False), + (None, False), + # last-wins: only the final value of each flag counts (Finding G2) + (["--spec-draft-ngl", "0", "--spec-draft-ngl", "-1"], False), # last = GPU + (["--spec-draft-ngl", "-1", "--spec-draft-ngl", "0"], True), # last = CPU + (["--spec-draft-device", "CUDA0", "--spec-draft-device", "none"], True), + (["--spec-draft-device", "none", "--spec-draft-device", "CUDA0"], False), + ], + ) + def test_draft_offloaded_to_cpu(self, args, expected): + assert _extra_args_draft_offloaded_to_cpu(args, env = {}) is expected + + def test_draft_offloaded_to_cpu_env(self): + # The child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT; an env-only CPU offload + # must drop the drafter from the budget too (review run3 #3). CLI wins. + assert ( + _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"}) + is True + ) + assert ( + _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "-1"}) + is False + ) + # CLI --spec-draft-ngl wins over the env (last-wins is CLI-only). + assert ( + _extra_args_draft_offloaded_to_cpu( + ["--spec-draft-ngl", "-1"], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"} + ) + is False + ) + assert _extra_args_draft_offloaded_to_cpu([], env = {}) is False + + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-draft-n-max", "4"], 4), + (["--spec-draft-n-max=6"], 6), + (["--spec-type", "draft-mtp", "--spec-draft-n-max", "3"], 3), + (["--spec-draft-n-max", "2", "--spec-draft-n-max", "5"], 5), # last wins + (["--spec-draft-n-max", "notanint"], None), + (["-c", "4096"], None), + (None, None), + (["--draft-max", "6"], 6), + (["--draft-max=4"], 4), + (["--spec-type", "draft-mtp", "--draft-max", "6"], 6), + (["--spec-draft-n-max", "2", "--draft-max", "5"], 5), + ], + ) + def test_spec_draft_n_max(self, args, expected): + assert _extra_args_spec_draft_n_max(args) == expected + + @pytest.mark.parametrize( + "args,expected", + [ + (["--model-draft", "/m/draft.gguf"], "/m/draft.gguf"), + (["--spec-draft-model", "/m/draft.gguf"], "/m/draft.gguf"), + (["-md", "/m/draft.gguf"], "/m/draft.gguf"), + (["--model-draft=/m/draft.gguf"], "/m/draft.gguf"), + (["--model-draft", "--spec-type"], None), + (["-c", "4096"], None), + (None, None), + ], + ) + def test_mtp_draft_path(self, args, expected): + # env={} isolates pure-CLI behavior from a polluted test environment. + assert _extra_args_mtp_draft_path(args, env = {}) == expected + + @pytest.mark.parametrize( + "args,expected", + [ + # HF draft repo flags are real llama-server flags; the budget must see them. + (["--spec-draft-hf", "big/repo:Q8_0"], "big/repo:Q8_0"), + (["-hfd", "big/repo"], "big/repo"), + (["-hfrd", "big/repo"], "big/repo"), + (["--hf-repo-draft=big/repo"], "big/repo"), + ], + ) + def test_mtp_draft_path_hf_flags(self, args, expected): + assert _extra_args_mtp_draft_path(args, env = {}) == expected + + def test_mtp_draft_path_env_fallback(self): + # The child honors LLAMA_ARG_SPEC_DRAFT_MODEL / _HF_REPO; CLI wins over env. + assert ( + _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "/m/e.gguf"}) + == "/m/e.gguf" + ) + assert _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"}) == "x/y" + assert ( + _extra_args_mtp_draft_path( + ["-md", "/m/cli.gguf"], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"} + ) + == "/m/cli.gguf" + ) + + @pytest.mark.parametrize( + "args,expected", + [ + (["--cache-type-k-draft", "q8_0"], ("q8_0", None)), + (["--spec-draft-type-k", "q4_0"], ("q4_0", None)), + (["-ctkd", "q8_0"], ("q8_0", None)), + (["--cache-type-v-draft", "q4_0"], (None, "q4_0")), # K stays f16, V only + (["--cache-type-k-draft", "q4_0", "--cache-type-v-draft", "q8_0"], ("q4_0", "q8_0")), + (["--cache-type-k-draft=q8_0"], ("q8_0", None)), + (["--cache-type-k", "q8_0"], (None, None)), # main type, not draft + (["-c", "4096"], (None, None)), + (None, (None, None)), + ], + ) + def test_draft_cache_types(self, args, expected): + assert _extra_args_draft_cache_types(args, env = {}) == expected + + def test_draft_cache_types_env_fallback(self): + # The child honors LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V per axis; CLI wins. + assert _extra_args_draft_cache_types( + [], env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0"} + ) == ("q8_0", None) + assert _extra_args_draft_cache_types( + [], + env = { + "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0", + "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V": "q4_0", + }, + ) == ("q8_0", "q4_0") + assert _extra_args_draft_cache_types( + ["-ctkd", "q4_0"], env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0"} + ) == ("q4_0", None) + + @pytest.mark.parametrize( + "args,expected", + [ + (["--ubatch-size", "1024"], 1024), + (["-ub", "4096"], 4096), + (["--ubatch-size=512"], 512), + (["--ubatch", "2048"], None), # not a real llama-server flag; ignore it + (["-c", "4096"], None), + (None, None), + ], + ) + def test_n_ubatch(self, args, expected): + assert _extra_args_n_ubatch(args, env = {}) == expected + + def test_n_ubatch_env_fallback(self): + # The child honors LLAMA_ARG_UBATCH; it must reach the compute-buffer reserve. + assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 4096 + assert ( + _extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) == 1024 + ) # CLI wins + assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None + + def test_env_main_cache_type_for_budget(self): + # The child inherits LLAMA_ARG_CACHE_TYPE_K/_V, but Studio emits no + # --cache-type when neither param nor extras set it -> a heavier env + # main KV (f32) must be adopted so the reserve matches the child. + assert _env_main_cache_type_for_budget(env = {}) is None + # f32 exceeds the f16 default -> adopt it (lower-cased so the launch + # re-emits it via _valid_cache_types). + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f32"}) == "f32" + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "F32"}) == "f32" + # Heavier of K/V (single knob; over-reserves the lighter axis). + assert ( + _env_main_cache_type_for_budget( + env = {"LLAMA_ARG_CACHE_TYPE_K": "f32", "LLAMA_ARG_CACHE_TYPE_V": "f16"} + ) + == "f32" + ) + # Quantized env types are <= f16 -> already over-reserved by the default. + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "q4_0"}) is None + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "q8_0"}) is None + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f16"}) is None + # Unknown env type self-neutralizes (treated as f16 by _kv_bytes_per_elem). + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "wat"}) is None + + def test_load_model_adopts_env_main_cache_type(self): + # Source-level: load_model budgets the heavier of asymmetric --cache-type + # extras, then (only when neither param nor extras set it) adopts the env + # main KV type, so the reserve covers a child that inherits a heavier + # LLAMA_ARG_CACHE_TYPE_*. Whitespace-stripped to survive formatter wraps. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_extra_args_main_cache_type_for_budget(extra_args)" in compact + assert "ifcache_type_kvisNone:" in compact + assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact + + def test_env_split_mode_is_tensor(self): + # The child inherits LLAMA_ARG_SPLIT_MODE, but Studio emits --split-mode + # only on its tensor branch -> a tensor env must flip the budget so the + # heavier per-device compute buffer is reserved (not layer overhead). + assert _env_split_mode_is_tensor(env = {}) is False + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "tensor"}) is True + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "Tensor"}) is True + # Other modes are not a runtime-heavier surprise -> not acted on. + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "layer"}) is False + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "row"}) is False + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "none"}) is False + + def test_effective_tensor_parallel_env_flip(self): + # Shared by load_model and both duplicate-load matchers, so they agree. + tensor_env = {"LLAMA_ARG_SPLIT_MODE": "tensor"} + # No extras, toggle off, tensor env -> flips on. + assert _effective_tensor_parallel(None, False, env = tensor_env) is True + # Extras override (any --split-mode) beats the env, even if non-tensor. + assert _effective_tensor_parallel(["--split-mode", "layer"], False, env = tensor_env) is False + # Explicit extras/toggle tensor stays on regardless of env. + assert _effective_tensor_parallel(["--split-mode", "tensor"], False, env = {}) is True + assert _effective_tensor_parallel(None, True, env = {}) is True + # One-directional: a non-tensor env never downgrades, and no env -> no flip. + assert _effective_tensor_parallel(None, False, env = {}) is False + assert ( + _effective_tensor_parallel(None, False, env = {"LLAMA_ARG_SPLIT_MODE": "layer"}) is False + ) + + def test_tensor_parallel_matches_loaded_env_downgrade(self): + # Env-only tensor matches a server that actually launched tensor, but a + # server load_model downgraded to layer (env scrubbed) must still match + # an identical request -- not reload forever (#6312). + tensor_env = {"LLAMA_ARG_SPLIT_MODE": "tensor"} + # Launched tensor: env-only request matches. + assert _tensor_parallel_matches_loaded(None, False, True, env = tensor_env) is True + # Downgraded to layer: same env-only request still matches (no reload loop). + assert _tensor_parallel_matches_loaded(None, False, False, env = tensor_env) is True + # No env: a plain request matches a layer server and mismatches a tensor one. + assert _tensor_parallel_matches_loaded(None, False, False, env = {}) is True + assert _tensor_parallel_matches_loaded(None, False, True, env = {}) is False + # Explicit tensor request stays strict: must have a tensor server. + assert _tensor_parallel_matches_loaded(None, True, False, env = {}) is False + assert _tensor_parallel_matches_loaded(None, True, True, env = {}) is True + # An explicit non-tensor --split-mode beats the env (no flip). + assert ( + _tensor_parallel_matches_loaded(["--split-mode", "layer"], False, True, env = tensor_env) + is False + ) + + def test_route_matcher_uses_tensor_parallel_matches_loaded(self): + # Fix: the route duplicate-load matcher must use the downgrade-aware + # helper, or an env-driven tensor server (or its layer downgrade) is + # needlessly reloaded (#6312). Read from disk (importing routes.inference + # drags in heavy deps). + routes_src = ( + Path(__file__).resolve().parent.parent / "routes" / "inference.py" + ).read_text() + start = routes_src.index("def _request_matches_loaded_settings") + end = routes_src.index("\ndef ", start + 1) + body = "".join(routes_src[start:end].split()) + assert ( + "_tensor_parallel_matches_loaded(effective_extra," + "request.tensor_parallel,llama_backend.tensor_parallel)" in body + ) + + def test_extra_args_main_cache_type_heavier_axis(self): + # Asymmetric --cache-type-k/-v must budget the heavier axis (extras win + # per axis at launch), not the last-wins single type that under-reserves. + H = _extra_args_main_cache_type_for_budget + assert H(["--cache-type-k", "f32", "--cache-type-v", "f16"]) == "f32" + assert H(["--cache-type-v", "f16", "--cache-type-k", "f32"]) == "f32" # order-free + assert H(["--cache-type-k=f32", "--cache-type-v=f16"]) == "f32" # = form + assert H(["-ctk", "q4_0", "-ctv", "q8_0"]) == "q8_0" # heavier quant + assert H(["--cache-type-k", "q8_0"]) == "q8_0" # single axis honored as-is + assert H(["-c", "4096"]) is None # no cache flags + assert H(None) is None + + def test_load_model_budgets_heavier_asymmetric_cache_axis(self): + # load_model must reserve from the heavier of asymmetric cache extras, or + # an f32 K against an f16 budget over-advertises context and can OOM. + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_extra_args_main_cache_type_for_budget(extra_args)" in load + + def test_load_model_tensor_drops_any_quantized_cache_axis(self): + # The heavier-by-bytes budget type can mask a quantized axis (an f16 + # budget hides a paired q4_0), so the tensor-safety drop must test each + # --cache-type-k/-v extra, not just cache_type_kv -- else the quantized + # axis survives into tensor mode and crashes the load (#6312). + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_ck_extra,_cv_extra=parse_cache_override_per_axis(extra_args)" in load + assert "forcin(cache_type_kv,_ck_extra,_cv_extra)" in load + assert "iftensor_paralleland_cache_non_tensor_safe:" in load + + def test_load_model_layer_downgrade_restores_original_cache_extras(self): + # Tensor mode strips asymmetric --cache-type-k/-v (it rejects quantized), + # but layer split supports them, so a downgrade must restore the ORIGINAL + # extras, not just the scalar heavier type (else q4_0/f16 silently becomes + # f16/f16 on the layer fallback) (#6312). + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_tensor_dropped_extra_args=list(extra_args)" in load + # Both tensor->layer downgrade points restore the saved originals. + assert load.count("strip_split_mode_only(_tensor_dropped_extra_argsif") == 2 + + def test_load_model_tensor_skips_reserve_for_cpu_drafter(self): + # A separate CPU-offloaded drafter (no embedded head) uses no GPU, so the + # tensor reserve must be suppressed like the layer path -- else tensor mode + # subtracts a phantom flat MTP reserve and under-advertises context (#6312). + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_mtp_will_engageandnot_draft_cpu_no_embedded" in load + assert "ifnot_mtp_reserves_gpu:" in load + assert "mtp_engaged=_mtp_reserves_gpu" in load + + def test_load_model_adopts_env_tensor_split_mode(self): + # load_model delegates the tensor decision to _effective_tensor_parallel, + # which flips to tensor only one-directionally: extras set no split mode, + # none is overridden, and the env selects tensor (an existing tensor plan + # is never downgraded). Whitespace-stripped to survive formatter wrapping. + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "tensor_parallel=_effective_tensor_parallel(extra_args,tensor_parallel)" in load + helper = "".join(inspect.getsource(_effective_tensor_parallel).split()) + assert "notresolved" in helper + assert "parse_split_mode_override(extra_args)isNone" in helper + assert "_env_split_mode_is_tensor(env)" in helper + + def test_load_model_does_not_emit_env_only_cache_type(self): + # Cluster C: an env-only (budget) cache type must not be re-emitted as + # --cache-type flags (that would rewrite an asymmetric K/V env). Emission + # is guarded by `not _cache_type_from_env`, set when the value came from + # _env_main_cache_type_for_budget(). Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact + assert "_cache_type_from_env=cache_type_kvisnotNone" in compact + assert "andnot_cache_type_from_env" in compact + + def test_load_model_clears_inherited_split_mode_on_layer(self): + # Cluster A: when the final decision is layer split, an inherited + # non-layer LLAMA_ARG_SPLIT_MODE (and paired LLAMA_ARG_TENSOR_SPLIT) must + # be popped from the child env so the child cannot run tensor/row/none + # against Studio's layer budget. Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert 'env.get("LLAMA_ARG_SPLIT_MODE")' in compact + assert '_inherited_sm!="layer"' in compact + assert 'env.pop("LLAMA_ARG_SPLIT_MODE",None)' in compact + assert 'env.pop("LLAMA_ARG_TENSOR_SPLIT",None)' in compact + + def test_load_model_clears_quantized_kv_env_for_tensor(self): + # Cluster B: tensor mode aborts on quantized KV. An inherited quantized + # LLAMA_ARG_CACHE_TYPE_K/_V must be popped from the child env so it cannot + # crash the tensor child (and matches the tensor-safe budget). + # Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert '("LLAMA_ARG_CACHE_TYPE_K","LLAMA_ARG_CACHE_TYPE_V")' in compact + assert "_ct_rawnotinself._TENSOR_PARALLEL_KV_TYPES" in compact + assert "env.pop(_ct_var,None)" in compact + + def test_load_model_clears_tensor_split_env_in_tensor_mode(self): + # review run3 #2: Studio owns the tensor split. When it emits no + # --tensor-split (even split), a stale inherited LLAMA_ARG_TENSOR_SPLIT must + # be cleared in the TENSOR branch too (not just the layer downgrade), or the + # child runs a split Studio didn't budget. The else (tensor) branch pops it. + src = inspect.getsource(LlamaCppBackend.load_model) + compact = "".join(src.split()) + # appears in both the layer branch and the tensor branch. + assert compact.count('env.pop("LLAMA_ARG_TENSOR_SPLIT",None)') >= 2 + + def test_load_model_layer_compute_buffer_fallback(self): + # review run3 #4: when GGUF dims are missing the compute-buffer estimate is + # 0; the layer path must still reserve the flat fallback (tensor buffer >= + # layer buffer), not fold 0, or it under-reserves at high --parallel. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "if_compute_buffer_pipeline<=0:" in compact + # The RHS expression only (no `_compute_buffer_pipeline=` prefix) so the + # match survives the formatter wrapping it in parens. It's unique within + # load_model (the other use of this constant is in MiB, no *1024*1024). + assert "self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB*1024*1024" in compact + + def test_load_model_passes_unsized_mtp_reserve_to_tensor_planner(self): + # review run3 #1/#5: a weights-only (KV-unsized) MTP reserve must flow into + # _plan_tensor_parallel as a flat cushion, else its binary search spends the + # unsized draft KV on context and OOMs. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_tp_unsized_mtp_reserve=" in compact + # Gated on _mtp_reserves_gpu so a CPU-offloaded drafter reserves nothing. + assert "(_mtp_reserves_gpuand_mtp_kv_unsized)" in compact + assert "mtp_flat_reserve_bytes=_tp_unsized_mtp_reserve" in compact + + def test_pool_budget_sums_per_gpu_usable(self): + # Finding #1: the multi-GPU pooled budget must sum each GPU's own usable + # budget (so an unknown-total GPU gets the free*frac cushion) rather than + # pooling free and total separately. The fit calls pass the precomputed + # budget as an absolute (budget_frac=1.0, total_mib=None) so fit and check + # agree. Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "def_pool_budget_mib(subset,frac):" in compact + assert "sum(max(0.0,_gpu_usable(g,frac))forginsubset)" in compact + # No revert to the pooled free/total form. + assert "def_pool_total(" not in compact + assert "budget_frac=1.0" in compact + + +# --------------------------------------------------------------------------- +# Regression: the reported Qwen3.6-27B MTP / 24 GB scenario +# --------------------------------------------------------------------------- + + +def test_qwen36_class_regression_picks_lower_ctx_with_mtp(): + """A 24 GB card that auto-picks a high context without MTP must pick a + strictly lower one once the MTP draft reserve is accounted for.""" + b = _make_backend() + b._can_estimate_kv = lambda: True + b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else int(n * 66_000)) + avail_mib = 24_000 + model = int(17.9 * GIB) # UD-Q4_K_XL weights + no_mtp = b._fit_context_to_vram(262144, avail_mib, model) + with_mtp = b._fit_context_to_vram( + 262144, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0, + ) + assert 0 < with_mtp < no_mtp + + +def test_mtp_draft_budget_prefers_user_extras_drafter(): + # A user --model-draft in extras is appended last and wins at launch, so the + # VRAM budget must size it first; then Studio's emitted mtp_draft_path (which + # overrides LLAMA_ARG_SPEC_DRAFT_MODEL), then the env drafter (load_model is too + # entangled to drive end-to-end; assert the precedence at the source level). + # Whitespace-stripped so the check survives any formatter line-wrapping. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + # CLI extras sized first (env={} so the env doesn't pre-empt Studio's drafter). + assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact + # Order: CLI extras, then Studio's mtp_draft_path, then the env drafter. + assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact + # The env must not be consulted before Studio's resolved drafter. + assert "_extra_args_mtp_draft_path(extra_args)ormtp_draft_path" not in compact diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index a9e0b98247..a2a505f479 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -209,10 +209,10 @@ class TestGgufVariantFileResolution: return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None] def fake_download( - *, repo_id, filename, token = None, + **_kwargs, ): downloaded.append(filename) return f"/fake/{repo_id}/{filename}" @@ -229,7 +229,7 @@ class TestGgufVariantFileResolution: ), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("huggingface_hub.hf_hub_download", fake_download), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), ): out = backend._download_gguf( hf_repo = "ggml-org/models", @@ -256,10 +256,10 @@ class TestGgufVariantFileResolution: return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None] def fake_download( - *, repo_id, filename, token = None, + **_kwargs, ): downloaded.append(filename) return f"/fake/{repo_id}/{filename}" @@ -269,7 +269,7 @@ class TestGgufVariantFileResolution: patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("huggingface_hub.hf_hub_download", fake_download), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), ): out = backend._download_gguf( hf_repo = "org/repo", @@ -639,17 +639,20 @@ class TestDownloadMmprojOfflineCacheFallback: raise OSError("offline") def fake_download( - *, repo_id, filename, token = None, + **kwargs, ): # Echo back so the test can verify the cache-resolved filename return f"/fake/cache/{repo_id}/{filename}" with ( patch("huggingface_hub.list_repo_files", boom_list), - patch("huggingface_hub.hf_hub_download", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_mmproj( hf_repo = "unsloth/vision-GGUF", @@ -675,17 +678,20 @@ class TestDownloadMmprojOfflineCacheFallback: captured = {} def fake_download( - *, repo_id, filename, token = None, + **kwargs, ): captured["filename"] = filename return f"/fake/{filename}" with ( patch("huggingface_hub.list_repo_files", boom_list), - patch("huggingface_hub.hf_hub_download", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): backend._download_mmproj( hf_repo = "unsloth/vision-GGUF", diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index b2b22f8934..2586076321 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -7,6 +7,7 @@ import os import sys import asyncio import json +import threading from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -27,20 +28,29 @@ from models.inference import ( from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) +from core.inference.api_monitor import ApiMonitor from routes.inference import ( _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, _clamp_finish_reason, + _cmpl_stream_event_out, _coalesce_consecutive_user_turns, _drop_empty_assistant_sentinels, _effective_max_tokens, _extract_content_parts, _friendly_error, _merge_user_content, + _monitor_openai_chunk, + _monitor_openai_sse_event, _openai_messages_for_gguf_chat, + _openai_passthrough_non_streaming, + _openai_passthrough_stream, _openai_stream_usage_chunk, + _proxy_to_external_provider, _set_or_prepend_system_message, + openai_completions, + openai_embeddings, openai_chat_completions, ) from state.tool_policy import reset_tool_policy @@ -325,7 +335,8 @@ class TestChatCompletionRequestToolFields: captured = {} - async def _fake_proxy(payload, request): + async def _fake_proxy(payload, request, current_subject): + assert current_subject == "test-user" captured["stream"] = payload.stream return JSONResponse({"choices": [], "object": "chat.completion"}) @@ -460,6 +471,7 @@ class TestChatCompletionRequestToolFields: supports_tools = False is_vision = False _is_audio = False + context_length = 4096 client = self._v1_client(monkeypatch, _GGUFBackend()) resp = client.post( @@ -473,13 +485,18 @@ class TestChatCompletionRequestToolFields: self._assert_unsupported_n(resp) def test_n_rejected_for_gguf_tools_passthrough_path(self, monkeypatch): + import routes.inference as inference_route + class _GGUFBackend: is_loaded = True model_identifier = "test-gguf" supports_tools = True is_vision = False _is_audio = False + context_length = 4096 + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) client = self._v1_client(monkeypatch, _GGUFBackend()) resp = client.post( "/v1/chat/completions", @@ -498,6 +515,10 @@ class TestChatCompletionRequestToolFields: }, ) self._assert_unsupported_n(resp) + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "n > 1 is not supported" in entry["error"] + assert monitor.active_count() == 0 def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: @@ -540,6 +561,8 @@ class TestChatCompletionRequestToolFields: "_detect_safetensors_features", lambda backend, chat_template: {"supports_tools": True}, ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend()) resp = client.post( "/v1/chat/completions", @@ -556,6 +579,10 @@ class TestChatCompletionRequestToolFields: body = resp.json() assert body["error"]["param"] == "confirm_tool_calls" assert "requires stream=true" in body["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "confirm_tool_calls requires stream=true" in entry["error"] + assert monitor.active_count() == 0 def test_multiturn_tool_loop_messages(self): req = ChatCompletionRequest( @@ -887,6 +914,35 @@ class TestOpenAICompatibilityHelpers: assert usage["completion_tokens"] == 7 assert usage["total_tokens"] == 7 + def test_completion_stream_monitor_reads_usage_before_client_strip(self, monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/completions", + method = "POST", + model = "m", + prompt = "hi", + context_length = 100, + ) + event = ( + b'data: {"id":"chatcmpl-test","choices":[{"text":"done","finish_reason":"stop"}],' + b'"usage":{"prompt_tokens":4,"completion_tokens":6,"total_tokens":10}}\n' + ) + + _monitor_openai_sse_event(monitor_id, event, context_length = 100) + out = _cmpl_stream_event_out(event, include_usage = False) + + assert out is not None + assert b'"usage"' not in out + [entry] = monitor.snapshot() + assert entry["reply"] == "done" + assert entry["prompt_tokens"] == 4 + assert entry["completion_tokens"] == 6 + assert entry["total_tokens"] == 10 + assert entry["context_usage"] == 0.1 + def test_developer_message_preserves_existing_system_prompt(self): payload = ChatCompletionRequest( messages = [ @@ -1168,6 +1224,10 @@ class TestGgufVisionMessages: class TestGgufVisionToolRouting: class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + async def is_disconnected(self): return False @@ -1203,9 +1263,12 @@ class TestGgufVisionToolRouting: is_vision = True, supports_tools = True, model_identifier = "gemma-4-12b-it-GGUF", + context_length = 4096, generate_chat_completion = _plain, generate_chat_completion_with_tools = _tools, ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) payload = ChatCompletionRequest( @@ -1258,9 +1321,12 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = True, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _plain, generate_chat_completion_with_tools = _tools, ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) payload = ChatCompletionRequest( @@ -1292,10 +1358,13 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = True, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _plain, generate_chat_completion_with_tools = _tools, ) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) payload = ChatCompletionRequest( model = "default", @@ -1316,6 +1385,62 @@ class TestGgufVisionToolRouting: ) assert exc.value.status_code == 400 assert "requires stream=true" in exc.value.detail["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "confirm_tool_calls requires stream=true" in entry["error"] + assert monitor.active_count() == 0 + + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): + import routes.inference as inf_mod + + calls = {"count": 0} + + def _generate(**_kwargs): + calls["count"] += 1 + text = f"reply {calls['count']}" + yield text + yield { + "type": "metadata", + "usage": { + "prompt_tokens": 3, + "completion_tokens": calls["count"], + "total_tokens": 3 + calls["count"], + }, + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + n = 2, + messages = [{"role": "user", "content": "two please"}], + ) + + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + body = json.loads(response.body) + + assert [c["message"]["content"] for c in body["choices"]] == ["reply 1", "reply 2"] + [entry] = monitor.snapshot() + assert entry["reply"] == "Choice 1:\nreply 1\n\nChoice 2:\nreply 2" + assert entry["completion_tokens"] == 3 + assert monitor.active_count() == 0 def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod @@ -1336,6 +1461,7 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = False, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _generate, ) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) @@ -1369,6 +1495,8 @@ class TestGgufVisionToolRouting: import routes.inference as inf_mod seen_seeds = [] + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) def _generate(**kwargs): seen_seeds.append(kwargs.get("seed")) @@ -1388,6 +1516,7 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = False, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _generate, ) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) @@ -1406,6 +1535,1081 @@ class TestGgufVisionToolRouting: assert seen_seeds == expected assert [choice["index"] for choice in body["choices"]] == [0, 1, 2] + assert body["usage"]["prompt_tokens"] == 5 + assert body["usage"]["completion_tokens"] == 21 + [entry] = monitor.snapshot() + assert entry["prompt_tokens"] == 5 + assert entry["completion_tokens"] == 21 + assert entry["total_tokens"] == 26 + + +class TestApiMonitorProviderAndCompletionStreams: + class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + async def is_disconnected(self): + return False + + def test_external_non_streaming_json_updates_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyExternalClient: + def __init__(self, **_kwargs): + pass + + async def stream_chat_completion(self, **kwargs): + assert kwargs["stream"] is False + yield json.dumps( + { + "choices": [{"message": {"content": "provider [DONE] reply"}}], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 4, + "total_tokens": 7, + }, + } + ) + + async def close(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "ExternalProviderClient", DummyExternalClient) + payload = ChatCompletionRequest( + model = "default", + external_model = "gpt-test", + provider_type = "openai", + provider_base_url = "https://api.openai.com/v1", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + response = await _proxy_to_external_provider(payload, self._Request()) + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + assert chunks[-1] == "data: [DONE]\n\n" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "provider [DONE] reply" + assert entry["prompt_tokens"] == 3 + assert entry["completion_tokens"] == 4 + assert entry["total_tokens"] == 7 + + asyncio.run(_run()) + + def test_external_stream_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyExternalClient: + def __init__(self, **_kwargs): + pass + + async def stream_chat_completion(self, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + await asyncio.sleep(3600) + + async def close(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "ExternalProviderClient", DummyExternalClient) + payload = ChatCompletionRequest( + model = "default", + external_model = "gpt-test", + provider_type = "openai", + provider_base_url = "https://api.openai.com/v1", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await _proxy_to_external_provider(payload, self._Request()) + iterator = response.body_iterator + first = await anext(iterator) + assert "hello" in first + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_preheader_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": True} + + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return None + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + response = await openai_completions(Request(), current_subject = "test") + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + assert chunks == [] + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_stream_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": True} + + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield b'data: {"choices":[{"text":"hello"}]}\n\n' + await asyncio.sleep(3600) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + + response = await openai_completions(Request(), current_subject = "test") + iterator = response.body_iterator + first = await anext(iterator) + assert b"hello" in first + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_non_streaming_post_error_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False} + + class FailingAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **_kwargs): + raise httpx.ConnectError("llama down") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FailingAsyncClient(), + ) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + with pytest.raises(httpx.ConnectError): + await openai_completions(Request(), current_subject = "test") + + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "Lost connection to the model server" in entry["error"] + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_monitor_openai_chunk_records_all_choice_replies(self, monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + _monitor_openai_chunk( + monitor_id, + { + "choices": [ + {"text": "first"}, + {"text": "second"}, + ], + "usage": { + "prompt_tokens": 2, + "completion_tokens": 5, + "total_tokens": 7, + }, + }, + 4096, + ) + + entry = monitor.get(monitor_id) + assert entry["reply"] == "Choice 1:\nfirst\n\nChoice 2:\nsecond" + assert entry["prompt_tokens"] == 2 + assert entry["completion_tokens"] == 5 + assert entry["context_length"] == 4096 + + def test_monitor_openai_chunk_records_tool_call_reply(self, monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + _monitor_openai_chunk( + monitor_id, + { + "choices": [ + { + "message": { + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"query":"weather"}', + }, + } + ] + } + } + ] + }, + 4096, + ) + + entry = monitor.get(monitor_id) + assert entry["reply"] == 'Tool call: lookup({"query":"weather"})' + + def test_embeddings_request_is_counted_active_and_completed(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/embeddings") + method = "POST" + + async def json(self): + return {"input": ["alpha", "beta"], "model": "embed"} + + class FakeAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **_kwargs): + assert monitor.active_count() == 1 + return httpx.Response( + 200, + json = { + "data": [{"embedding": [0.1]}], + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeAsyncClient(), + ) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + response = await openai_embeddings(Request(), current_subject = "test") + + assert response.status_code == 200 + [entry] = monitor.snapshot() + assert entry["endpoint"] == "/v1/embeddings" + assert entry["status"] == "completed" + assert entry["prompt_preview"] == "alpha\nbeta" + assert entry["prompt_tokens"] == 4 + assert entry["total_tokens"] == 4 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + await asyncio.sleep(3600) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + first = await anext(iterator) + assert "hello" in first + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class CancellingAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **_kwargs): + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: CancellingAsyncClient(), + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + assert chunks == ['data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'] + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + +class TestApiMonitorSafetensorsUsage: + class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + def test_non_streaming_safetensors_records_usage(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyBackend: + active_model_name = "safe-model" + models = {"safe-model": {"context_length": 2048}} + + def generate_chat_response(self, *, stats_holder, **_kwargs): + stats_holder["stats"] = { + "usage": { + "prompt_tokens": 8, + "completion_tokens": 5, + "total_tokens": 13, + } + } + yield "safe reply" + + def reset_generation_state(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, + "_detect_safetensors_features", + lambda *_args, **_kwargs: {"supports_tools": False}, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + body = json.loads(response.body) + + assert body["choices"][0]["message"]["content"] == "safe reply" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "safe reply" + assert entry["prompt_tokens"] == 8 + assert entry["completion_tokens"] == 5 + assert entry["total_tokens"] == 13 + assert entry["context_length"] == 2048 + + asyncio.run(_run()) + + def test_non_streaming_safetensors_tool_cancel_records_cancelled(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + reset_tool_policy() + + class DummyBackend: + active_model_name = "safe-model" + models = {"safe-model": {"context_length": 2048}} + + def generate_chat_response(self, **_kwargs): + raise AssertionError("plain safetensors path should not be used") + + def generate_chat_completion_with_tools( + self, *, cancel_event, stats_holder, **_kwargs + ): + stats_holder["stats"] = { + "usage": { + "prompt_tokens": 8, + "completion_tokens": 5, + "total_tokens": 13, + } + } + yield {"type": "content", "text": "partial"} + cancel_event.set() + yield {"type": "content", "text": "ignored"} + + def reset_generation_state(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, + "_detect_safetensors_features", + lambda *_args, **_kwargs: {"supports_tools": True}, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + enable_tools = True, + enabled_tools = ["web_search"], + cancel_id = "safe-cancel", + ) + + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + body = json.loads(response.body) + + assert body["choices"][0]["message"]["content"] == "partial" + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "partial" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_non_streaming_safetensors_tool_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + reset_tool_policy() + reset_called = False + + class DummyBackend: + active_model_name = "safe-model" + models = {"safe-model": {"context_length": 2048}} + + def generate_chat_response(self, **_kwargs): + raise AssertionError("plain safetensors path should not be used") + + def generate_chat_completion_with_tools(self, **_kwargs): + yield {"type": "content", "text": "unused"} + + def reset_generation_state(self): + nonlocal reset_called + reset_called = True + + async def fake_to_thread(*_args, **_kwargs): + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod.asyncio, "to_thread", fake_to_thread) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, + "_detect_safetensors_features", + lambda *_args, **_kwargs: {"supports_tools": True}, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + enable_tools = True, + enabled_tools = ["web_search"], + cancel_id = "safe-cancel", + ) + + with pytest.raises(asyncio.CancelledError): + await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert reset_called is True + + asyncio.run(_run()) + + +class TestApiMonitorAudioInput: + def _patch_audio_backend(self, monkeypatch, chunks): + import routes.inference as inf_mod + + class DummyAudioBackend: + active_model_name = "audio-model" + models = { + "audio-model": { + "has_audio_input": True, + "audio_type": "audio-input", + } + } + + def generate_audio_input_response(self, **_kwargs): + yield from chunks + + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False), + ) + monkeypatch.setattr( + inf_mod, + "get_inference_backend", + lambda: DummyAudioBackend(), + ) + monkeypatch.setattr( + inf_mod, + "_decode_audio_base64", + lambda _payload: object(), + ) + return inf_mod + + def test_audio_input_non_streaming_records_active_monitor(self, monkeypatch): + async def _run(): + inf_mod = self._patch_audio_backend(monkeypatch, ["hello", " world"]) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "describe this audio")], + audio_base64 = "ZmFrZQ==", + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + response = await openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + body = json.loads(response.body) + + assert body["choices"][0]["message"]["content"] == "hello world" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hello world" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_audio_input_streaming_records_monitor_reply(self, monkeypatch): + async def _run(): + inf_mod = self._patch_audio_backend(monkeypatch, ["hello", " world"]) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + async def is_disconnected(): + return False + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "describe this audio")], + audio_base64 = "ZmFrZQ==", + stream = True, + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + is_disconnected = is_disconnected, + ) + + response = await openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk) + + assert chunks[-1] == "data: [DONE]\n\n" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hello world" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_non_gguf_tts_auto_route_records_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyTtsBackend: + active_model_name = "tts-model" + models = { + "tts-model": { + "is_audio": True, + "audio_type": "snac", + } + } + + async def fake_generate_audio( + _payload, + _request, + current_subject = None, + ): + return inf_mod.JSONResponse( + content = { + "choices": [ + { + "message": { + "content": "[Generated audio]", + } + } + ] + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend()) + monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "say hello")], + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + response = await inf_mod.openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == ( + "[Generated audio]" + ) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["model"] == "tts-model" + assert entry["reply"] == "[Generated audio]" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_non_gguf_tts_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyTtsBackend: + active_model_name = "tts-model" + models = { + "tts-model": { + "is_audio": True, + "audio_type": "snac", + } + } + + async def fake_generate_audio( + _payload, + _request, + current_subject = None, + ): + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend()) + monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "say hello")], + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + with pytest.raises(asyncio.CancelledError): + await inf_mod.openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["model"] == "tts-model" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_gguf_tts_auto_route_records_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_generate_audio( + _payload, + _request, + current_subject = None, + ): + return inf_mod.JSONResponse( + content = { + "choices": [ + { + "message": { + "content": "[Generated audio]", + } + } + ] + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + _is_audio = True, + model_identifier = "gguf-tts", + context_length = 2048, + ), + ) + monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "say hello")], + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + await inf_mod.openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["model"] == "gguf-tts" + assert entry["context_length"] == 2048 + assert entry["reply"] == "[Generated audio]" + assert monitor.active_count() == 0 + + asyncio.run(_run()) # ===================================================================== diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 89155c2daf..0bea355668 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -40,6 +40,7 @@ from fastapi import HTTPException from fastapi.responses import JSONResponse from pydantic import ValidationError +from core.inference.api_monitor import ApiMonitor from models.inference import ( ChatMessage, ResponsesFunctionCallInputItem, @@ -781,6 +782,129 @@ class TestResponsesNonStreamingAdapter: assert "" not in body["output"][1]["content"][0]["text"] assert "" not in body["output"][1]["content"][0]["text"] + def test_monitor_records_translated_visible_text(self, monkeypatch): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + assert request.state.skip_api_monitor is True + return JSONResponse( + content = { + "model": "test-model", + "choices": [{"message": {"content": "plananswer"}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/responses"), + method = "POST", + ) + + async def run(): + response = await _responses_non_streaming(payload, messages, request) + return json.loads(response.body.decode()) + + body = asyncio.run(run()) + + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][1]["content"][0]["text"] == "answer" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "answer" + assert entry["prompt_tokens"] == 2 + assert entry["completion_tokens"] == 3 + assert request.state.skip_api_monitor is False + + def test_monitor_records_tool_only_reply(self, monkeypatch): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + assert request.state.skip_api_monitor is True + return JSONResponse( + content = { + "model": "test-model", + "choices": [ + { + "message": { + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"query":"weather"}', + }, + } + ], + } + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + payload = ResponsesRequest( + input = "hi", + tools = [{"type": "function", "name": "lookup"}], + ) + messages = [ChatMessage(role = "user", content = "hi")] + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/responses"), + method = "POST", + ) + + async def run(): + response = await _responses_non_streaming(payload, messages, request) + return json.loads(response.body.decode()) + + body = asyncio.run(run()) + + assert body["output"][0]["type"] == "function_call" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == 'Tool call: lookup({"query":"weather"})' + assert request.state.skip_api_monitor is False + + def test_cancelled_chat_completion_finalizes_monitor(self, monkeypatch): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + assert request.state.skip_api_monitor is True + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + payload = ResponsesRequest(input = "hi") + messages = [ChatMessage(role = "user", content = "hi")] + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/responses"), + method = "POST", + ) + + async def run(): + with pytest.raises(asyncio.CancelledError): + await _responses_non_streaming(payload, messages, request) + + asyncio.run(run()) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert request.state.skip_api_monitor is False + def test_literal_think_tags_remain_visible_without_reasoning_request(self, monkeypatch): body = self._run_with_message(monkeypatch, {"content": "show x tags"}) @@ -939,6 +1063,275 @@ class TestResponsesStreamAdapter: assert completed["response"]["output"][0]["content"][0]["text"] == "plan" assert completed["response"]["output"][1]["content"][0]["text"] == "33" + def test_usage_only_chunk_updates_monitor(self, monkeypatch): + import routes.inference as inf_mod + + chunks = [ + {"choices": [{"delta": {"content": "33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + asyncio.run(run()) + + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "33" + assert entry["prompt_tokens"] == 2 + assert entry["completion_tokens"] == 3 + assert entry["total_tokens"] == 5 + assert entry["context_length"] == 4096 + + def test_function_call_chunk_updates_monitor_reply(self, monkeypatch): + import routes.inference as inf_mod + + chunks = [ + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "function": { + "name": "lookup", + "arguments": '{"query":"weather"}', + }, + } + ] + } + } + ] + } + ] + self._install_stream_mock(monkeypatch, chunks) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert self._payloads(lines, "response.output_item.done")[-1]["item"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == 'Tool call: lookup({"query":"weather"})' + + def test_preheader_cancel_finalizes_monitor(self, monkeypatch): + import routes.inference as inf_mod + + self._install_stream_mock(monkeypatch, []) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + + async def fake_send(*_args, **_kwargs): + return None + + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + asyncio.run(run()) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + def test_stream_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + await asyncio.sleep(3600) + + self._install_stream_mock(monkeypatch, []) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + iterator = response.body_iterator + first = "" + for _ in range(8): + first = await anext(iterator) + if "hello" in first: + break + else: + pytest.fail("stream did not emit text delta") + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_final_visible_text_updates_monitor(self, monkeypatch): + import routes.inference as inf_mod + + class FakeExtractor: + def __init__(self, **_kwargs): + pass + + def feed( + self, + _content, + _reasoning_content = None, + ): + return "", "" + + def finish(self): + return "", "tail" + + self._install_stream_mock(monkeypatch, [{"choices": [{"delta": {"content": ""}}]}]) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_ResponsesReasoningExtractor", FakeExtractor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert self._payloads(lines, "response.output_text.delta")[-1]["delta"] == "plan" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plan" + def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch): chunks = [ {"choices": [{"delta": {"content": "show pool_usable # old behavior over-spent the cushion + + +def test_tensor_unknown_total_keeps_fraction_cushion(): + # A two-column nvidia-smi probe yields total 0. The planner must fall back to + # free*frac (keep the 5% cushion), like _select_gpus/_gpu_usable, not raw free, + # or it over-advertises context exactly where the PR is hardening the budget. + b = _kv_seeded_backend() + gpus = [(0, 20000), (1, 20000)] + MIB = 1024 * 1024 + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + model = int(18 * _GB) + ec_zero, *_ = b._plan_tensor_parallel(gpus, model, 131072, total_by_idx = {0: 0, 1: 0}) + ec_none, *_ = b._plan_tensor_parallel(gpus, model, 131072) + assert ec_zero == ec_none # total 0 == total absent: both use free*frac + pool_free = sum(f for _, f in gpus) + foot = (model + b._estimate_kv_cache_bytes(ec_zero, None)) / MIB + len(gpus) * reserve + assert foot <= pool_free * _CTX_FIT_VRAM_FRACTION + 2 # within free*frac, not raw free + + +def test_tensor_reserve_scales_with_ubatch(): + # A user --ubatch override must enlarge the per-device reserve -> less ctx room. + b = _kv_seeded_backend() + b._vocab_size = 152064 # enable the deterministic compute-buffer estimate + gpus = [(0, 16000), (1, 16000)] + model = int(18 * _GB) + small_ub, *_ = b._plan_tensor_parallel(gpus, model, 131072, n_ubatch = 512) + big_ub, *_ = b._plan_tensor_parallel(gpus, model, 131072, n_ubatch = 4096) + assert big_ub < small_ub + + +def test_plan_tensor_carries_unsized_mtp_flat_reserve(): + # review run3 #1/#5: with a weights-only (KV-unsized) MTP reserve, the planner + # gets a non-None mtp_overhead_fn but must still subtract the flat unsized-KV + # cushion, or its binary search spends it on context. Passing the reserve must + # pick a strictly smaller context than passing 0. + b = _kv_seeded_backend() + gpus = [(0, 14000), (1, 14000)] # tight pool so the context is actually capped + model = int(8 * _GB) + weights_only = lambda c: 3 * _GB # noqa: E731 -- constant drafter weights, no KV term + ctx_no_flat, *_ = b._plan_tensor_parallel( + gpus, + model, + 131072, + mtp_engaged = True, + mtp_overhead_fn = weights_only, + mtp_flat_reserve_bytes = 0, + ) + ctx_flat, *_ = b._plan_tensor_parallel( + gpus, + model, + 131072, + mtp_engaged = True, + mtp_overhead_fn = weights_only, + mtp_flat_reserve_bytes = 2 * _GB, + ) + assert 0 < ctx_flat < ctx_no_flat + + +def test_tensor_admission_drops_gpu_below_usable_budget(): + # A partly-used big card can clear the buffer reserve on raw free yet have no + # usable budget left (free - 0.05*total). Admit by usable budget: GPU 0 here is + # 6000 free on an 80 GB card -> usable 1904 < flat reserve 5120, so it's dropped + # (leaving <2 -> no split). Without total_by_idx, raw free 6000 >= 5120 admits it. + b = _kv_seeded_backend() + gpus = [(0, 6000), (1, 40000)] + totals = {0: 81920, 1: 81920} + _ec, _mac, gi, ts = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192, total_by_idx = totals) + assert gi == [1] and ts is None # GPU 0 excluded on usable budget + _ec2, _mac2, gi_raw, _ts2 = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192) + assert gi_raw == [0, 1] # raw free would have admitted both + + +def test_load_model_tensor_admission_and_capacity_gate_use_usable_budget(): + # load_model is too entangled (subprocess + GPU probe) to drive end-to-end, so + # assert at the source level that the tensor prefilter admits on the usable + # budget (_gpu_usable), not raw free, and downgrades to layer split when the + # pooled budget can't hold weights + per-device compute buffers. + src = inspect.getsource(LlamaCppBackend.load_model) + assert "_gpu_usable(g) >= reserve_mib" in src # admit by usable budget + assert "g[1] >= reserve_mib" not in src # not raw free + assert "_tp_weight_budget_mib" in src # pooled-weight capacity gate + assert "falling back to layer split" in src # downgrade on overcommit + # The gate's required footprint must include the non-shrinkable MTP reserve, + # not weights alone, or a separate-drafter MTP load can still overcommit. + assert "_tp_mtp_floor" in src + assert "model_size + _tp_mtp_floor" in src + + +def test_load_model_tensor_floor_keeps_flat_reserve_for_weights_only(): + # Tensor mode has no --fit valve, so a weights-only drafter (KV unsized) must + # keep the flat reserve as the draft-KV cushion, not just the byte weights + # (Finding H1, the tensor analog of the layer-split _mtp_kv_unsized handling). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + # byte-only floor used only when KV is sizable (not the weights-only case) + assert "mtp_overhead_fnisnotNoneandnot_mtp_kv_unsized" in compact + # weights-only / dims-unavailable: flat reserve, never below the byte floor + assert "_tp_mtp_floor=max(" in compact + + +def test_load_model_reserves_pipeline_per_device_overhead(): + # Layer split must reserve the fixed per-device overhead per EXTRA device so a + # tight multi-GPU split can't pin a context that OOMs a device (Finding A); k=1 + # adds nothing. + assert LlamaCppBackend._PIPELINE_PER_DEVICE_OVERHEAD_MIB > 0 + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "def_subset_model_size(n_gpus:int)->int:" in compact + assert "max(0,n_gpus-1)*_pipeline_overhead_bytes" in compact + assert "_subset_model_size(n_gpus)" in compact # used in the layer-split fit + + +def test_load_model_restores_quantized_kv_on_tensor_downgrade(): + # A quantized KV dropped for the tensor attempt must be restored if tensor + # downgrades to layer split (Finding D); captured once, restored at both the + # GPU-count and capacity-gate downgrades. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_tensor_dropped_cache_type_kv=cache_type_kv" in compact # captured pre-null + assert compact.count("cache_type_kv=_tensor_dropped_cache_type_kv") >= 2 # restored diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py new file mode 100644 index 0000000000..393a74daaf --- /dev/null +++ b/studio/backend/tests/test_torchao_select.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for _select_torchao_spec in install_python_stack.py. + +torchao's C++ extensions are built against one exact torch release, so the +installer must pick the torchao version matching the torch installed in the +venv (otherwise the cpp kernels are skipped). This pins that mapping. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +# install_python_stack.py lives at repo_root/studio/install_python_stack.py +_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py" + + +def _load_module(monkeypatch): + """(Re-)import install_python_stack and return it (mirrors test_pytorch_mirror).""" + sys.modules.pop("install_python_stack", None) + monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent)) + import install_python_stack + + return install_python_stack + + +@pytest.mark.parametrize( + "torch_version, expected", + [ + # torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0, + # independent of the local +cuXXX/+rocm/+cpu suffix or patch level. + ("2.10.0+cu130", "torchao==0.16.0"), + ("2.10.0+rocm6.4", "torchao==0.16.0"), + ("2.10.0+cpu", "torchao==0.16.0"), + ("2.10.1", "torchao==0.16.0"), + ("2.10.0", "torchao==0.16.0"), + # Pre-release / dev / rc builds: the minor is cleaned of non-digits. + ("2.10.0rc1", "torchao==0.16.0"), + ("2.10.0.dev20250804+cu130", "torchao==0.16.0"), + ("2.10rc1", "torchao==0.16.0"), + # torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0. + ("2.11.0+cu130", "torchao==0.17.0"), + ("2.11.0", "torchao==0.17.0"), + ("2.12.0", "torchao==0.17.0"), + # torch <=2.9 keeps today's pin (already a correct match for 2.9.0). + ("2.9.0+cu128", "torchao==0.14.0"), + ("2.9.1", "torchao==0.14.0"), + ("2.8.0", "torchao==0.14.0"), + ("2.4.0", "torchao==0.14.0"), + # Unparseable / missing / non-2.x major -> conservative default. + (None, "torchao==0.14.0"), + ("", "torchao==0.14.0"), + ("garbage", "torchao==0.14.0"), + ("2", "torchao==0.14.0"), + ("3.0.0", "torchao==0.14.0"), + ], +) +def test_select_torchao_spec(monkeypatch, torch_version, expected): + mod = _load_module(monkeypatch) + assert mod._select_torchao_spec(torch_version) == expected + + +def test_default_spec_matches_table(monkeypatch): + """The default/floor stays the historical pin so older torch is unchanged.""" + mod = _load_module(monkeypatch) + assert mod._TORCHAO_DEFAULT_SPEC == "torchao==0.14.0" + assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC diff --git a/studio/backend/tests/test_training_xet_fallback.py b/studio/backend/tests/test_training_xet_fallback.py new file mode 100644 index 0000000000..b4a3864334 --- /dev/null +++ b/studio/backend/tests/test_training_xet_fallback.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Parent-side training Xet->HTTP fallback: a model-load stall respawns the +worker once with Xet disabled, preserving the DB run row. Driven via +_handle_event with a fake spawn context; no GPU, no network, no real subprocess. +""" + +from __future__ import annotations + +import contextlib +import logging +import queue +import sys +import threading +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub the heavy module-level imports of core/training/training.py so it imports +# under CPU-only/no-network, then restore them (see the restore loop below). +_SAVED: dict = {} + + +def _stub(name, mod): + _SAVED[name] = sys.modules.get(name) + sys.modules[name] = mod + + +_lg = _types.ModuleType("loggers") +_lg.get_logger = lambda name: logging.getLogger(name) +_stub("loggers", _lg) +_stub("structlog", _types.ModuleType("structlog")) +_mpl = _types.ModuleType("matplotlib") +_plt = _types.ModuleType("matplotlib.pyplot") +_plt.Figure = type("Figure", (), {}) # referenced in a class-def annotation +_mpl.pyplot = _plt +_stub("matplotlib", _mpl) +_stub("matplotlib.pyplot", _plt) +_hw = _types.ModuleType("utils.hardware") +_hw.prepare_gpu_selection = lambda *a, **k: (None, None) +_stub("utils.hardware", _hw) +_npl = _types.ModuleType("utils.native_path_leases") +_npl.native_path_secret_removed_for_child_start = lambda: contextlib.nullcontext() +_npl.run_without_native_path_secret = lambda fn: fn +_stub("utils.native_path_leases", _npl) +_pth = _types.ModuleType("utils.paths") +_pth.outputs_root = lambda *a, **k: "/tmp/outputs" +_stub("utils.paths", _pth) + +import core.training.training as training_mod +from core.training.training import TrainingBackend + +# Restore every stubbed module so this file never pollutes the shared session: a +# leaked bare ``structlog`` (no ``get_logger``) would break every later module +# that logs at import. training_mod already bound the stubs it needs at runtime. +for _name in ( + "loggers", + "structlog", + "matplotlib", + "matplotlib.pyplot", + "utils.hardware", + "utils.native_path_leases", + "utils.paths", +): + _prev = _SAVED.get(_name) + if _prev is None: + sys.modules.pop(_name, None) + else: + sys.modules[_name] = _prev + + +@pytest.fixture(autouse = True) +def _stub_worker_module(): + """Stub ``core.training.worker`` so the respawn's lazy import of the + torch-heavy worker is never required.""" + prev = sys.modules.get("core.training.worker") + stub = _types.ModuleType("core.training.worker") + stub.run_training_process = lambda **kwargs: None + sys.modules["core.training.worker"] = stub + yield + if prev is None: + sys.modules.pop("core.training.worker", None) + else: + sys.modules["core.training.worker"] = prev + + +class _FakeProc: + def __init__(self, **kwargs): + self._alive = True + self.pid = 4321 + self.kwargs = kwargs + + def start(self): + pass + + def is_alive(self): + return self._alive + + def terminate(self): + self._alive = False + + def kill(self): + self._alive = False + + def join(self, timeout = None): + self._alive = False + + +class _FakeQueue: + def put(self, *a, **k): + pass + + def get(self, *a, **k): + raise queue.Empty + + +class _FakeCtx: + def __init__(self): + self.spawned: list = [] + + def Queue(self): + return _FakeQueue() + + def Process(self, **kwargs): + self.spawned.append(kwargs) + return _FakeProc(**kwargs) + + +def _backend_mid_load(): + b = TrainingBackend() + b._last_full_config = {"model_name": "org/model", "disable_xet": False, "hf_token": "tok"} + b._in_model_load = True + b._xet_fallback_used = False + proc = _FakeProc() + b._proc = proc + return b, proc + + +def test_stall_during_load_arms_respawn_and_terminates_worker(): + b, proc = _backend_mid_load() + b._handle_event({"type": "stall", "message": "no progress for 180s"}) + assert b._needs_xet_respawn is True + assert b._xet_fallback_used is True + assert proc.is_alive() is False, "stalled worker must be terminated" + + +def test_respawn_uses_disable_xet_and_preserves_run_row(monkeypatch): + b, _ = _backend_mid_load() + b._handle_event({"type": "stall", "message": "x"}) + + fake_ctx = _FakeCtx() + monkeypatch.setattr(training_mod, "_CTX", fake_ctx) + monkeypatch.setattr(b, "_pump_loop", lambda: None) # neutralize the new pump + created = {"n": 0} + finalized = {"n": 0} + monkeypatch.setattr( + b, "_ensure_db_run_created", lambda: created.__setitem__("n", created["n"] + 1) + ) + monkeypatch.setattr( + b, "_finalize_run_in_db", lambda **k: finalized.__setitem__("n", finalized["n"] + 1) + ) + + b._respawn_worker_disable_xet() + + assert len(fake_ctx.spawned) == 1, "respawn must start exactly one worker" + cfg = fake_ctx.spawned[0]["kwargs"]["config"] + assert cfg["disable_xet"] is True, "respawned worker must run with Xet disabled" + assert cfg["model_name"] == "org/model" + assert created["n"] == 0, "respawn must not recreate the DB run row" + assert finalized["n"] == 0, "a successful respawn must not finalize the run as error" + + +def test_second_stall_surfaces_error_without_respawn(): + b, proc = _backend_mid_load() + b._xet_fallback_used = True # HTTP fallback already spent + b._handle_event({"type": "stall", "message": "stalled again over http"}) + assert b._needs_xet_respawn is False + assert b._progress.error and "stalled" in b._progress.error.lower() + assert proc.is_alive() is False + + +def test_model_load_completed_disarms_recovery(): + b, _ = _backend_mid_load() + b._handle_event({"type": "model_load_completed"}) + assert b._in_model_load is False + # A stall after the load finished is not a transport stall to recover from. + b._handle_event({"type": "stall", "message": "post-load"}) + assert b._needs_xet_respawn is False + + +def test_model_load_started_arms_recovery_window(): + b = TrainingBackend() + assert b._in_model_load is False + b._handle_event({"type": "model_load_started"}) + assert b._in_model_load is True + + +def test_child_should_disable_xet_truth_table(): + from utils.hf_xet_fallback import child_should_disable_xet + + assert child_should_disable_xet({"disable_xet": True}) is True + assert child_should_disable_xet({"disable_xet": False}) is False + assert child_should_disable_xet({}) is False diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index c66d56528a..cc7d04ca0a 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -88,6 +88,19 @@ class TestGetDevice: ): assert _reset_and_detect() == DeviceType.CUDA + @needs_torch + def test_detect_survives_device0_probe_failure(self, capsys): + # is_available() True but the device-0 name probe raises: startup must + # still resolve CUDA rather than crash. + with ( + patch("utils.hardware.hardware._has_torch", return_value = True), + patch("torch.cuda.is_available", return_value = True), + patch("torch.cuda.device_count", return_value = 1), + patch("torch.cuda.get_device_properties", side_effect = RuntimeError("probe")), + ): + assert _reset_and_detect() == DeviceType.CUDA + assert "" in capsys.readouterr().out + @needs_mlx def test_returns_mlx_when_on_apple_silicon_with_mlx(self): with ( @@ -303,6 +316,103 @@ class TestLogGpuMemory: assert "No GPU available" in captured.out +# ========== CUDA_DEVICE_ORDER pinning ========== + + +class TestCudaDeviceOrder: + """Importing the hardware module pins CUDA_DEVICE_ORDER=PCI_BUS_ID when unset, + but setdefault keeps an explicit user override, so nvidia-smi indices, torch + ordinals, and CUDA_VISIBLE_DEVICES agree on a mixed-GPU host.""" + + @staticmethod + def _order_after_fresh_import(preset): + # Fresh interpreter so the module-level setdefault runs against a clean env. + import os, subprocess, sys + from pathlib import Path + + env = os.environ.copy() + backend = str(Path(__file__).resolve().parents[1]) + existing = env.get("PYTHONPATH", "") + # Avoid a trailing os.pathsep (empty entry -> cwd on sys.path) when unset. + env["PYTHONPATH"] = (backend + os.pathsep + existing) if existing else backend + if preset is None: + env.pop("CUDA_DEVICE_ORDER", None) + else: + env["CUDA_DEVICE_ORDER"] = preset + out = subprocess.run( + [ + sys.executable, + "-c", + "import os, utils.hardware.hardware; print(os.environ.get('CUDA_DEVICE_ORDER'))", + ], + env = env, + capture_output = True, + text = True, + check = True, + ) + return out.stdout.strip().splitlines()[-1] + + def test_import_pins_pci_bus_id_when_unset(self): + assert self._order_after_fresh_import(None) == "PCI_BUS_ID" + + def test_import_respects_explicit_user_override(self): + assert self._order_after_fresh_import("FASTEST_FIRST") == "FASTEST_FIRST" + + +# ========== _print_cuda_device_list() ========== + + +class TestPrintCudaDeviceList: + """The startup console lists every CUDA GPU with its index, not just + device 0, so a multi-GPU host shows the full available set.""" + + @needs_torch + def test_lists_all_devices_when_multi_gpu(self, capsys): + props = [ + MagicMock(name = "p0"), + MagicMock(name = "p1"), + ] + props[0].name = "NVIDIA GeForce RTX 5090" + props[1].name = "NVIDIA RTX PRO 6000 Blackwell Workstation Edition" + with ( + patch("torch.cuda.device_count", return_value = 2), + patch("torch.cuda.get_device_properties", side_effect = lambda i: props[i]), + ): + _hw_module._print_cuda_device_list(is_rocm = False) + out = capsys.readouterr().out + assert "[0] NVIDIA GeForce RTX 5090" in out + assert "[1] NVIDIA RTX PRO 6000 Blackwell Workstation Edition" in out + assert "CUDA_DEVICE_ORDER=" in out + + @needs_torch + def test_silent_on_single_gpu(self, capsys): + with patch("torch.cuda.device_count", return_value = 1): + _hw_module._print_cuda_device_list(is_rocm = False) + assert capsys.readouterr().out == "" + + @needs_torch + def test_never_raises_on_probe_failure(self, capsys): + with patch("torch.cuda.device_count", side_effect = RuntimeError("no cuda")): + _hw_module._print_cuda_device_list(is_rocm = False) + assert capsys.readouterr().out == "" + + @needs_torch + def test_rocm_label_omits_cuda_device_order(self, capsys): + # CUDA_DEVICE_ORDER governs CUDA only, so the ROCm listing must not claim it. + props = [MagicMock(), MagicMock()] + props[0].name = "AMD Instinct MI300X" + props[1].name = "AMD Instinct MI300X" + with ( + patch("torch.cuda.device_count", return_value = 2), + patch("torch.cuda.get_device_properties", side_effect = lambda i: props[i]), + ): + _hw_module._print_cuda_device_list(is_rocm = True) + out = capsys.readouterr().out + assert "ROCm devices (2):" in out + assert "CUDA_DEVICE_ORDER" not in out + assert "[0] AMD Instinct MI300X" in out + + # ========== format_error_message() ========== diff --git a/studio/backend/tests/test_validate_model_error.py b/studio/backend/tests/test_validate_model_error.py new file mode 100644 index 0000000000..f752f8eada --- /dev/null +++ b/studio/backend/tests/test_validate_model_error.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""routes/inference.py::validate_model surfaces actionable RuntimeError/ValueError +messages (e.g. "llama-server binary not found - run setup.sh") instead of a blank +"Invalid model", while keeping unexpected exceptions generic so internals never +leak to the client. +""" + +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +pytest.importorskip("fastapi") + +from fastapi import HTTPException # noqa: E402 + +import routes.inference as inf # noqa: E402 +from models.inference import ValidateModelRequest # noqa: E402 + + +def _provoke( + monkeypatch, + exc: BaseException, + *, + native: bool = False, +) -> HTTPException: + """Drive validate_model so from_identifier raises ``exc``; return the + HTTPException it converts that into.""" + monkeypatch.setattr( + inf, + "_resolve_model_identifier_for_request", + lambda request, operation: ("org/repo", "org/repo", native), + ) + + def _raise(*_args, **_kwargs): + raise exc + + monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(_raise)) + + req = ValidateModelRequest(model_path = "org/repo") + with pytest.raises(HTTPException) as excinfo: + asyncio.run(inf.validate_model(req, current_subject = "tester")) + return excinfo.value + + +def test_runtime_error_surfaces_actionable_message(monkeypatch): + err = RuntimeError( + "llama-server binary not found - cannot load GGUF models. " + "Run setup.sh to build it, or set LLAMA_SERVER_PATH." + ) + http = _provoke(monkeypatch, err) + assert http.status_code == 400 + assert "llama-server binary not found" in http.detail + assert http.detail != "Invalid model" + + +def test_value_error_not_supported_is_wrapped(monkeypatch): + http = _provoke(monkeypatch, ValueError("architecture FooBar is not supported")) + assert http.status_code == 400 + assert "not supported yet" in http.detail.lower() + # Original cause is preserved for context. + assert "FooBar" in http.detail + + +def test_unexpected_exception_stays_generic(monkeypatch): + # A non-user-facing exception type must NOT have its message surfaced. + http = _provoke(monkeypatch, KeyError("secret-internal-detail")) + assert http.status_code == 400 + assert http.detail == "Invalid model" + assert "secret-internal-detail" not in http.detail + + +def test_empty_runtime_error_falls_back_to_generic(monkeypatch): + # A RuntimeError with no message should not produce an empty 400 detail. + http = _provoke(monkeypatch, RuntimeError("")) + assert http.status_code == 400 + assert http.detail == "Invalid model" diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py index 88a1a28d14..a4ec3f3fb5 100644 --- a/studio/backend/tests/test_windows_gpu_detection_mock.py +++ b/studio/backend/tests/test_windows_gpu_detection_mock.py @@ -211,6 +211,22 @@ class TestWindowsGpuDetectionAfter5106Fix: gpus = LlamaCppBackend._get_gpu_free_memory() assert gpus == [(1, 24576)], gpus + def test_get_gpu_memory_parses_three_and_two_column(self, monkeypatch): + """Total is parsed when present; a legacy two-column line or a non-integer + total ("N/A") yields total 0 (back-compat) rather than dropping the GPU, + which would silently spill to CPU.""" + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + with _mock_nvidia_smi_run("0, 22805, 24576\n"): + assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 24576)] + with _mock_nvidia_smi_run("0, 22805\n"): + assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 0)] + # A non-integer total must keep the GPU (total 0), not drop it. + with _mock_nvidia_smi_run("0, 22805, N/A\n"): + assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 0)] + # A bad free still skips that line (free is required). + with _mock_nvidia_smi_run("0, N/A, 24576\n1, 22805, 24576\n"): + assert LlamaCppBackend._get_gpu_memory() == [(1, 22805, 24576)] + def test_windows_install_dir_has_all_three_cudart_dlls(self, tmp_path): """All three bundle DLLs must land in install_dir/build/bin/ Release; any missing one breaks ggml-cuda.dll's PE import chain.""" diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 86dfa8a93f..08a470f11f 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -35,6 +35,21 @@ from typing import Optional, Dict, Any logger = get_logger(__name__) +# ── GPU index ordering ────────────────────────────────────────────────────── +# CUDA defaults to CUDA_DEVICE_ORDER=FASTEST_FIRST, numbering GPUs by compute +# performance. nvidia-smi -- and every free-VRAM probe in Studio -- numbers GPUs +# by PCI bus id instead. On a mixed-GPU host (e.g. an RTX 5090 alongside an RTX +# PRO 6000) the two orderings disagree, so an index picked from nvidia-smi data +# ("the emptiest card is GPU 1") gets written into CUDA_VISIBLE_DEVICES and then +# reinterpreted by CUDA against FASTEST_FIRST -- landing the model on a different +# physical GPU than the one selected. Pinning PCI_BUS_ID makes torch, nvidia-smi, +# and CUDA_VISIBLE_DEVICES share a single index space, matching what users see in +# `nvidia-smi -L`. Set at import (before any torch.cuda call latches the order +# at context creation) and inherited by child processes, since the llama-server +# and spawn workers copy os.environ. setdefault so an explicit user override wins. +os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") + + # ========== Device Enum ========== @@ -91,6 +106,40 @@ def _has_mlx() -> bool: return False +def _print_cuda_device_list(is_rocm: bool) -> None: + """List every visible CUDA/ROCm GPU with its index at startup. + + The "Hardware detected" banner names only device 0, which hides the other + cards on a multi-GPU host. This lists the full visible set in CUDA-ordinal + order, matching `nvidia-smi -L` when no CUDA_VISIBLE_DEVICES mask is set + (under a mask the indices are visible ordinals, not physical PCI ids). + CUDA_DEVICE_ORDER governs only CUDA, so it is shown for CUDA but not ROCm. + No-ops on single-GPU hosts and never raises -- it is purely informational. + """ + try: + import torch + + count = torch.cuda.device_count() + if count <= 1: + return + if is_rocm: + header = f"ROCm devices ({count}):" + else: + order = os.environ.get("CUDA_DEVICE_ORDER", "default") + header = f"CUDA devices ({count}, CUDA_DEVICE_ORDER={order}):" + lines = [header] + for i in range(count): + try: + name = torch.cuda.get_device_properties(i).name + except Exception as e: + logger.debug("CUDA device %d property probe failed: %s", i, e) + name = "" + lines.append(f" [{i}] {name}") + print("\n".join(lines)) + except Exception: + return # purely informational; never disrupt startup + + def detect_hardware() -> DeviceType: """ Detect the best compute device and set the module-level DEVICE global. @@ -112,7 +161,11 @@ def detect_hardware() -> DeviceType: if torch.cuda.is_available(): DEVICE = DeviceType.CUDA CHAT_ONLY = False - device_name = torch.cuda.get_device_properties(0).name + try: + device_name = torch.cuda.get_device_properties(0).name + except Exception as e: + logger.debug("CUDA device 0 property probe failed: %s", e) + device_name = "" # Distinguish ROCm from CUDA for display only (DeviceType stays CUDA). # AMD SDK wheels don't set torch.version.hip, so fall back to __version__. @@ -123,6 +176,7 @@ def detect_hardware() -> DeviceType: print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}") else: print(f"Hardware detected: CUDA -- {device_name}") + _print_cuda_device_list(IS_ROCM) return DEVICE # --- XPU: Intel GPU --- diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py new file mode 100644 index 0000000000..80f42f92d1 --- /dev/null +++ b/studio/backend/utils/hf_xet_fallback.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall. + +Xet (``hf_xet``) is the fast default but can hang with no progress and no +exception, and a blocked native thread cannot be killed. Keep Xet primary; fall +back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET`` +is read at import time, so the fallback runs in a fresh ``spawn`` child (not a +thread) that sets the env before importing ``huggingface_hub``. Cached files +short-circuit with no child; deterministic errors (401/403/404/disk-full) and +cancellation propagate without a fallback. Mirrors the safetensors inference +recovery in core/inference/{orchestrator,worker}.py. +""" + +from __future__ import annotations + +import multiprocessing as mp +import os +import queue +import signal +import sys +import threading +import time +from typing import Any, Callable, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +_CTX = mp.get_context("spawn") + +# Defaults match the existing inference watchdog and hub shutdown deadline. +DEFAULT_HEARTBEAT_INTERVAL = 30.0 +DEFAULT_STALL_TIMEOUT = 180.0 +DEFAULT_GRACE_PERIOD = 10.0 +_POLL_INTERVAL = 0.5 + + +class DownloadStallError(RuntimeError): + """Raised when no download progress is observed for too long. + + Canonical home; orchestrator.py re-imports it so all paths share one type. + """ + + +def child_should_disable_xet(config: dict) -> bool: + """Single source of truth for the per-worker Xet env flip.""" + return bool(config.get("disable_xet")) + + +def get_hf_download_state( + repo_ids: Optional[list[str]] = None, *, repo_type: str = "model" +) -> Optional[tuple[int, bool]]: + """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache. + + Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete`` + is not mistaken for full-size progress. ``None`` means the state could not be + measured, so callers skip stall logic for that tick. + """ + try: + from hub.utils.hf_cache_state import ( + blob_bytes_present, + has_active_incomplete_blobs, + hf_cache_root, + iter_active_repo_cache_dirs, + ) + + if hf_cache_root() is None: + return (0, False) + + total = 0 + has_incomplete = False + for repo_id in repo_ids or []: + # Skip local paths: HF IDs never start with / . ~ or contain "\". + if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id: + continue + for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + continue + for blob in blobs_dir.iterdir(): + try: + if blob.is_file(): + total += blob_bytes_present(blob) + except OSError: + pass + if has_active_incomplete_blobs(repo_type, repo_id): + has_incomplete = True + return (total, has_incomplete) + except Exception as e: + logger.debug("Failed to determine HF download state: %s", e) + return None + + +def start_watchdog( + *, + repo_ids: list[str], + on_stall: Callable[[str], None], + repo_type: str = "model", + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + stall_timeout: float = DEFAULT_STALL_TIMEOUT, + xet_disabled: bool = False, + on_heartbeat: Optional[Callable[[str], None]] = None, +) -> threading.Event: + """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a + ``*.incomplete`` is present AND the on-disk size is unchanged for + *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so + post-download init is never misread as a stall. Returns a stop event the + caller sets when the download phase ends. + """ + stop = threading.Event() + transport = "https" if xet_disabled else "xet" + fired = False + + def _beat() -> None: + nonlocal fired + state = get_hf_download_state(repo_ids, repo_type = repo_type) + last_size = state[0] if state is not None else 0 + last_change = time.monotonic() + + while not stop.wait(interval): + state = get_hf_download_state(repo_ids, repo_type = repo_type) + now = time.monotonic() + + if state is None: + if on_heartbeat is not None: + on_heartbeat(f"Downloading ({transport} transport)...") + continue + + current_size, has_incomplete = state + if current_size != last_size: + last_size = current_size + last_change = now + + # Reset unless .incomplete confirms an active download, so model init + # and lock waits are not counted as a stall. + if not has_incomplete: + last_change = now + elif now - last_change >= stall_timeout: + if not fired: + fired = True + on_stall( + f"Download appears stalled ({transport} transport) " + f"-- no progress for {int(now - last_change)}s" + ) + return + + if on_heartbeat is not None: + on_heartbeat(f"Downloading ({transport} transport)...") + + threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start() + return stop + + +def _download_child_entry( + *, + repo_id: str, + filename: str, + token: Optional[str], + repo_type: str, + disable_xet: bool, + result_queue: Any, +) -> None: + """Spawn-child entrypoint: download one file and report the result. + + Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub, + forms its own process group so the parent can kill the whole transfer, and + never logs the token or signed URLs. + """ + if hasattr(os, "setsid"): + try: + os.setsid() + except OSError: + pass + + if disable_xet: + os.environ["HF_HUB_DISABLE_XET"] = "1" + # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse + # partials a sequential resume cannot safely continue). + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + + # Test-only fault injection (never set in production): stall the Xet attempt + # so the watchdog + HTTP fallback can be exercised against a real repo. + if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1": + import time as _t + try: + from huggingface_hub.constants import HF_HUB_CACHE + + blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs") + os.makedirs(blobs, exist_ok = True) + with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh: + fh.write(b"\0" * 4096) + except OSError: + pass + while True: + _t.sleep(3600) + + try: + from huggingface_hub import hf_hub_download + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + repo_type = repo_type, + token = token, + ) + result_queue.put({"ok": True, "path": path}) + except BaseException as e: # noqa: BLE001 - report every failure to the parent + error = f"{type(e).__name__}: {e}" + try: + from hub.utils.download_registry import scrub_secrets + error = scrub_secrets(error, hf_token = token) + except Exception: + pass + result_queue.put({"ok": False, "error": error}) + + +def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None: + """Kill *proc* and its whole process group (Xet may spawn helper procs). + + The child calls ``os.setsid()`` so its pgid equals its pid; signal via + ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a + group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*. + """ + pid = proc.pid + + def _signal_group(sig: int) -> None: + if pid is not None and hasattr(os, "killpg"): + try: + os.killpg(pid, sig) + return + except (ProcessLookupError, PermissionError, OSError): + pass + # Windows or pre-setsid: best effort on the single process. + try: + proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill() + except Exception: + pass + + _signal_group(getattr(signal, "SIGTERM", signal.SIGINT)) + proc.join(timeout = grace_period) + if proc.is_alive(): + _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM)) + proc.join(timeout = 5.0) + + +def _run_download_attempt( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str, + disable_xet: bool, + cancel_event: Optional[threading.Event], + stall_timeout: float, + interval: float, + grace_period: float, + on_status: Optional[Callable[[str], None]], +) -> tuple[str, Optional[str]]: + """Run one download in a spawn child supervised by the no-progress watchdog. + + Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or + ``("error", message)``. This is the seam tests monkeypatch to avoid spawning. + """ + result_queue: Any = _CTX.Queue() + proc = _CTX.Process( + target = _download_child_entry, + kwargs = dict( + repo_id = repo_id, + filename = filename, + token = token, + repo_type = repo_type, + disable_xet = disable_xet, + result_queue = result_queue, + ), + daemon = True, + ) + proc.start() + + stalled = threading.Event() + stop_watchdog = start_watchdog( + repo_ids = [repo_id], + on_stall = lambda msg: stalled.set(), + repo_type = repo_type, + interval = interval, + stall_timeout = stall_timeout, + xet_disabled = disable_xet, + on_heartbeat = on_status, + ) + + result: Optional[dict] = None + try: + while proc.is_alive(): + if cancel_event is not None and cancel_event.is_set(): + _terminate_process_group(proc, grace_period) + return ("cancelled", None) + if stalled.is_set(): + _terminate_process_group(proc, grace_period) + return ("stall", None) + try: + result = result_queue.get(timeout = _POLL_INTERVAL) + break + except queue.Empty: + continue + else: + # Process exited; drain any result it enqueued. + try: + result = result_queue.get_nowait() + except queue.Empty: + result = None + finally: + stop_watchdog.set() + proc.join(timeout = grace_period) + + if result is None: + return ( + "error", + f"download process for '{repo_id}/{filename}' exited " + f"(code={proc.exitcode}) without a result", + ) + if result.get("ok"): + return ("ok", result["path"]) + return ("error", result.get("error") or "unknown download error") + + +def hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + cancel_event: Optional[threading.Event] = None, + repo_type: str = "model", + stall_timeout: float = DEFAULT_STALL_TIMEOUT, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + grace_period: float = DEFAULT_GRACE_PERIOD, + on_status: Optional[Callable[[str], None]] = None, +) -> str: + """Download a single file with Xet primary and HTTP as a stall-only fallback. + + Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if + *cancel_event* is set, re-raises a deterministic child error unchanged (no + fallback), and raises ``DownloadStallError`` only if BOTH transports stall. + """ + # Finalized blob already cached: return it with no child and no network. + try: + from huggingface_hub import try_to_load_from_cache + cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) + if isinstance(cached, str) and os.path.exists(cached): + return cached + except Exception as e: + logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) + + if cancel_event is not None and cancel_event.is_set(): + raise RuntimeError("Cancelled") + + disable_xet = False + for attempt in range(2): + if disable_xet: + # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume + # over a sparse Xet/hf_transfer partial silently corrupts the blob. + try: + from hub.utils.download_registry import prepare_cache_for_transport + prepare_cache_for_transport(repo_type, repo_id, "http") + except Exception as e: + logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e) + + kind, payload = _run_download_attempt( + repo_id, + filename, + token, + repo_type = repo_type, + disable_xet = disable_xet, + cancel_event = cancel_event, + stall_timeout = stall_timeout, + interval = interval, + grace_period = grace_period, + on_status = on_status, + ) + + if kind == "ok": + return payload # type: ignore[return-value] + if kind == "cancelled": + raise RuntimeError("Cancelled") + if kind == "error": + # Deterministic failure: the other transport would fail identically. + raise RuntimeError(payload) + # kind == "stall" + if attempt == 0 and not disable_xet: + logger.warning( + "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1", + repo_id, + filename, + ) + if on_status is not None: + on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP") + disable_xet = True + continue + raise DownloadStallError( + f"Download stalled for '{repo_id}/{filename}' even with " + f"HF_HUB_DISABLE_XET=1 -- check your network connection" + ) + + # Unreachable: the loop either returns or raises on each attempt. + raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'") diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index 87d0d2ec01..7d077bfa3b 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -33,6 +33,8 @@ _INSTALL_MARKER_NAME = "UNSLOTH_PREBUILT_INFO.json" _marker_cache: dict[str, Optional[dict]] = {} _release_memo: dict[str, tuple[float, Optional[str]]] = {} +# Newest-release asset sizes (name -> bytes), memoized like the tag (24h TTL). +_assets_memo: dict[str, tuple[float, dict[str, int]]] = {} def _cache_dir() -> Path: @@ -180,6 +182,113 @@ def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optio return latest +def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]: + """Asset name -> size (bytes) for the newest published release of `repo`, + selected exactly like _fetch_latest_release_tag. None on any failure.""" + import urllib.error + import urllib.request + + url = f"https://api.github.com/repos/{repo}/releases?per_page=30" + headers = { + "Accept": "application/vnd.github+json", + "User-Agent": "unsloth-studio-freshness-check", + } + token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") + if token: + headers["Authorization"] = f"Bearer {token}" + req = urllib.request.Request(url, headers = headers) + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + data = json.loads(resp.read().decode("utf-8")) + except ( + urllib.error.URLError, + urllib.error.HTTPError, + OSError, + json.JSONDecodeError, + ) as exc: + logger.debug("freshness asset fetch failed", repo = repo, error = str(exc)) + return None + if not isinstance(data, list): + return None + published = [ + r + for r in data + if isinstance(r, dict) + and not r.get("draft") + and not r.get("prerelease") + and isinstance(r.get("tag_name"), str) + and r.get("tag_name") + ] + if not published: + return None + newest = max(published, key = lambda r: r.get("published_at") or "") + assets: dict[str, int] = {} + for a in newest.get("assets") or []: + name, size = a.get("name"), a.get("size") + if isinstance(name, str) and isinstance(size, int): + assets[name] = size + return assets + + +def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]: + """Newest-release asset sizes for `repo`, memoized (24h TTL). None when + offline and never fetched. In-memory only -- a restart simply re-fetches.""" + if not repo: + return None + now = time.time() + if not force_refresh: + memo = _assets_memo.get(repo) + if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS: + return memo[1] + assets = _fetch_latest_release_assets(repo) + if assets is None: + memo = _assets_memo.get(repo) + return memo[1] if memo else None + _assets_memo[repo] = (now, assets) + return assets + + +def update_download_size_bytes( + marker: Optional[dict], + latest_tag: Optional[str], + repo: Optional[str], + *, + force_refresh: bool = False, +) -> Optional[int]: + """Download size of the latest-release asset matching this host's installed + bundle (same platform/arch/runtime suffix as the installed asset). None when + there is no marker asset, the latest assets can't be read, or no match.""" + if not marker or not latest_tag or not repo: + return None + installed_asset = marker.get("asset") + if not isinstance(installed_asset, str): + return None + # Tag-independent platform suffix: accept the fork's "app-*" bundles and the + # upstream ggml-org "ubuntu-*"/"win-*" prebuilts ("windows" before "win"). + m = re.search(r"-((?:linux|ubuntu|windows|win|macos|darwin)-.*)$", installed_asset) + if not m: + return None + suffix = m.group(1) + # Upstream ubuntu/win assets live in the marker's binary_repo, not the fork + # publish repo; try the publish repo first, then it. + repos = [repo] + binary_repo = marker.get("binary_repo") + if isinstance(binary_repo, str) and binary_repo and binary_repo != repo: + repos.append(binary_repo) + want = f"app-{latest_tag}-{suffix}" + for r in repos: + assets = latest_release_assets(r, force_refresh = force_refresh) + if not assets: + continue + if want in assets: + return assets[want] + # Tag formatting can vary (mix suffixes); fall back to the platform suffix. + for name, size in assets.items(): + if name.endswith(suffix): + return size + return None + + def _parse_installed_at(value: object) -> Optional[datetime]: if not isinstance(value, str) or not value: return None @@ -313,6 +422,7 @@ def reset_caches(*, drop_disk: bool = False) -> None: open (off) instead of pointing at the just-replaced build.""" _marker_cache.clear() _release_memo.clear() + _assets_memo.clear() if drop_disk: import shutil diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index c2c6c67432..22ffc58b24 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -37,9 +37,11 @@ from utils.llama_cpp_freshness import ( _INSTALL_MARKER_NAME, check_prebuilt_freshness, latest_published_release, + latest_release_assets, parse_base_build, read_install_marker, reset_caches, + update_download_size_bytes, ) logger = structlog.get_logger(__name__) @@ -291,6 +293,18 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: update_available = False # Display the mix tag when that's what makes it newer; otherwise the base. latest = release_tag if latest_is_mix else base_tag + # Size of the resolved prebuilt, so source builds show it like the marker + # path. Fails open to None (offline / asset absent from the release). + update_size_bytes = None + if update_available: + asset_name = res.get("asset") + if isinstance(asset_name, str) and asset_name: + try: + assets = latest_release_assets(res.get("repo"), force_refresh = force_refresh) + if assets: + update_size_bytes = assets.get(asset_name) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: source-build size lookup failed", error = str(exc)) with _job_lock: job = dict(_job) return { @@ -303,6 +317,7 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: "installed_at_utc": None, "age_days": None, "source_build": True, + "update_size_bytes": update_size_bytes, "job": job, } @@ -345,6 +360,20 @@ def get_update_status(*, force_refresh: bool = False) -> dict: # (see llama_cpp_freshness.is_behind). update_available = bool(freshness.get("has_marker") and freshness.get("behind")) + # Size of the prebuilt that Update would download, for the banner. Only when + # an update is offered; fails open to None (offline / no matching asset). + update_size_bytes = None + if update_available: + try: + update_size_bytes = update_download_size_bytes( + marker, + latest, + freshness.get("published_repo") or repo, + force_refresh = force_refresh, + ) + except Exception as exc: # pragma: no cover - network defensive + logger.debug("llama update: size lookup failed", error = str(exc)) + with _job_lock: job = dict(_job) @@ -358,6 +387,7 @@ def get_update_status(*, force_refresh: bool = False) -> dict: "installed_at_utc": freshness.get("installed_at_utc"), "age_days": freshness.get("age_days"), "source_build": False, + "update_size_bytes": update_size_bytes, "job": job, } diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index a2912cc843..c24ec28e1d 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -50,6 +50,10 @@ _CACHE_MAX_ENTRIES = 4096 # keyed by (file cache key, wanted key). None = key absent / file unreadable. _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} +# Native training context length (``{arch}.context_length``). None = absent / +# unreadable. Lets the UI show the real context ceiling before a model loads. +_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {} + def _cache_key(path: str) -> Optional[_CacheKey]: try: @@ -138,6 +142,92 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: return out +def read_gguf_context_length(path: str) -> Optional[int]: + """Return the GGUF's native training context length (``{arch}.context_length``), + or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size). + Lets the UI populate the context slider before the model is loaded.""" + key = _cache_key(path) + if key is None: + return None + with _CACHE_LOCK: + if key in _CONTEXT_CACHE: + return _CONTEXT_CACHE[key] + result = _parse_gguf_context_length(path) + with _CACHE_LOCK: + while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE))) + except StopIteration: + break + _CONTEXT_CACHE[key] = result + return result + + +def _parse_gguf_context_length(path: str) -> Optional[int]: + # The context key is architecture-namespaced (``llama.context_length`` etc.), + # so we learn the key only after reading ``general.architecture``. GGUF writes + # general.* before arch.* keys, matching the loader's own parser. + ctx_key: Optional[str] = None + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20: # 1 MB sanity bound + break + kbytes = f.read(klen) + if len(kbytes) < klen: + break + key = kbytes.decode("utf-8", "replace") + vt_bytes = f.read(4) + if len(vt_bytes) < 4: + break + vtype = struct.unpack(" 1 << 22: # 4 MB sanity bound + break + sbytes = f.read(slen) + if len(sbytes) < slen: + break + ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length" + elif ctx_key is not None and key == ctx_key and vtype in (4, 10): + width = 4 if vtype == 4 else 8 + n_bytes = f.read(width) + if len(n_bytes) < width: + break + value = struct.unpack(" 0 else None + else: + if not _skip_gguf_value(f, vtype): + break + except (struct.error, UnicodeDecodeError): + break + except OSError as e: + logger.debug(f"read_gguf_context_length: cannot open {path}: {e}") + return None + except Exception as e: + logger.debug(f"read_gguf_context_length: parse failure on {path}: {e}") + return None + return None + + # Strings (8) and arrays (9) are handled inline. _FIXED_VTYPE_SIZES: Dict[int, int] = { 0: 1, # uint8 diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 39e04843a9..29644d04cf 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -335,7 +335,7 @@ function TauriWrapper({ children }: { children: ReactNode }) { export function AppProvider({ children }: AppProviderProps) { return ( - + {children} diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 4792e78c4d..7fb9380037 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -119,6 +119,7 @@ function RootLayout() { chatRuntime.setActiveThreadId(null); chatRuntime.setActiveProjectId(null); chatRuntime.setIncognito(false); + if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel(); void navigate({ to: "/chat", search: { new: crypto.randomUUID() }, @@ -135,6 +136,7 @@ function RootLayout() { chatRuntime.setActiveProjectId(null); chatRuntime.setActiveThreadId(null); chatRuntime.setIncognito(false); + if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel(); }, [isChatRoute]); return ( diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx index dcd6617ec8..c623ef9848 100644 --- a/studio/frontend/src/app/routes/hub.tsx +++ b/studio/frontend/src/app/routes/hub.tsx @@ -14,6 +14,9 @@ const ModelsPage = lazy(() => export interface ModelsSearch { tab?: "discover" | "downloaded"; + model?: string; + section?: "trending" | "latest" | "finetune"; + kind?: "models" | "datasets"; } export const Route = createRoute({ @@ -22,8 +25,19 @@ export const Route = createRoute({ beforeLoad: () => requireAuth(), component: ModelsPage, validateSearch: (search: Record): ModelsSearch => { + const next: ModelsSearch = {}; const raw = search.tab; - if (raw === "discover" || raw === "downloaded") return { tab: raw }; - return {}; + if (raw === "discover" || raw === "downloaded") next.tab = raw; + const model = search.model; + if (typeof model === "string" && model.length > 0) next.model = model; + const section = search.section; + if (section === "trending" || section === "latest" || section === "finetune") { + next.section = section; + } + const kind = search.kind; + if (kind === "models" || kind === "datasets") { + next.kind = kind; + } + return next; }, }); diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 66accdce3a..8255d49de8 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -84,7 +84,7 @@ import { } from "@/components/ui/tooltip"; import { Tooltip as TooltipPrimitive } from "radix-ui"; import { HugeiconsIcon } from "@hugeicons/react"; -import { ChevronDown, MoreHorizontalIcon, Moon } from "lucide-react"; +import { ChevronDown, Moon } from "lucide-react"; import { Link, useNavigate, useRouterState } from "@tanstack/react-router"; import { archiveChatItem, @@ -226,7 +226,7 @@ function NavItem({ onClick={onClick} isActive={active} data-tour={dataTour} - className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:mx-auto" + className="sidebar-nav-btn h-[33px] rounded-[14px] gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:mx-auto" > {label} @@ -293,6 +293,7 @@ export function AppSidebar() { }; const isRecipesRoute = pathname.startsWith("/data-recipes"); + const isExportRoute = pathname === "/export" || pathname.startsWith("/export/"); const { displayTitle, avatarDataUrl } = useEffectiveProfile(); const { projects } = useChatProjects(); @@ -334,10 +335,14 @@ export function AppSidebar() { undefined : undefined; - // Training runs + // Training runs: surfaced as sidebar "Recents" on Train, Recipes, and Export, + // falling back to chat recents when there are no runs yet. + const trainingRecentsRoute = isStudioRoute || isRecipesRoute || isExportRoute; const { items: runItems } = useTrainingHistorySidebarItems( - !chatOnly && isStudioRoute, + !chatOnly && trainingRecentsRoute, ); + const showTrainingRecents = + !chatOnly && trainingRecentsRoute && runItems.length > 0; const activeJobId = useTrainingRuntimeStore((s) => s.jobId); const currentRunViewActive = useTrainingRuntimeStore((s) => s.currentRunViewActive); const selectedHistoryRunId = useTrainingRuntimeStore((s) => s.selectedHistoryRunId); @@ -667,7 +672,7 @@ export function AppSidebar() { ? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; const buttonClass = cn( - "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium", + "sidebar-nav-btn h-[33px] cursor-pointer rounded-[14px] pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium", // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the // title with the nav items above. variant === "project" ? "pl-[39px]" : "pl-3", @@ -921,7 +926,7 @@ export function AppSidebar() { @@ -1278,14 +1286,14 @@ export function AppSidebar() {
diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index efdef37e72..40fc8b8da6 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -4,6 +4,13 @@ "use client"; import { ArtifactCard, useChatRuntimeStore } from "@/features/chat"; +import { + getCodeFence, + isFullHtmlDocument, + isHtmlFence, + isRenderableRenderHtmlToolPart, + isSvgFence, +} from "@/features/chat/artifacts/html-fences"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { preprocessLaTeX } from "@/lib/latex"; import { openLink } from "@/lib/open-link"; @@ -45,62 +52,16 @@ const STREAMDOWN_COMPONENTS = { }; const COPY_RESET_MS = 2000; const MERMAID_SOURCE_RE = /```mermaid\s*([\s\S]*?)```/i; -const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/; const ACTION_PANEL_CLASS = "pointer-events-auto flex shrink-0 items-center gap-1"; const ACTION_BUTTON_CLASS = "flex size-8 cursor-pointer items-center justify-center rounded-[10px] text-chat-icon-fg transition-all hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover disabled:cursor-not-allowed disabled:opacity-50"; -type CodeFence = { - language: string | null; - source: string; -}; - -type ToolCallPartLike = { - type?: string; - toolName?: string; - args?: unknown; - result?: unknown; -}; - -function isRenderableRenderHtmlToolPart(part: unknown): boolean { - const toolPart = part as ToolCallPartLike; - if (toolPart.type !== "tool-call" || toolPart.toolName !== "render_html") { - return false; - } - if ( - typeof toolPart.result === "string" && - toolPart.result.startsWith("Error:") - ) { - return false; - } - if ( - typeof toolPart.result === "string" && - toolPart.result.startsWith("Rendered HTML canvas") - ) { - return true; - } - const args = toolPart.args as { code?: unknown } | undefined; - return typeof args?.code === "string" && args.code.trim().length > 0; -} - function getMermaidSource(blockContent: string): string | null { const source = blockContent.match(MERMAID_SOURCE_RE)?.[1]?.trim(); return source && source.length > 0 ? source : null; } -function getCodeFence(blockContent: string): CodeFence | null { - const match = blockContent.trimEnd().match(CODE_FENCE_RE); - if (!match) { - return null; - } - - return { - language: match[1]?.trim() || null, - source: match[2], - }; -} - function getCodeFilename(language: string | null) { const extByLanguage: Record = { bash: "sh", @@ -131,28 +92,6 @@ function getCodeFilename(language: string | null) { return `snippet.${ext}`; } -function isSvgFence(codeFence: CodeFence): boolean { - const lang = codeFence.language?.toLowerCase() ?? ""; - if (lang === "svg") return true; - if (lang === "xml" || lang === "html") { - const trimmed = codeFence.source.trimStart(); - // Match followed by ]/i.test(trimmed); -} - const UNSAFE_SVG_RE = /]|on\w+\s*=|javascript:|]|]|]|]/i; @@ -285,15 +224,13 @@ function CodeBlockActions({ ); } -// DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in -// thread.tsx) and has the HTML canvas feature on by default, so a full-HTML answer -// (e.g. a playable game) renders as an interactive card without the global toggle. +// Collapse a full-HTML answer in place into an artifact card. Diffusion keeps the +// raw code visible instead (the trailing MessageHtmlArtifacts appends its card). function StreamdownBlock(props: BlockProps) { const shouldCollapseHtmlArtifacts = useChatRuntimeStore( (state) => - state.artifactsEnabled || - state.collapseHtmlArtifacts || - state.loadedIsDiffusion, + (state.artifactsEnabled || state.collapseHtmlArtifacts) && + !state.loadedIsDiffusion, ); const messageHasRenderableRenderHtmlTool = useAuiState(({ message }) => message.parts.some(isRenderableRenderHtmlToolPart), diff --git a/studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx b/studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx new file mode 100644 index 0000000000..7555287211 Binary files /dev/null and b/studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx differ diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 80090bfd63..d17892e1bd 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -8,6 +8,8 @@ import { PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; +import { InfoHint } from "@/components/ui/info-hint"; +import { Switch } from "@/components/ui/switch"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { usePlatformStore } from "@/config/env"; import { isCustomProviderType } from "@/features/chat/external-providers"; @@ -108,6 +110,11 @@ interface ModelSelectorProps { activeGgufVariant?: string | null; onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; + /** When provided, renders a persisted "Load on selection" toggle in the + * popover. Off → picking a model stages it for a deferred, configured load + * instead of loading immediately. */ + loadOnSelection?: boolean; + onLoadOnSelectionChange?: (value: boolean) => void; onFoldersChange?: () => void; onPickLocalModel?: () => void | Promise; onModelsChange?: (deletedModel?: DeletedModelRef) => void; @@ -210,6 +217,8 @@ function ModelSelectorContent({ onPickLocalModel, onModelsChange, deleteDisabled, + loadOnSelection, + onLoadOnSelectionChange, className, dataTour, }: { @@ -223,6 +232,8 @@ function ModelSelectorContent({ onPickLocalModel?: () => void; onModelsChange?: (deletedModel?: DeletedModelRef) => void; deleteDisabled?: boolean; + loadOnSelection?: boolean; + onLoadOnSelectionChange?: (value: boolean) => void; className?: string; dataTour?: string; }) { @@ -378,6 +389,37 @@ function ModelSelectorContent({
) : null} + {onLoadOnSelectionChange ? ( +
+
+
+
+ Load on selection + +
+
+ On: load the model + immediately after selection. +
+
+ Off: configure options + first, then click Load model. +
+
+ +
+ + Local GGUF models only + +
+ +
+
+ ) : null} ); } @@ -395,6 +437,8 @@ export function ModelSelector({ onPickLocalModel, onModelsChange, deleteDisabled, + loadOnSelection, + onLoadOnSelectionChange, variant = "outline", size = "default", className, @@ -513,6 +557,8 @@ export function ModelSelector({ onPickLocalModel={onPickLocalModel ? handlePickLocalModel : undefined} onModelsChange={onModelsChange} deleteDisabled={deleteDisabled} + loadOnSelection={loadOnSelection} + onLoadOnSelectionChange={onLoadOnSelectionChange} className={contentClassName} dataTour={contentDataTour} /> diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 3165057021..6f11ec96f7 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1357,15 +1357,14 @@ export function HubModelPicker({ <> LM Studio {lmStudioModels.map((m) => { + const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); const isGguf = isGgufRepo(m.id) || isGgufRepo(m.display_name); const optionKey = makeModelOptionKey("lm-studio", m.id); return (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index 4cc5d779ce..6a86e4f7ed 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -31,6 +31,9 @@ export interface ModelSelectorChangeMeta { ggufVariant?: string; isDownloaded?: boolean; expectedBytes?: number; + /** Direct local .gguf file picked without a variant (custom folder / LM + * Studio). Marks it as a GGUF source for the deferred-load staging flow. */ + isGguf?: boolean; } export interface DeletedModelRef { diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index 3eaa21a19e..a6c326c64e 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -20,7 +20,9 @@ import { } from "@assistant-ui/react"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { type VariantProps, cva } from "class-variance-authority"; -import { CheckIcon, ChevronDownIcon, CopyIcon, LightbulbIcon } from "lucide-react"; +import { ChevronDownIcon, CopyIcon, LightbulbIcon } from "lucide-react"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { HugeiconsIcon } from "@hugeicons/react"; import { type CSSProperties, type ComponentProps, @@ -293,7 +295,7 @@ function ReasoningCopyButton({ startIndex, endIndex }: { startIndex: number; end aria-label="Copy reasoning" > {copied ? ( - + ) : ( )} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index ee3754bdf6..a65a157707 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -11,6 +11,7 @@ import { } from "@/components/assistant-ui/generated-image-overlay-context"; import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources"; @@ -161,6 +162,7 @@ import { useState, } from "react"; import { create } from "zustand"; +import { extractTaggedText, updateThreadMessage } from "@/features/chat/utils/update-thread-message"; // True while a file is dragged anywhere over the chat page, so the composer // can show its "Drop files here" affordance. @@ -3194,40 +3196,138 @@ const DiffusionCanvas: FC = () => { ); }; +/** + * AssistantMessage handles the display and inline-editing of AI responses. + * + * It utilizes a "Tagged Text" system ( and tags) to allow users + * to edit structured reasoning and tool outputs within a plain-text textarea + * while preserving the underlying data schema and tool-call metadata. + */ const AssistantMessage: FC = () => { + const aui = useAui(); + const messageId = useAuiState(({ message }) => message.id); + const messageContent = useAuiState(({ message }) => message.content); + const incognito = useChatRuntimeStore((s) => s.incognito); + + // Use global store for editing state to ensure a single source of truth + const editingId = useChatRuntimeStore((s) => s.editingMessageId); + const setEditingId = useChatRuntimeStore((s) => s.setEditingMessageId); + const isEditing = editingId === messageId; + + const textareaRef = useRef(null); + + // Auto-grow textarea height based on content + const adjustHeight = () => { + const el = textareaRef.current; + if (el) { + el.style.height = "auto"; + el.style.height = `${el.scrollHeight}px`; + } + }; + + useEffect(() => { + if (isEditing) setTimeout(adjustHeight, 0); + }, [isEditing]); + + const handleSave = async () => { + const finalText = textareaRef.current?.value || ""; + + // Prioritize the specific thread item ID, then fallback to the global active thread ID + const remoteId = aui.threadListItem().getState().remoteId + || useChatRuntimeStore.getState().activeThreadId; + + if (!remoteId || remoteId === "" || remoteId === "/") { + toast.error("Save failed: No thread ID found."); + setEditingId(null); + return; + } + + try { + await updateThreadMessage({ + thread: { + export: () => aui.thread().export(), + import: (data) => aui.thread().import(data) + }, + messageId, + remoteId, + newText: finalText, + isIncognito: incognito, + }); + } catch (error) { + console.error("UI: Error during save:", error); + toast.error("Failed to save message edits."); + } finally { + setEditingId(null); + } + }; + return (
- - - - - - - + {isEditing ? ( +
+